MEMTrack

Deep learning based automated detection and tracking of bacteria in complex environments such as Collagen.

Project Description

Tracking microrobots is a challenging task, considering their minute size and high speed. As the field progresses towards developing microrobots for biomedical applications and studying them in physiologically relevant or in vivo environments, this challenge is exacerbated by the dense surrounding environments with feature size and shape comparable to those of microrobots. To address this challenge, we developed Motion Enhanced Multi-level Tracker (MEMTrack), a robust pipeline for detecting and tracking micro-motors in bright-field microscopy videos using synthetic motion features, deep learning-based object detection, and a modified Simple Online and Real-time Tracking (SORT) algorithm with interpolation for tracking. Our object detection approach combines different models based on the object’s motion pattern. We trained and validated our model using bacterial micro- motors in the tissue-like collagen environment and tested it in collagen and liquid (aqueous) media. We demonstrate that MEMTrack can accurately predict and track even the most challenging bacterial micro-motors missed by skilled human annotators, achieving precision and recall of 77% and 48% in collagen and 94% and 35% in liquid media, respectively. We also show that MEMTrack is able to accurately quantitate the average speed of bacterial micromotors with no statistically significant difference from the laboriously produced manual tracking data. Our proposed pipeline not only represents a significant contribution to the field of microrobot image analysis and tracking using computer vision but also opens the potential of applying deep learning methods in vision-based control of microrobots for various applications, including disease diagnosis and treatment.

Paper Link: https://onlinelibrary.wiley.com/doi/full/10.1002/aisy.202300590

Paper Arxiv Link: https://arxiv.org/abs/2310.09441 (For updated results and information)

Google Colab Tutorial: https://colab.research.google.com/drive/1PltZ8q_AEmVxoSoh5gog288GcPfOgvZy?usp=sharing

Getting Started

Installation

To get started with this project, follow these steps:

1. Clone the repository to your local machine using Git:

   git clone https://github.com/sawhney-medha/MEMTrack.git
   

2. Navigate to the project directory:

   cd MEMTrack
   

3. Create environemnt and install the required dependencies using pip :

    python3.8 -m venv memtrack_env
    source memtrack_env/bin/activate
    cd MEMTrack
    
    pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu118 (Check PyTorch version according to CUDA version/drivers)
    python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'
    pip install -r requirements.txt
   
   

4. Download models.zip containing pre-trained and store it in MEMTrack/models

Download Pretrained Models

Reproducing results on sample test data

• Download Sample Data
• Run Data Loading (Adding Data and Preprocessing Data) and Feature Generation Scripts for sample data.
• Follow Inference, Tracking and Evaluation instructions below.
• /src/Automated-ConfScoreThresholds.ipynb Notebook has the code for generating confidence score thresholds for trained models on the validation set.
• src/CollagenPRPlot.ipynb Notebook generates the results on our test set and creates the plot to visualize it.

Inference on custom data using pre-trained models

• Generate Data and store in the same format as mentioned in Adding Data
• Run Data Loading (Adding Data and Preprocessing Data) and Feature Generation Scripts for custom data.
• Ensure custom data is in the mentioned format, else write your own code to have the same output format as the preprocessing code.
• Since inference test data will not have a Raw Data.csv, the code will automaticallly load data assuming no ground truth annotations were present in the data.
• Run Inference script and follow instructions in Inference from Trained Object Detector Model. The coco_instances json file will have a list of all predicted bounding boxes for every frame.
• Run Tracking scripts to generate tracklets

Evaluation on custom data using pre-trained models

• Follow same instructions as Inference but have an updated RawData.csv with ground truth annotations.
• For evaluation follow instructions as mentioned in Evaluation and Analysis

Data Preparation

Adding Data

1. Data Collection

   • We recorded bacteria (i.e., the micromotor in bacteria-based biohybrid microrobots) swimming behavior in collagen, as a tissue surrogate, and in an aqueous environment. In order to generate training, validation, and test datasets for MEMTrack, the microscopy videos acquired in collagen and in aqueous media were imported into ImageJ software. MTrackJ plugin was used to label all bacteria in each frame of each video manually, and their x and y coordinates were recorded.
   • Any data can be used as long as a list of frames along with their respective x and y coordinates are available.
   • To evaluate the results, track_ids would also be required for the annotated bacteria.
   • To be able to train different models based on different motilites, motility sub population labels would also be required. If these are not available, one can train a single model for all annotated bacteria.

2. Organize Data Directory:

   • The final data should be packed in a zip file with the following structure:

     ├── video_zip_file/
     │   ├── Images without Labels/
     │   │   ├── sample_name1.jpg
     │   │   ├── sample_name2.jpg
     │   │   └── ...
     │   ├── Raw Data.csv
     

Download Sample Data

3. Data Format:

• RawData.csv should have the following columns:
  • "TID" --> Track Id
  • "PID" --> Picture Id/Frame Number
  • "x [pixel]"" --> x coordinate
  • "y [pixel]" --> y coordinate
  • "subpopulation" -->containing "N"/"L"/"M"/"H" for subpopulations (Non Motile/Low Motility/Mid Motility/High Motility)

Preprocessing Data

4. Preprocessing Code:

   • Run the DataPreparation_Motility.ipynb notebook located in the MEMTrack/src directory and update the path variables (Cell 3) according to your directory structure.
   • The code needs to be run for every video that needs to be loaded to data/
   • The videomap.txt will automatically be generated after the preprocessing code

5. Preprocessed Data Directory:

   • Expected Directory Structure after running the preprocessing code:

     ├── MEMTRack/
     │   ├──data/
     │   │   ├── videomap.txt
     │   │   ├── collagen/
     |   |   |     ├── video1/
     |   |   |     |    ├── frame1/
     |   |   |     |         ├── annotations_motility_no/
     |   |   |     |         |    ├── 0.txt
     |   |   |     |         |    ├── 1.txt
     |   |   |     |         |    └── ...
     |   |   |     |         ├── annotations_motility_low/
     |   |   |     |         ├── annotations_motility_mid/
     |   |   |     |         ├── annotations_motility_high/
     |   |   |     |         ├── bacteria/
     |   |   |     |         |    ├── 1/
     |   |   |     |         |    |    ├── xy_coord/
     |   |   |     |         |    |          ├── 0.txt
     |   |   |     |         |    |          ├── 1.txt
     |   |   |     |         |    |          └── ...
     |   |   |     |         |    |
     |   |   |     |         |    ├── 2/
     |   |   |     |         |    └── ...
     |   |   |     |         ├── images/
     |   |   |     |              ├── 0.tif
     |   |   |     |              ├── 1.tif
     |   |   |     |              └── ...
     |   |   |     ├── video2/
     |   |   |     └── ...
     │   │   └── ...
     │   ├── src/
     

Data Usage

7. Feature Generation Code:

   • Run the DataFeaturePreparation.ipynb notebook located in the MEMTrack/src directory and update the path variables (Cell 3) according to your directory structure.
   • Also update path variables and choose feature generation method at the end of the notebook (Cell 18 and 19 for Training Data and Cell 20 for Test/Inefrence) to generate and store features that would be used for training.
   • The notebook provides multiple ways to generate features, the one recommended based on experiments on Collagen data is: "optical_flow_median_back". This generates 3 channels for each frame: 1. Original Frame, 2. Consecutive Optical Flow Vector and 3. Difference from Median Background.
   • Similarly, "optical flow median back" with "optical_flow_prior"=x variable is optical flow from xth previous frame. "diff_from_max_absolute_consecutive_frame_diff" creates a feature for difference from the "max consecutive frame diff" feature, with a frame diff prior for xth frame diff
   • The train/test/val split can be provided in this code as dict of video numbers that have been loaded and accordingly their fearture sets will be generated.

8. Final Data Directory:

   • Expected Directory Structure after feature generation code:

        ├── MEMTRack/
        │   ├──data/
        │   ├──data_features/
        │   │   ├── exp_name/
        |   |   |     ├── data_features_set/
        |   |   |     |    ├── train/
        |   |   |     |    |     ├── annotations_motility_no/
        |   |   |     |    |     |    ├── 0.txt
        |   |   |     |    |     |    ├── 1.txt
        |   |   |     |    |     |    └── ...
        |   |   |     |    |     ├── annotations_motility_low/
        |   |   |     |    |     ├── annotations_motility_mid/
        |   |   |     |    |     ├── annotations_motility_high/
        |   |   |     |    |     ├── images/
        |   |   |     |    |          ├── 0.tif
        |   |   |     |    |          ├── 1.tif
        |   |   |     |    |          └── ...
        |   |   |     |    |     ├── images_feature/
        |   |   |     |    |          ├── 0.tif
        |   |   |     |    |          ├── 1.tif
        |   |   |     |    |          └── ...
        |   |   |     |    ├── test/
        |   |   |     |    ├── val/
        |   |   |     ├── data_feature_video1/test/
        |   |   |     ├── data_feature_video2/test/
        |   |   |     └── ...
        │   │   └── ...
        │   ├── src/
     
     
     

The following sections describes the training, inference, tracking and evaluation procedures. The codebase is built using Python, PyTorch and Detectron 2.0.

Training Object Detector Model

• Run the training script /scripts/train.sh to start training from the MEMTRack/ root directory.

• Update exp_name, data_path (feature directory) and output_dir paths as approriate.

• The training parameters such as learning rate, epochs, etc. can be updated from the bash script.

• There are two python scripts for training: src/trainBacteriaRetinanetMotionData_Motility.py and src/trainBacteriaRetinanetMotionData_Motility_Val_loss.py. The only difference is that the latter generates loss plots for validation set during training that the Detectron2 code does not automatically generate, which can be visualized in the src/Train_Val_Loss_Curve.ipynb notebook. The code for it has been taken from the following repository: https://gist.github.com/ortegatron/c0dad15e49c2b74de8bb09a5615d9f6b

• The training script saves regular checkpoints along with the final model, which are stored in the output_dir specifed in the bash script.

• Post trainng save the trained model checkpoint in /src/models/

   bash scripts/train.sh
  

Inference from Trained Object Detector Model

• Run the inference script /scripts/inference.sh from the MEMTrack/ root directory.

• Update exp_name, data_path (feature directory) and model_dir (directory with trained models) paths as approriate.

• The inference.sh scripts calls the /src/inferenceBacteriaRetinanet_Motility_v2.py script, the paramters for which can be updated in the bash script. The output from inference is a json file containing object predictions, the json file is saved in the output_dir of the model.

   bash scripts/inference.sh
  

Tracking

• Run the tracking script /scripts/tracking.sh from the MEMTrack/ root directory.

• The tracking scripts generates data using the object detector predictions in the format expected by the SORT algorithm and then implements tracking.

• The script also calls a script to generate output videos from the tracked data.

• The parameters to update for trackiing code generation such as confidence thresholds, or tracking parameters like min_age can be updated in the bash script.

  bash scripts/tracking.sh
  

Evaluation and Analysis

• Run the evaluation script /scripts/test_set_eval.sh from the MEMTrack/ root directory.

• The script will automatically generate test_results_.txt for every step of the MEMTrack pipline and stor the True Postives, False Positives and False Negatives to generate results on test sets.

    bash scripts/test_set_eval.sh

• Can also just call /scripts/run.sh to automate the inference to evaluation process.

• Finally run src/CombinedTestSet-StepWiseEvaluation.ipynb notebook to generate precision and recall values for an entire test set along with plots to visualize results

Debugging

• All known issues are consistently updated in debug.md. Please refer to it before raising an issue.

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgement

This research was supported in part by NSF grants CBET-2133739, CBET-1454226 and 4-VA grant to Dr. Bahareh Behkam, and NSF grant IIS-2107332 to Dr. Anuj Karpatne. Access to computing resources was provided by the Advanced Research Computing (ARC) Center at Virginia Tech.