Surgical Instrument Segmentation in Endoscopic Spine Surgery — Code Release

This repository accompanies the manuscript "Deep Learning Architectures for Surgical Instrument Segmentation in Endoscopic Spine Surgery" and provides the code needed to reproduce all training and evaluation results.

Repository: https://github.com/grotyx/endo_instruments

The release covers seven segmentation pipelines:

1. U-Net
2. Attention U-Net
3. U-Net++
4. SegFormer-B0
5. DeepLabV3+ (ResNet50, ImageNet-pretrained)
6. nnU-Net v2 (5-fold cross-validation, external code)
7. MedSAM2-Tiny + YOLOv8-nano — fully automatic two-stage foundation-model pipeline

A SAM2.1-Large variant of the foundation-model pipeline is included as a robustness check (scripts/train_medsam2.py --variant large) and is reported in Supplementary Table 5 of the manuscript.

Repository layout

endo_instruments/
├── README.md
├── requirements.txt
├── LICENSE                        MIT
├── configs/
│   ├── config_variant_A.yaml      biportal-only training (Variant A)
│   └── config_variant_B.yaml      combined biportal + uniportal training (Variant B)
├── sam2_configs/                  Hydra configs for the SAM2 model
│   ├── sam2.1_hiera_t512.yaml     MedSAM2-Tiny @ 512×512
│   └── sam2.1_hiera_l512.yaml     SAM2.1-Large @ 512×512
├── models/                        architecture definitions
│   ├── medsam2_wrapper.py
│   ├── unet.py
│   ├── attention_unet.py
│   ├── unet_plus.py
│   ├── segformer.py
│   └── deeplabv3_plus.py
├── utils/                         data loading, losses, metrics, augmentation
│   ├── dataset.py
│   ├── losses.py
│   ├── metrics.py
│   └── augmentation.py
├── scripts/
│   ├── train_medsam2.py           MedSAM2-Tiny / SAM2.1-Large fine-tuning
│   ├── train_yolo.py              YOLOv8-nano detector training (Variants A & B)
│   ├── evaluate_medsam2.py        GT-bbox oracle evaluation
│   └── auto_bbox_pipeline.py      YOLO + MedSAM2 end-to-end (deployment)
└── data/
    └── README_data.md             dataset layout and access instructions

Requirements

The pipeline was developed and tested with:

• Python 3.12
• PyTorch 2.10 + CUDA 12.8
• NVIDIA RTX 4090 (24 GB)

Install dependencies:

pip install -r requirements.txt

MedSAM2 / SAM2 setup

git clone https://github.com/bowang-lab/MedSAM2.git
cd MedSAM2
pip install -e .

After installation, copy the two Hydra configs from sam2_configs/ into the cloned MedSAM2/sam2/configs/ directory (sam2.1_hiera_t512.yaml ships with MedSAM2; sam2.1_hiera_l512.yaml is a 512-resolution variant of Meta's sam2.1_hiera_l.yaml provided here).

Pretrained checkpoints

Download into weights/medsam2/:

• MedSAM2_latest.pt — wanglab/MedSAM2 on Hugging Face
• sam2.1_hiera_large.pt — https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt
• yolov8n.pt — auto-downloaded by Ultralytics on first use

Datasets

See data/README_data.md for the expected directory layout and dataset access instructions.

Reproducing the paper

All commands assume the repository root as the working directory.

# 1. Prepare YOLO-format bbox annotations from binary masks
#    (convert your mask PNGs to YOLO .txt label files and build data.yaml)
#    Use standard Ultralytics tooling or derive bboxes from mask contours.

# 2. Train the YOLO detectors (Variants A and B)
python scripts/train_yolo.py --variant both --epochs 100 --patience 20

# 3. Train MedSAM2-Tiny for both variants
python scripts/train_medsam2.py --config configs/config_variant_A.yaml --variant tiny --epochs 50 --patience 10
python scripts/train_medsam2.py --config configs/config_variant_B.yaml --variant tiny --epochs 50 --patience 10

# (optional) Train SAM2.1-Large variant
python scripts/train_medsam2.py --config configs/config_variant_A.yaml --variant large --epochs 50 --patience 10

# 4. Evaluate under GT-bbox prompts (oracle upper bound)
python scripts/evaluate_medsam2.py --segmentor tiny \
    --decoder_ckpt checkpoints/variant_A/medsam2_tiny_best.pth \
    --out_dir results/variant_A/evaluation

# 5. Run fully automatic YOLO + MedSAM2 pipeline
python scripts/auto_bbox_pipeline.py --segmentor tiny --detector A \
    --decoder_ckpt checkpoints/variant_A/medsam2_tiny_best.pth \
    --test internal \
    --out results/variant_A/auto_bbox/medsam2_tiny__yolo_A__internal.csv
# Repeat for temporal, uniportal_v2, and Variant B with detector B

The six conventional baselines (U-Net, Attention U-Net, U-Net++, SegFormer-B0, DeepLabV3+, nnU-Net) are not retrained by this script set; any standard implementation trained with the hyperparameters in configs/config_variant_*.yaml (AdamW lr=1e-4, cosine annealing warm restarts, composite Dice+BCE+Focal loss, 150 epochs, early stopping patience=15) will reproduce the reported numbers.

Citation

Citation will be added on publication. The MedSAM2 model used here is from:

• Ma, J., et al. MedSAM2: Segment Anything in 3D Medical Images and Videos. Wang Lab, Toronto. https://github.com/bowang-lab/MedSAM2
• Ravi, N., et al. SAM 2: Segment Anything in Images and Videos. arXiv 2024;2408.00714.

The YOLOv8 detector is from:

• Ultralytics YOLOv8. https://github.com/ultralytics/ultralytics

License

MIT License. See LICENSE. Third-party dependencies (SAM2/MedSAM2: Apache-2.0; Ultralytics: AGPL-3.0; nnU-Net: Apache-2.0) retain their own licenses.