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

TL;DR

This work tackles distribution shifts in medical imaging by enabling online adaptation of interactive segmentation models guided by clinician inputs. It introduces OAIMS, combining Post-Interaction and Mid-Interaction adaptation with a Click-Centered Gaussian loss to reinforce responses around user prompts. The method uses a U-Net with guidance maps, pretrained with simulated clicks and a Dice-Focal objective, and adapts online using pseudo-ground-truth from user refinements. Experiments across fundus and brain MRI datasets show consistent gains over state-of-the-art adaptation methods, with notable improvements under large distribution shifts and minimal latency.

Abstract

Interactive segmentation uses real-time user inputs, such as mouse clicks, to iteratively refine model predictions. Although not originally designed to address distribution shifts, this paradigm naturally lends itself to such challenges. In medical imaging, where distribution shifts are common, interactive methods can use user inputs to guide models towards improved predictions. Moreover, once a model is deployed, user corrections can be used to adapt the network parameters to the new data distribution, mitigating distribution shift. Based on these insights, we aim to develop a practical, effective method for improving the adaptive capabilities of interactive segmentation models to new data distributions in medical imaging. Firstly, we found that strengthening the model's responsiveness to clicks is important for the initial training process. Moreover, we show that by treating the post-interaction user-refined model output as pseudo-ground-truth, we can design a lean, practical online adaptation method that enables a model to learn effectively across sequential test images. The framework includes two components: (i) a Post-Interaction adaptation process, updating the model after the user has completed interactive refinement of an image, and (ii) a Mid-Interaction adaptation process, updating incrementally after each click. Both processes include a Click-Centered Gaussian loss that strengthens the model's reaction to clicks and enhances focus on user-guided, clinically relevant regions. Experiments on 5 fundus and 4 brain-MRI databases show that our approach consistently outperforms existing methods under diverse distribution shifts, including unseen imaging modalities and pathologies. Code and pretrained models will be released upon publication.

Figures (6)

• Figure 1: Method overview. For Pretraining, the model is trained with simulated clicks, provided as additional input channels besides the image. During Inference and adaptation, images arrive sequentially. For each image, the user iteratively provides $T$ clicks to correct the segmentation, until the final prediction $P_{\text{final}}=\! P_T$ is obtained. Mid-interaction adaptation: After each corrective click $c_t$, the model's output $P_t^{\text{initial}}$ is used as pseudo-label compared with the pre-correction output $P_{t-1}$ via the DF and CCG losses, to update model parameters. The updated model then produces refined output $P_t$, which is then shown to the user, ending iteration $t$. Post-interaction adaptation: Once the final corrected segmentation $P_T$ is obtained, it is used as pseudo-label to first fine-tune the model using a localization click (Stage 1), and then to fine-tune using multiple correction clicks, generated from areas where the prediction of Stage 1 disagrees with $P_T$ (Stage 2).
• Figure 2: Dice-score performance for automatic and interactive models. All networks, except the in‑distribution baseline, are trained on REFUGE (fundus) or BRATS‑FLAIR (brain MRI). The x-axis represents the number of clicks. Horizontal lines mark automatic performance of automatic models in cross‑distribution and in‑distribution settings. Curves show the interactive segmentation model with and without the CCG loss; clicks significantly improve performance in all test cases, with the CCG loss providing additional gains, especially for large distribution gaps (e.g., BRATS‑T1 ).
• Figure 3: Visualizations on BRATS Databases demonstrating adaptation to different modalities and pathologies (Trained on BRATS Flair for BRATS T1 and BRATS T2, Trained on BRATS Flair/T1/T1c for ATLAS
• Figure 4: Illustration of the databases used. The distribution across different datasets and modalities can be visually observed.
• Figure 5: Illustration of how the predicted segmentation from our OAIMS (PI+MI) method evolves as more interaction clicks are provided. The examples are brain MRI images (ATLAS and BRATS T1). From left to right: input image, ground truth, and predictions with 1, 2, 3, 5, and 10 clicks. The results show progressive refinement of the segmentation as more clicks are provided.