One-Shot Medical Video Object Segmentation via Temporal Contrastive Memory Networks

TL;DR

This paper tackles one-shot medical video object segmentation, where a single annotated first frame must guide segmentation across an entire video. It introduces a temporal contrastive memory network comprising an image encoder, a mask encoder, a memory bank with a temporal contrastive objective, and a decoder to propagate masks, without needing fine-tuning on new structures. Evaluated on a diverse, multi-modal medical video dataset (colonoscopy and echocardiography), the approach achieves state-of-the-art results for both seen and unseen structures, demonstrating strong generalization from a single exemplar. By reducing annotation burden and enabling generalizable segmentation, it points to practical tools for clinicians and adaptable memory-based strategies for medical video analysis; future work includes memory-bank optimization and dynamic frame selection.

Abstract

Video object segmentation is crucial for the efficient analysis of complex medical video data, yet it faces significant challenges in data availability and annotation. We introduce the task of one-shot medical video object segmentation, which requires separating foreground and background pixels throughout a video given only the mask annotation of the first frame. To address this problem, we propose a temporal contrastive memory network comprising image and mask encoders to learn feature representations, a temporal contrastive memory bank that aligns embeddings from adjacent frames while pushing apart distant ones to explicitly model inter-frame relationships and stores these features, and a decoder that fuses encoded image features and memory readouts for segmentation. We also collect a diverse, multi-source medical video dataset spanning various modalities and anatomies to benchmark this task. Extensive experiments demonstrate state-of-the-art performance in segmenting both seen and unseen structures from a single exemplar, showing ability to generalize from scarce labels. This highlights the potential to alleviate annotation burdens for medical video analysis. Code is available at https://github.com/MedAITech/TCMN.

Figures (2)

• Figure 1: Overview of our proposed model for one-shot medical video object segmentation. Our model consists of four main components: an image encoder, a mask encoder, a temporal contrastive memory bank, and a decoder. We feed the first frame and its corresponding mask annotation into the image encoder and mask encoder, respectively, to generate key and value feature maps. These are stored in the memory bank as initial memories. For subsequent frames, we input them into the image encoder to generate new key features and utilize the memory bank to produce segmentations. Concurrently, we feed the predicted segmentations to the mask encoder to generate new value features, which, along with the newly generated key features, are stored in the memory bank as updated memories. Our designed memory bank optimizes the memory information through temporal contrastive learning. This process is repeated until the entire video is segmented.
• Figure 2: Qualitative results. The first two rows show segmentation results on classes that our model has learned during training, while the last two rows show results on classes that the model has not seen during training. The green masks indicate annotation masks of the first frames within videos. The red masks represent predicted segmentation masks obtained by our model on subsequent frames, and white lines indicate the corresponding ground truth contours for these frames.