This repository contains notebooks and resources used to train a state-of-the-art military vehicle tracker. Its main focus is on building a dataset of relevant images and annotations to fine-tune pre-trained object detection models, namely a Yolov8 model. The yolo-tracking library is used to provide the multi-object tracker algorithm.
After a first pass at training such a model from images available from object detection datasets, we explore other options to improve the performance of our model. We use scraping tools to collect more images of military vehicles from Google images. We also extend the classes to be able to discriminate between different types of vehicles: Armoured Fighting Vehicle (AFV), Armoured Personnel Carrier (APC), Military Engineering Vehicle (MEV) and Light armoured vehicle (LAV). We provide a sample annotated dataset to test performance improvement from extending our training data.
We also explore using diffusion models and the dreambooth method to generate new training images in different scenes and conditions.
- The adomvi directory contains jupyter notebooks to create a dataset of military vehicles, use this dataset to finetune a Yolov8 model for object detection, and to run object tracking on video inputs.
- The resources directory contains video samples for vehicle detection task.
You can install the project locally using poetry with
poetry install
You can run the notebooks from this project in Google Colab to benefit from GPU acceleration:
- 01 - Train a YOLOv8 model with a custom dataset:
- 02 - Run tracking using the trained model on a sample video:
- 03 - Scrape images from google to extend the training dataset:
- 04 - Fine tune Dreambooth to generate images of a tank:
Some sample results of tracking different types of military vehicles (AFV, APC, MEV, LAV) using a finetuned yolov8-large model.