Boundary Loss in Highly Unbalanced Segmentation
Purchased at Istockphoto.com. Copyright.
Semantic Image Segmentation and Imbalanced Tasks
Semantic image segmentation—consisting of automatically drawing the contour of an object in an image—is one of the most active topics of research in the medical imaging community. The resulting segmentation can then be used in many downstream applications, such as diagnosis, surgery planning or the follow-up of many diseases.
One difficulty that arises in many medical applications, such as brain tumor segmentation, is the imbalance between the healthy part of the image (the “background”) and the lesions—being several orders of magnitude smaller, it is akin to finding a needle in a haystack.
Figure 1: Illustration of the imbalance problem: the vast majority of pixels belong to the background, whereas we are interested in segmenting the small brain lesions.
Nowadays, semantic segmentation is mostly automated using artificial neural networks—mathematical constructs that mimic a brain model. Those networks are “trained” with a series of examples that contain both the input image and the contours of the objects it should predict (the “ground truth”). In the case of imbalanced tasks, the usual methods to train them become unreliable, and networks tend to predict all pixels as background making the lesions disappear. Several previous works have tried to improve the way training is done for imbalanced tasks, by changing the “weight” of each kind of pixel—making the lesions more “important”—but results are still unsatisfying.
Our Contribution: Minimizing the Distance between Boundaries
Our contribution—inspired by traditional works in computer vision—is to handle the training, not as a series of pixel predictions, but rather as predicting a contour that should match the ground-truth contour. As such, we attempt to minimize the distance between both contours until they overlap perfectly.
Minimizing the distance between contours (q) is not affected by the aforementioned imbalance; only what happens at the boundary matters. The way this distance is minimized at training involves reformulating the distance calculation as a sum of independent pixel multiplications (details are available in the paper). The end result is that our way of computing this distance is very efficient. This new boundary loss can then be combined to existing losses, to improve and rectify their shortcomings.
Results
Combining our boundary loss to other losses proved to both improve absolute results—helping to recover small lesions as shown in the examples—and to make the training more stable. This is a very useful property, as practitioners need to know when to stop training their networks and when to continue—difficult to decide when it is unstable.
Adoption by the Community
Due to its efficiency, good results, and ease of use the loss has been adopted by a lot of practitioners in the community, as shown by the number of citations, high-number of stars on GitHub, and the active community around it.
Additional Information
More technical details can be found in the full paper, describing different settings across several datasets:
Kervadec H.; Bouchtiba, J.; Desrosiers, C.; Granger, E.; Dolz, J.; Ben Ayed, I. 2021. Boundary Loss for Highly Unbalanced Segmentation. Oral presentation at Medical Imaging with Deep Learning (MIDL) 2019, London. Runner-up for best paper award. Invited for a deep learning special issue in Medical Image Analysis (MEDIA), vol 67.
The code is available online and is free for reuse and modification: https://github.com/LIVIAETS/boundary-loss
