MedShift: Implicit Conditional Transport for X-Ray Domain Adaptation

Probleemstelling

Het trainen van robuuste medische beeldanalysemodellen vereist grote datasets. Synthetische data biedt een schaalbare oplossing, maar er bestaat een aanzienlijke domeinkloof (domain gap) tussen synthetische en echte klinische X-ray beelden. Deze kloof beperkt de generaliseerbaarheid van modellen die op synthetische data zijn getraind naar de echte wereld. De belangrijkste verschillen die deze kloof veroorzaken zijn:

1. Attenuatieprofielen: Synthetische beelden repliceren de complexe verzwakking van X-stralen door heterogene anatomische structuren (bot, luchtholtes, zacht weefsel) vaak niet volledig.
2. Ruiskenmerken: Echte X-ray-systemen introduceren gestructureerde ruis, verstrooiing en compressie-artefacten die in simulatoren ontbreken.
3. Weefselrepresentatie: Er zijn verschillen in de weergave van zacht weefsel en contrastdynamiek, vooral op de grenzen van botstructuren.

Bestaande methoden voor beeld-naar-beeld vertaling (zoals GANs of Diffusiemodellen) vereisen vaak gepaarde datasets (die zeldzaam zijn in de medische wereld) of moeten voor elk domeinpaar apart worden getraind, wat schaalbaarheid beperkt.

Methodologie: MedShift

De auteurs stellen MedShift voor, een unificerend, class-conditional generatief model dat ongepaarde beeldvertaling tussen meerdere domeinen mogelijk maakt.

• Kernarchitectuur: MedShift combineert Flow Matching (FM) en Schrödinger-bruggen. In plaats van een stochastisch proces (zoals bij diffusiemodellen), parameteriseert het model een tijdsafhankelijke vectorveld dat monsters deterministisch transporteert van een basisverdeling naar de doelverdeling via een gewone differentiaalvergelijking (ODE).
• Latente Ruimte: Het model opereert in een gedeelde, domein-agnostische latente ruimte die wordt geleerd via een vooraf getrainde VAE (Variational Autoencoder). Dit verlaagt de rekenkosten aanzienlijk terwijl semantische details behouden blijven.
• Inference Proces (Twee-staps):
  1. Encodering: Een bronbeeld (bijv. synthetisch) wordt via tijd-integratie achterwaarts geëncodeerd naar een tussenliggende latente representatie $z_\tau$ (waarbij $\tau \in (0,1)$). Deze staat ligt op een gedeeld manifold dat voor alle domeinen ongeveer uitgelijnd is.
  2. Vertaling: Vervolgens wordt voorwaarts geïntegreerd van $\tau$ naar 1, nu geconditioneerd op het label van het doel-domein (bijv. "echt"). Dit genereert het vertaalde beeld.
• Flexibiliteit: Het model gebruikt Classifier-Free Guidance (CFG). Tijdens inferentie kan de gebruiker de parameter $\tau$ en de CFG-schaal aanpassen om de balans te sturen tussen het behoud van de oorspronkelijke anatomische structuur en het overnemen van de stijl (ruis, contrast) van het doel-domein, zonder het model opnieuw te hoeven trainen.

Belangrijkste Bijdragen

1. MedShift Model: Een nieuw model gebaseerd op Flow Matching en Schrödinger-bruggen voor hoge-resolutie, ongepaarde beeldvertaling. Het leert één gedeelde transportkaart voor meerdere domeinen.
2. X-DigiSkull Dataset: De publicatie van een nieuwe dataset met uitgelijnde synthetische en echte schedel-X-rays.
   • Synthetisch: Genereerd met de Mentice VIST® simulator (Low/High dosering).
   • Echt: Opgehaald met een Philips Azurion IGT-systeem op een fysieke schedelfantoom (Low/Normal/Exposure dosering).
   • De dataset bevat meerdere hoeken en doseringen, wat een robuust benchmark biedt voor domeinadaptatie.
3. Benchmarking: Een uitgebreide evaluatie van state-of-the-art generatieve modellen (GANs, Normalizing Flows, Diffusiemodellen) op deze dataset.

Resultaten

De experimenten tonen aan dat MedShift een superieure balans biedt tussen realisme en structurele integriteit:

• Prestatie: MedShift presteert consistent beter dan bestaande methoden zoals CycleGAN-Turbo, Z-STAR en SDEdit op de gemiddelde rangorde van realisme (CFID, Coverage) en structuurbehoud (LPIPS, SSIM).
• Aanpassingsvermogen:
  • Bij een hoge $\tau$ (bijv. 0.6) behoudt MedShift de anatomische structuur uitstekend (vergelijkbaar met Hierarchy Flow) maar met veel betere domein-aanpassing.
  • Bij een lage $\tau$ (bijv. 0.3) bereikt het hoge realisme (vergelijkbaar met CycleGAN) maar met minder anatomische vervormingen dan concurrenten.
• Efficiëntie: MedShift is aanzienlijk lichter dan diffusie-gebaseerde modellen. Het gebruikt een aangepaste U-Net die zes keer kleiner is dan de Stable Diffusion U-Nets die door concurrenten worden gebruikt. Dit resulteert in een lagere VRAM-gebruik en snellere trainingstijden, terwijl het toch meerdere inferentiestappen vereist (in tegenstelling tot single-pass GANs).
• Kwaliteit: De vertaalde beelden tonen realistischere randen, subtielere intensiteitsvariaties binnen het bot en hersteld zacht weefsel (zoals de hoofdhuid), wat ontbreekt in de synthetische bron.

Betekenis en Toekomstperspectief

MedShift biedt een schaalbare en generaliseerbare oplossing voor domeinadaptatie in de medische beeldvorming. Het lost het probleem op van het ontbreken van gepaarde data en de noodzaak van domeinspecifieke modellen. Door de flexibiliteit om tijdens inferentie te kiezen tussen structuurbehoud of stijltransformatie, is het ideaal voor klinische toepassingen zoals chirurgische navigatie en training.

Toekomstig werk richt zich op het verbeteren van de inferentie-efficiëntie (bijv. via modeldistillatie), het uitbreiden naar multi-class scenario's (zoals doseringsstandaardisatie) en het integreren van extra conditionele inputs (zoals ruimtelijke maskers) voor nog meer controle. De code en dataset zijn openbaar beschikbaar gesteld.