Semantic segmentation pipeline built on top of NVIDIA SegFormer, fine-tuned for detecting wires and obstacles in agricultural imagery. Supports both binary (background/wire) and multilabel (background/wire/obstacle) segmentation.
The model uses a pre-trained SegFormer encoder-decoder from HuggingFace Transformers, trained with PyTorch Lightning and tracked via Weights & Biases. Training uses weighted Dice loss, Albumentations-based augmentation, and cosine annealing LR scheduling at 512x512 resolution.
- Python >= 3.13
- CUDA-capable GPU (tested with CUDA 12.6)
Dependencies are managed with uv:
git clone https://github.com/adimac13/image-segmentation.git
cd image-segmentation
uv syncThe pipeline expects a COCO-format dataset. The dataset directory should contain an images/ folder with PNG files and a JSON annotation file.
Category mapping:
| ID | Category |
|---|---|
| 1 | apple |
| 2 | background |
| 3 | obstacle |
| 4 | leaves |
| 5 | wire |
If your dataset is split across multiple batches, use data/json_concatenation.py to merge annotations into a single file.
Set the DATASET_PATH environment variable to point to your dataset root before running.
Training:
python -m train.segformer_binary
python -m train.segformer_multiclassInference (loads a checkpoint, runs prediction on a random image, displays the result):
python -m inference.segformer_binary
python -m inference.segformer_multiclassThe losses/ module provides a custom Dice loss supporting binary and multiclass modes with per-class weighting for handling class imbalance.