Yolov5 training with Custom Dataset
By: Ashish Veeragoni
Tags: Computer Vision, Object Detection, Yolov5

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

yolov5 model comparison

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.

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:



$ %cd yolov5

$ !pip install -r requirements.txt # install dependencies

$ !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

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

# Weights & Biases (optional)

$ %pip install -q wandb

$ !wandb login

batch images

metrics

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

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