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

2025-07-21

AI Summary

Existing medical image segmentation methods uniformly penalize all pixel-level errors, ignoring semantic hierarchical relationships among anatomical or pathological labels—thereby limiting fine-grained, multi-class segmentation performance. To address this, we propose a label-hierarchy-tree-based semantic-aware segmentation framework. We design two novel tree-structured semantic loss functions that explicitly incorporate inter-class hierarchical semantics into loss computation for the first time. Furthermore, we integrate a sparse-label training mechanism to enhance discriminative capability for fine-grained classes under limited annotation budgets. Our method is architecture-agnostic and compatible with mainstream segmentation backbones. It achieves state-of-the-art performance on two clinically relevant tasks: whole-brain parcellation in cranial MRI and hyperspectral surgical scene understanding in neurosurgery. Empirical results demonstrate significant improvements in multi-class segmentation accuracy and clinical utility.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Multi-class/Multi-label Learning & Extreme Classification
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for heterogeneous, signed, attributed, multi-relational, temporal, higher-order, and annotated Web-related graphs

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.

Problem

• Exploiting inter-class semantics in medical image segmentation
• Addressing errors in high-cardinality label spaces
• Enhancing segmentation with hierarchical label organization

Innovation

• Tree-based semantic loss functions for segmentation
• Hierarchical label organization for rich classes
• Integration with sparse annotation training