Upload folder using huggingface_hub

09fee2a verified 2 days ago

8.44 kB

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

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

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

# 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

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

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

@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 for the 3D Mamba visual backbone.
VoxTell & SegVol for multimodal prompting inspirations.
TotalSegmentator for CT dataset annotations.
Research funded by LPPM Universitas Negeri Surabaya (UNESA).