You Point, I Learn: Online Adaptation of Interactive Segmentation Models for Handling Distribution Shifts in Medical Imaging

1. Problem Statement

Medical image segmentation models often suffer from distribution shifts when deployed in real-world scenarios (e.g., different scanners, clinics, or imaging modalities). While deep learning models excel on training data, their performance degrades significantly on unseen target distributions.

• The Gap: Existing interactive segmentation models allow clinicians to correct predictions via clicks, but they typically treat these interactions as static inputs without updating the model's parameters. Consequently, the model does not "learn" from the corrections to adapt to the new data distribution.
• The Challenge: How to design a framework where a model can continuously and effectively adapt to new data distributions using sparse, real-time user corrections (clicks) without requiring complex regularizers or suffering from catastrophic forgetting.

2. Methodology: OAIMS (Online Adaptation for Interactive Medical-image Segmentation)

The authors propose OAIMS, a framework that integrates two distinct online adaptation phases with a novel loss function.

A. Core Component: Click-Centered Gaussian (CCG) Loss

To ensure the model reacts appropriately to user inputs, the authors introduce the CCG Loss.

• Mechanism: Instead of penalizing errors only on the specific clicked pixel or the entire image, this loss applies a Gaussian-weighted penalty to the region surrounding a click.
• Class Limitation: Crucially, the loss is applied only to pixels that share the same class label as the click (e.g., a foreground click only penalizes surrounding foreground pixels that are predicted as background). This prevents the model from overfitting to incorrect class assumptions in surrounding areas.
• Formula: It combines Cross-Entropy with a Gaussian kernel $G_c(i,j)$ centered at the click, weighted by an indicator function $I_c$ that restricts the loss to the correct class region.

B. Two-Stage Adaptation Framework

The method operates during the inference of a sequence of images, utilizing two complementary adaptation processes:

1. Mid-Interaction Adaptation (Online, Iterative)

• Timing: Occurs after every user click during the refinement of a single image.
• Process:
  • The user provides a corrective click $c_t$.
  • The model generates an initial prediction $P_t^{initial}$ based on the current weights.
  • The previous prediction $P_{t-1}$ is treated as the pseudo-ground-truth.
  • The model is updated by minimizing the CCG Loss and Dice-Focal (DF) Loss between $P_{t-1}$ and $P_t^{initial}$.
  • The updated weights produce a refined prediction $P_t$, which is shown to the user for the next iteration.
• Benefit: Allows the model to correct itself immediately within the current image and improves the quality of the final mask used for the next stage.

2. Post-Interaction Adaptation (Batch, Image-level)

• Timing: Occurs after the user has finished refining an image and the final mask $P_{final}$ is obtained.
• Assumption: The final user-corrected mask is of sufficient quality to serve as a pseudo-ground-truth for the new distribution.
• Process (Two Stages):
  • Stage 1 (Localization): Fine-tune the model using the first localization click and the final mask $P_{final}$ as the target, using DF Loss. This improves the model's initial guess for subsequent images.
  • Stage 2 (Correction): Generate artificial correction clicks ($\hat{C}$) by identifying regions where the Stage 1 prediction disagrees with $P_{final}$ (False Positives/Negatives). Fine-tune the model using these synthetic clicks and the final mask, employing both DF Loss and CCG Loss.
• Benefit: Solidifies the model's adaptation to the new distribution for all future images in the sequence.

3. Key Contributions

1. Novel Adaptation Framework: Proposes a practical, model-agnostic online adaptation method (OAIMS) that combines Mid-Interaction and Post-Interaction learning to handle distribution shifts in medical imaging.
2. Click-Centered Gaussian (CCG) Loss: Introduces a specialized loss function that focuses learning on the user-guided regions while respecting class boundaries, preventing overfitting to sparse labels.
3. Pseudo-Ground-Truth Strategy: Demonstrates that user-refined outputs can effectively serve as pseudo-labels for online learning without complex regularizers, simplifying the adaptation pipeline.
4. Synthetic Click Generation: A novel Stage 2 in Post-Interaction adaptation that automatically generates correction clicks from prediction errors to maximize the utility of the final mask.

4. Experimental Results

The method was evaluated on 5 Fundus databases (REFUGE2, G1020, GS1-Drishti, GAMMA, PAPILA) and 4 Brain MRI databases (BRATS2023, ATLAS, WMH, TBI) involving various pathologies and modalities.

• Performance under Distribution Shift:
  • OAIMS consistently outperformed state-of-the-art (SOTA) methods like IA+SA, TSCA, and Med-SA (Medical SAM Adapter).
  • Brain MRI: On BRATS T1 (a significant modality shift from Flair), OAIMS with both adaptations (PI+MI) achieved a Dice score improvement of >10% compared to the base interactive model and significantly outperformed TSCA.
  • Fundus: On large distribution shifts (e.g., REFUGE to G1020), OAIMS surpassed previous methods even with fewer clicks.
• Efficiency:
  • The method requires minimal computational overhead. Updates take 0.05s (Mid-Interaction) and 0.09s (Post-Interaction) on an NVIDIA A5000 GPU, and remain imperceptible (<0.5s) even on a CPU.
• Robustness:
  • Overfitting: Unlike TSCA, OAIMS did not exhibit overfitting even when tested with 50 clicks per image.
  • Noisy Inputs: The method remained robust even when 40% of user clicks were intentionally incorrect (wrong class), significantly outperforming non-adaptive baselines.
  • Extreme Shifts: In scenarios with extreme domain gaps (e.g., training on Flair, testing on T1), OAIMS recovered high performance (80% Dice) where baselines failed (<25% Dice).

5. Significance

• Clinical Applicability: The framework enables a single model to adapt continuously to new scanners, modalities, or pathologies in real-time, reducing the need for retraining or managing multiple specialized models.
• Human-in-the-Loop Efficiency: By leveraging user corrections not just for immediate refinement but for model learning, the system reduces the number of clicks required to achieve high-quality segmentation in new environments.
• Generalizability: The approach is architecture-agnostic (demonstrated on U-Net) and can theoretically extend to other input types (scribbles, bounding boxes) and medical domains.

In summary, OAIMS transforms interactive segmentation from a static correction tool into a dynamic, self-improving system that effectively bridges the gap between training data distributions and real-world clinical deployment.