Repository files navigation Fine-tuning a YOLO26-Segmentation model with custom dataset.
Torch Version: 2.12.0.dev20260326+cu128
Torchvision Version: 0.26.0.dev20260326+cu128
CUDA Device: NVIDIA GeForce RTX 5050 Laptop GPU
segTrain/ : Base dataset, "YOLO with Images" export from Label-Studio .
segData/ : Cleaned up dataset, running SplitSet.ipynb turns segTrain/ into segData/. This one has 3 labels (0, 1, 2) -> (battery, knight, bottle)
dataset.yaml : Need for training, edit to your own needs.
runs/ : training output.
output/ : testing output from ONNXInference.ipynb and TRTInference.ipynb.
SplitSet.ipynb : Start here, splits your dataset into training and validation (segData).
Training.ipynb : Fine-tunes yolo26m-seg.pt with dataset (dataset.yaml), produces segModel.pt.
Testing.ipynb : Test segModel.pt on images, produces processed images in output/.
Conversion.ipynb : Convert to ONNX (segModel.onnx) and TensorRT (segModel.engine).
ONNXInference.ipynb : Tests segModel.onnx on images, has custom NMS due to ONNX export not allowing inbuilt nms, produces images in output/.
Bonus: Custom neck/choke detection code leftover from my arm code.
TRTInference.ipynb : Tests segModel.engine on images, produces images in output/.
About
YOLO 26 Segmentation Fine-tuning pipeline.
Topics
Resources
Stars
Watchers
Forks
You can’t perform that action at this time.