Upload folder using huggingface_hub
Browse files
README.md
CHANGED
|
@@ -20,33 +20,72 @@ library_name: promptm-unet
|
|
| 20 |
[](https://github.com/kiuyha/PromptM-UNet)
|
| 21 |
[](https://opensource.org/licenses/Apache-2.0)
|
| 22 |
|
| 23 |
-
**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
|
| 24 |
|
| 25 |
* **GitHub Repository**: [https://github.com/kiuyha/PromptM-UNet](https://github.com/kiuyha/PromptM-UNet)
|
| 26 |
-
* **Target Organ
|
| 27 |
|
| 28 |
---
|
| 29 |
|
| 30 |
## 🌟 Key Features
|
| 31 |
|
| 32 |
-
* **Linear
|
| 33 |
-
* **Ultra-Lightweight Footprint**: Active visual backbone is only **~1.87 Million parameters** (
|
| 34 |
* **Zero-Overhead Language Conditioning**: Decoupled text embedding caching ensures frozen language models (CLIP, BioBERT, Sentence-BERT) consume 0 extra GPU memory during training/inference.
|
| 35 |
* **Dual-Resolution Spatial Zoom**: SegVol-inspired dual-scale pipeline (3mm global context + 1.5mm high-resolution target crop).
|
| 36 |
-
* **Clinical Performance**: Achieves **
|
| 37 |
|
| 38 |
---
|
| 39 |
|
| 40 |
-
## 📊 Benchmark Results
|
| 41 |
|
| 42 |
-
|
| 43 |
|
| 44 |
-
|
| 45 |
-
|
| 46 |
-
|
| 47 |
-
|
|
| 48 |
-
|
|
| 49 |
-
| **
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 50 |
|
| 51 |
---
|
| 52 |
|
|
|
|
| 20 |
[](https://github.com/kiuyha/PromptM-UNet)
|
| 21 |
[](https://opensource.org/licenses/Apache-2.0)
|
| 22 |
|
| 23 |
+
**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.
|
| 24 |
|
| 25 |
* **GitHub Repository**: [https://github.com/kiuyha/PromptM-UNet](https://github.com/kiuyha/PromptM-UNet)
|
| 26 |
+
* **Target Organ**: Spleen (TotalSegmentator CT Cohort)
|
| 27 |
|
| 28 |
---
|
| 29 |
|
| 30 |
## 🌟 Key Features
|
| 31 |
|
| 32 |
+
* **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.
|
| 33 |
+
* **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).
|
| 34 |
* **Zero-Overhead Language Conditioning**: Decoupled text embedding caching ensures frozen language models (CLIP, BioBERT, Sentence-BERT) consume 0 extra GPU memory during training/inference.
|
| 35 |
* **Dual-Resolution Spatial Zoom**: SegVol-inspired dual-scale pipeline (3mm global context + 1.5mm high-resolution target crop).
|
| 36 |
+
* **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.
|
| 37 |
|
| 38 |
---
|
| 39 |
|
| 40 |
+
## 📊 Benchmark & Ablation Results
|
| 41 |
|
| 42 |
+
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)):
|
| 43 |
|
| 44 |
+
### 1. Fusion Strategy Ablation
|
| 45 |
+
*Evaluated using Multiplication operation, CLIP text encoder, and 50:50 BCE supervision:*
|
| 46 |
+
|
| 47 |
+
| Fusion Strategy | Peak Epoch | Peak DSC | Peak NSD | Final DSC (Epoch 30) | Final NSD (Epoch 30) | Peak Val VRAM |
|
| 48 |
+
| :--- | :---: | :---: | :---: | :---: | :---: | :---: |
|
| 49 |
+
| **Early Fusion** | **19** | **`0.8622`** | **`0.7928`** | `0.8196` | `0.7514` | **< 4.3 GB** |
|
| 50 |
+
| **All-Stage Fusion** | 24 | `0.8619` | `0.7871` | `0.7737` | `0.6848` | < 4.3 GB |
|
| 51 |
+
| **Late Fusion** | 12 | `0.8576` | `0.7675` | `0.7347` | `0.6095` | < 4.3 GB |
|
| 52 |
+
|
| 53 |
+
### 2. Loss Function Ablation
|
| 54 |
+
*Evaluated across compound loss formulations:*
|
| 55 |
+
|
| 56 |
+
| Loss Configuration | Peak Epoch | Peak DSC | Peak NSD | Final DSC (Epoch 30) | Final NSD (Epoch 30) |
|
| 57 |
+
| :--- | :---: | :---: | :---: | :---: | :---: |
|
| 58 |
+
| **Focal Loss (50:50)** | **20** | **`0.8627`** | `0.7827` | `0.8045` | `0.7195` |
|
| 59 |
+
| **Standard BCE (50:50)** | 24 | `0.8619` | **`0.7871`** | `0.7737` | `0.6848` |
|
| 60 |
+
| **Dynamic Loss (80:20 \(\to\) 65:35)** | 27 | `0.8537` | `0.7668` | `0.8280` | `0.7299` |
|
| 61 |
+
| **Focal Loss (80:20)** | 20 | `0.7827` | `0.6821` | `0.7358` | `0.6393` |
|
| 62 |
+
|
| 63 |
+
### 3. Fusion Operation Ablation
|
| 64 |
+
*Evaluated using 50:50 BCE supervision:*
|
| 65 |
+
|
| 66 |
+
| Operation | Peak Epoch | Peak DSC | Peak NSD | Final DSC (Epoch 30) | Final NSD (Epoch 30) |
|
| 67 |
+
| :--- | :---: | :---: | :---: | :---: | :---: |
|
| 68 |
+
| **Multiplication** | **24** | **`0.8619`** | **`0.7871`** | `0.7737` | `0.6848` |
|
| 69 |
+
| **Concatenation** | 26 | `0.8381` | `0.7647` | `0.8120` | `0.7387` |
|
| 70 |
+
| **Hybrid** | 28 | `0.8152` | `0.7301` | `0.7962` | `0.7041` |
|
| 71 |
+
|
| 72 |
+
### 4. Text Encoder Ablation
|
| 73 |
+
*Evaluated with Early Fusion:*
|
| 74 |
+
|
| 75 |
+
| Text Encoder | Peak Epoch | Peak DSC | Peak NSD | Final DSC (Epoch 30) | Final NSD (Epoch 30) |
|
| 76 |
+
| :--- | :---: | :---: | :---: | :---: | :---: |
|
| 77 |
+
| **CLIP** | **19** | **`0.8622`** | **`0.7928`** | `0.8196` | `0.7514` |
|
| 78 |
+
| **BioBERT** | 19 | `0.8466` | `0.7584` | `0.7968` | `0.7148` |
|
| 79 |
+
|
| 80 |
+
### 5. Multi-Tier Clinical Prompt Robustness
|
| 81 |
+
*Evaluated on Best Early Fusion + CLIP checkpoint across prompt complexity tiers:*
|
| 82 |
+
|
| 83 |
+
| Prompt Tier | Clinical Description | Peak DSC |
|
| 84 |
+
| :--- | :--- | :---: |
|
| 85 |
+
| **Tier \(N\)** | Organ Name (*e.g., "spleen"*) | `0.8620` |
|
| 86 |
+
| **Tier \(NS\)** | Name + Synonym (*e.g., "spleen, lien"*) | `0.8621` |
|
| 87 |
+
| **Tier \(NL\)** | Name + Location (*e.g., "spleen in left upper quadrant"*) | **`0.8624`** |
|
| 88 |
+
| **Tier \(NSL\)** | Name + Synonym + Location (*e.g., "spleen, lien in upper left abdomen"*) | `0.8622` |
|
| 89 |
|
| 90 |
---
|
| 91 |
|