skip to content W3C SVG

Yolov5 training with Custom Dataset

Publish date: Dec 23, 2020
Tags:

joey

intro

I recently came across wildlife projects using AI and how they are helping in protecting endangered species and thought of making a model to help detecting tigers🐅 in the wild. Scroll down, too see how i trained my object detection model.

​

If you are active in computer vision, you may have heard about yolov5. There’s some controversy around its naming, you can read details from here. Ultralytics team put a great effort in open-sourcing this model 👏👏👏

​

yolo

I decided to use yolov5s model, which is light weight version and claiming better fps on edge devices.

data

Regarding data, I googled about tiger dataset and got to know about Amur Tiger Re-identification in the Wild (ATRW) dataset.

data-set website: https://cvwc2019.github.io/challenge.html

paper: https://arxiv.org/abs/1906.05586

Train set and Validation set consists of 2485 and 277 images respectively. And converted the data-set which is labelled in voc (.xml) format to yolo (.txt) format using this code: link

If your dataset had unlabelled images, use tools like CVAT, makesense.ai or labelme to annotate them

data directory looks like this:

├── train
│   ├── images
│   └── labels
└── valid
    ├── images
    └── labels

Prepare data.yaml file, before proceed to training like this:

train: ../train/images
val: ../valid/images

nc: 1
names: ['tiger']

here nc refers to number of classes.

train

Zip the entire folder along with yaml file and uploaded to google drive, so that easy to download in colab. Based on your luck and timing you may get P100 gpu in google colab, use it to train the model.

  • Clone repo and Install required dependecies:
$ !git clone https://github.com/ultralytics/yolov5  # clone repo
$ %cd yolov5
$ !pip install -r requirements.txt  # install dependencies

  • download the dataset from gdrive and unzip it.

    $ !gdown --id <your dataset download-id> ##download from drive
$ !unzip data.zip # unzip your dataset

go through this link to get familiar with gdown

  • start training by selecting input image size, batch size and setting number of epochs

    $ !python train.py --img 640 --batch 16 --epochs 300 --data data.yaml --weights yolov5s.pt

All training results are saved to runs/train/ directory

You can go through this colab notebook for more info.

Model Performance Metrics

You can setup weights and bias account and start seeing model metrics and visualise batch images while it is training. It is integrated with yolov5, so that its easy for you to setup.

# Weights & Biases (optional)
$ %pip install -q wandb  
$ !wandb login

​ batch images

wb1

​ metrics

wb2

Onnx model conversion

Install onnx tools

$ !pip install onnx>=1.7.0  # for ONNX export

Export model to onnx

$ !python models/export.py --weights yolov5s.pt --img 640 --batch 1  # export at 640x640 with batch size 1
  • you can see more about model export on this thread :link:
  • About pytorch to tensorflow model conversion :link:

inference

for inference on batch of images:

!python detect.py --source data/images --weights yolov5s.pt --conf 0.25

for video inference:

!python detect.py --source file.mp4 --weights yolov5s.pt --conf 0.25

tiger

final thoughts

Successfully trained the model, now looking forward to use it on edge device like raspberry pi. Last november PyTorch announced their officials builds for Arm64 devices. (nightly)