VISCERA-SAM: Fine-Tuning UltraSam for Visceral Organ Segmentation in Ultrasound

Fine-tuning UltraSam on the FASS (Fetal Abdominal Structures Segmentation) dataset for visceral organ segmentation (artery, liver, stomach, vein) with box-prompt and mask refinement. This project introduces additional loss functions (boundary loss, Hausdorff DT loss), largest connected component (LCC) post-processing, and geometric augmentation as part of an ablation study.

Based on: UltraSam by Meyer et al. (IJCARS 2025)

---

Installation

Prerequisites

• OS: Windows 10/11 or Linux
• GPU: NVIDIA GPU with CUDA 11.8 support (tested on RTX 4060 Ti)
• Conda: Anaconda or Miniconda

Setup environment

Clone the repository and create the conda environment:

git clone https://github.com/MarkusKhoa/VISCERA-SAM.git
cd VISCERA-SAM

Option A — From environment.yml (recommended):

conda env create -f environment.yml
conda activate VisceraSam

Option B — Manual setup:

conda create --name VisceraSam python=3.8 -y
conda activate VisceraSam

pip install torch==2.0.0 torchvision==0.15.1 --index-url https://download.pytorch.org/whl/cu118
pip install -U openmim
mim install mmengine
mim install "mmcv==2.1.0"
mim install mmdet
mim install mmpretrain
pip install tensorboard scipy scikit-image SimpleITK tqdm

Download pre-trained weights

# UltraSam base checkpoint (required for fine-tuning)
wget -O ./UltraSam.pth "https://s3.unistra.fr/camma_public/github/ultrasam/UltraSam.pth"

# SAM ViT-B mapped weights (used by default configs)
# See weights/ directory for conversion scripts

Set PYTHONPATH

# Linux / macOS
export PYTHONPATH=$PYTHONPATH:.

# Windows (cmd)
set PYTHONPATH=%PYTHONPATH%;.

---

Project Structure

VISCERA-SAM/
├── configs/
│   ├── _base_/                  # Base configs (model, dataset, runtime)
│   └── UltraSAM/
│       ├── UltraSAM_full/       # Original UltraSam configs
│       └── ablation/            # Ablation study configs (FASS)
│           ├── baseline_fass.py
│           ├── ablation1_boundary_hausdorff.py
│           ├── ablation2_lcc.py
│           └── ablation3_geometric_aug.py
├── endosam/
│   ├── datasets/                # Custom transforms, evaluation metrics
│   ├── hooks/                   # Training hooks
│   ├── models/                  # SAM model, heads, backbones
│   │   └── dense_heads/
│   │       └── sam_mask_decoder.py   # SAMHead + HausdorffDTLoss
│   └── visualization/           # Visualization utilities
├── scripts/
│   ├── run_with_notify.py       # Run experiment + email notification
│   └── run_all_ablations.py     # Batch runner for all ablations
├── datasets/                    # Dataset converters and tools
├── weights/                     # Weight download & conversion scripts
├── sample_dataset/              # Sample data for quick test
├── environment.yml              # Conda environment specification
└── README.md

---

Quick Start (Inference)

Run inference on the sample dataset:

conda activate VisceraSam
export PYTHONPATH=$PYTHONPATH:.

mim test mmdet configs/UltraSAM/UltraSAM_full/UltraSAM_box_refine.py \
    --checkpoint UltraSam.pth \
    --cfg-options \
        test_dataloader.dataset.data_root="sample_dataset" \
        test_dataloader.dataset.ann_file="sample_coco_MMOTU2D.json" \
        test_dataloader.dataset.data_prefix.img="sample_images" \
        test_evaluator.ann_file="sample_dataset/sample_coco_MMOTU2D.json" \
    --work-dir ./work_dir/example \
    --show-dir ./show_dir

Predicted masks will be saved in ./show_dir.

---

Training on FASS

Dataset preparation

1. Download the FASS dataset
2. Place images in work_dir/datasets/images/ and the COCO annotation at work_dir/datasets/annotations/FASS__coco.json

Train the baseline

mim train mmdet configs/UltraSAM/ablation/baseline_fass.py \
    --work-dir ./work_dir/ablation/baseline

Test a checkpoint

mim test mmdet configs/UltraSAM/ablation/baseline_fass.py \
    --checkpoint ./work_dir/ablation/baseline/best_coco_segm_mAP_iter_30000.pth \
    --work-dir ./work_dir/ablation/baseline

---

Ablation Studies

Four experiments are defined under configs/UltraSAM/ablation/:

Config	Description
baseline_fass.py	Baseline: focal + dice + IoU loss, LCC off, RandomFlip only
ablation1_boundary_hausdorff.py	Baseline + boundary loss + Hausdorff DT loss
ablation2_lcc.py	Baseline + LCC post-processing (test-only, reuses baseline checkpoint)
ablation3_geometric_aug.py	Baseline + geometric augmentation (RandomResize + RandomCrop)

Run all ablations with email notification

# Set up email notifications (optional)
# Linux
export GMAIL_APP_PASSWORD=your_16_char_app_password

# Windows
set GMAIL_APP_PASSWORD=your_16_char_app_password

# Run all experiments sequentially
python scripts/run_all_ablations.py

Each experiment sends an email notification upon completion with a summary of results.

Run individually

python scripts/run_with_notify.py \
    --name "Baseline" \
    --mode train \
    --config configs/UltraSAM/ablation/baseline_fass.py \
    --work-dir work_dir/ablation/baseline

---

Evaluation Metrics

All experiments report the following metrics:

Metric	Log Key	Description
Bbox mAP	coco/bbox_mAP	COCO bounding box mean Average Precision
Segm mAP	coco/segm_mAP	COCO segmentation mean Average Precision
Dice	segm/Dice_mean	Mean Dice coefficient (per-class also reported)
mIoU	segm/IoU_mean	Mean Intersection over Union (per-class also reported)
Hausdorff	segm/Hausdorff_mean	Mean Hausdorff distance in pixels (lower is better)

Where to find results

• Text logs: work_dir/ablation/<name>/<timestamp>/<timestamp>.log
• TensorBoard: tensorboard --logdir work_dir/ablation/ --port 6006
• Best checkpoint: work_dir/ablation/<name>/best_coco_segm_mAP_iter_*.pth

---

Note: Create the repository on GitHub first at https://github.com/new (name: VISCERA-SAM, private or public as desired). Do not initialize with a README.

---

Original UltraSam

This project builds on top of UltraSam. The original US-43d dataset and full UltraSam training pipeline are documented below.

Click to expand US-43d dataset table

Dataset	Link
105US	researchgate
AbdomenUS	kaggle
ACOUSLIC	grand-challenge
ASUS	onedrive
AUL	zenodo
brachial plexus	github
BrEaST	cancer imaging archive
BUID	qamebi
BUS_UC	mendeley
BUS_UCML	mendeley
BUS-BRA	github
BUS (Dataset B)	mmu
BUSI	HomePage
CAMUS	insa-lyon
CardiacUDC	kaggle
CCAUI	mendeley
DDTI	github
EchoCP	kaggle
EchoNet-Dynamic	github
EchoNet-Pediatric	github
FALLMUD	kalisteo
FASS	mendeley
Fast-U-Net	github
FH-PS-AOP	zenodo
GIST514-DB	github
HC	grand-challenge
kidneyUS	github
LUSS_phantom	Leeds
MicroSeg	zenodo
MMOTU-2D	github
MMOTU-3D	github
MUP	zenodo
regPro	HomePage
S1	ncbi
Segthy	TUM
STMUS_NDA	mendeley
STU-Hospital	github
TG3K	github
Thyroid US Cineclip	standford
TN3K	github
TNSCUI	grand-challenge
UPBD	HomePage
US nerve Segmentation	kaggle

---

---

License

This project is licensed under CC BY-NC-SA 4.0, following the original UltraSam license.