Recreating IRIS: Universal Medical Image Segmentation via In-Context Learning

class IRISModel(nn.Module):
    """
    IRIS model for universal medical image segmentation.
    Combines 3D UNet encoder, task encoding, and query-based decoder.
    """
    def __init__(self, in_channels=1, base_channels=32, embed_dim=512, 
                 num_tokens=10, num_classes=1, num_heads=8):
        super().__init__()

        # Image encoder (shared for reference and query)
        self.encoder = Encoder3D(in_channels, base_channels, num_blocks_per_stage=2)

        # Task encoding from reference examples
        self.task_encoder = TaskEncodingModule(
            encoder_channels[-1], embed_dim, num_tokens
        )

        # Query-based decoder with cross-attention
        self.decoder = QueryBasedDecoderFixed(
            encoder_channels, embed_dim, num_classes, num_heads
        )

    def forward(self, query_image, reference_image=None, reference_mask=None):
        # Extract features from query image
        query_features = self.encoder(query_image)

        # Generate task embedding from reference
        if reference_image is not None:
            ref_features = self.encoder(reference_image)
            task_embedding = self.task_encoder(ref_features[-1], reference_mask)

        # Decode with task guidance
        segmentation = self.decoder(query_features, task_embedding)
        return segmentation
        

🔧 Task Encoding Module

The task encoding module consists of two parallel streams:

• Foreground Feature Encoding: High-resolution processing to preserve fine anatomical details
• Contextual Feature Encoding: Learnable query tokens with cross-attention for global context

DICE Coefficient Formula

*** QuickLaTeX cannot compile formula:
  DICE = 2 × |P ∩ G| / (|P| + |G|) 

*** Error message:
Unicode character ∩ (U+2229)
leading text: $  DICE = 2 × |P ∩

Why DICE for Medical Images?

• Class Imbalance Robust: Focuses on overlap rather than accuracy, crucial when organs occupy small portions of images
• Overlap-Centric: Emphasizes intersection between prediction and ground truth
• Range [0,1]: 0 = no overlap, 1 = perfect overlap
• Differentiable: Can be used as a loss function for backpropagation

Visual Representation of DICE Calculation

✅ Core Capabilities

• End-to-end Training: Fixed channel mismatch issues for complete gradient flow
• Multiple Inference Modes: One-shot, ensemble, and memory bank strategies
• 3D Volume Processing: Native support for medical 3D imaging
• Multi-modal Support: Works across CT, MRI, PET imaging modalities
• Flexible Task Encoding: Adapts to different anatomical structures

Primary Metrics

• DICE Coefficient: Spatial overlap measure (0-1, higher better)
• Surface DICE: Boundary-based evaluation for surface accuracy
• Hausdorff Distance: Maximum surface distance error
• IoU (Jaccard Index): Intersection over Union measure

# Training Setup
BATCH_SIZE = 4 # Limited by GPU memory for 3D volumes
LEARNING_RATE = 1e-4 # Adam optimizer with Lamb variant
MAX_ITERATIONS = 1000 # Episodic training iterations
LOSS_FUNCTION = DiceLoss + CrossEntropyLoss # Combined loss

# Install dependencies
pip install -e .

# Download datasets
abts25_download_data

# Quick training test
python simple_train.py

# Full training on AMOS22 dataset
python train_amos22.py –data_dir src/data/amos –batch_size 4 –max_iterations 1000

# Evaluation
python evaluate_amos22.py –model_path checkpoints/model.pth

# Compute metrics
abts25_compute_metrics PREDICTIONS_FOLDER -num_processes 4

# Load trained model
model = IRISModel(in_channels=1, embed_dim=512, num_classes=1)
model.load_state_dict(torch.load(‘checkpoints/model.pth’))

# Initialize inference engine
inference = IRISInferenceFixed(model, device=’cuda’)

# One-shot inference with reference
result = inference.one_shot_inference(
query_image=query_volume, # (1, 1, D, H, W)
reference_image=ref_volume, # (1, 1, D, H, W)
reference_mask=ref_mask, # (1, 1, D, H, W)
apply_sigmoid=True,
threshold=0.5
)

# Get segmentation prediction
prediction = result[‘prediction’] # Binary mask
probabilities = result[‘probabilities’] # Confidence scores

🏆 Key Achievements

• Universal Framework: Single model works across 12+ medical imaging datasets
• No Fine-tuning Required: Adapts to new tasks through reference examples
• State-of-the-art Performance: Competitive with task-specific models
• Computational Efficiency: Task embeddings can be precomputed and reused
• Clinical Applicability: Supports real-world medical imaging workflows

• In-context learning enables flexible adaptation without retraining
• Task encoding decouples context definition from inference
• DICE loss is crucial for handling class imbalance in medical images
• 3D architectures are essential for volumetric medical data
• Cross-attention mechanisms effectively integrate task-specific guidance

Repository: mister-weeden/abts25

Paper: “Show and Segment: Universal Medical Image Segmentation via In-Context Learning”

For technical questions or collaboration opportunities, please open an issue on GitHub.