Skip to content

Alexledev/Yolo26Seg

Repository files navigation

Fine-tuning a YOLO26-Segmentation model with custom dataset.

License

GPU

PyTorch ONNX TensorRT

Torch Version: 2.12.0.dev20260326+cu128

Torchvision Version: 0.26.0.dev20260326+cu128

CUDA Device: NVIDIA GeForce RTX 5050 Laptop GPU

Guide:

Folders:

  • 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.

Notebooks/Running:

  • 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/.

Releases

No releases published

Packages

 
 
 

Contributors