Kiuyha commited on
Commit
09fee2a
·
verified ·
1 Parent(s): ae3ab30

Upload folder using huggingface_hub

Browse files
Files changed (1) hide show
  1. README.md +52 -13
README.md CHANGED
@@ -20,33 +20,72 @@ library_name: promptm-unet
20
  [![GitHub Repository](https://img.shields.io/badge/GitHub-Repository-black?logo=github)](https://github.com/kiuyha/PromptM-UNet)
21
  [![License: Apache 2.0](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](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 (v1.0)**: 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.787+$ Normalized Surface Distance (NSD)** on TotalSegmentator CT scans.
37
 
38
  ---
39
 
40
- ## 📊 Benchmark Results
41
 
42
- Evaluated on the official TotalSegmentator spleen validation/test cohort across multi-tier clinical text prompts ($N$: Name, $NS$: Name+Synonym, $NL$: Name+Location, $NSL$: Name+Synonym+Location):
43
 
44
- | Configuration | Parameters | Peak DSC | Peak NSD | Inference VRAM |
45
- | :--- | :---: | :---: | :---: | :---: |
46
- | **PromptM-UNet (Bottleneck Fusion, 50:50 BCE)** | **1.87M** | **`0.8622`** | **`0.7928`** | **< 4.3 GB** |
47
- | **PromptM-UNet (All-Stage Fusion, 50:50 BCE)** | **1.91M** | **`0.8619`** | **`0.7871`** | **< 4.3 GB** |
48
- | **LightM-UNet Baseline (Vision-Only)** | 1.87M | `0.7820` | `0.6950` | < 4.3 GB |
49
- | **SegVol (Foundation Model)** | 181M | `0.8590` | `0.7710` | > 24 GB |
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
50
 
51
  ---
52
 
 
20
  [![GitHub Repository](https://img.shields.io/badge/GitHub-Repository-black?logo=github)](https://github.com/kiuyha/PromptM-UNet)
21
  [![License: Apache 2.0](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](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