| --- |
| license: apache-2.0 |
| tags: |
| - medical |
| - 3d-segmentation |
| - mamba |
| - state-space-model |
| - vision-language |
| - prompt-guided-segmentation |
| - ct-scan |
| - spleen |
| - totalsegmentator |
| - lightm-unet |
| pipeline_tag: image-segmentation |
| library_name: promptm-unet |
| --- |
| |
| # PromptM-UNet: Efficient Text-Prompted 3D Medical Image Segmentation with Mamba SSM |
|
|
| [](https://github.com/kiuyha/PromptM-UNet) |
| [](https://opensource.org/licenses/Apache-2.0) |
|
|
| **PromptM-UNet** is an ultra-lightweight, multi-modal 3D medical image segmentation framework that integrates **Residual Vision Mamba (SSM)** with **multi-stage language conditioning**. It enables accurate, prompt-guided volumetric CT segmentation while maintaining linear \(O(N)\) computational complexity and a minimal memory footprint feasible for standard consumer GPUs. |
|
|
| * **GitHub Repository**: [https://github.com/kiuyha/PromptM-UNet](https://github.com/kiuyha/PromptM-UNet) |
| * **Target Organ**: Spleen (TotalSegmentator CT Cohort) |
|
|
| --- |
|
|
| ## 🌟 Key Features |
|
|
| * **Linear \(O(N)\) Volumetric Modeling**: Built on 3D Residual Vision Mamba (RVM) blocks from LightM-UNet, providing global receptive field modeling without quadratic \(O(N^2)\) Transformer memory explosion. |
| * **Ultra-Lightweight Footprint**: Active visual backbone is only **~1.87 Million parameters** (\(16\times\) smaller than standard 3D nnU-Net, \(100\times\) smaller than SegVol). |
| * **Zero-Overhead Language Conditioning**: Decoupled text embedding caching ensures frozen language models (CLIP, BioBERT, Sentence-BERT) consume 0 extra GPU memory during training/inference. |
| * **Dual-Resolution Spatial Zoom**: SegVol-inspired dual-scale pipeline (3mm global context + 1.5mm high-resolution target crop). |
| * **Clinical Performance**: Achieves **\(0.862+\) Dice Similarity Coefficient (DSC)** and **\(0.792+\) Normalized Surface Distance (NSD)** on TotalSegmentator CT scans with under **4.3 GB** validation VRAM. |
|
|
| --- |
|
|
| ## 📊 Benchmark & Ablation Results |
|
|
| The following results are evaluated on the official TotalSegmentator spleen validation cohort across 30 training epochs (from [`notebooks/ablation_promptm_unet.ipynb`](notebooks/ablation_promptm_unet.ipynb)): |
|
|
| ### 1. Fusion Strategy Ablation |
| *Evaluated using Multiplication operation, CLIP text encoder, and 50:50 BCE supervision:* |
|
|
| | Fusion Strategy | Peak Epoch | Peak DSC | Peak NSD | Final DSC (Epoch 30) | Final NSD (Epoch 30) | Peak Val VRAM | |
| | :--- | :---: | :---: | :---: | :---: | :---: | :---: | |
| | **Early Fusion** | **19** | **`0.8622`** | **`0.7928`** | `0.8196` | `0.7514` | **< 4.3 GB** | |
| | **All-Stage Fusion** | 24 | `0.8619` | `0.7871` | `0.7737` | `0.6848` | < 4.3 GB | |
| | **Late Fusion** | 12 | `0.8576` | `0.7675` | `0.7347` | `0.6095` | < 4.3 GB | |
|
|
| ### 2. Loss Function Ablation |
| *Evaluated across compound loss formulations:* |
|
|
| | Loss Configuration | Peak Epoch | Peak DSC | Peak NSD | Final DSC (Epoch 30) | Final NSD (Epoch 30) | |
| | :--- | :---: | :---: | :---: | :---: | :---: | |
| | **Focal Loss (50:50)** | **20** | **`0.8627`** | `0.7827` | `0.8045` | `0.7195` | |
| | **Standard BCE (50:50)** | 24 | `0.8619` | **`0.7871`** | `0.7737` | `0.6848` | |
| | **Dynamic Loss (80:20 \(\to\) 65:35)** | 27 | `0.8537` | `0.7668` | `0.8280` | `0.7299` | |
| | **Focal Loss (80:20)** | 20 | `0.7827` | `0.6821` | `0.7358` | `0.6393` | |
|
|
| ### 3. Fusion Operation Ablation |
| *Evaluated using 50:50 BCE supervision:* |
|
|
| | Operation | Peak Epoch | Peak DSC | Peak NSD | Final DSC (Epoch 30) | Final NSD (Epoch 30) | |
| | :--- | :---: | :---: | :---: | :---: | :---: | |
| | **Multiplication** | **24** | **`0.8619`** | **`0.7871`** | `0.7737` | `0.6848` | |
| | **Concatenation** | 26 | `0.8381` | `0.7647` | `0.8120` | `0.7387` | |
| | **Hybrid** | 28 | `0.8152` | `0.7301` | `0.7962` | `0.7041` | |
|
|
| ### 4. Text Encoder Ablation |
| *Evaluated with Early Fusion:* |
|
|
| | Text Encoder | Peak Epoch | Peak DSC | Peak NSD | Final DSC (Epoch 30) | Final NSD (Epoch 30) | |
| | :--- | :---: | :---: | :---: | :---: | :---: | |
| | **CLIP** | **19** | **`0.8622`** | **`0.7928`** | `0.8196` | `0.7514` | |
| | **BioBERT** | 19 | `0.8466` | `0.7584` | `0.7968` | `0.7148` | |
|
|
| ### 5. Multi-Tier Clinical Prompt Robustness |
| *Evaluated on Best Early Fusion + CLIP checkpoint across prompt complexity tiers:* |
|
|
| | Prompt Tier | Clinical Description | Peak DSC | |
| | :--- | :--- | :---: | |
| | **Tier \(N\)** | Organ Name (*e.g., "spleen"*) | `0.8620` | |
| | **Tier \(NS\)** | Name + Synonym (*e.g., "spleen, lien"*) | `0.8621` | |
| | **Tier \(NL\)** | Name + Location (*e.g., "spleen in left upper quadrant"*) | **`0.8624`** | |
| | **Tier \(NSL\)** | Name + Synonym + Location (*e.g., "spleen, lien in upper left abdomen"*) | `0.8622` | |
|
|
| --- |
|
|
| ## 🚀 Quickstart & Inference |
|
|
| ### 1. Installation |
|
|
| ```bash |
| # Clone the repository |
| git clone https://github.com/kiuyha/PromptM-UNet.git |
| cd PromptM-UNet |
| |
| # Install dependencies and package |
| pip install -e . |
| ``` |
|
|
| ### 2. Python Inference |
|
|
| ```python |
| import torch |
| import yaml |
| from promptm_unet.models.PromptMUNet import PromptMUNet |
| |
| # Load configuration |
| with open("configs/default.yml", "r") as f: |
| config = yaml.safe_load(f) |
| |
| # Instantiate model |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model = PromptMUNet(config).to(device) |
| |
| # Load checkpoint |
| checkpoint = torch.load("best_model.pth", map_location=device) |
| model.load_state_dict(checkpoint["model_state_dict"]) |
| model.eval() |
| |
| # Dummy 3D input: (Batch, Channels, Depth, Height, Width) |
| dummy_ct = torch.randn(1, 1, 64, 128, 128).to(device) |
| prompt = ["spleen segmentation in abdominal CT scan"] |
| |
| # Forward pass |
| with torch.no_grad(): |
| prediction = model(dummy_ct, prompt) # Output: (1, 1, 64, 128, 128) |
| probabilities = torch.sigmoid(prediction) |
| binary_mask = (probabilities > 0.5).cpu().numpy() |
| |
| print(f"Predicted spleen mask shape: {binary_mask.shape}") |
| ``` |
|
|
| ### 3. CLI Training & Evaluation |
|
|
| ```bash |
| # Run training with default 50:50 compound supervision |
| accelerate launch -m promptm_unet.cli train \ |
| --path.raw_data_dir "/path/to/totalsegmentator_raw" \ |
| --path.prepro_data_dir "/path/to/preprocessed_data" \ |
| --training.batch_size 2 \ |
| --training.epochs 30 \ |
| --loss.dice_weight 0.5 \ |
| --loss.bce_weight 0.5 |
| |
| # Run multi-tier evaluation |
| python -m promptm_unet.cli test \ |
| --checkpoint "./experiments/checkpoints/best_model.pth" |
| ``` |
|
|
| --- |
|
|
| ## 🏗️ Model Architecture |
|
|
| ``` |
| 3D CT Volume (1.5mm / 3.0mm) |
| │ |
| ▼ |
| ┌───────────────┐ |
| │ Visual Encoder│ (Residual Vision Mamba - RVM Blocks) |
| └───────┬───────┘ |
| │ |
| ▼ |
| ┌───────────────┐ ┌──────────────────┐ |
| │ Bottleneck │ ◄────► │ Text Projection │ ◄── Frozen CLIP / BioBERT / SBERT |
| └───────┬───────┘ └──────────────────┘ (512D / 768D / 384D) |
| │ ▲ |
| ▼ │ (Multi-Stage Hadamard Product) |
| ┌───────────────┐ │ |
| │ Visual Decoder│ ────────────────┘ |
| └───────┬───────┘ |
| │ |
| ▼ |
| Deep Supervision Heads ──► High-Resolution Spleen Binary Mask |
| ``` |
|
|
| --- |
|
|
| ## 📜 Citation & License |
|
|
| This project is licensed under the **Apache License 2.0**. |
|
|
| If you find this work useful in your research, please cite our repository: |
|
|
| ```bibtex |
| @misc{promptm_unet2026, |
| title={PromptM-UNet: Efficient Text-Prompted 3D Medical Image Segmentation Using State-Space Model Mamba for Spleen}, |
| author={Shridhara, Ketut and Setyawan, Ivan Andika and Al-Habib, Hasanuddin}, |
| year={2026}, |
| publisher={Universitas Negeri Surabaya}, |
| howpublished={\url{https://github.com/kiuyha/PromptM-UNet}} |
| } |
| ``` |
|
|
| --- |
|
|
| ## 🤝 Acknowledgements |
| * [LightM-UNet](https://github.com/Mr-TalhaIlyas/LightM-UNet) for the 3D Mamba visual backbone. |
| * [VoxTell](https://arxiv.org/abs/2511.11450) & [SegVol](https://arxiv.org/abs/2311.13385) for multimodal prompting inspirations. |
| * [TotalSegmentator](https://github.com/wasserth/TotalSegmentator) for CT dataset annotations. |
| * Research funded by **LPPM Universitas Negeri Surabaya (UNESA)**. |
|
|
