Deep-Worm-Tracker

This repository contains the code to Deep-Worm-Tracker a realtime Deep Learning based C. elegans worm tracker. It is a combination of Deep object detection (YOLO) and tracking (Strong SORT) models.

Features

A large scale annotated dataset containing 3000 worm images is made available for training the object detection model.
A first of its kind C. elegans re-id dataset containing 32 worm identities is used for training the tracking model.
Annotated images account for background variability like worm trails, eggs, change in magnification, dust particles and marker prints used for demacrating quadrants in actual chemotaxis assays.
Training time reduced to just 9 to 26 min based on network dimensions.
Fast inference speed of 8 to 15 ms based on the YOLO object detection model.
Added functionality to segment and skeletonize tracked worms.
Individual worm trajectories are also highlighted.

Steps to Installation

It is reccomended to create separate conda environments for installing individual packages for the object detection (training YOLO model), tracking (training the torchreid model) and combined final model (for getting final track results). The steps would look something like:

Clone the repository:

git clone https://github.com/knaticat/Deep-Worm-Tracker.git

Setup environments:

   #for running the Deep-Worm-Tracker using pretrained model weights (Quick start)
   cd Deep-Worm-Tracker
   conda create -n yolostrong
   pip install -r requirements.txt
   
   #for training yolo model
   cd yolov5
   conda create -n yolo
   pip install -r requirements.txt
   
   #for training torchreid model
   cd strong_sort
   conda create -n torchreid
   pip install -r requirements.txt

Running the tracker after cloning the Deep-Worm-Tracker repository:

   $ python3 track.py --source <can be video, webcam, image file> --yolo-weights <path to weights file stored in worm_object_weights> 
   --strong-sort-weights <one of mobilenetv2_x1_0_worm.pt, mobilenetv2_x1_4_worm.pt, osnet_ain_x0_5_worm.pt, osnet_x0_5_worm.pt, osnet_x0_25_worm.pt>
   --img <network dimension> --do-segment --do-skeleton --show-track --show-id-black --show-vid --save-vid

Training YOLO model

The step by step instructions for training a YOLO model with a new dataset can be found in YOLO training . The worm detection dataset is available at worm-data which can be extracted and training resumed. A combination of DarkMark and Roboflow were used for annotating the dataset images. It's reccomended to use Roboflow for future annotations on newer image data for the user friendly interface.

Training Strong SORT reid model

We have made available a re-id dataset for C. elegans that contains 32 worm identities and is structured similar to the MARS dataset . We use torchreid library for training different reid models. Pretrained checkpoint files for different reid models can be found here . The available pretrained models are: mobilenetv2_x1_0_worm.pt, mobilenetv2_x1_4_worm.pt, osnet_ain_x0_5_worm.pt, osnet_x0_5_worm.pt, osnet_x0_25_worm.pt. These model weight files are automatically downloaded by the script. More details on training Strong SORT model can be found in the strong_sort folder.

YOLOv4-DeepSORT training

The entire dataset for training the YOLOv4-DeepSORT model remain the same. The individual code and configuration files used in this study can be found here .

Acknowledgements

This work was possible due to the contributions from:

Yolov5-StrongSORT_OSNet by mikel-brostrom
YOLOv5 by ultralytics
Deep-Person-Reid by KaiyangZhou
YOLOv4 by AlexeyAB
DeepSORT by nwojke
YOLOv4-DeepSORT by theAIGuysCode
DarkMark by stephanecharette

Cite the work

@article{BANERJEE2023106024,
title = {Deep-worm-tracker: Deep learning methods for accurate detection and tracking for behavioral studies in C. elegans},
journal = {Applied Animal Behaviour Science},
volume = {266},
pages = {106024},
year = {2023},
issn = {0168-1591},
doi = {https://doi.org/10.1016/j.applanim.2023.106024},
url = {https://www.sciencedirect.com/science/article/pii/S016815912300196X},
author = {Shoubhik Chandan Banerjee and Khursheed Ahmad Khan and Rati Sharma},
keywords = {, Real time tracking, Deep learning, Tracker},
abstract = {Accurate detection and tracking of model organisms such as C. elegans worms remains a fundamental task in behavioral studies. Traditional Machine Learning and Computer Vision methods produce poor detection results and suffer from repeated ID switches during tracking under occlusions and noisy backgrounds. Considering this, we propose Deep-Worm-Tracker, an end-to-end Deep Learning (DL) model, which is a combination of You Only Look Once (YOLOv5) object detection model and Strong Simple Online Real Time Tracking (Strong SORT) tracking backbone that is highly accurate and provides tracking results in real-time inference speeds. Present literature has few solutions to track animals under occlusions and even fewer publicly available large-scale animal re-ID datasets. Thus, we also provide a worm re-ID dataset to minimize worm ID switches, which, to the best of our knowledge, is the first of its kind for C. elegans. We are able to track worms at a mean Average Precision (mAP@0.5) >98% within just 9 min of training time with inference speeds of 9–15 ms for worm detection and on average 27 ms for worm tracking. Our tracking results show that Deep-Worm-Tracker is well suited for ethological studies involving C. elegans.}
}

somacoder/Deep-Worm-Tracker