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.
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
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:
│ ├── images
│ └── labels
data.yaml file, before proceed to training like this:
val: ../valid/imagesnc: 1
nc refers to number of classes.
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
You can go through this colab notebook for more info.
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
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:
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