Follow the link below to see how the model was trained using Colab.
How to use in Colab:
1. Clone the required repositories, install required packages and download pretrained weights:
!git clone https://github.com/ultralytics/yolov5
!wget https://github.com/Stroma-Vision/machine-learning-challenge/releases/download/v0.1/challenge.zip
!unzip challenge.zip
!rm challenge.zip
%cd yolov5
!pip install -r requirements.txt
!pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu nvidia-tensorrt
!pip install onnxruntime-gpu
!wget https://github.com/sarpalgan/nuts-bolts/releases/download/v1.0.0/best.pt
%cd data
!wget https://github.com/sarpalgan/nuts-bolts/releases/download/v1.0.0/custom.yaml
%cd ..
2. Run the following code to export best.py PyTorch model to ONNX and TensorRT formats:
!python export.py --device 0 --weights best.pt --include engine --data data/data.yaml
3. Make detection with one of the following code:
# Detect with PyTorch
!python detect.py --name pt --weights best.pt \
--source /content/challenge/images/test/test.mp4 \
--data data/custom.yaml --conf-thres 0.85
# Detect with ONNX
!python detect.py --name onnx --weights best.onnx \
--source /content/challenge/images/test/test.mp4 \
--data data/custom.yaml --conf-thres 0.85
#Detect with TensorRT
!python detect.py --name engine --weights best.engine \
--source /content/challenge/images/test/test.mp4 \
--data data/custom.yaml --conf-thres 0.85 --device 0
4. Results will be saved into ./yolov5/runs/detect/model.extension
model.extension
pt
(PyTorch)onnx
(ONNX)engine
(TensorRT)