OOD-SEG: Exploiting out-of-distribution detection techniques for learning image segmentation from sparse multi-class positive-only annotations

1. Probleemstelling

Het artikel adresseert twee fundamentele uitdagingen in de medische beeldsegmentatie met diepe neurale netwerken:

• Tekort aan gedetailleerde annotaties: Het verkrijgen van volledige, pixel-voor-pixel segmentatielabels is tijdrovend, kostbaar en vereist domeinexpertise. In de praktijk zijn annotaties vaak "spaars" (sparse) en beperken ze zich tot positieve voorbeelden (bijv. specifieke organen of weefsels), zonder annotaties voor de achtergrond of andere niet-gemarkeerde gebieden.
• Gebrek aan Out-of-Distribution (OOD) detectie: Traditionele segmentatiemodellen gaan ervan uit dat testdata uit dezelfde verdeling komt als de trainingsdata (In-Distribution of ID). Ze kunnen echter niet detecteren wanneer een pixel behoort tot een onbekende klasse of achtergrond (OOD). Dit leidt tot foutieve voorspellingen in veilige klinische omgevingen, waar het systeem moet kunnen waarschuwen voor onzekere situaties.

Bestaande methoden voor "Weakly Supervised Learning" (WSL) behandelen ongemarkeerde pixels vaak ten onrechte als achtergrond (negatief), wat leidt tot ruis in de labels. Het artikel stelt dat dit een Positive-Unlabelled (PU) leerprobleem is, waarbij ongemarkeerde pixels zowel positief als negatief kunnen zijn.

2. Methodologie: OOD-SEG Framework

De auteurs stellen een nieuw raamwerk voor, OOD-SEG, dat PU-learning combineert met technieken voor OOD-detectie.

• Leerparadigma: Het model wordt getraind uitsluitend op spaars geannoteerde, positieve pixels van meerdere klassen (ID-classes). Er worden geen labels gebruikt voor de achtergrond. Ongemarkeerde pixels worden niet als achtergrond behandeld, maar als potentiële OOD-data.
• Architectuur: De basis is een U-Net met een EfficientNet-B4 encoder. Tijdens het trainingsstadium wordt alleen de Cross-Entropy Loss berekend op de geannoteerde positieve pixels.
• Inferentie en OOD-detectie: Tijdens de inferentie wordt een pixel-voor-pixel scorefunctie ($S$) gebruikt om de waarschijnlijkheid te bepalen of een pixel tot een ID-klasse behoort of tot de OOD-klasse (achtergrond/onbekend).
  • Als de maximale score boven een drempelwaarde ($\tau$) ligt, wordt de pixel toegewezen aan de beste ID-klasse.
  • Is de score lager dan $\tau$, dan wordt de pixel gemarkeerd als OOD (achtergrond).
• Geïntegreerde OOD-methoden: Het framework is modulair en kan verschillende OOD-detectiemethoden integreren:
  1. Baseline: Softmax-score van het netwerk.
  2. ODIN: Gebruikt temperatuurschaling (temperature scaling) om de betrouwbaarheid van de scores te kalibreren.
  3. Mahalanobis: Berekent de afstand van de features tot de klassenspecifieke verdelingen (Gaussian) in de feature space.
  4. GODIN (Generalised ODIN): Een aangepaste, convolutie-gebaseerde versie van de GODIN-methode die leert een temperatuurschaling en een similariteitsmaat per pixel, specifiek ontworpen voor segmentatietaken.

3. Evaluatiestrategie

Een belangrijke innovatie is de evaluatiemethode, aangezien er geen bestaande benchmarks zijn voor OOD-detectie in medische segmentatie met spaarse annotaties.

• Twee-niveau Cross-validatie: In plaats van een apart OOD-testset te gebruiken (wat lastig is om te verkrijgen), houden de auteurs tijdens de training bepaalde geannoteerde klassen achter. Deze "gehouden" klassen fungeren als OOD-data voor die specifieke fold.
  • Subject-partities (SP): Zorgen voor geen overlap van patiënten tussen train- en testset.
  • Klasse-partities (CP): Zorgen ervoor dat bepaalde klassen niet worden getraind en dus als OOD worden behandeld.
• Metingen: Er worden twee drempel-onafhankelijke metrieken gebruikt:
  • AUROC: Gebaseerd op True Positive Rate (voor ID-classes) en False Positive Rate (voor OOD-classes).
  • AUPR: Area Under the Precision-Recall curve.

4. Belangrijkste Resultaten

De methode is getest op drie datasets: twee hyperspectrale datasets (Heiporspectral en ODSI-DB) en één RGB-laparoscopische dataset (DSAD).

• Prestatie: De ODIN-methode leverde over het algemeen de beste prestaties op in termen van AUROC en AUPR. Dit suggereert dat goed gekalibreerde confidence-scores cruciaal zijn voor het onderscheiden van OOD-data.
• Vergelijking:
  • De Mahalanobis-methode presteerde minder goed, waarschijnlijk door de hoge dimensionaliteit van de feature space die leidt tot slecht geconditioneerde covariantiematrices.
  • De GODIN-methode (met convolutielagen) was effectief, maar ODIN bleef iets beter.
• Kwaliteit: Visuele inspectie toonde aan dat het framework in staat is om OOD-pixels (achtergrond/onbekend) correct te identificeren zonder de nauwkeurigheid voor de bekende (ID) klassen te verstoren.
• Robuustheid: Het framework bleek robuust over verschillende beeldmodi (HSI en RGB) en verschillende anatomische contexten.

5. Bijdragen en Significantie

De belangrijkste bijdragen van dit werk zijn:

1. Nieuw Raamwerk: De eerste formulering van multi-class medische beeldsegmentatie met spaarse positieve-only annotaties als een PU-learning probleem, opgelost met pixel-voor-pixel OOD-detectie.
2. Evaluatieprotocol: Een innovatieve tweeniveau cross-validatiestrategie die de noodzaak elimineert voor extra OOD-testsets, wat een grote barrière voor onderzoek in dit domein wegneemt.
3. Technische Innovatie: Een nieuwe convolutie-gebaseerde aanpassing van de GODIN-methode, die de toepasbaarheid van classificatie-gebaseerde OOD-methoden uitbreidt naar segmentatietaken.

Significantie:
Dit werk biedt een praktische oplossing voor het veilig implementeren van geautomatiseerde segmentatiesystemen in de kliniek. Door het vermogen om onbekende weefsels of achtergrond te detecteren (zonder dat deze expliciet zijn geannoteerd), verhoogt het de veiligheid en betrouwbaarheid van AI-systemen. Het stelt ook een nieuwe standaard voor de evaluatie van OOD-detectie in medische beeldvorming, wat essentieel is voor de volgende generatie diagnostische hulpmiddelen.