LesionDiffusion: Towards Text-controlled General Lesion Synthesis

Probleemstelling

Voltooid toezicht (fully-supervised) methoden voor het herkennen van laesies in medische beeldvorming, zoals CT-scans, kampen met grote uitdagingen. Deze methoden zijn sterk afhankelijk van grote, handmatig geannoteerde datasets, die vanwege privacyzorgen en de hoge kosten van annotatie door radiologen schaars en duur zijn om te verzamelen.

Bestaande oplossingen voor synthetische laesiegeneratie (bijvoorbeeld met GANs of eerdere diffusiemodellen) hebben drie belangrijke beperkingen:

1. Beperkte schaalbaarheid: Ze trainen vaak per laesietype apart, wat generalisatie naar nieuwe organen of laesietypen bemoeilijkt.
2. Gebrek aan controle: Er is weinig fijnmazige controle over de attributen van de laesie (zoals vorm, dichtheid of locatie).
3. Complexiteit: Bestaande modellen zijn vaak ontworpen voor solide, ronde laesies en hebben moeite met het genereren van complexe structuren (zoals holle laesies bij maag- of darmkanker) of realistische texturen.

Methodologie: LesionDiffusion

De auteurs stellen LesionDiffusion voor, een framework voor tekstgestuurde, generatieve synthese van laesies in 3D CT-beelden. Het systeem genereert zowel de laesie zelf als het bijbehorende segmentatiemasker, geleid door gestructureerde medische rapporten.

Het framework bestaat uit twee hoofdcomponenten en een dataset:

1. Dataset:
De auteurs hebben een nieuwe dataset samengesteld van 1.505 geannoteerde CT-scans met gepaarde laesiemaskers en gestructureerde rapporten. Deze dekt 14 verschillende laesietypen over 8 organen (o.a. long, lever, nieren, pancreas, maag, darmen). De gestructureerde rapporten bevatten 10 attributen, onderverdeeld in:

• Masker-attributen: Vorm, locatie, relatie met omliggende organen.
• Beeld-attributen: Dichtheid, textuur, versterkingsstatus.

2. Architectuur:
Het proces verloopt in twee fasen:

• Fase I: Lesion Mask Synthesis Network (LMNet)

  • Doel: Genereren van een nauwkeurig laesiemasker op basis van de tekstuele attributen (vorm, locatie, orgaan).
  • Techniek: Een diffusiemodel dat gebruikmaakt van een b-box-gewogen verliesfunctie (bounding box weighted loss). Dit dwingt het model om zich te focussen op het specifieke gebied binnen de geschatte bounding box.
  • Conditioning: Tekstuele attributen worden ingebed via een vooraf getrainde tekst-encoder (BiomedCLIP) en via cross-attention modules aan het diffusiemodel doorgegeven.

• Fase II: Lesion Inpainting Network (LINet)

  • Doel: Het "inpainten" van de laesie in de CT-scan om realistische texturen en dichtheden te creëren.
  • Techniek: Gezien het hoge geheugengebruik van 3D CT-data, wordt een Latent Diffusion Model gebruikt. Een paar auto-encoders (ECT, DCT) comprimeren de 3D-beelden naar een latentere ruimte.
  • Proces: Het diffusieproces vindt plaats alleen binnen het gebied dat door het gegenereerde masker is gedefinieerd. Het model gebruikt beeld-attributen (zoals "heterogene dichtheid") om de laesie realistisch te reconstrueren.

3. Inference (Toepassing):
Tijdens inferentie kan een gebruiker een gestructureerd rapport genereren (handmatig of via een LLM). Op basis hiervan worden de bounding box en het masker gegenereerd (via LMNet), waarna de daadwerkelijke laesie in het beeld wordt ingevuld (via LINet).

Belangrijkste Bijdragen

1. Eerste tekstgestuurde, algemene framework: Het introduceert het eerste model dat tekstuele instructies gebruikt om generieke laesies in 3D CT te synthetiseren, in plaats van gespecialiseerde modellen per ziektebeeld.
2. Gestructureerde rapporten als conditionering: Door gestructureerde medische rapporten te gebruiken als input, biedt het model fijnmazige controle over zowel de vorm als de visuele eigenschappen van de laesie.
3. Nieuwe dataset: Publicatie van een groot, divers dataset met 1.505 scans, 14 laesietypen en gestructureerde annotaties, wat een knelpunt in de huidige literatuur wegneemt.
4. Twee-staps architectuur: Een effectieve scheiding tussen maskergeneratie (LMNet) en beeldinpainting (LINet) die complexiteit en realisme verbetert.

Resultaten

De prestaties zijn geëvalueerd op basis van de verbetering van segmentatiemodellen (nnUNet-v2) die getraind zijn met synthetische data.

• Verbeterde Segmentatie: LesionDiffusion leidt tot een aanzienlijke stijging in de Dice Similarity Coefficient (DSC). Modellen getraind met LesionDiffusion-data presteren beter dan die getraind met data van concurrenten zoals DiffTumor.
• Generalisatie naar onbekende taken: Een cruciaal resultaat is de generalisatie naar onbekende organen en laesietypen die niet in de trainingsset zaten (bijv. hersenbloedingen/hematencephalon).
  • Zelfs zonder hersen-scans in de training, presteerde het model getraind op LesionDiffusion-data aanzienlijk beter dan bestaande methoden.
  • Na fijnafstemming (fine-tuning) met echte data, overtrof het model dat getraind was op LesionDiffusion-data zelfs modellen die alleen op echte data waren getraind (met +2.4% DSC).
• Controleerbaarheid: Er is een sterke correlatie aangetoond tussen de tekstuele prompt (bijv. "ronde vorm" vs. "irreguliere vorm") en de gegenereerde morfologische eigenschappen (convexiteit en sfericiteit).

Significantie

LesionDiffusion biedt een schaalbare oplossing voor het tekort aan geannoteerde medische data. Door de mogelijkheid om willekeurige laesies met specifieke attributen te genereren, kan het worden gebruikt als een krachtige data-augmentatiestrategie. Dit verbetert niet alleen de nauwkeurigheid van bestaande diagnostische modellen, maar maakt het ook mogelijk om modellen te trainen voor zeldzame ziektebeelden waarvoor weinig tot geen echte data beschikbaar is. Het paper markeert een stap in de richting van een fundamenteel model voor controllable laesiegeneratie in de medische beeldvorming.