Dual-stage 3D medical image segmentation integrating learnable prompt generation and memory attention

Abstract

Although MedSAM2 achieves 3D medical image segmentation through a memory attention mechanism, its performance declines significantly when manually designed prompts are replaced by automatically generated ones—particularly in the context of rare cases or multi-object segmentation scenarios. Current automatic prompt generation methods often extract prompt cues directly from image features, which typically lack rich spatiotemporal context and semantic information, resulting in suboptimal performance. To overcome these limitations, we propose DSSAM2-LAPG, a dual-stage 3D medical image segmentation network that integrates a Learnable Automatic Prompt-space Generator (LAPG) with memory attention. The LAPG acts as a trainable mapper that transforms raw 3D image features into a semantically-rich and spatially-aligned prompt embedding space. In the preliminary stage, this mapper utilizes learnable object tokens (concept embeddings) to dynamically interact with image features, generating coarse but reliable spatial priors $M_{\text {coarse}}$ and initial semantic prompts. In the refinement stage, a memory attention mechanism integrates these priors with historical context from a support memory to precisely delineate boundaries and ensure 3D consistency. This integrated approach specifically addresses the failure of existing methods in capturing 3D context, encoding rare-object semantics, and providing instance-aware guidance. Experimental results demonstrate that DSSAM2-LAPG achieves Dice score improvements of 7.2%, 6.0% and 4.1% on the private XYCH-cervical dataset, public dataset CCTH-Cervical and Multi-Organ BTCV datasets, respectively, compared to the strong baseline MedSAM2, and all without requiring any manual prompts. Our code is available https://github.com/liufangcoca-515/APG/tree/main.