---
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

[![GitHub Repository](https://img.shields.io/badge/GitHub-Repository-black?logo=github)](https://github.com/kiuyha/PromptM-UNet)
[![License: Apache 2.0](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](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)**.