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

2025-03-25

AI Summary

Medical image segmentation suffers from poor generalization across imaging modalities, anatomical structures, and tasks; existing deep learning methods require task-specific training or fine-tuning, hindering zero-shot adaptation to unseen categories or data distributions. To address this, we propose Iris, a reference-image-guided in-context learning framework. Iris introduces a lightweight architecture that decouples task encoding from inference, incorporates reference image–label pair embedding distillation, and employs multi-strategy contextual reasoning—including in-context tuning. It supports one-shot segmentation, contextual ensemble, object-level retrieval, and automatically discovers cross-dataset anatomical correspondences. Evaluated on 12 in-distribution datasets, Iris matches the performance of dedicated models; on 7 unseen out-of-distribution datasets, it significantly surpasses state-of-the-art methods. Notably, Iris achieves, for the first time, zero-shot cross-modal anatomical relationship discovery without any labeled data from target domains.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Multimodal Learning
• Intelligent Robots: Multimodal Perception & Sensor Fusion

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs

Abstract

Medical image segmentation remains challenging due to the vast diversity of anatomical structures, imaging modalities, and segmentation tasks. While deep learning has made significant advances, current approaches struggle to generalize as they require task-specific training or fine-tuning on unseen classes. We present Iris, a novel In-context Reference Image guided Segmentation framework that enables flexible adaptation to novel tasks through the use of reference examples without fine-tuning. At its core, Iris features a lightweight context task encoding module that distills task-specific information from reference context image-label pairs. This rich context embedding information is used to guide the segmentation of target objects. By decoupling task encoding from inference, Iris supports diverse strategies from one-shot inference and context example ensemble to object-level context example retrieval and in-context tuning. Through comprehensive evaluation across twelve datasets, we demonstrate that Iris performs strongly compared to task-specific models on in-distribution tasks. On seven held-out datasets, Iris shows superior generalization to out-of-distribution data and unseen classes. Further, Iris's task encoding module can automatically discover anatomical relationships across datasets and modalities, offering insights into medical objects without explicit anatomical supervision.

Problem

• Universal medical image segmentation across diverse tasks
• Generalization to unseen classes without fine-tuning
• Automatic discovery of anatomical relationships across datasets

Innovation

• In-context learning for medical segmentation
• Lightweight context task encoding module
• Decouples task encoding from inference