Label-set Loss Functions for Partial Supervision: Application to Fetal Brain 3D MRI Parcellation (2107.03846v2)

Abstract: Deep neural networks have increased the accuracy of automatic segmentation, however, their accuracy depends on the availability of a large number of fully segmented images. Methods to train deep neural networks using images for which some, but not all, regions of interest are segmented are necessary to make better use of partially annotated datasets. In this paper, we propose the first axiomatic definition of label-set loss functions that are the loss functions that can handle partially segmented images. We prove that there is one and only one method to convert a classical loss function for fully segmented images into a proper label-set loss function. Our theory also allows us to define the leaf-Dice loss, a label-set generalization of the Dice loss particularly suited for partial supervision with only missing labels. Using the leaf-Dice loss, we set a new state of the art in partially supervised learning for fetal brain 3D MRI segmentation. We achieve a deep neural network able to segment white matter, ventricles, cerebellum, extra-ventricular CSF, cortical gray matter, deep gray matter, brainstem, and corpus callosum based on fetal brain 3D MRI of anatomically normal fetuses or with open spina bifida. Our implementation of the proposed label-set loss functions is available at https://github.com/LucasFidon/label-set-loss-functions

• The paper's main contribution is the formulation of label-set loss functions, including the innovative leaf-Dice loss, tailored for partially supervised medical image segmentation.
• It develops a theoretical framework converting traditional loss functions to work with incomplete annotations, ensuring predicted distributions align with partially labeled ground truth.
• Empirical tests on fetal brain 3D MRI data demonstrate enhanced segmentation accuracy with improved Dice scores and reduced Hausdorff distances compared to conventional methods.

Label-set Loss Functions for Partial Supervision: Insights and Applications in Fetal Brain 3D MRI Parcellation

The paper "Label-set Loss Functions for Partial Supervision: Application to Fetal Brain 3D MRI Parcellation" presents significant theoretical contributions and practical applications in the domain of medical image analysis, specifically targeting the challenges posed by partially annotated datasets in the field of fetal brain MRI segmentation. Recognizing the constraints associated with manually segmenting extensive medical imaging datasets, the authors propose novel methodologies to leverage partial supervision effectively.

Theoretical Contributions

The authors introduce a theoretical framework to address the limitations of existing loss functions when dealing with partial supervision in image segmentation tasks. The core contribution is the definition and axiomatic establishment of "label-set loss functions." These functions are specifically crafted to ensure compatibility between label-sets and individual leaf-labels, accounting for partially segmented images where some regions of interest might not be annotated. The paper proves that there is a singular method to convert classical loss functions designed for fully segmented images into suitable label-set loss functions for partially segmented images.

Crucially, the authors propose the "leaf-Dice loss," a generalization of the Dice loss, tailored for scenarios with missing labels. This loss function is grounded in the theoretical framework of label-set loss functions and shows substantial improvements over previous methods. It ensures that predicted leaf-label probability distributions align with the partially labeled ground truth, leveraging the provided annotations efficiently while accounting for the uncertainty of unannotated regions.

Empirical Validation with Fetal Brain MRI Data

The paper's authors apply their theoretical advancements in the domain of fetal brain 3D MRI segmentation, a task critical for understanding fetal brain development and assisting in the diagnosis of central nervous system pathologies. They employ a dataset comprising fetal brain MRI scans from multiple clinical centers, including scans of fetuses with open spina bifida, to validate their approach. Utilizing 3D U-Net architectures, the paper demonstrates that models trained with label-set loss functions markedly outperform those utilizing traditional fully supervised losses, as evidenced by improved Dice scores and reduced Hausdorff distances for various brain tissues.

Implications and Future Directions

The research findings have significant implications for both theoretical and practical aspects of artificial intelligence in medical imaging. By providing a rigorous method for employing partially supervised training, the paper offers a pathway to harness the vast arrays of unannotated medical imaging data without necessitating exhaustive manual annotation. This has practical implications for clinical research where obtaining fully annotated datasets is often infeasible.

The successful application of the leaf-Dice loss function suggests that similar approaches might be developed for other domains within medical imaging and beyond, potentially leading to refined models capable of handling incomplete and noisy datasets. Future research could explore extending this framework to other imaging modalities and tasks, as well as enhancing the computational efficiency and scalability of such approaches when integrated into large-scale, automated systems.

In conclusion, this paper makes a significant contribution to the field of partially supervised deep learning. It lays a robust theoretical foundation while delivering practical solutions with demonstrable improvements in fetal brain segmentation tasks. These advancements open avenues for enhanced automation and accuracy in medical image analysis, suggesting a promising direction for future AI developments in healthcare.