Yolov5 training with Custom Dataset

detection of tigers using yolov5

Introduction

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 👏👏👏

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 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:

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

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:

> download the dataset from gdrive and unzip it.

> go through this link to get familiar with gdown.

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

All training results are saved to runs/train/ directory

You can go through this colab notebook for more info.

Model Performance

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.

** run this step before you start training

Onnx model conversion

  • Install onnx tools
  • Export model to onnx
  • 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:

for video inference:

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store