MagicSeg: Open-World Segmentation Pretraining via Counterfactural Diffusion-Based Auto-Generation

Probleemstelling

Open-wereld semantische segmentatie (open-world semantic segmentation) heeft tot nu toe sterk afhankelijk gestaan van uitgebreide datasets met afbeelding-tekstparen. Een groot nadeel hiervan is het gebrek aan fijnmazige, pixel-level annotaties voor voldoende categorieën. Het verzamelen van dergelijke data is economisch onhaalbaar vanwege de enorme investeringen in menselijke arbeid en tijd die nodig zijn voor handmatige annotatie. Bestaande synthetische methoden (zoals GANs of eerdere diffusion-modellen) hebben vaak nog steeds een deel van echte, gelabelde data nodig of genereren ruisachtige maskers via attentie-mechanismen, wat niet ideaal is voor open-wereld scenario's waar een breed scala aan categorieën vereist is.

Methodologie: MagicSeg

MagicSeg introduceert een nieuw, gestroomlijnd pipeline om datasets voor open-wereld segmentatie automatisch te genereren. De aanpak bestaat uit twee hoofdfasen:

1. Counterfactual Diffusion-based Generation (Datageneratie)

In plaats van alleen positieve voorbeelden te genereren, bouwt MagicSeg een dataset op die bestaat uit vier elementen: tekst, afbeelding, masker en een counterfactual afbeelding.

• Tekstgeneratie: Uit een grote vocabulaire van categorieën (1205 categorieën) worden willekeurig 1-2 namen geselecteerd. Een Large Language Model (ChatGPT) genereert hierop rijke, gedetailleerde beschrijvingen die de achtergrond en context divers maken, zodat de gegenereerde afbeeldingen niet te monotoon zijn.
• Afbeeldinggeneratie: Een diffusion model (Stable Diffusion) genereert de afbeeldingen op basis van deze gedetailleerde prompts.
• Counterfactual Pairs: Om het model te leren objecten te lokaliseren zonder afhankelijk te zijn van perfecte maskers, wordt de categorie-naam in de tekst vervangen door "nothing" (bijv. "een hond" wordt "niets"). Dit resulteert in een counterfactual afbeelding die qua achtergrond identiek is, maar het doelobject mist.
• Automatische Annotatie: Voor de originele afbeeldingen worden objectmaskers gegenereerd door eerst een open-vocabulary detector (Grounding DINO) te gebruiken om bounding boxes te vinden, en vervolgens een interactief segmentatiemodel (SAM - Segment Anything Model) om de precieze pixel-maskers te extraheren.

2. Training van het Open-World Segmentatiemodel

Het gegenereerde dataset wordt gebruikt om een CLIP-gebaseerd segmentatiemodel (ZegCLIP) te trainen met twee innovatieve strategieën:

• Categorie Random Sampling: Omdat de dataset veel meer categorieën bevat dan er per afbeelding voorkomen, zou het trainen op alle categorieën leiden tot een onbalans in de gradienten en hoge rekenkosten. MagicSeg selecteert per iteratie een willekeurige subset van categorieën (bijv. 100) om maskers voor te voorspellen. Dit voorkomt overfitting op specifieke klassen en vermindert de rekenlast.
• Counterfactual Contrastive Training: Het model wordt getraind met een contrastieve loss tussen de originele afbeelding en de counterfactual afbeelding. Doel is om het model een class-agnostic objectlokalisatie-vaardigheid bij te brengen. Het model leert dat het object in de eerste afbeelding wel aanwezig is en in de tweede niet, zelfs zonder dat het masker perfect is. Dit helpt ruis in de pseudo-maskers te mitigeren.

De totale loss functie combineert Focal Loss, Dice Loss (voor de maskers) en de Cosine Similarity Loss (voor de counterfactual contrastive training).

Belangrijkste Bijdragen

1. MagicSeg Pipeline: Een volledig geautomatiseerde methode voor het genereren van open-wereld segmentatiedata met pixel-level maskers, tekstbeschrijvingen en counterfactual afbeeldingen, zonder afhankelijkheid van grote hoeveelheden echte gelabelde data.
2. Nieuwe Trainingstechnieken:
   • Een strategie voor willekeurige steekproeven van categorieën om gradient-onbalans en overfitting op te lossen.
   • Counterfactual contrastive training om het model robuuster te maken tegen ruis in maskers en om objectlokalisatie zonder expliciete pixel-annotaties te leren.
3. Dataset: Een synthetische dataset van 380.000 paren met 1205 categorieën, die diversiteit en controle over de gegenereerde inhoud biedt.

Resultaten

MagicSeg is geëvalueerd op meerdere benchmarks en presteert state-of-the-art (SOTA) voor methoden die gebruikmaken van synthetische data of open-wereld training:

• PASCAL VOC: 62,9% mIoU (versus 59,1% bij de beste eerdere open-wereld methode).
• PASCAL Context: 26,7% mIoU.
• COCO: 40,2% mIoU (een aanzienlijke verbetering ten opzichte van eerdere methoden zoals 32,0%).
• Zero-Shot Point Prediction (OpenImages-v7): 45,6% p-mIoU, wat aantoont dat het model goed generaliseert naar ongeziene categorieën.

De resultaten tonen aan dat MagicSeg de prestatie-kloof tussen open-wereld methoden (vaak zonder maskers) en methoden met handmatige annotaties aanzienlijk verkleint. In sommige gevallen (zoals op COCO) overtreft het zelfs methoden die specifiek zijn ontworpen voor die dataset.

Betekenis en Impact

MagicSeg demonstreert dat hoogwaardige, schaalbare datasets voor open-wereld visuele taken volledig synthetisch kunnen worden gegenereerd. Door gebruik te maken van de kracht van diffusion-modellen in combinatie met LLM's voor tekst en vision-language modellen voor maskering, wordt de dure en tijdrovende stap van handmatige annotatie omzeild. De introductie van counterfactual afbeeldingen biedt een nieuwe manier om modellen te trainen om objecten te "zien" door contrast, wat de robuustheid van open-wereld segmentatie significant verbetert. Dit opent de deur voor het trainen van modellen op een veel bredere scala aan concepten dan ooit tevoren mogelijk was met alleen echte data.