Jugular Vein Segmentation using SegFormer

A PyTorch implementation of SegFormer-based semantic segmentation for Near-Infrared (NIR) jugular vein images.

The repository provides:

• 5-Fold Cross Validation training
• Test-Time Augmentation (TTA)
• Ensemble inference across folds
• Automatic threshold sweeping
• Video inference
• Multiple segmentation losses

---

Features

• SegFormer (MiT-B0–B5 backbones)
• Dice Loss
• Focal Loss
• Dice + Focal Loss
• Tversky Loss
• Focal Tversky Loss
• CLAHE preprocessing
• Albumentations data augmentation
• Ensemble inference
• Automatic threshold optimisation
• IoU and Dice evaluation

---

Repository Structure

.
├── configs.py
├── dataset.py
├── train.py
├── infer.py
├── infer_video.py
├── evaluate.py
├── model.py
├── loss.py
├── requirements.txt
├── README.md
│
├── checkpoints/
├── data/
│   ├── train/
│   ├── val/
│   └── test/
│
└── reports/

---

Dataset Structure

Expected directory layout:

data/
│
├── train/
│   ├── images/
│   └── masks/
│
├── val/
│   ├── images/
│   └── masks/
│
└── test/
    ├── images/
    └── masks/

Images may be

• JPG
• PNG
• JPEG
• TIFF

Masks must be

PNG

with the same filename stem as the image.

Example

001.jpg
001.png

---

Installation

Clone the repository

git clone https://github.com/<username>/<repository>.git

cd <repository>

Create a virtual environment (optional)

python -m venv venv

Activate it.

Install the appropriate PyTorch build for your system from:

https://pytorch.org/get-started/locally/

For example, with CUDA 12.1:

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

Then install the remaining dependencies:

pip install -r requirements.txt

---

Training

Configure hyperparameters inside

configs.py

Start training

python train.py

Training performs 5-Fold Cross Validation and stores the best checkpoint for each fold.

Saved checkpoints

checkpoints/

best_model_fold1.pth
best_model_fold2.pth
...
best_model_fold5.pth

---

Single Image Inference

from infer import infer

infer(
    image_path="data/test/images/sample.png",
    fold=1,
    output_path="prediction.png",
)

---

Folder Inference

from infer import infer_directory

infer_directory(
    input_dir="data/test/images",
    output_dir="predictions",
    ensemble=True,
)

---

Video Inference

from infer_video import infer_video

infer_video(
    video_path="input.mp4",
    output_path="output.mp4",
    ensemble=True,
    overlay=True,
)

---

Threshold Evaluation

Run

python evaluate.py

The script automatically

• loads all fold checkpoints
• performs ensemble inference
• sweeps multiple thresholds
• computes IoU and Dice
• identifies the optimal threshold
• saves masks and overlay images
• exports CSV reports

Generated outputs are written under

reports/

---

Configuration

Most settings can be modified from

configs.py

including

• SegFormer backbone
• image size
• learning rate
• batch size
• number of folds
• loss function
• threshold
• CLAHE
• Test-Time Augmentation

---

Supported Backbones

nvidia/mit-b0
nvidia/mit-b1
nvidia/mit-b2
nvidia/mit-b3
nvidia/mit-b4
nvidia/mit-b5

---

Supported Losses

dice
focal
dice_focal
tversky
tversky_focal

---

Outputs

Training

checkpoints/
training_log.txt
best_model_fold1.pth
...

Evaluation

reports/

threshold_report.txt
threshold_sweep.csv
threshold_sweep_summary.csv
masks/
overlays/

---

License

This project is released under the MIT License.