Label tree semantic losses for rich multi-class medical image segmentation

TL;DR

This work tackles the problem of richly labeled medical image segmentation by leveraging explicit label hierarchies. It introduces two semantically informed losses: a Wasserstein distance-based loss on label space and a tree-weighted semantic cross-entropy loss, both integrated with a sparse positive-only annotation framework via out-of-distribution detection. Empirical results demonstrate state-of-the-art performance on whole-brain parcellation from MRI and neurosurgical hyperspectral imaging, with notable gains for small classes and improved semantic discrimination between related tissues. The approach offers a scalable way to encode label semantics into segmentation models, with practical implications for preoperative planning and intraoperative navigation.

Abstract

Rich and accurate medical image segmentation is poised to underpin the next generation of AI-defined clinical practice by delineating critical anatomy for pre-operative planning, guiding real-time intra-operative navigation, and supporting precise post-operative assessment. However, commonly used learning methods for medical and surgical imaging segmentation tasks penalise all errors equivalently and thus fail to exploit any inter-class semantics in the labels space. This becomes particularly problematic as the cardinality and richness of labels increases to include subtly different classes. In this work, we propose two tree-based semantic loss functions which take advantage of a hierarchical organisation of the labels. We further incorporate our losses in a recently proposed approach for training with sparse, background-free annotations to extend the applicability of our proposed losses. Extensive experiments are reported on two medical and surgical image segmentation tasks, namely head MRI for whole brain parcellation (WBP) with full supervision and neurosurgical hyperspectral imaging (HSI) for scene understanding with sparse annotations. Results demonstrate that our proposed method reaches state-of-the-art performance in both cases.

Figures (7)

• Figure 1: The neuro-anatomical label hierarchy of Mindboggle dataset. From left to right, the hierarchy progresses from coarse object categories to specific classes. Rich annotations correspond to leaf node classes. The colour coding matches the ground-truth mask at each level. of this tree is included as supplementary material (Figure \ref{['fig:hierarchy_mindboggle_large']}).
• Figure 2: Visual comparison of the baseline loss and the proposed Wasserstein-based loss on the AOMIC dataset. Each column shows the predicted segmentation masks at progressively finer levels of the label hierarchy. The white arrow marks the challenging class non-WM hypointensities, which the Wasserstein-based loss segments correctly, whereas the baseline fails to capture it..
• Figure 3: Qualitative result on top-level classes. We show the result of same image using different methods at confidence threshold $\tau_m$. Baseline results at $\tau_0=0$ are added to represent result without outlier detection.
• Figure 4: Confusion matrices for the WBP and HSI tasks. For WBP, the evaluation is on 10 small classes of the AOMIC dataset. Class names from top-left to bottom-right: Left-Inf-Lat-Vent, Left-vessel, Left-choroid-plexus, Right-vessel, Right-choroid-plexus, 5th-Ventricle, WM-hypointensities, non-WM-hypointensities, Optic-Chiasm, ctx-lh-unknown, ctx-rh-unknown. For HSI, the evaluation is on top-level nodes. Class names from top-left to bottom-right: Other, Out-of-focus Area, Vascular, Normal Tissue, Abnormal Tissue and Surgical Equipment.
• Figure 5: Sensitivity analysis result by changing scaling parameter $\kappa$ on HSI dataset. On each trial model is evaluated on a four-fold cross-validation.