uploading MEMTrack codebase

.gitignore

@@ -6,3 +6,8 @@ flagged/
 frames/
 diff_frames/
 *.wmv
+data/
+raw_data/
+DataFeatures/
+app_trial.py
+models*

MEMTrack/.gitignore

@@ -0,0 +1,13 @@
+#ignore models
+models/*
+
+#zip files
+*.zip
+
+#ignore data
+data/*
+!data/sample_videomap.txt
+/DataFeatures/*
+
+#ipynb checkpoints
+*/.ipynb_checkpoints

MEMTrack/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Medha Sawhney
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

MEMTrack/README.md

@@ -0,0 +1,228 @@
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) ![Python Version](https://img.shields.io/badge/python-3.8-blue) ![PyTorch Version](https://img.shields.io/badge/pytorch-1.9.0-blue) ![Detectron2 Version](https://img.shields.io/badge/Detectron2-v0.5.1-blue) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1PltZ8q_AEmVxoSoh5gog288GcPfOgvZy?usp=sharing) [![arXiv](https://img.shields.io/badge/arXiv-1234.56789-b31b1b.svg)](https://arxiv.org/abs/2310.09441)
+
+
+
+
+# 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:
+
+   ```bash
+   git clone https://github.com/sawhney-medha/MEMTrack.git
+   
+2. Navigate to the project directory:
+
+   ```bash
+   cd MEMTrack
+
+3. Create environemnt and install the required dependencies using pip:
+
+   ```bash
+    python3.8 -m venv memtrack_env
+    source memtrack_env/bin/activate
+    cd MEMTrack
+    pip install -r requirements.txt
+    python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'
+
+   
+4. Download models.zip containing pre-trained and store it in *MEMTrack/models*
+
+**[Download Pretrained Models](https://drive.google.com/file/d/1agsLD5HV_VmDNpDhjHXTCAVmGUm2IQ6p/view?usp=sharing)**
+
+### Reproducing results on sample test data
+- **[Download Sample Data](https://drive.google.com/file/d/1YHgT--xyMrQYeyaBYiIxIfLdTu2nW7nQ/view?usp=sharing)**
+- Run Data Loading ([Adding Data](#adding-data) and [Preprocessing Data](#preprocessing-data)) and [Feature Generation](#data-usage) Scripts for sample data.
+- Follow [Inference](#inference-from-trained-object-detector-model), [Tracking](#tracking) and [Evaluation](#evaluation-and-analysis) 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](#adding-data)
+-  Run Data Loading ([Adding Data](#adding-data) and [Preprocessing Data](#preprocessing-data)) and [Feature Generation](#data-usage) 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](#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](#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](https://drive.google.com/file/d/1YHgT--xyMrQYeyaBYiIxIfLdTu2nW7nQ/view?usp=sharing)** 
+
+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
+
+6. **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.
+
+
+6. **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
+   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
+   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
+    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_<step>.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
+    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](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](https://me.vt.edu/people/faculty/behkam-bahareh.html), and NSF grant IIS-2107332 to [Dr. Anuj Karpatne](https://people.cs.vt.edu/karpatne/). Access to computing resources was provided by the Advanced Research Computing (ARC) Center at Virginia Tech.

MEMTrack/debug.md

@@ -0,0 +1,19 @@
+# Debugging known Issues
+
+1. **Error with OpenCV functions**
+   OpenCV(4.5.1) C:\Users\appveyor\AppData\Local\Temp\1\pip-req-build-1drr4hl0\opencv\modules\highgui\src\window.cpp:651: error: (-2:Unspecified error) The function is not implemented. Rebuild the library with Windows, GTK+ 2.x or Cocoa support.
+
+   ```bash
+   pip uninstall opencv-python-headless -y
+   pip install opencv-python --upgrade
+   
+2. **Numpy error**
+ImportError: lap requires numpy, please "pip install numpy"
+
+   ```bash
+   pip install numpy
+
+3. **Installing pycocotools**
+Gthub link fails
+   ```bash
+   pip install pycocotools>=2.0.2

MEMTrack/requirements.txt

@@ -0,0 +1,34 @@
+pip==22.0.3
+numpy
+filterpy==1.4.5
+json5==0.9.6
+jsonlines==3.0.0
+jsonpatch==1.32
+jsonpointer==2.1
+jsonschema==3.2.0
+jupyter-client==7.0.3
+jupyter-core==4.8.1
+jupyter-server==1.11.0
+jupyterlab==3.1.13
+jupyterlab-pygments==0.1.2
+jupyterlab-server==2.8.1
+jupyterlab-widgets==1.0.2
+lap==0.4.0
+matplotlib==3.3.4
+motmetrics==1.2.5
+natsort==7.1.1
+numpy==1.19.5
+opencv-python==4.5.5.64
+opencv-python-headless==4.5.4.60
+openpyxl
+pandas==1.1.5
+plotly==5.11.0
+scikit-image==0.17.2
+scikit-learn==0.24.2
+scipy==1.5.4
+seaborn==0.11.2
+torch==1.9.1
+torchfile==0.1.0
+torchmetrics==0.5.1
+torchvision==0.10.1
+tqdm==4.62.3

MEMTrack/scripts/inference.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+set -e
+exp_name="collagen_motility_inference" # exp name
+data_path="/home/medha/BacteriaDetectionTracking/MEMTrack/DataFeatures/" # feature directory path
+data_path+=$exp_name
+echo $data_path 
+
+
+#path to saved models
+low_motility_model_path="/home/medha/BacteriaDetectionTracking/MEMTrack/models/motility/low/collagen_optical_flow_median_bkg_more_data_90k/" 
+wiggle_motility_model_path="/home/medha/BacteriaDetectionTracking/MEMTrack/models/motility/wiggle/collagen_optical_flow_median_bkg_more_data_90k/" 
+mid_motility_model_path="/home/medha/BacteriaDetectionTracking/MEMTrack/models/motility/mid/collagen_optical_flow_median_bkg_more_data_90k/" 
+high_motility_model_path="/home/medha/BacteriaDetectionTracking/MEMTrack/models/motility/high/collagen_optical_flow_median_bkg_more_data_90k/"  
+
+
+#update test video numbers from video map
+for video_num in  2 
+do
+   
+   #To genearate testing files for all motilities
+      python src/inferenceBacteriaRetinanet_Motility_v2.py  --output_dir $low_motility_model_path   --annotations_test "All" --video $video_num   --test_dir $data_path 
+
+      python  src/inferenceBacteriaRetinanet_Motility_v2.py  --output_dir $mid_motility_model_path  --annotations_test "Motility-mid" --video $video_num --test_dir $data_path  
+      
+      python  src/inferenceBacteriaRetinanet_Motility_v2.py  --output_dir $high_motility_model_path --annotations_test "Motility-high" --video $video_num --test_dir $data_path 
+      
+      python  src/inferenceBacteriaRetinanet_Motility_v2.py  --output_dir $low_motility_model_path  --annotations_test "Motility-low" --video $video_num --test_dir $data_path  
+done

MEMTrack/scripts/run.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+set -e 
+
+#RUN ONLY INFERENCE WITH EVAL(GT DATA)
+bash scripts/inference.sh 
+#run with plot and plt gt
+bash scripts/tracking.sh 
+bash scripts/test_set_eval.sh
+
+# RUN ONLY INFERENCE NO EVAL(NO GT DATA)
+ bash scripts/inference_no_eval.sh
+# #run only with plot
+# bash scripts/tracking.sh 

MEMTrack/scripts/test_set_eval.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+set -e
+
+exp_name="collagen_motility_optical_flow_median_bkg_more_data/" # exp name
+source_path="/data/medha/Bacteria" 
+data_path="$source_path/DataFeatures/$exp_name" # path toi feature directory
+echo $data_path 
+
+video_map_path="$source_path/Data/videomap.txt" # path to video map
+echo $video_map_path
+
+test_results_path="$data_path/test_set_results*.txt" #path to test result output files
+
+rm -rf $test_results_path
+
+for video_num in  141 148 153 160 167 170  
+do  
+     python src/evaluation_step_wise_motility.py --video_map_path $video_map_path --data_path $data_path --video $video_num
+done

MEMTrack/scripts/tracking.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+set -e
+
+exp_name="collagen_motility_optical_flow_median_bkg_more_data/"
+data_path="/data/medha/Bacteria/DataFeatures/"
+data_path+=$exp_name
+echo $data_path  
+video_map_path="/home/medha/BacteriaDetectionTracking/MEMTrack/data/videomap.txt"
+data_root_path="/home/medha/BacteriaDetectionTracking/MEMTrack/data/"
+
+for video_num in  141 148 153 160 167 170  
+do
+     python src/GenerateTrackingData.py --filter_thresh 0.3 --video_num $video_num --data_path $data_path 
+     python src/Tracking.py --video_num $video_num  --data_path $data_path  
+     #plot predictions and gt
+     python src/TrackingAnalysis.py --video_num $video_num --data_feature_path $data_path --video_map_path $video_map_path --data_root_path $data_root_path #--plot --plot_gt 
+     # python src/GenerateVideo.py --video_num $video_num --fps 1 --data_path $data_path  
+     # python src/GenerateVideo.py --video_num $video_num --fps 60 --data_path $data_path  
+done
+

MEMTrack/scripts/train.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+set -e
+
+#test dir
+exp_name="collagen_motility_optical_flow_consecutive_median_bkg" #exp name
+data_path="/DataFeatures/" #path to data features dir
+data_path+="$exp_name/data_feature_optical_flow_median_back_2pyr_18win_background_img/"
+echo $data_path 
+
+
+output_dir="/alldata/medha/CleanCodeData/training-output/" #output dir to store models
+output_dir+=$exp_name
+
+ 
+low_output_dir="$output_dir/low_90k_00125"
+high_output_dir="$output_dir/high_90k_00125"
+mid_output_dir="$output_dir/mid_90k_00125"
+wiggle_output_dir="$output_dir/wiggle_90k_00125"
+
+
+python src/trainBacteriaRetinanetMotionData_Motility.py --source_path $data_path --output_dir $low_output_dir      --annotations_train "Motility-low" --annotations_test "Motility-low" --bbox_size 31 --lr "0.00125" --epochs "90000"
+
+python src/trainBacteriaRetinanetMotionData_Motility.py --source_path $data_path --output_dir $high_output_dir      --annotations_train "Motility-high" --annotations_test "Motility-high" --bbox_size 31 --lr "0.00125" --epochs "90000"
+
+python src/trainBacteriaRetinanetMotionData_Motility.py --source_path $data_path --output_dir $mid_output_dir      --annotations_train "Motility-mid" --annotations_test "Motility-mid" --bbox_size 31 --lr "0.00125" --epochs "90000"
+
+python src/trainBacteriaRetinanetMotionData_Motility.py --source_path $data_path --output_dir $wiggle_output_dir      --annotations_train "Motility-wiggle" --annotations_test "Motility-wiggle" --bbox_size 31 --lr "0.00125" --epochs "90000"

MEMTrack/src/GenerateTrackingData.py

@@ -0,0 +1,248 @@
+import json
+import cv2
+import os
+import shutil
+import argparse
+import itertools
+from natsort import natsorted
+
+def combine_predictions(predictions_json_path, op_path):
+    #print(predictions_json_path)
+    all_preds = []
+    tracking_info_predictions_path = op_path
+    for path, diff in predictions_json_path:
+        if os.path.exists(path):
+            predictions_json = open(path)
+            predictions = json.load(predictions_json)
+            for pred in predictions:
+                pred["diff"] = diff
+                all_preds.append(pred)
+    all_preds.sort(key = lambda json: json['image_id'])
+    with open(tracking_info_predictions_path, "w") as track_info_file:
+        json.dump(all_preds, track_info_file)
+        
+def filter_boxes_size(tracking_info_predictions_path, op_path):
+    tracking_info_predictions_filetered_path = op_path
+    # filter boxes by size to remove extremely big boxes
+    combined_predictions_json = open(tracking_info_predictions_path)
+    combined_predictions = json.load(combined_predictions_json)
+    combined_predictions_filtered = []
+    for pred in combined_predictions:
+        width = pred["bbox"][2]
+        # print("w", width)
+        # print("h", height)
+        height = pred["bbox"][3]
+        if (width > 35) or (height > 35):
+            # print(width)
+            # print(height)
+            #print("removed")
+            #tracking_data.remove(tracked_box)
+            continue
+        else:
+            combined_predictions_filtered.append(pred)
+    with open(tracking_info_predictions_filetered_path, "w") as track_info_file:
+        json.dump(combined_predictions_filtered, track_info_file)
+    specific_file = tracking_info_predictions_filetered_path.split("filtered.json")[0] + "filter_box_size" +".json"
+    shutil.copy(tracking_info_predictions_filetered_path, specific_file)
+
+def filter_conf_score(tracking_info_predictions_path, op_path):
+    tracking_info_predictions_filetered_path = op_path
+    # filter boxes by size to remove extremely big boxes
+    combined_predictions_json = open(tracking_info_predictions_path)
+    combined_predictions = json.load(combined_predictions_json)
+    combined_predictions_filtered = []
+    for pred in combined_predictions:
+        score = pred["score"]
+        # if pred["diff"] == "Motility-low" and score < 0.2:
+        #     continue
+        if pred["diff"] == "Motility-wiggle" and score < 0.99:
+            continue
+        elif pred["diff"] == "Motility-mid" and score < 0.99:
+            continue
+        elif pred["diff"] == "Motility-high" and score < 0.99:
+            continue
+        else:
+            combined_predictions_filtered.append(pred)
+    with open(tracking_info_predictions_filetered_path, "w") as track_info_file:
+        json.dump(combined_predictions_filtered, track_info_file)
+    specific_file = tracking_info_predictions_filetered_path.split("filtered.json")[0] + "filter_conf_score" +".json"
+    shutil.copy(tracking_info_predictions_filetered_path, specific_file)
+
+def nms_filter(tracking_info_predictions_path, op_path, iou_thresh_nms):
+    tracking_info_predictions_filetered_path = op_path
+    # filter boxes using nms to remove near duplicate boxes with lower confidence score
+    combined_predictions_json = open(tracking_info_predictions_path)
+    combined_predictions = json.load(combined_predictions_json)
+    nms_filtered_preds = []
+    for image_id, preds in itertools.groupby(combined_predictions, key = lambda k:k["image_id"]):
+        #print(image_id)
+        img_preds = []
+        for pred in preds:
+            img_preds.append(pred)
+        filtered_preds = nms_helper(img_preds, iou_thresh_nms)
+        nms_filtered_preds.extend(filtered_preds)
+    with open(tracking_info_predictions_filetered_path, "w") as track_info_file:
+        json.dump(nms_filtered_preds, track_info_file)
+    specific_file = tracking_info_predictions_filetered_path.split("filtered.json")[0] + "filter_nms" +".json"
+    shutil.copy(tracking_info_predictions_filetered_path, specific_file)
+
+    
+def nms_helper(combined_predictions, iou_thresh_nms):
+    iou_thresh_nms=0.70
+    final_preds = []
+    #get iou of all boxes wrt all boxes
+    iou_mat = get_iou_mat_image(combined_predictions, combined_predictions)
+    #get matching box pairings
+    matching_boxes_iou = get_matching_boxes_iou(combined_predictions, iou_mat, iou_thresh = iou_thresh_nms)
+    #run nms while loop only if overlapping boxes exists
+    while(len(matching_boxes_iou)>0):
+        #sort the list acc to confidence score
+        sorted_bbox_list = sorted(combined_predictions, key=lambda k: k["score"])
+        #print(sorted_bbox_list[0])
+        #print("largest:",sorted_bbox_list[-1])
+        #get iou of all boxes wrt to box with max conf score
+        iou_mat = get_iou_mat_image(sorted_bbox_list[:-1], [sorted_bbox_list[-1]])
+        preds_temp = []
+        for index, iou in enumerate(iou_mat):
+            if iou[0]<iou_thresh_nms:
+                #print(sorted_bbox_list[index])
+                preds_temp.append(sorted_bbox_list[index])
+        final_preds.append(sorted_bbox_list[-1]) 
+        combined_predictions = preds_temp
+        iou_mat = get_iou_mat_image(combined_predictions, combined_predictions)
+        matching_boxes_iou = get_matching_boxes_iou(combined_predictions, iou_mat, iou_thresh = iou_thresh_nms)
+    return final_preds  
+
+#Calculate IoU of 2 bounding boxes
+def bb_intersection_over_union(boxA, boxB):
+    # determine the (x, y)-coordinates of the intersection rectangle
+    xA = max(boxA[0], boxB[0])
+    yA = max(boxA[1], boxB[1])
+    xB = min(boxA[2], boxB[2])
+    yB = min(boxA[3], boxB[3])
+
+     # compute the area of intersection rectangle
+    interArea = abs(max((xB - xA, 0)) * max((yB - yA), 0))
+    if interArea == 0:
+        return 0
+
+    # compute the area of both the prediction and ground-truth
+    # rectangles
+    boxAArea = abs((boxA[2] - boxA[0]) * (boxA[3] - boxA[1]))
+    boxBArea = abs((boxB[2] - boxB[0]) * (boxB[3] - boxB[1]))
+
+    # compute the intersection over union by taking the intersection
+    # area and dividing it by the sum of prediction + ground-truth
+    # areas - the interesection area
+    iou_box = interArea / float(boxAArea + boxBArea - interArea)
+
+    # return the intersection over union value
+    return iou_box
+
+def get_iou_mat_image(preds_list1, preds_list2 ):
+    iou_mat = []
+    for pred1 in preds_list1:
+        box1 = pred1["bbox"]
+        iou_row = []
+        for pred2 in preds_list2:
+            box2 = pred2["bbox"]
+            box1_xyxy = [box1[0], box1[1], box1[0] + box1[2], box1[1] + box1[3]]
+            box2_xyxy = [box2[0], box2[1], box2[0] + box2[2], box2[1] + box2[3]]
+            iou_boxes = bb_intersection_over_union(box1_xyxy, box2_xyxy)
+            iou_row.append(iou_boxes)
+        iou_mat.append(iou_row)
+    return iou_mat
+
+def get_matching_boxes_iou(preds_list, iou_mat, iou_thresh ):
+    matching_iou_boxes = []
+    for i in range(0, len(preds_list)):
+        iou_row_max = max(iou_mat[i])
+        iou_row_max_pred_id = iou_mat[i].index(iou_row_max)
+        # print((iou_row_max))
+        # print(iou_row_max_pred_id)
+        #print(iou_row_max)
+        if iou_row_max>iou_thresh:
+            matching_iou_boxes.append([i, iou_row_max_pred_id, iou_row_max])
+    # print(matching_iou_boxes)  
+    #print("Number of matching IOU Ground truth and Predicted boxes: " , len(matching_iou_boxes))
+    return matching_iou_boxes
+
+def gen_tracking_data_(video_num, predictions_json_path=None, op_path=None):
+    # predictions_json_path = f"../NewData-60FPS-Center/data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/test/coco_instances_results.json"
+    # predictions_json_path = f"../NewData-60FPS-Center/video{video_num}_feature_optical_flow_median_back_2pyr_18win/test/coco_instances_results.json"
+    predictions_json = open(predictions_json_path)
+    predictions = json.load(predictions_json)
+    try:
+        print(predictions[0])
+    except:
+        print("No predictions")
+    tracking_info_predictions_path = op_path + f"./video{video_num}_predictions.json"
+    instance_dict = {}
+    tracking_info_dict = {}
+    instances_image = []
+    count = 0
+    prev_image_id = predictions[0]['image_id']
+    prev_image_name = ""
+    image_id = 0
+    #print(predictions[-12])
+    for image_id, preds in itertools.groupby(predictions, key = lambda k:k["image_id"]):
+        #print(image_id)
+        instances_image = []
+        image_name = str(image_id) + ".tif"
+        for prediction in preds:
+            score = prediction['score']
+            if  score > 0.0: 
+                x,y,w,h = prediction['bbox']
+                x2 = x + w
+                y2 = y + h
+                instance_bbox = [x, y, x2, y2]
+                instace_dict = {'bbox': instance_bbox, 'labels': prediction['category_id'], 'scores':prediction['score'], 'diff': prediction['diff']}
+                instances_image.append(instace_dict)
+        tracking_info_dict[image_name] = instances_image
+    with open(tracking_info_predictions_path, "w") as track_info_file:
+        json.dump(tracking_info_dict, track_info_file)
+        
+
+def gen_tracking_data(video_num, data_path=None, custom_test_dir=None, filter_thresh=0):
+    combined_pred_path = "coco_instances_results_combined.json"
+    combined_pred_filtered_path = "coco_instances_results_combined_filtered.json"
+
+    video_path = f"data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/test/"
+    paths = []
+    if custom_test_dir:
+        data_path =  custom_test_dir
+        video_path = ""
+    #print(os.path.join(data_path, video_path, f"coco_instances_results_{diff}.json"))
+
+    for diff in ["Motility-mid", "Motility-wiggle", "Motility-high" ]:
+        paths.append((os.path.join(data_path, video_path, f"coco_instances_results_{diff}.json"), diff))
+
+    op_path0 = os.path.join(data_path, video_path, combined_pred_path)
+    op_path1 = os.path.join(data_path, video_path, combined_pred_filtered_path)
+
+    combine_predictions(paths, op_path0)
+
+    filter_boxes_size(op_path0, op_path=op_path1)
+    filter_conf_score(op_path1, op_path=op_path1)
+    nms_filter(op_path1, op_path=op_path1, iou_thresh_nms=filter_thresh)
+
+    gen_tracking_data_(video_num, op_path1, os.path.join(data_path, video_path))
+
+
+if __name__ == "__main__":
+
+    ap = argparse.ArgumentParser(description='Training')
+    ap.add_argument('--filter_thresh', default="0.0", type=float, metavar='THRESH')
+    ap.add_argument('--video_num', default="", type=str, metavar='VIDEO')
+    ap.add_argument('--data_path', default="19", type=str, metavar='PATH')
+    ap.add_argument('--custom_test_dir', type=str, metavar='CELL PATH')
+    args = ap.parse_args()
+    filter_thresh = args.filter_thresh
+    video_num = args.video_num
+    data_path =  args.data_path 
+    custom_test_dir = args.custom_test_dir
+    
+    gen_tracking_data(video_num, data_path, custom_test_dir, filter_thresh)
+
+
+    

MEMTrack/src/GenerateVideo.py

@@ -0,0 +1,48 @@
+import cv2
+import os
+import argparse
+from natsort import natsorted
+
+def create_video(data_dir, image_dir, video_name,fps):
+    # choose codec according to format needed
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v') 
+    #print(data_dir)
+    img_sample = cv2.imread(os.path.join(image_dir,"0.png"))
+    #print(img_sample.shape)
+    height, width, channels = img_sample.shape
+    
+    video = cv2.VideoWriter(data_dir + video_name + ".mp4", fourcc, fps, (width, height))
+    #data_dir = "./Data/video3/"
+    #image_dir = os.path.join(_dir, "images")
+    for frame in natsorted(os.listdir(image_dir)):
+        #print(frame)
+        img = cv2.imread(os.path.join(image_dir, frame))
+        video.write(img)
+    video.release()
+
+def gen_tracking_video(video_num, fps=60, data_path=None, custom_test_dir=None):
+    if custom_test_dir:
+        data_dir = custom_test_dir
+    else:
+        data_dir = data_path + f"/data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/"
+    image_dir = data_dir + "/test/tracklets-filtered/"
+    video_name = f'video{video_num}-tracklets-filtered-{fps}'
+    create_video(data_dir, image_dir, video_name,fps)
+    return os.path.join(data_dir, video_name)+ ".mp4"
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser(description='Training')
+    ap.add_argument('--video_num', default="", type=str, metavar='VIDEO')
+    ap.add_argument('--fps', default=1, type=int, metavar='FPS')
+    ap.add_argument('--data_path', default="19", type=str, metavar='PATH')
+    ap.add_argument('--custom_test_dir', type=str, metavar='CELL PATH')
+    args = ap.parse_args()
+    
+    video_num = args.video_num
+    fps = args.fps
+    data_path = args.data_path
+    custom_test_dir = args.custom_test_dir
+    
+    
+    gen_tracking_video(video_num, fps=fps, data_path=data_path, custom_test_dir=custom_test_dir)
+    

MEMTrack/src/LossEvalHook.py

@@ -0,0 +1,70 @@
+from detectron2.engine.hooks import HookBase
+from detectron2.evaluation import inference_context
+from detectron2.utils.logger import log_every_n_seconds
+from detectron2.data import DatasetMapper, build_detection_test_loader
+import detectron2.utils.comm as comm
+import torch
+import time
+import datetime
+import logging
+import numpy as np
+
+class LossEvalHook(HookBase):
+    def __init__(self, eval_period, model, data_loader):
+        self._model = model
+        self._period = eval_period
+        self._data_loader = data_loader
+    
+    def _do_loss_eval(self):
+        # Copying inference_on_dataset from evaluator.py
+        total = len(self._data_loader)
+        num_warmup = min(5, total - 1)
+            
+        start_time = time.perf_counter()
+        total_compute_time = 0
+        losses = []
+        for idx, inputs in enumerate(self._data_loader):            
+            if idx == num_warmup:
+                start_time = time.perf_counter()
+                total_compute_time = 0
+            start_compute_time = time.perf_counter()
+            if torch.cuda.is_available():
+                torch.cuda.synchronize()
+            total_compute_time += time.perf_counter() - start_compute_time
+            iters_after_start = idx + 1 - num_warmup * int(idx >= num_warmup)
+            seconds_per_img = total_compute_time / iters_after_start
+            if idx >= num_warmup * 2 or seconds_per_img > 5:
+                total_seconds_per_img = (time.perf_counter() - start_time) / iters_after_start
+                eta = datetime.timedelta(seconds=int(total_seconds_per_img * (total - idx - 1)))
+                log_every_n_seconds(
+                    logging.INFO,
+                    "Loss on Validation  done {}/{}. {:.4f} s / img. ETA={}".format(
+                        idx + 1, total, seconds_per_img, str(eta)
+                    ),
+                    n=5,
+                )
+            loss_batch = self._get_loss(inputs)
+            losses.append(loss_batch)
+        mean_loss = np.mean(losses)
+        self.trainer.storage.put_scalar('validation_loss', mean_loss)
+        comm.synchronize()
+
+        return losses
+            
+    def _get_loss(self, data):
+        # How loss is calculated on train_loop 
+        metrics_dict = self._model(data)
+        metrics_dict = {
+            k: v.detach().cpu().item() if isinstance(v, torch.Tensor) else float(v)
+            for k, v in metrics_dict.items()
+        }
+        total_losses_reduced = sum(loss for loss in metrics_dict.values())
+        return total_losses_reduced
+        
+        
+    def after_step(self):
+        next_iter = self.trainer.iter + 1
+        is_final = next_iter == self.trainer.max_iter
+        if is_final or (self._period > 0 and next_iter % self._period == 0):
+            self._do_loss_eval()
+        self.trainer.storage.put_scalars(timetest=12)

MEMTrack/src/Tracking.py

@@ -0,0 +1,114 @@
+import json
+import matplotlib 
+import matplotlib.pylab as plt
+from contextlib import suppress
+from sort import Sort
+import collections
+from pprint import pprint
+from sort import *
+import os
+import cv2
+import numpy as np
+import shutil
+import argparse
+from natsort import natsorted
+
+def track_bacteria(video_num, max_age, max_interpolation, op_path=None, custom_test_dir=None, data_path=None):
+    #mot_tracker = Sort(max_age = 25, min_hits=0, iou_threshold=0.05, max_interpolation=25)
+    mot_tracker = Sort(max_age = max_age, min_hits=0, iou_threshold=0.05, max_interpolation=max_interpolation)
+    #video_num = 29
+    if custom_test_dir:
+        op_path = custom_test_dir
+        video_num = ""
+    else:
+        op_path = data_path +  f"/data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/test/"
+    
+    jsonpath = op_path + f"./video{video_num}_predictions.json"
+    save_path = op_path + f"./video{video_num}-tracklets/"
+    tracking_predictions_path = op_path +f"./video{video_num}_tracking_predictions.json"
+    img_path = op_path + "/images/"
+    
+   
+
+    with open(jsonpath) as data_file:
+        data = json.load(data_file)
+    odata = collections.OrderedDict(sorted(data.items()))
+    print(jsonpath)
+    
+    
+    print(img_path)
+    shutil.rmtree(save_path, ignore_errors=True)
+    os.makedirs(save_path, exist_ok=True)
+
+    tracking_predictions = []
+
+    for key in natsorted(odata.keys()): 
+        arrlist = []
+        det_img = cv2.imread(os.path.join(img_path, key))
+        height, width, channels = det_img.shape
+        overlay = det_img.copy()
+        det_result = data[key] 
+
+        for info in det_result:
+            bbox = info['bbox']
+            #add filter bbox for size
+            labels = info['labels']
+            scores = info['scores']
+            templist = bbox+[scores]       
+            arrlist.append(templist)
+
+        track_bbs_ids = mot_tracker.update(np.array(arrlist))
+
+        mot_imgid = key.replace('.tif','')
+        newname = save_path + mot_imgid + '_mot.jpg'
+        #print(mot_imgid)
+
+        for j in range(track_bbs_ids.shape[0]):  
+            ele = track_bbs_ids[j, :]
+            x = int(ele[0])
+            y = int(ele[1])
+            x2 = int(ele[2])
+            y2 = int(ele[3])
+            x_cen = x + int((x2-x)/2)
+            y_cen = y + int((y2-y)/2)
+            track_label = str(int(ele[4]))
+            if x_cen>= width:
+                continue
+            if y_cen>= height:
+                continue
+            #cv2.rectangle(det_img, (x, y), (x2, y2), (0, 255, 255), 4)
+            # cv2.line(det_img, (x, y), (x2, y2), (0, 255, 255), 1, )
+            # cv2.line(det_img, (x2, y), (x, y2), (0, 255, 255), 1)
+            cv2.drawMarker(det_img, (x_cen + 5, y_cen + 5),(0,255,255), markerType=cv2.MARKER_CROSS, 
+            markerSize=10, thickness=1, line_type=cv2.LINE_AA)
+            cv2.putText(det_img, '#'+track_label, (x, y-4), 0,0.4,(0,255,255),thickness=1)
+            w = x2 - x
+            h = y2 - y
+            bbox = [x, y, w, h]
+            instance_prediction = {'image_id': int(mot_imgid) , 'category_id': 0, 'bbox': bbox, 'score':-1, 'track_label': track_label}
+            tracking_predictions.append(instance_prediction)
+
+        cv2.imwrite(newname,det_img)
+        
+        
+    with open(tracking_predictions_path, "w") as track_preds_file:
+        json.dump(tracking_predictions, track_preds_file)
+    
+if __name__ == "__main__":
+
+    ap = argparse.ArgumentParser(description="")
+    ap.add_argument('--video_num', default="", type=str, metavar='VIDEO')
+    ap.add_argument('--data_path', default="19", type=str, metavar='PATH')
+    ap.add_argument('--custom_test_dir', default="", type=str, metavar='CELL PATH')
+    ap.add_argument('--max_age', type=int, metavar='CELL PATH', default=25)
+    ap.add_argument('--max_interpolation', type=int, metavar='CELL PATH', default=25)
+    
+    args = ap.parse_args()
+    video_num = args.video_num
+    custom_test_dir = args.custom_test_dir
+    data_path = args.data_path
+
+    max_age = args.max_age
+    max_interpolation = args.max_interpolation
+     
+    track_bacteria(video_num, max_age, max_interpolation, custom_test_dir=custom_test_dir, data_path=data_path)

MEMTrack/src/TrackingAnalysis.py

@@ -0,0 +1,464 @@
+import json
+import seaborn as sns
+import matplotlib.pyplot as plt
+import numpy as np
+import cv2
+import os
+import glob
+import shutil
+import pandas as pd
+import argparse
+from natsort import natsorted
+
+def filter_boxes_size(tracking_data):
+    # filter boxes by size to remove extremely big boxes
+    tracking_data_filtered = []
+    for tracked_box in tracking_data:
+        width = tracked_box["bbox"][2]
+        # print("w", width)
+        # print("h", height)
+        height = tracked_box["bbox"][3]
+        if (width > 35) or (height > 35):
+            # print(width)
+            # print(height)
+            #print("removed")
+            #tracking_data.remove(tracked_box)
+            continue
+        else:
+            tracking_data_filtered.append(tracked_box)
+    return tracking_data_filtered
+
+
+def groupby_trackid(tracking_data):
+    #generate dict to contains tracked boxes per track id
+    track_id_details = {}
+    for tracked_box in tracking_data:
+        #print(tracked_box["track_label"])
+        track_label = tracked_box["track_label"]
+        #print("track_label", track_label) 
+        if track_label not in track_id_details.keys():
+            #print("adding to dict")
+            track_id_details[track_label] = []
+        track_id_details[track_label].append(tracked_box)
+    print("----------------------------------------------------------")
+    print("sample track id list in dict: ", track_id_details[track_label][:2])
+    print("----------------------------------------------------------")
+    print("dict track id keys: ", track_id_details.keys())
+    print("----------------------------------------------------------")
+    return track_id_details
+
+def analyze_track_preds(track_id_details, filter_length=False, min_con_length=11):   
+    #get avg frames per track
+    #get num of fragments
+    #get fragments length
+    if filter_length == True:
+        track_id_filtered_length = track_id_details.copy()
+    num_frames = []
+    num_fragments = []
+    fragment_lengths = []
+    skipped_frames = []
+    for track_id in track_id_details.keys():
+        num_frames_track = 0
+        num_fragments_track = 0
+        prev_image_id = 0
+        prev_fragment_length = 0 
+        new_image = True
+        frags_len_track = []
+        skip_frames_track = []
+        skipped_frames_track = 0
+        #print("new track: ", track_id)
+        tracked_boxes = track_id_details[track_id]
+        #print(tracked_boxes)
+        for tracked_box in tracked_boxes:
+            #print("box", tracked_box)
+            image_id = tracked_box["image_id"]
+            num_frames_track += 1
+            #print("num frames: ", num_frames_track)
+            if new_image == True:
+                new_image = False
+                num_fragments_track +=1
+                prev_image_id = image_id
+                fragment_start = image_id
+
+                #print("num fragments: ", num_fragments_track)
+            else:
+                if prev_image_id + 1 == image_id:
+                    prev_fragment_length = prev_image_id - fragment_start + 1
+                    prev_image_id +=1
+
+                else:
+                    #print("image_id: ", image_id)
+                    #print("prev_image_id: ", prev_image_id)
+                    skipped_frames_track = image_id-prev_image_id-1
+                    skip_frames_track.append(skipped_frames_track)
+                    prev_fragment_length = prev_image_id - fragment_start + 1
+                    frags_len_track.append(prev_fragment_length)
+                    num_fragments_track +=1
+                    fragment_start = image_id
+                    prev_image_id = image_id
+                    #print("num fragments: ", num_fragments_track)
+        prev_fragment_length = prev_image_id - fragment_start + 1
+        frags_len_track.append(prev_fragment_length)
+        if num_fragments_track == 1:
+            skip_frames_track.append(skipped_frames_track)
+        num_frames.append(num_frames_track)
+        if num_frames_track< min_con_length and filter_length is True:
+            track_id_filtered_length.pop(track_id, None)
+            #print("removing")
+        num_fragments.append(num_fragments_track)
+        fragment_lengths.append(frags_len_track)
+        skipped_frames.extend(skip_frames_track)  
+    if filter_length is True:
+        return num_frames, num_fragments, fragment_lengths, skipped_frames, track_id_filtered_length
+    return num_frames, num_fragments, fragment_lengths, skipped_frames
+
+def groupby_imageid(tracking_data):
+    #generate dict to contains tracked boxes per track id
+    image_id_details = {}
+    for track_label in tracking_data:
+        for tracked_box in tracking_data[track_label]: 
+            image_id = tracked_box["image_id"]
+            if image_id not in image_id_details.keys():
+            #print("adding to dict")
+                image_id_details[image_id] = []
+            image_id_details[image_id].append(tracked_box)
+    try:
+        
+        print("----------------------------------------------------------")
+        print("sample image id list in dict: ", image_id_details[image_id][:2])
+        print("----------------------------------------------------------")
+        print("dict image id keys: ", image_id_details.keys())
+        print("----------------------------------------------------------")
+    except:
+        print("Error in print statements")
+    return image_id_details
+
+
+
+def analyse_tracking(video_num, min_track_length=60, custom_test_dir=None, data_feature_path=None, 
+                     video_map_path="data/videomap.txt", data_root_path="", plot=True, plot_gt=False):
+    #tracking-results:
+    if custom_test_dir:
+        op_path = custom_test_dir
+        video_num= ""
+    else:
+        op_path = data_feature_path + f"/data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/test/"
+    tracking_predictions_path = op_path + f"./video{video_num}_tracking_predictions.json"
+
+    video_map_path = video_map_path
+    video_map = open(video_map_path,'r',)
+    header = [x.strip() for x in (video_map.readline().strip()).split(",")]
+    video_num_id = header.index("video_num")
+    strain_id = header.index("strain")
+    strain_map = {}
+    for line in video_map.readlines():
+        line_details = [x.strip() for x in line.split(",")]
+        video_id = int(line_details[video_num_id])
+        strain = line_details[strain_id]
+        strain_map[video_id] = strain
+    print(strain_map)
+    
+    with open(tracking_predictions_path) as data_file:
+        tracking_data = json.load(data_file)
+    print("Total tracking data: ", len(tracking_data))
+
+    tracking_data_filtered = tracking_data
+
+
+    track_id_details = groupby_trackid(tracking_data_filtered)
+    num_frames, num_fragments, fragment_lengths, skipped_frames, track_id_filtered_length =analyze_track_preds(track_id_details,filter_length=True,min_con_length=min_track_length)
+    
+    try:
+        print("Number of unique track ids: ", len(num_frames))
+        print("----------------------------------------------------------")
+        print("Average number of frames per track id: ", round(np.mean(num_frames)))
+        print("Maximum number of frames per track id: ", np.max(num_frames))
+        print("Minimum number of frames per track id: ", np.min(num_frames))
+        print("Median number of frames per track id: ", np.median(num_frames))
+        print("----------------------------------------------------------")
+        print("Average number of fragments per track id: ", round(np.mean(num_fragments)))
+        print("Maximum number of fragments per track id: ", np.max(num_fragments))
+        print("Minimum number of fragments per track id: ", np.min(num_fragments))
+        print("Median number of fragments per track id: ", np.median(num_fragments))
+        print("----------------------------------------------------------")
+        print("Average number of skipped frames per fragment: ", round(np.mean(skipped_frames)))
+        print("Maximum number of skipped frames per fragment: ", np.max(skipped_frames))
+        print("Minimum number of skipped frames per fragment: ", np.min(skipped_frames))
+    except:
+        print("Error in print statements")
+
+    num_frames, num_fragments, fragment_lengths, skipped_frames = analyze_track_preds(track_id_filtered_length)
+    
+    try:
+
+        print("Number of unique track ids: ", len(num_frames))
+        print("----------------------------------------------------------")
+        print("Average number of frames per track id: ", round(np.mean(num_frames)))
+        print("Maximum number of frames per track id: ", np.max(num_frames))
+        print("Minimum number of frames per track id: ", np.min(num_frames))
+        print("Median number of frames per track id: ", np.median(num_frames))
+        print("----------------------------------------------------------")
+        print("Average number of fragments per track id: ", round(np.mean(num_fragments)))
+        print("Maximum number of fragments per track id: ", np.max(num_fragments))
+        print("Minimum number of fragments per track id: ", np.min(num_fragments))
+        print("Median number of fragments per track id: ", np.median(num_fragments))
+        print("----------------------------------------------------------")
+        print("Average number of skipped frames per fragment: ", round(np.mean(skipped_frames)))
+        print("Maximum number of skipped frames per fragment: ", np.max(skipped_frames))
+        print("Minimum number of skipped frames per fragment: ", np.min(skipped_frames))
+    except:
+        print("Error in print statements")
+
+    all_preds = []
+    if custom_test_dir:
+        final_preds_path = custom_test_dir
+    else:
+        final_preds_path = data_feature_path + f"/data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/test/"
+    final_preds_path += "coco_instances_results_final.json"
+    for track_id in track_id_filtered_length.keys():
+        for pred in track_id_filtered_length[track_id]:
+            all_preds.append(pred)
+    with open(final_preds_path, "w") as track_info_file:
+            json.dump(all_preds, track_info_file)
+        
+
+
+    image_id_filtered = groupby_imageid(track_id_filtered_length)
+
+    if plot :
+        #create images with labels plot
+        data_dir = data_root_path
+        data_sub_dirs = glob.glob(f'{data_dir}/*')  
+        #print("data dirs: ", data_sub_dirs)
+        video_dirs = glob.glob(f'{data_dir}/*/*')  
+        for video in video_dirs:
+                #print(video)
+                if str(video_num) in video:
+                    print(video)
+                    video_path = video
+                    break
+
+        img_path =  video_path + f"/frame1/images/"
+        if custom_test_dir:
+            save_path = custom_test_dir +"tracklets-filtered/"
+        else:
+            save_path = data_feature_path + f"/data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/test/tracklets-filtered/" 
+        shutil.rmtree(save_path, ignore_errors=True)
+        os.makedirs(save_path, exist_ok=True)
+
+        if plot_gt:
+            ground_truth_json_path = op_path + "/test_All.json"
+            ground_truth_json_path2 = op_path + "/test_Easy+Hard.json"
+            if os.path.exists(ground_truth_json_path):
+                gt_json = open(ground_truth_json_path)
+                ground_truth = json.load(gt_json)
+            else: # os.path.exists(ground_truth_json_path2):
+                gt_json = open(ground_truth_json_path)
+                ground_truth = json.load(gt_json)
+            # else:
+            #     ground_truth_json_path = args.data_path + f"/data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/test/test_Hard.json"
+            #     gt_json = open(ground_truth_json_path)
+            #     ground_truth = json.load(gt_json)
+
+            #need frame id, bacteria track id and bbox
+            data_path = video_path + f"/frame1/bacteria/"
+            bacteria_dirs = os.listdir(data_path)
+            print(bacteria_dirs)
+
+            ground_truth_data_bacteria_track_specific = []
+            for bacteria in bacteria_dirs:
+                #print(bacteria)
+                for coord_txt in os.listdir(os.path.join(data_path,bacteria,"xy_coord")):
+                    #print(coord_txt)
+                    if not coord_txt.endswith(".ipynb_checkpoints"):
+                        frame_id = int(coord_txt.split(".txt")[0])
+                        coord_file = open(os.path.join(data_path,bacteria,"xy_coord", coord_txt),'r')
+                        line = coord_file.readlines()
+                        #print(line)
+                        if len(line)>0:
+                            x = float(line[0].split(" ")[0])
+                            y = float(line[0].split(" ")[1])
+                        # print(line)
+                        # print(x)
+                        # print(y)
+                        #print(frame_id)
+                            width = height = 30
+                            factor_w = 1
+                            factor_h = 1
+                            x1 = int(x*factor_w - (width // 2))
+                            y1 = int(y*factor_h - (height // 2))
+                            w=h=30
+                            bbox = [x1,y1,w,h]
+                            entry = {'image_id': frame_id, 'track_label':int(bacteria), 'bbox':  bbox}
+                            ground_truth_data_bacteria_track_specific.append(entry)
+            track_id_ground_truth = groupby_trackid(ground_truth_data_bacteria_track_specific)
+            ground_truth = groupby_imageid(track_id_ground_truth)
+
+
+    if plot :
+        for image_id in image_id_filtered:
+            #print(image_id)
+            newname = save_path + str(image_id) + '.png'
+            det_img = cv2.imread(os.path.join(img_path,str(image_id))+".tif")
+            det_img_gt_only = det_img.copy()
+            det_img_p_only = det_img.copy()
+
+            height, width, channels = det_img.shape
+            #print (height, width, channels)
+            # plotting filtered predictions
+            for tracked_box in image_id_filtered[image_id]:
+                bbox = tracked_box["bbox"]
+                x = bbox[0]
+                y = bbox[1]
+                w = bbox[2]
+                h = bbox[3]
+                x_cen = x + int(w/2)
+                y_cen = y + int(h/2)
+                track_label = tracked_box["track_label"]
+                cv2.drawMarker(det_img, (x_cen , y_cen),(0,255,255), markerType=cv2.MARKER_CROSS, 
+                markerSize=5, thickness=1, line_type=cv2.LINE_AA)
+
+                # print("Y:",y_cen)
+                # print("X:",x_cen)
+                if x_cen>660:
+                    print("here")
+                if y_cen>490:
+                    print("here")
+                cv2.putText(det_img, '#'+track_label, (x, y-6), 0,0.6,(0,255,255),thickness=1)
+            if plot_gt:
+                #plotting ground truth  
+                try:
+                    for ground_truth_box in ground_truth[image_id]:
+                        bbox = ground_truth_box["bbox"]
+                        x = bbox[0]
+                        y = bbox[1]
+                        w = bbox[2]
+                        h = bbox[3]
+                        x_cen = x + int(w/2)
+                        y_cen = y + int(h/2)
+                        if x_cen>660:
+                            print(image_id)
+                            print("here")
+                            print("Y:",y_cen)
+                            print("X:",x_cen)
+                        if y_cen>490:
+                            print(image_id)
+                            print("here")
+                            print("Y:",y_cen)
+                            print("X:",x_cen)
+                        track_label = ground_truth_box["track_label"]
+                        if y_cen>490:
+                            print(track_label)
+                        cv2.drawMarker(det_img, (x_cen, y_cen),(255,255,0), markerType=cv2.MARKER_TILTED_CROSS, 
+                        markerSize=5, thickness=1, line_type=cv2.LINE_AA)
+                        cv2.putText(det_img, '#'+str(track_label), (x, y+4), 0,0.6,(255,255,0),thickness=1)
+                except:
+                    pass
+            cv2.imwrite(newname,det_img)
+
+  
+     #image id wise csv
+    bacteria_count = 0
+    tracked_raw_data = pd.DataFrame(columns = ['Nr', 'TID', 'PID', 'x [pixel]', 'y [pixel]'])
+    for image_id in image_id_filtered.keys():
+        #print(image_id)
+        pid = image_id
+        for tracked_box in image_id_filtered[image_id]:
+            tracked_bacteria_data = []
+            tid = tracked_box["track_label"]
+            #print(tid)
+            bbox = tracked_box["bbox"]
+            x = bbox[0]
+            y = bbox[1]
+            w = bbox[2]
+            h = bbox[3]
+            x_cen = x + w/2
+            y_cen = y + h/2
+            bacteria_count +=1
+            tracked_bacteria_data ={"Nr" : bacteria_count,'TID' : tid, 'PID' : pid, 
+                                    'x [pixel]' : x_cen, 'y [pixel]' : y_cen } 
+            tracked_raw_data = tracked_raw_data.append(tracked_bacteria_data, ignore_index=True)
+
+    #track id wise csa raw data
+    bacteria_count = 0
+    tracked_raw_data = pd.DataFrame(columns = ['Nr', 'TID', 'PID', 'x [pixel]', 'y [pixel]'])
+    for track_id in natsorted(track_id_filtered_length.keys()):
+        tid = track_id
+        for tracked_box in track_id_filtered_length[track_id]:
+            tracked_bacteria_data = []
+            pid = tracked_box["image_id"]
+            #print(pid)
+            bbox = tracked_box["bbox"]
+            x = bbox[0]
+            y = bbox[1]
+            w = bbox[2]
+            h = bbox[3]
+            x_cen = x + w/2
+            y_cen = y + h/2
+            bacteria_count +=1
+            tracked_bacteria_data ={"Nr" : bacteria_count,'TID' : tid, 'PID' : pid, 
+                                    'x [pixel]' : x_cen, 'y [pixel]' : y_cen } 
+            tracked_raw_data = tracked_raw_data.append(tracked_bacteria_data, ignore_index=True)
+
+    if len(str(video_num))>0:        
+        tracked_raw_data.to_csv(op_path+ f'TrackedRawData_{min_track_length}_video{video_num}_{strain_map[video_num]}.csv', index=False)
+    else:
+         tracked_raw_data.to_csv(op_path+ f'TrackedRawData_{min_track_length}_video{video_num}.csv', index=False)
+
+    #track id wise csa raw data
+    bacteria_count = 0
+    tracked_raw_data = pd.DataFrame(columns = ['Nr', 'TID', 'PID', 'x [pixel]', 'y [pixel]'])
+    for track_id in natsorted(track_id_details.keys()):
+        tid = track_id
+        for tracked_box in track_id_details[track_id]:
+            tracked_bacteria_data = []
+            pid = tracked_box["image_id"]
+            #print(pid)
+            bbox = tracked_box["bbox"]
+            x = bbox[0]
+            y = bbox[1]
+            w = bbox[2]
+            h = bbox[3]
+            x_cen = x + w/2
+            y_cen = y + h/2
+            bacteria_count +=1
+            tracked_bacteria_data ={"Nr" : bacteria_count,'TID' : tid, 'PID' : pid, 
+                                    'x [pixel]' : x_cen, 'y [pixel]' : y_cen } 
+            tracked_raw_data = tracked_raw_data.append(tracked_bacteria_data, ignore_index=True)
+        
+    print("video_num:", video_num)
+    if len(str(video_num))>0:   
+        tracked_raw_data.to_csv(op_path+ f'TrackedRawData_all_video{video_num}_{strain_map[video_num]}.csv', index=False)
+    else:
+         tracked_raw_data.to_csv(op_path+ f'TrackedRawData_all_video{video_num}.csv', index=False)
+            
+    return op_path         
+            
+            
+            
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser(description='Training')
+    ap.add_argument('--video_num', default="", type=str, metavar='VIDEO')
+    ap.add_argument('--data_feature_path', default="19", type=str, metavar='PATH')
+    ap.add_argument('--data_root_path', default="19", type=str, metavar='PATH')
+    ap.add_argument('--plot_gt', action='store_true')
+    ap.add_argument('--plot', action='store_true')
+    ap.add_argument('--custom_test_dir', type=str, metavar='CELL PATH', default="")
+    ap.add_argument('--min_track_len', default = 60, type=int)
+    ap.add_argument('--video_map_path', default="data/videomap.txt", type=str, metavar='PATH')
+    
+    args = ap.parse_args()
+    
+    video_num = args.video_num
+    custom_test_dir = args.custom_test_dir
+    min_track_length = args.min_track_len
+    data_feature_path = args.data_feature_path
+    video_map_path = args.video_map_path 
+    data_root_path = args.data_root_path
+    plot = args.plot
+    plot_gt = args.plot_gt
+    
+    analyse_tracking(video_num, min_track_length=min_track_length, custom_test_dir=custom_test_dir,
+                     data_feature_path=data_feature_path, 
+                     video_map_path=video_map_path, data_root_path=data_root_path,
+                     plot=plot, plot_gt=plot_gt)

MEMTrack/src/data_feature_gen.py

@@ -0,0 +1,628 @@
+import zipfile
+import openpyxl
+import os
+import tqdm
+import csv
+import cv2
+import shutil
+import PIL
+import glob
+import pandas as pd
+import numpy as np
+from natsort import natsorted
+from PIL import Image
+import argparse
+
+
+def create_video(data_dir):
+    # choose codec according to format needed
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v') 
+    #print(data_dir)
+    img_sample = cv2.imread(os.path.join(data_dir,"images/0.tif"))
+    #print(img_sample.shape)
+    height, width, channels = img_sample.shape
+    
+    video = cv2.VideoWriter(data_dir + 'video.mp4', fourcc, 1, (width, height))
+    #data_dir = "./Data/video3/"
+    image_dir = os.path.join(data_dir, "images")
+    for frame in natsorted(os.listdir(image_dir)):
+        #print(frame)
+        img = cv2.imread(os.path.join(image_dir, frame))
+        video.write(img)
+
+    cv2.destroyAllWindows()
+    video.release()
+    
+    
+    
+def get_background(file_path, mean=True, sample=False):
+    if not os.path.exists(file_path):
+        create_video(file_path.rsplit("/",1)[0] +"/frame1")
+    cap = cv2.VideoCapture(file_path)
+    #print(cap.read())
+    # we will randomly select 50 frames for the calculating the median
+    #frame_indices = cap.get(cv2.CAP_PROP_FRAME_COUNT) * np.random.uniform(size=20)
+    frame_indices = list(range(0,int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) ))
+    print(len(frame_indices))
+    # we will store the frames in array
+    frames = []
+    for idx in frame_indices:
+        # set the frame id to read that particular frame
+        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
+        ret, frame = cap.read()
+        #print(ret)
+        frames.append(frame)
+    if mean:
+         # calculate the mean
+        background_frame = np.mean(frames, axis=0).astype(np.uint8)
+    else:
+        # calculate the median
+        background_frame = np.median(frames, axis=0).astype(np.uint8)
+    if sample==True:
+        background_frame = cv2.imread("./Control_2_b0t5306c0x0-660y0-492.tiff")
+        #background_frame = cv2.imread("./RBS 2_1_b0t2791c0x0-660y0-492.tiff")
+    return background_frame
+
+
+def get_absolute_consec_frame_diff_feature(video_path, max_consec_frame_diff=True, num_prior_frames=1):
+    cap = cv2.VideoCapture(video_path)
+    frames = []
+    differences = []
+    prev_frames = []
+    while True:
+        ret, frame = cap.read()
+        #print(frame)
+        if not ret:
+            break
+        #storing old frames till prior frame number
+        if num_prior_frames>0:
+            prev_frames.append(frame)
+            num_prior_frames -= 1
+            continue
+        #retrieving the previous xth frame
+        prev = prev_frames[0]
+        prev_frames.pop(0)
+        #computing frame diff between current and one previous
+        consecutive_diff = np.abs(frame - prev)
+        frames.append(frame) #creating frame list
+        #creating consecutive frame diff list
+        differences.append(consecutive_diff)
+    #creating consecutive frame diff features by taking the max at every pixel along the frame diff list
+    if max_consec_frame_diff:
+        max_abs_consec_diff_feature = np.max(differences, axis=0)
+    else:
+         max_abs_consec_diff_feature = np.min(differences, axis=0)
+    return max_abs_consec_diff_feature
+
+
+
+def get_diff_from_absolute_consec_frame_diff_feature(image_path, frame_diff_feature, frame=False):
+    if frame==False:
+        image = PIL.Image.open(image_path).convert('L')
+    if frame == True:
+        #print("getting diff from frame")
+        image = Image.fromarray(image_path).convert('L')
+    L = image.getchannel(0)
+    frame_diff_feature = PIL.Image.fromarray(frame_diff_feature).convert('L')
+    frame_diff_feature.getbands()
+    L1 = frame_diff_feature.getchannel(0)
+    diff = PIL.ImageChops.difference(L, L1)
+    return diff
+
+
+def get_absolute_all_frame_diff_feature(video_path, max_feature=True):
+    cap = cv2.VideoCapture(video_path)
+    frames = []
+    while True:
+        ret, frame = cap.read()
+        #print(frame)
+        if not ret:
+            break
+        frames.append(frame)
+    
+    features = []
+    count = 0
+    for index in tqdm.tqdm(range(len(frames))):
+        differences = []
+        count+=1
+        frame = frames[index]
+        for index_1 in (range(len(frames))):
+            if index == index_1:
+                #print(count)
+                continue
+            frame1 = frames[index_1]
+            differences.append(np.abs(frame - frame1))
+        if max_feature:
+            max_diff_feature = np.max(differences, axis=0)
+            features.append(max_diff_feature)
+        else:
+            #print("min")
+            min_diff_feature = np.min(differences, axis=0)
+            features.append(min_diff_feature)
+    return features    
+
+
+
+#updated to include optical floiw computtaion from x frames prior
+def gen_dense_optical_flow_data(method, video_path, params=[], to_gray=False, median=False, median_frame=None, num_frames_prior=1):
+    frames_optical_flow = []
+    frames_orignal = []
+    # Read the video and first x frames
+    cap = cv2.VideoCapture(video_path)
+    #print("fps",cap.get(cv2.CAP_PROP_FPS))
+    old_frames = []
+    for i in range(num_frames_prior):
+        ret, old_frame = cap.read()
+        # crate HSV & make Value a constant
+        hsv = np.zeros_like(old_frame)
+        hsv[..., 1] = 255
+        if to_gray:
+            old_frame = cv2.cvtColor(old_frame, cv2.COLOR_BGR2GRAY)
+        old_frames.append(old_frame)
+
+    #to compute optical flow from the median background    
+    if median == True:
+        old_frame = median_frame
+        old_frame = cv2.cvtColor(old_frame, cv2.COLOR_GRAY2BGR)
+        
+  
+
+    
+    while True:
+        # Read the next frame
+        ret, new_frame = cap.read()
+        frame_copy = new_frame
+        if not ret:
+            break
+        
+        old_frame = old_frames[0]
+        if median == True:
+            old_frame = median_frame
+
+        # Preprocessing for exact method
+        if to_gray:
+            new_frame = cv2.cvtColor(new_frame, cv2.COLOR_BGR2GRAY)
+                
+            
+        # Calculate Optical Flow
+      
+        flow = method(old_frame, new_frame, None, *params)
+
+        # Encoding: convert the algorithm's output into Polar coordinates
+        mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
+        # Use Hue and Value to encode the Optical Flow
+        hsv[..., 0] = ang * 180 / np.pi / 2
+        hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
+
+        # Convert HSV image into BGR for demo
+        bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
+        #plt.imshow(frame_copy)
+        #plt.imshow(bgr)
+        # plt.imshow(v,cmap='gray', vmin=0, vmax=255)
+        # plt.show()
+    
+        # Update the previous frame
+        old_frames.append(new_frame)
+        old_frames.pop(0)
+        frames_orignal.append(frame_copy)
+        frames_optical_flow.append(bgr)
+    return frames_orignal, frames_optical_flow
+
+
+
+def get_background_diff_feature(image_path, background):
+    image = PIL.Image.open(image_path).convert('L')
+    L = image.getchannel(0)
+    background_img = PIL.Image.fromarray(background)
+    background_img.getbands()
+    L1 = background_img.getchannel(0)
+    diff = PIL.ImageChops.difference(L, L1)
+    return diff
+    
+def get_prev_frame_diff_feature(image_path, prev_image_path):
+    image = PIL.Image.open(image_path).convert('L')
+    L = image.getchannel(0)
+    image = PIL.Image.open(prev_image_path).convert('L')
+    L1 = image.getchannel(0)
+    diff = PIL.ImageChops.difference(L, L1)
+    return diff
+
+def create_feature_image(image_path, background, prev_image_path, prev_image_path2=None):
+    prev_image_diff = get_prev_frame_diff_feature(image_path, prev_image_path)
+    image = PIL.Image.open(image_path).convert('L')
+    L = image.getchannel(0)
+    if prev_image_path2 is None:
+        background_diff = get_background_diff_feature(image_path, background)
+        newImagediff = PIL.Image.merge("RGB", [L, background_diff, prev_image_diff])
+    if prev_image_path2 is not None:
+        prev_image_diff2 = get_prev_frame_diff_feature(image_path, prev_image_path2)
+        newImagediff = PIL.Image.merge("RGB", [L, prev_image_diff2, prev_image_diff])
+        
+    return newImagediff
+
+
+
+def create_feature_image_optical_flow(frame, optical_flow, pure=False, background=None, optical_flow2=None, final_channel=False):
+    frame = Image.fromarray(frame).convert('L')
+    L = frame.getchannel(0)
+    flow = PIL.Image.fromarray(optical_flow)
+    #flow.save("flow.png")
+    hsv_optical_flow = cv2.cvtColor(optical_flow, cv2.COLOR_BGR2HSV)
+    h, s, v = cv2.split(hsv_optical_flow)
+    flow = PIL.Image.fromarray(v)
+    #flow.save("flow_hsv.png")
+    
+    #print(v.shape)
+    last_channel = np.zeros([v.shape[0], v.shape[1]],dtype=np.uint8)
+    last_channel[:] = 255
+    v = Image.fromarray(v).convert('L')
+    v = v.getchannel(0)
+    
+    
+    #print(last_channel)
+    last_channel = Image.fromarray(last_channel).convert('L')
+    last_channel = last_channel.getchannel(0)
+
+    gray_optical_flow = Image.fromarray(cv2.cvtColor(optical_flow, cv2.COLOR_BGR2GRAY)).convert('L').getchannel(0)
+    
+    feature_image = PIL.Image.merge("RGB", [L, v, last_channel])
+    feature_image = PIL.Image.merge("RGB", [L, gray_optical_flow, last_channel])
+    
+    if pure == False:
+        if background is None:
+            gray_optical_flow2 = Image.fromarray(cv2.cvtColor(optical_flow2, cv2.COLOR_BGR2GRAY)).convert('L').getchannel(0)
+            feature_image = PIL.Image.merge("RGB", [L, gray_optical_flow, gray_optical_flow2])
+            #print(here)
+        else:
+            if final_channel == False:
+                background_img = PIL.Image.fromarray(background)
+                L1 = background_img.getchannel(0)
+                feature_image = PIL.Image.merge("RGB", [L, v, L1]) #-->org code
+                # diff = PIL.ImageChops.difference(L, L1)
+                # feature_image = PIL.Image.merge("RGB", [L, v, diff]) #-->up code
+                #print("here adding background")
+                #print("background")
+            if final_channel==True:
+                #print("final channel tru, just adding diff")
+                feature_image = PIL.Image.merge("RGB", [L, v, background]) 
+            
+       
+        
+    # print(feature_image.size)
+    # print(feature_image.mode)
+        
+    return feature_image
+
+
+
+def save_data(image_dest, count, feature_image, frame_num, data_dir, video, minivideo, count_test, count_train, train):
+    #save image feature
+    feature_image.save(image_dest +"images_feature/"+str(count)+".tif")
+
+    text_file = str(frame_num) +".txt"
+    annotations_easy_source = os.path.join(data_dir, video, minivideo, "annotations_easy",
+                                           text_file)
+    annotations_easy_hard_source = os.path.join(data_dir, video, minivideo, 
+                                            "annotations_easy_hard", text_file)
+    
+    annotations_very_hard_source = os.path.join(data_dir, video, minivideo, "annotations_veryhard",
+                                           text_file)
+    annotations_easy_hard_veryhard_source = os.path.join(data_dir, video, minivideo, 
+                                            "annotations_easy_hard_veryhard", text_file)
+
+
+    shutil.copy(annotations_easy_source, image_dest +"annotation_easy/" +str(count)+".txt")
+    shutil.copy(annotations_easy_hard_source, image_dest +"annotation_easy_hard/" 
+                + str(count)+".txt")
+    if os.path.exists(annotations_very_hard_source):
+        shutil.copy(annotations_very_hard_source, image_dest +"annotation_veryhard/" +str(count)+".txt")
+    if os.path.exists(annotations_easy_hard_veryhard_source):
+        shutil.copy(annotations_easy_hard_veryhard_source, image_dest +"annotation_easy_hard_veryhard/" 
+                + str(count)+".txt")
+    #create hard only annotation
+    text_file_easy_hard = open(annotations_easy_hard_source, 'r')
+    xy_coords_easy_hard  = text_file_easy_hard.readlines()
+
+    text_file_easy = open(annotations_easy_source, 'r')
+    xy_coords_easy  = text_file_easy.readlines()
+
+    xy_coords_hard = [coord for coord in xy_coords_easy_hard if coord not in xy_coords_easy ] 
+    text_file_hard = open(image_dest +"annotation_hard/" +str(count)+".txt", 'w')
+    for coord in xy_coords_hard:
+        text_file_hard.write(coord)
+    text_file_hard.close()
+
+    annotations_low_motility_source = os.path.join(data_dir, video, minivideo, "annotations_motility_low",
+                                           text_file)
+    annotations_high_motility_source = os.path.join(data_dir, video, minivideo, "annotations_motility_high",
+                                           text_file)
+    annotations_mid_motility_source = os.path.join(data_dir, video, minivideo, "annotations_motility_mid",
+                                           text_file)
+    annotations_wiggle_motility_source = os.path.join(data_dir, video, minivideo, "annotations_motility_wiggle",
+                                           text_file)
+    
+    annotations_sticking_motile_source = os.path.join(data_dir, video, minivideo, "annotations_sticking_motile",
+                                           text_file)
+    annotations_sticking_stick_source = os.path.join(data_dir, video, minivideo, "annotations_sticking_stick",
+                                           text_file)
+    annotations_sticking_non_motile_source = os.path.join(data_dir, video, minivideo, "annotations_sticking_non_motile",
+                                           text_file)
+    
+    shutil.copy(annotations_low_motility_source, image_dest +"annotation_motility_low/" +str(count)+".txt")
+    shutil.copy(annotations_high_motility_source, image_dest +"annotation_motility_high/" +str(count)+".txt")
+    shutil.copy(annotations_mid_motility_source, image_dest +"annotation_motility_mid/" +str(count)+".txt")
+    shutil.copy(annotations_wiggle_motility_source, image_dest +"annotation_motility_wiggle/" +str(count)+".txt")
+    
+    shutil.copy(annotations_sticking_stick_source, image_dest +"annotation_sticking_stick/" +str(count)+".txt")
+    shutil.copy(annotations_sticking_motile_source, image_dest +"annotation_sticking_motile/" +str(count)+".txt")
+    shutil.copy(annotations_sticking_non_motile_source, image_dest +"annotation_sticking_non_motile/" +str(count)+".txt")
+    
+    
+    if train == True:
+        count_train += 1
+    else:
+        count_test += 1
+        
+    return count_test, count_train
+    
+    
+    
+#get backgorund frame for every mini video
+#skip the first frame in every mini video
+#store image in train images set
+#store image+background diff + prev image diff in train images feature set
+#similarly for test
+
+def create_data(data_dir, dest_dir, trainfolder, train_video, testfolder, test_video, valfolder, val_video, method="background", num_prev=None, mean=False, sample=False, test_only = False, params=None, optical_flow_prior=1, frame_diff_prior=1):
+    os.makedirs(dest_dir, exist_ok=True)
+    data_dir_types = ["/images/", "/images_feature/", "/annotation_easy/", "/annotation_hard/", "/annotation_easy_hard/", "/annotation_easy_hard_veryhard/" , "/annotation_veryhard/" , "/annotation_motility_low", "/annotation_motility_wiggle", "/annotation_motility_mid", "/annotation_motility_high", "/annotation_sticking_stick", "/annotation_sticking_motile", "/annotation_sticking_non_motile"]
+    for video in test_video:
+        for dir_type in data_dir_types:
+            os.makedirs(os.path.join(dest_dir, testfolder, video) + dir_type, exist_ok=True)
+            
+    for video in val_video:
+        for dir_type in data_dir_types:
+            if test_only == True:
+                continue
+            os.makedirs(os.path.join(dest_dir, valfolder, video) + dir_type, exist_ok=True)
+            
+    for dir_type in data_dir_types:
+            os.makedirs(os.path.join(dest_dir, testfolder) + dir_type, exist_ok=True)
+            if test_only == True:
+                continue
+            os.makedirs(os.path.join(dest_dir, valfolder) + dir_type, exist_ok=True)
+            os.makedirs(dest_dir+trainfolder + dir_type, exist_ok=True)
+    
+    count_train = 0
+    count_test = 0
+    count_test_all = 0
+    count_val_all = 0
+
+    for video in natsorted(os.listdir(data_dir)):
+        
+        if not video.startswith('.') and os.path.isdir(os.path.join(data_dir,video))==True:
+            if test_only == True and video not in test_video:
+                continue
+            if video not in train_video + test_video + val_video:
+                continue
+            for minivideo in natsorted(os.listdir(os.path.join(data_dir,video))) :
+                if not minivideo.startswith('.') and os.path.isdir(os.path.join(data_dir,video,minivideo)) == True:
+                    video_path = os.path.join(data_dir,video) + "/" + minivideo + "video.mp4"
+                    print(video_path)
+
+                    if method == "background":
+                        # get the background model
+                        background = get_background(video_path, mean, sample=sample)
+                        # convert the background model to grayscale format
+                        background = cv2.cvtColor(background, cv2.COLOR_BGR2GRAY)
+                        skip_frame_list = [0]
+                    if method == "nprev":
+                        skip_frame_list = list(range(num_prev))
+                    if method in ["diff_from_max_absolute_consecutive_frame_diff", "max_absolute_consecutive_frame_diff"]:
+                        max_absolute_consecutive_frame_diff = get_absolute_consec_frame_diff_feature(video_path, max_consec_frame_diff=True)
+                    if method in ["max_absolute_all_frame_diff"]:
+                        absolute_all_frame_diff = get_absolute_all_frame_diff_feature(video_path, max_feature=True)
+                    if method in ["min_absolute_all_frame_diff"]:
+                        absolute_all_frame_diff = get_absolute_all_frame_diff_feature(video_path, max_feature=False)
+                        
+                        
+                        
+                    if video not in (test_video + val_video): 
+                        #print(video)
+                        train=True
+                    else:
+                        train=False
+                        count_test = 0
+                        if video in test_video:
+                            testfolder_video = os.path.join(testfolder, video)
+                        else:
+                            testfolder_video = os.path.join(valfolder, video)
+                    
+                        
+                    if method in ["optical_flow", "optical_flow_median_back", "optical_flow_from_median_frame", 
+                                  "optical_flow_combined", "diff_from_max_absolute_consecutive_frame_diff", 
+                                  "max_absolute_consecutive_frame_diff", "min_absolute_all_frame_diff", 
+                                  "max_absolute_all_frame_diff"]:
+                        #print(method)
+                        method_flow = cv2.calcOpticalFlowFarneback
+                        #params = [0.5, 3, 15, 3, 5, 1.2, 0]  # default Farneback's algorithm parameters
+                        # params = [0.5, 4, 18, 3, 5, 1.2, 0]  
+                        #params = [0.5, 2, 18, 3, 5, 1.2, 0]  
+                        #params = [0.5, 2, 20, 3, 5, 1.2, 0]  
+                        
+                        background = get_background(video_path, mean, sample=sample)
+                        # convert the background model to grayscale format
+                        background = cv2.cvtColor(background, cv2.COLOR_BGR2GRAY)
+                        
+                        if method == "optical_flow_from_median_frame":
+                            frames_org, frames_optical_flow = gen_dense_optical_flow_data(method_flow, video_path, params, to_gray=True, median=True, median_frame=background)
+                        else :
+                            frames_org, frames_optical_flow = gen_dense_optical_flow_data(method_flow, video_path, params, to_gray=True,num_frames_prior=optical_flow_prior)
+                            
+                        if method == "optical_flow_combined":
+                            m_frames_org, m_frames_optical_flow = gen_dense_optical_flow_data(method_flow, video_path, params, to_gray=True, median=True, median_frame=background)
+                            c_frames_org, c_frames_optical_flow = gen_dense_optical_flow_data(method_flow, video_path, params, to_gray=True)
+    #                     print(len(frames_org))
+    #                     print(frames_org[0].shape)
+    #                     print(len(frames_optical_flow))
+    #                     print(frames_optical_flow[0].shape)
+    
+                        if method in ["optical_flow", "optical_flow_from_median_frame"]:
+                            pure = True
+                        else:
+                            pure = False
+                                              
+
+                        print(len(frames_org))        
+                        for i, frame in enumerate(frames_org):
+                            #save frame in images
+                            if train==True:
+                                image_dest = dest_dir + "/" +trainfolder +"/"
+                                count = count_train
+                            else:
+                                #print(video)
+                                image_dest = dest_dir +"/" + testfolder_video + "/"
+                                count = count_test
+
+                            #print(frame.shape)
+                            img = Image.fromarray(frame, "RGB")
+                            img.save(image_dest +"images/"+str(count)+".tif")
+                            if method == "optical_flow_combined":
+                                feature_image = create_feature_image_optical_flow(frame, c_frames_optical_flow[i], pure=pure, optical_flow2=m_frames_optical_flow[i])
+                            elif method == "diff_from_max_absolute_consecutive_frame_diff":
+                                diff_from_max_absolute_consecutive_frame_diff = get_diff_from_absolute_consec_frame_diff_feature(frame_diff_feature=max_absolute_consecutive_frame_diff, image_path=frame, frame=True)
+                                feature_image = create_feature_image_optical_flow(frame, frames_optical_flow[i], pure, diff_from_max_absolute_consecutive_frame_diff, final_channel=True)
+                            elif method == "max_absolute_consecutive_frame_diff":
+                                feature_image = create_feature_image_optical_flow(frame, frames_optical_flow[i], pure, max_absolute_consecutive_frame_diff, final_channel=False)
+                            elif method in ["max_absolute_all_frame_diff", "min_absolute_all_frame_diff"]:
+                                feature_image = create_feature_image_optical_flow(frame, frames_optical_flow[i], pure, absolute_all_frame_diff[i+1], final_channel=False)
+                            else:
+                                feature_image = create_feature_image_optical_flow(frame, frames_optical_flow[i], pure, background)
+                            frame_num = i+1
+                            
+                            count_test, count_train = save_data(image_dest, count, feature_image, frame_num, data_dir, video, minivideo, count_test, 
+                                                                                count_train, train)
+                            if train == False:
+                                if video in test_video:
+                                    count = count_test_all
+                                    image_dest_all = dest_dir +"/" + testfolder + "/"
+                                    count_test_all, count_train = save_data(image_dest_all, count, feature_image, frame_num, data_dir, video, minivideo, count_test_all, 
+                                                                                 count_train, train)
+                                else:
+                                    count = count_val_all
+                                    image_dest_all = dest_dir +"/" + valfolder + "/"
+                                    count_val_all, count_train = save_data(image_dest_all, count, feature_image, frame_num, data_dir, video, minivideo, count_val_all, 
+                                                                                 count_train, train)
+                                img.save(image_dest_all +"images/"+str(count)+".tif")
+                                
+                            
+                    else:
+                        # print("in else")
+                        # print(method)
+                        for frame in natsorted(os.listdir(os.path.join(data_dir, video, minivideo, "images"))):
+                            frame_num = int(frame.split(".tif")[0])
+                            #skip first frame
+                            if frame_num not in skip_frame_list:
+                                #save frame in images
+                                images_source = os.path.join(data_dir, video, minivideo, "images", frame)
+                                if train==True:
+                                    image_dest = dest_dir + "/" +trainfolder +"/"
+                                    count = count_train
+                                    #print(count)
+                                else:
+                                    #print(video)
+                                    image_dest = dest_dir +"/" + testfolder_video + "/"
+                                    image_dest_all = dest_dir +"/" + testfolder + "/"
+                                    count = count_test
+
+                                shutil.copy(images_source, image_dest +"images/"+str(count)+".tif")
+
+                                #create new image
+                                prev_frame = str(frame_num-1) +".tif"
+                                prev_image = os.path.join(data_dir, video, minivideo, "images", prev_frame) 
+                                if method == "background":
+                                    feature_image = create_feature_image(image_path=images_source, background=background,prev_image_path=prev_image)
+                                if method == "nprev":
+                                    background=None
+                                    prev_frame2 = str(frame_num-num_prev) +".tif"
+                                    prev_image2 = os.path.join(data_dir, video, minivideo, "images", prev_frame2) 
+                                    feature_image = create_feature_image(image_path=images_source, background=background, 
+                                                                     prev_image_path=prev_image, prev_image_path2=prev_image2)
+
+
+
+                                count_test, count_train = save_data(image_dest, count, feature_image, frame_num, data_dir, video, minivideo, count_test, count_train, train)
+                                if train == False:
+                                    if video in test_video:
+                                        count = count_test_all
+                                        image_dest_all = dest_dir +"/" + testfolder + "/"
+                                        count_test_all, count_train = save_data(image_dest_all, count, feature_image, frame_num, data_dir, video, minivideo, count_test_all, 
+                                                                                     count_train, train)
+                                    else:
+                                        count = count_val_all
+                                        image_dest_all = dest_dir +"/" + valfolder + "/"
+                                        count_val_all, count_train = save_data(image_dest_all, count, feature_image, frame_num, data_dir, video, minivideo, count_val_all, 
+                                                                                     count_train, train)
+                                    shutil.copy(images_source, image_dest_all +"images/"+str(count)+".tif")
+                                    
+                            
+
+def create_train_data(target_data_sub_dir, dest_sub_dir, exp_name, train_video,  test_video, val_video, feature_method="optical_flow_median_back"):
+    feature_dir = "data_feature_optical_flow_median_back_2pyr_18win_background_img/"
+    data_dir = target_data_sub_dir
+    dest_dir = os.path.join(dest_sub_dir, exp_name, feature_dir)
+    trainfolder = "train"
+    testfolder = "test"
+    valfolder = "val"
+    params = [0.5, 2, 18, 3, 5, 1.2, 0] 
+
+    create_data(data_dir, dest_dir, trainfolder, train_video, testfolder, test_video, valfolder, val_video, 
+            method=feature_method, params=params)
+
+    
+    
+def create_test_data(target_data_sub_dir, dest_sub_dir, exp_name, test_video_list, feature_method="optical_flow_median_back"):
+    params = [0.5, 2, 18, 3, 5, 1.2, 0] 
+    for video in test_video_list:
+        print(video)
+        data_dir = target_data_sub_dir
+        dest_dir = os.path.join(dest_sub_dir, exp_name , f"data_{video}_feature_optical_flow_median_back_2pyr_18win_background_img/")
+        trainfolder = "train"
+        testfolder = "test"
+        valfolder = "val"
+        val_video = []
+        train_video = []
+        test_video = [f"{video}"]
+        params = [0.5, 2, 18, 3, 5, 1.2, 0]  
+
+        create_data(data_dir, dest_dir, trainfolder, train_video, testfolder, test_video, valfolder, val_video, 
+                    method=feature_method, test_only = True, params=params)
+    
+
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="Feature Preparation")
+    parser.add_argument("--target_data_sub_dir", default = "MEMTrack/data/collagen/", help="Path to the folder to process")
+    parser.add_argument("--dest_sub_dir",  default ="MEMTrack/DataFeatures/", help="Path to the folder to process")
+    parser.add_argument("--exp_name", default ="collagen_motility_inference",  help="Path to the folder to process")
+    parser.add_argument("--feature_method",  default="optical_flow_median_back", help="Path to the folder to process")
+    # optical flow median back with optical_Flow_prior is optical flow from xth previous frame
+    # "diff_from_max_absolute_consecutive_frame_diff" creates a feature for diff from the "max consecutive frame diff" feature, with a frame diff prior for xth frame diff
+    parser.add_argument('--train_video', type=str, nargs='+', help='a list of strings', default=[])
+    parser.add_argument('--val_video',type=str, nargs='+', help='a list of strings', default=[])
+    parser.add_argument('--test_video', type=str, nargs='+', help='a list of strings', default=[])
+
+    args = parser.parse_args(args)
+    target_data_sub_dir = args.target_data_sub_dir
+    dest_sub_dir  = args.dest_sub_dir
+    exp_name = args.exp_name
+    feature_method = args.feature_method
+    train_video = args.train_video
+    test_video = args.test_video
+    val_video = args.val_video
+    
+    create_train_data(target_data_sub_dir, dest_sub_dir, exp_name, train_video,  test_video, val_video, feature_method)
+    create_test_data(target_data_sub_dir, dest_sub_dir, exp_name, test_video_list, feature_method)
+    

MEMTrack/src/data_prep_utils.py

@@ -0,0 +1,531 @@
+import zipfile
+import openpyxl
+import os
+import csv
+import sys
+import cv2
+import shutil
+import PIL
+import glob
+import pandas as pd
+import numpy as np
+from natsort import natsorted
+from PIL import Image
+import argparse
+
+def load_map(video_map_path):
+    # Initialize an empty dictionary
+    video_dict = {}
+
+    # Open and read the text file
+    with open(video_map_path, 'r') as file:
+        next(file)
+        for line in file:
+            #print(line)
+            # Split each line into ID and name using the comma as a separator
+            video_id, video_name = line.strip().split(',')
+
+            # Convert the ID to an integer if needed
+            video_id = int(video_id)
+
+            # Add the key-value pair to the dictionary
+            video_dict[video_id] = video_name
+
+    # Print the loaded dictionary
+#     print(video_dict)
+    return video_dict
+
+#find video number from video map
+def find_next_video_num(data_dir):
+    #data_dir = "../Data/"
+    data_sub_dirs = glob.glob(f'{data_dir}/*')  
+    print("data dirs: ", data_sub_dirs)
+    video_dirs = glob.glob(f'{data_dir}/*/*')  
+    max_id = 0
+    for video in video_dirs:
+        #print(video)
+        if "ipynb" not in video:
+            video_id = int(video.split("video")[1])
+            #print(video_id)
+            max_id = max(video_id, max_id)
+    print("last video num: ",max_id)
+    return max_id +1
+
+def video_in_videomap(video_map_path, folder):
+   # load videomap
+    video_list = [value.strip() for value in load_map(video_map_path).values()]
+    video_list = [value.lstrip("_") for value in video_list]
+    
+    if os.path.basename(folder.split(".zip")[0]) in video_list:
+        return True
+        
+def add_video_to_videomap(video_map_path, video_dir, final_data_dir):
+    video_num = find_next_video_num(final_data_dir)
+    print("curr video num: ",video_num)
+    
+    if not os.path.exists(video_map_path):
+            video_map = open(video_map_path,'a+')
+            video_map.write("video_num, strain")
+            video_map.close()
+    
+    video_map = open(video_map_path,'a+')
+    video_strain = os.path.basename(video_dir.strip().strip("/"))
+    video_map.write("\n")
+    video_map.write(f"{video_num}, {video_strain}")
+    video_map.close()
+    
+    print("Added to videomap")
+    return video_num
+    
+    
+def create_annotations(video_dir, csv_file_name, inference_mode, video_num, target_data_sub_dir ):
+    count = 0
+    images_source = os.path.join(video_dir, "Images without Labels")
+    files = natsorted(os.listdir(images_source))
+#     files[:5]
+    data_dir = video_dir
+    all_files = os.listdir(data_dir)
+    if not inference_mode:
+        try:
+            #csv_file_name = list(filter(lambda f: f.endswith('Tracked.csv'), all_files))[0] if list(filter(lambda f: f.endswith('.csv'), all_files)) else list(filter(lambda f: f.endswith('.xlsx'), all_files))[0]
+            #csv_file_name = "Raw Data.csv"
+            csv_file = os.path.join(data_dir, csv_file_name)
+#             print(csv_file_name)
+            if ".csv" not in csv_file_name:
+                read_file = pd.read_excel(csv_file, engine='openpyxl')
+                read_file.to_csv (csv_file.split("xlsx")[0]+".csv", 
+                              index = None, header=True)
+                csv_file = csv_file.split("xlsx")[0]+".csv"
+#             print(csv_file)
+            file = open(csv_file)
+        except:
+            print("No Raw Data present, assuming Inference mode only")
+            inference_mode=True
+
+    if not inference_mode:
+        csvreader = csv.reader(file)
+        header = []
+        header = next(csvreader)
+        # header = header[0].split("\t")
+        print(header)
+        header = [x.lower() for x in header]
+        img_id = header.index('pid') #  header.index('t [Frame]') 
+        x_id = header.index('x [pixel]')
+        y_id = header.index('y [pixel]')
+        try:
+            sticking_id = header.index('sticking')
+            subpop_id = header.index('subpopulation')
+        except: 
+            sticking_id = None
+            subpop_id = None
+        try:
+            diff_id = header.index('e/h')
+
+        except:
+            diff_id = None
+            print("No diff id found, assuming all easy")
+#         print(img_id)
+#         print(x_id)
+#         print(y_id)
+#         print(sticking_id)
+#         print(subpop_id)
+
+        rows = []
+        for row in csvreader:
+                rows.append(row)
+        #rows
+        file.close()
+        
+    dest_dir = os.path.join(target_data_sub_dir, f"video{video_num}/")
+#     print("dest_dir: ", dest_dir)
+    images_target =os.path.join(dest_dir, "frame1/images/")
+    annotations_easy = os.path.join(dest_dir, "frame1/annotations_easy/")
+    annotations_easy_hard = os.path.join(dest_dir, "frame1/annotations_easy_hard/")
+    annotations_hard = os.path.join(dest_dir, "frame1/annotations_hard/")
+    annotations_very_hard = os.path.join(dest_dir, "frame1/annotations_veryhard/")
+    annotations_easy_hard_veryhard = os.path.join(dest_dir, "frame1/annotations_easy_hard_veryhard/")
+
+    annotations_motility_high = os.path.join(dest_dir, "frame1/annotations_motility_high/")
+    annotations_motility_low = os.path.join(dest_dir, "frame1/annotations_motility_low/")
+    annotations_motility_wiggle = os.path.join(dest_dir, "frame1/annotations_motility_wiggle/")
+    annotations_motility_mid = os.path.join(dest_dir, "frame1/annotations_motility_mid/")
+
+    annotations_sticking_motile = os.path.join(dest_dir, "frame1/annotations_sticking_motile/")
+    annotations_sticking_non_motile = os.path.join(dest_dir, "frame1/annotations_sticking_non_motile/")
+    annotations_sticking_stick = os.path.join(dest_dir, "frame1/annotations_sticking_stick/")
+
+
+    os.makedirs(images_target, exist_ok=True)
+    os.makedirs(annotations_easy_hard, exist_ok=True)
+    os.makedirs(annotations_easy_hard_veryhard, exist_ok=True)
+    os.makedirs(annotations_hard, exist_ok=True)
+    os.makedirs(annotations_easy, exist_ok=True)
+    os.makedirs(annotations_very_hard, exist_ok=True)
+
+
+    os.makedirs(annotations_sticking_stick, exist_ok=True)
+    os.makedirs(annotations_sticking_motile, exist_ok=True)
+    os.makedirs(annotations_sticking_non_motile, exist_ok=True)
+
+    os.makedirs(annotations_motility_high, exist_ok=True)
+    os.makedirs(annotations_motility_wiggle, exist_ok=True)
+    os.makedirs(annotations_motility_mid, exist_ok=True)
+    os.makedirs(annotations_motility_low, exist_ok=True)
+
+    for i,image in enumerate(files):
+        #copy and rename images
+        image_name = str(count) +".tif"
+        shutil.copy(os.path.join(images_source, image), os.path.join(images_target,image_name))
+#         print(os.path.join(images_source, image))
+#         print(os.path.join(images_target,image_name))
+        #create annoatations txt file
+        #image_id = int(files_track[i].split(")")[1].split(".tif")[0])
+        txt_file = open(annotations_easy_hard + str(count) +".txt",'w')
+        txt_file_easy = open(annotations_easy + str(count) +".txt",'w')
+        txt_file_hard = open(annotations_hard + str(count) +".txt",'w')
+        txt_file_very_hard = open(annotations_very_hard + str(count) +".txt",'w')
+        txt_file_all = open(annotations_easy_hard_veryhard + str(count) +".txt",'w')
+
+        txt_file_motility_high = open(annotations_motility_high + str(count) +".txt",'w')
+        txt_file_motility_low = open(annotations_motility_low + str(count) +".txt",'w')
+        txt_file_motility_wiggle = open(annotations_motility_wiggle  + str(count) +".txt",'w')
+        txt_file_motility_mid = open(annotations_motility_mid + str(count) +".txt",'w')
+
+        txt_file_sticking_stick = open(annotations_sticking_stick + str(count) +".txt",'w')
+        txt_file_sticking_motile = open(annotations_sticking_motile + str(count) +".txt",'w')
+        txt_file_sticking_non_motile = open(annotations_sticking_non_motile + str(count) +".txt",'w')
+
+        if not inference_mode:
+            for row in rows:
+                # print(image_id)
+                # print( row[img_id])
+                ##print(row)
+                if int(row[img_id])-1 == int(count):#PID starts from 1
+                    # print(image_id)
+                    # print( row[img_id])
+                    txt_file_all.write(row[x_id])
+                    txt_file_all.write(" ")
+                    txt_file_all.write(row[y_id])
+                    txt_file_all.write("\n")
+                    try:
+                        if row[diff_id]=="E":
+                            txt_file_easy.write(row[x_id])
+                            txt_file_easy.write(" ")
+                            txt_file_easy.write(row[y_id])
+                            txt_file_easy.write("\n")
+
+                            txt_file.write(row[x_id])
+                            txt_file.write(" ")
+                            txt_file.write(row[y_id])
+                            txt_file.write("\n")
+
+                        if row[diff_id]=="H":
+                            txt_file_hard.write(row[x_id])
+                            txt_file_hard.write(" ")
+                            txt_file_hard.write(row[y_id])
+                            txt_file_hard.write("\n")
+
+                            txt_file.write(row[x_id])
+                            txt_file.write(" ")
+                            txt_file.write(row[y_id])
+                            txt_file.write("\n")                
+
+                        if row[diff_id]=="VH":
+                            txt_file_very_hard.write(row[x_id])
+                            txt_file_very_hard.write(" ")
+                            txt_file_very_hard.write(row[y_id])
+                            txt_file_very_hard.write("\n")
+
+                        if row[subpop_id]=="L":
+                            txt_file_motility_low.write(row[x_id])
+                            txt_file_motility_low.write(" ")
+                            txt_file_motility_low.write(row[y_id])
+                            txt_file_motility_low.write("\n")
+                        if row[subpop_id]=="W":
+                           #print("wiggle")
+                            txt_file_motility_wiggle.write(row[x_id])
+                            txt_file_motility_wiggle.write(" ")
+                            txt_file_motility_wiggle.write(row[y_id])
+                            txt_file_motility_wiggle.write("\n")
+                        if row[subpop_id]=="M":
+                            txt_file_motility_mid.write(row[x_id])
+                            txt_file_motility_mid.write(" ")
+                            txt_file_motility_mid.write(row[y_id])
+                            txt_file_motility_mid.write("\n")
+                        if row[subpop_id]=="H":
+                            txt_file_motility_high.write(row[x_id])
+                            txt_file_motility_high.write(" ")
+                            txt_file_motility_high.write(row[y_id])
+                            txt_file_motility_high.write("\n")
+
+                        if row[sticking_id]=="S":
+                            txt_file_sticking_stick.write(row[x_id])
+                            txt_file_sticking_stick.write(" ")
+                            txt_file_sticking_stick.write(row[y_id])
+                            txt_file_sticking_stick.write("\n")
+                        if row[sticking_id]=="M":
+                            txt_file_sticking_motile.write(row[x_id])
+                            txt_file_sticking_motile.write(" ")
+                            txt_file_sticking_motile.write(row[y_id])
+                            txt_file_sticking_motile.write("\n")
+                        if row[sticking_id]=="NM":
+                            txt_file_sticking_non_motile.write(row[x_id])
+                            txt_file_sticking_non_motile.write(" ")
+                            txt_file_sticking_non_motile.write(row[y_id])
+                            txt_file_sticking_non_motile.write("\n")
+                    except:
+                        txt_file_easy.write(row[x_id])
+                        txt_file_easy.write(" ")
+                        txt_file_easy.write(row[y_id])
+                        txt_file_easy.write("\n")
+
+            txt_file.close()
+        count = count+1
+    print("Annotations processed")
+    return inference_mode
+    
+    
+
+#run only once to generate bacteria data
+def generate_bacteria_data(file, video_dir):
+    # Generate Bacteria Tracks specific data for Bacteria Analysis
+
+    csvreader = csv.reader(file)
+    header = []
+    header = next(csvreader)
+    # header = header[0].split("\t")
+    print(header)
+    header = [x.lower() for x in header]
+    img_id = header.index('pid') #  header.index('t [Frame]') 
+    x_id = header.index('x [pixel]')
+    y_id = header.index('y [pixel]')
+    try:
+        diff_id = header.index('e/h')
+    except:
+        diff_id = None
+        print("No diff id found, assuming all easy")
+    track_id = header.index('tid')
+    print(img_id)
+    print(x_id)
+    print(y_id)
+    print(track_id)
+
+    rows = []
+    for row in csvreader:
+            rows.append(row)
+    print(rows[:6])
+    file.close()
+    
+    tid_visited = []
+    #video_dir = "./video9_feature_optical_flow_median_back_2pyr_18win/test/"
+    bacteria_folder = "bacteria"
+    bacteria_easy_hard_state_file = "easy_hard_veryhard"
+    bacteria_coords = "xy_coord"
+    count = 0
+    max_pid = 0
+    prev_tid = 0
+
+    for row in rows:
+        pid = int(row[img_id])-1
+        if max_pid < pid:
+            max_pid = pid
+
+    for row in rows:
+        tid = row[track_id]
+        if tid not in tid_visited:
+            tid_visited.append(tid)
+
+    #         if count<(max_pid-1) and count>0:
+    #             # print(row)
+    #             # print(rows[i+1])
+    #             print(count)
+    #             print(max_pid)
+    #             for i in range((max_pid - count-1)):
+    #                 txt_file.write(str(count))
+    #                 txt_file.write(" ")
+    #                 txt_file.write("NotPresent")
+    #                 txt_file.write("\n")
+
+    #                 coord_file = open(os.path.join(video_dir, bacteria_folder, str(prev_tid), bacteria_coords, str(count)) +".txt",'w')
+    #                 coord_file.close()
+    #                 count = count+1
+            # txt_file.close()
+            # coord_file.close()
+            count = 0
+            os.makedirs(os.path.join(video_dir, bacteria_folder, str(tid), bacteria_coords), exist_ok=True)
+            #os.makedirs(os.path.join(video_dir, bacteria_folder, str(tid)), exist_ok=True)
+            try:
+                os.remove(os.path.join(video_dir, bacteria_folder, str(tid), bacteria_easy_hard_state_file) +".txt")
+                #os.remove(os.path.join(video_dir, bacteria_folder, str(tid), bacteria_coords, str(count)) +".txt")
+            except OSError:
+                pass
+
+        txt_file = open(os.path.join(video_dir, bacteria_folder, str(tid), bacteria_easy_hard_state_file) +".txt",'a')
+        pid = int(row[img_id]) - 1
+        if int(pid) == 0: #for optical flow since first frame is skipped
+            continue
+        if pid-2>count: # pid-1 because 1 is skipped
+            # print(count)
+            # print(pid)
+            for i in range((pid - count-1)):
+                txt_file.write(str(count))
+                txt_file.write(" ")
+                txt_file.write("NotPresent")
+                txt_file.write("\n")
+
+                coord_file = open(os.path.join(video_dir, bacteria_folder, str(tid), bacteria_coords, str(count)) +".txt",'w')
+                coord_file.close()
+                count = count+1
+
+        txt_file.write(str(count))
+        txt_file.write(" ")
+        try:
+            txt_file.write(row[diff_id])
+        except:
+            txt_file.write("E")
+            #print("No diff id found, assuming all easy")
+        txt_file.write("\n")
+
+
+
+        coord_file = open(os.path.join(video_dir, bacteria_folder, str(tid), bacteria_coords, str(count)) +".txt",'a')
+        coord_file.write(row[x_id])
+        coord_file.write(" ")
+        coord_file.write(row[y_id])
+        coord_file.write("\n")
+
+
+        count = count+1
+
+    if count<(max_pid-1) and count>0:
+            # print(row)
+            # print(rows[i+1])
+            print(count)
+            print(max_pid)
+            for i in range((max_pid - count)):
+                txt_file.write(str(count))
+                txt_file.write(" ")
+                txt_file.write("NotPresent")
+                txt_file.write("\n")
+
+                coord_file = open(os.path.join(video_dir, bacteria_folder, str(tid), bacteria_coords, str(count)) +".txt",'w')
+                coord_file.close()
+                count = count+1
+    txt_file.close()
+    coord_file.close()
+
+    
+def create_video(data_dir):
+    # choose codec according to format needed
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v') 
+    #print(data_dir)
+    img_sample = cv2.imread(os.path.join(data_dir,"images/0.tif"))
+    #print(img_sample.shape)
+    height, width, channels = img_sample.shape
+    
+    video = cv2.VideoWriter(data_dir + 'video.mp4', fourcc, 1, (width, height))
+    #data_dir = "./Data/video3/"
+    image_dir = os.path.join(data_dir, "images")
+    for frame in natsorted(os.listdir(image_dir)):
+        #print(frame)
+        img = cv2.imread(os.path.join(image_dir, frame))
+        video.write(img)
+    video.release()
+
+def get_background(file_path, mean=True, sample=False):
+    cap = cv2.VideoCapture(file_path)
+    #print(cap.read())
+    # we will randomly select 50 frames for the calculating the median
+    #frame_indices = cap.get(cv2.CAP_PROP_FRAME_COUNT) * np.random.uniform(size=20)
+    frame_indices = list(range(0,int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) ))
+    print(len(frame_indices))
+    # we will store the frames in array
+    frames = []
+    for idx in frame_indices:
+        # set the frame id to read that particular frame
+        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
+        ret, frame = cap.read()
+        #print(ret)
+        frames.append(frame)
+    if mean:
+         # calculate the mean
+        background_frame = np.mean(frames, axis=0).astype(np.uint8)
+    else:
+        # calculate the median
+        background_frame = np.median(frames, axis=0).astype(np.uint8)
+    if sample==True:
+        background_frame = cv2.imread("./Control_2_b0t5306c0x0-660y0-492.tiff")
+        #background_frame = cv2.imread("./RBS 2_1_b0t2791c0x0-660y0-492.tiff")
+    return background_frame
+
+
+def process_data(folder, src, final_data_dir, out_sub_dir, videomap="videomap.txt", csv_file_name="Raw Data.csv", inference_mode=False, unzip=False):
+   
+    target_data_sub_dir = os.path.join(final_data_dir, out_sub_dir)
+    print("target_data_sub_dir: ", target_data_sub_dir)
+    os.makedirs(target_data_sub_dir, exist_ok=True)
+    
+    video_map_path = os.path.join(final_data_dir, videomap )
+    print("video_map_path: ", video_map_path)
+    
+    video_dir = os.path.join(src, folder.split(".zip")[0])
+    
+    if unzip:
+        # unzip data
+        zip_data_dir = os.path.join(src, folder)
+        with zipfile.ZipFile(zip_data_dir, 'r') as zip_ref:
+            zip_ref.extractall(src)
+    
+    if os.path.exists(video_map_path) and video_in_videomap(video_map_path, folder):
+        print("video alread processed: ", folder.split(".zip")[0])
+        raise Exception("video alread processed")
+        return
+    
+    # add video to videomap
+    video_num = add_video_to_videomap(video_map_path, video_dir, final_data_dir)
+    inference_mode = create_annotations(video_dir, csv_file_name, inference_mode, video_num, target_data_sub_dir)
+   
+    if not inference_mode:
+        target_video_dir = f"{target_data_sub_dir}/video{video_num}/frame1"
+        csv_file = os.path.join(video_dir, csv_file_name)
+        file = open(csv_file)
+        generate_bacteria_data(file, target_video_dir)
+        file.close()
+    
+    data_path = target_data_sub_dir
+    test_video = [f"video{video_num}"]
+    for video in natsorted(test_video):
+        if not video.startswith('.') and os.path.isdir(os.path.join(data_path,video))==True:
+            for minivideo in natsorted(os.listdir(os.path.join(data_path,video))):
+                if not minivideo.startswith('.') and os.path.isdir(os.path.join(data_path,video,minivideo)) == True:
+                    #print(minivideo)
+                    create_video(os.path.join(data_path,video,minivideo))
+    
+    print(f"Video {video_num} processed")
+    return video_num
+    
+    
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="Process folders.")
+    parser.add_argument("--src_dir", default = "DataAll/Collagen/", help="Path to the folder to process")
+    parser.add_argument("--folder",  default ="video1", help="Path to the folder to process")
+    parser.add_argument("--out_dir", default ="MEMTrack/data/",  help="Path to the folder to process")
+    parser.add_argument("--out_sub_dir",   help="Path to the folder to process")
+    parser.add_argument("--csv_file_name", help="Path to the folder to process", default= "Raw Data.csv")
+    parser.add_argument("--videomap", default="videomap.txt")
+    parser.add_argument("--no_gt", action="store_true")
+    parser.add_argument("--unzip", action="store_true", help="Unzip folder")
+
+
+    args = parser.parse_args(args)
+    src = args.src_dir
+    final_data_dir  = args.out_dir
+    out_sub_dir = args.out_sub_dir
+    videomap = args.videomap
+    csv_file_name = args.csv_file_name
+    inference_mode = args.no_gt
+    folder= args.folder
+    unzip = args.unzip    
+    
+    process_data(folder, src, final_data_dir, out_sub_dir, videomap, csv_file_name, inference_mode, unzip)

MEMTrack/src/evaluation_step_wise_motility.py

@@ -0,0 +1,755 @@
+import random
+import os
+import cv2
+import json
+import matplotlib
+import argparse
+from matplotlib import image, patches
+import matplotlib.pyplot as pyplot
+import seaborn as sns
+import numpy as np
+import seaborn as sns
+import pandas as pd
+import matplotlib.pyplot as plt
+
+ap = argparse.ArgumentParser(description='Inference')
+ap.add_argument('--video_map_path', type=str, metavar='PATH')
+ap.add_argument('--data_path',  type=str, metavar='PATH')
+ap.add_argument('--video', type=str, metavar='PATH')
+ap.add_argument('--all_labels_file',  type=str, metavar='PATH')
+args = ap.parse_args()
+
+def bb_intersection_over_union(boxA, boxB):
+    # determine the (x, y)-coordinates of the intersection rectangle
+    xA = max(boxA[0], boxB[0])
+    yA = max(boxA[1], boxB[1])
+    xB = min(boxA[2], boxB[2])
+    yB = min(boxA[3], boxB[3])
+
+     # compute the area of intersection rectangle
+    interArea = abs(max((xB - xA, 0)) * max((yB - yA), 0))
+    if interArea == 0:
+        return 0
+
+    # compute the area of both the prediction and ground-truth
+    # rectangles
+    boxAArea = abs((boxA[2] - boxA[0]) * (boxA[3] - boxA[1]))
+    boxBArea = abs((boxB[2] - boxB[0]) * (boxB[3] - boxB[1]))
+
+    # compute the intersection over union by taking the intersection
+    # area and dividing it by the sum of prediction + ground-truth
+    # areas - the interesection area
+    iou_box = interArea / float(boxAArea + boxBArea - interArea)
+
+    # return the intersection over union value
+    return iou_box
+
+def get_iou_mat_image(detectron_train_file_bbox_list, coco_results_file_bbox_list ):
+    # iou_mat = []
+    # for gt_box in detectron_train_file_bbox_list:
+    #     iou_row = []
+    #     for pred_box in coco_results_file_bbox_list:
+    #         gt_box_xyxy = [gt_box[0], gt_box[1], gt_box[0] + gt_box[2], gt_box[1] + gt_box[3]]
+    #         pred_box_xyxy = [pred_box[0], pred_box[1], pred_box[0] + pred_box[2], pred_box[1] + pred_box[3]]
+    #         # print(pred_box_xyxy)
+    #         iou_boxes = bb_intersection_over_union(gt_box_xyxy, pred_box_xyxy)
+    #         iou_row.append(iou_boxes)
+    #     iou_mat.append(iou_row)
+    iou_mat = np.zeros((len(detectron_train_file_bbox_list), len(coco_results_file_bbox_list)))
+    for i, gt_box in enumerate(detectron_train_file_bbox_list):
+        for j, pred_box in enumerate(coco_results_file_bbox_list):
+            gt_box_xyxy = [gt_box[0], gt_box[1], gt_box[0] + gt_box[2], gt_box[1] + gt_box[3]]
+            pred_box_xyxy = [pred_box[0], pred_box[1], pred_box[0] + pred_box[2], pred_box[1] + pred_box[3]]
+            iou_mat[i, j] = bb_intersection_over_union(gt_box_xyxy, pred_box_xyxy)
+    return iou_mat
+
+def get_matching_boxes_iou(detectron_train_file_bbox_list, iou_mat, iou_thresh = 0.75):
+    matching_iou_boxes = []
+    iou_mat = iou_mat.tolist()
+    for i in range(0, len(detectron_train_file_bbox_list)):
+        if iou_mat[i]:
+            iou_row_max = max(iou_mat[i])
+            iou_row_max_pred_id = iou_mat[i].index(iou_row_max)
+           # print("iou_mat:", iou_mat)           
+            # print((iou_row_max))
+            # print(iou_row_max_pred_id)
+            #print(iou_row_max)
+            if iou_row_max>iou_thresh:
+                matching_iou_boxes.append([i, iou_row_max_pred_id, iou_row_max])
+    # print(matching_iou_boxes)  
+    #print("Number of matching IOU Ground truth and Predicted boxes: " , len(matching_iou_boxes))
+    return matching_iou_boxes
+
+def compute_mse_matching_boxes(matching_iou_boxes, detectron_train_file_bbox_list, coco_results_file_bbox_list):
+    # Compute MSE of fiber intersection based on matching iou-boxes
+    mse = []
+    iou_sum = 0
+    for matching_bbox in matching_iou_boxes:
+           # ground truth
+            x = detectron_train_file_bbox_list[matching_bbox[0]][0]
+            y = detectron_train_file_bbox_list[matching_bbox[0]][1]
+            w = detectron_train_file_bbox_list[matching_bbox[0]][2]
+            h = detectron_train_file_bbox_list[matching_bbox[0]][3]
+            center_x_gt = x + (w/2)
+            center_y_gt = y + (h/2)
+
+            # predicted
+            x = coco_results_file_bbox_list[matching_bbox[1]][0]
+            y = coco_results_file_bbox_list[matching_bbox[1]][1]
+            w = coco_results_file_bbox_list[matching_bbox[1]][2]
+            h = coco_results_file_bbox_list[matching_bbox[1]][3]
+            center_x_pred = x + (w / 2)
+            center_y_pred = y + (h / 2)  
+
+            mse.append(((center_x_gt - center_x_pred)**2 + (center_y_gt - center_y_pred)** 2)**0.5)
+            iou_sum += matching_bbox[2]
+    
+   # np.array(matching_iou_boxes)[2].mean()
+    if len(matching_iou_boxes)>0:
+        return np.array(mse).mean(), iou_sum/len(matching_iou_boxes), mse
+    else:
+        return np.array(mse).mean(), -1, mse
+
+def nms_predicted_bounding_boxes(pred_bbox_list, pred_bbox_scores, iou_thresh_nms=0.95):
+    #NMS of predicted boxes
+    #print(len(pred_bbox_list))
+    final_tracking_bbox_list =[]
+    iou_mat = get_iou_mat_image(pred_bbox_list, pred_bbox_list)
+    
+    matching_pred_boxes_iou = get_matching_boxes_iou(pred_bbox_list, iou_mat, iou_thresh = iou_thresh_nms)
+    while(len(matching_pred_boxes_iou)>0):
+        sorted_bbox_list = sorted(zip(pred_bbox_scores, pred_bbox_list))
+        iou_mat = get_iou_mat_image(pred_bbox_list, [pred_bbox_list[-1]])
+        #print(iou_mat)
+        pred_bbox_list_temp = []
+        pred_bbox_scores_temp = []
+        for index, iou in enumerate(iou_mat):
+            if iou[0]<iou_thresh_nms:
+                pred_bbox_list_temp.append(pred_bbox_list[index])
+                pred_bbox_scores_temp.append(pred_bbox_scores[index])
+        # matching_pred_boxes_iou = get_matching_boxes_iou(pred_bbox_list, iou_mat, iou_thresh = iou_thresh_nms)
+        # print(matching_pred_boxes_iou)
+        final_tracking_bbox_list.append(pred_bbox_list[-1]) #add highest scored bbox to final list
+        # matching_pred_boxes_index=[]
+        # for bbox in matching_pred_boxes_iou:
+        #     matching_pred_boxes_index.append(bbox[0])
+        # print("hello",matching_pred_boxes_index)
+      
+        # for index, bbox in enumerate(pred_bbox_list):
+        #     if index not in matching_pred_boxes_index:
+        #         pred_bbox_list_temp.append(pred_bbox_list[index])
+        #         pred_bbox_scores_temp.append(pred_bbox_scores[index])
+        pred_bbox_list = pred_bbox_list_temp
+        pred_bbox_scores = pred_bbox_scores_temp
+        iou_mat = get_iou_mat_image(pred_bbox_list, pred_bbox_list)
+        matching_pred_boxes_iou = get_matching_boxes_iou(pred_bbox_list, iou_mat, iou_thresh = iou_thresh_nms)
+        
+      
+    #print(final_tracking_bbox_list + (pred_bbox_list))
+    # print(len(final_tracking_bbox_list))
+    # print(len(pred_bbox_list))
+    #print(len(final_tracking_bbox_list + (pred_bbox_list)))
+    return final_tracking_bbox_list + (pred_bbox_list)
+
+    
+def get_statistics(data_path, images_path, coco_eval_path, test_json_path, train_file_name="train.json", detectron_train_file_name = "boardetect_train_coco_format.json",iou_thresh=0.75, bacteria_tp_path="./bacteria-detections/", coco_file = "coco_instances_results.json",print_details=True, store_results_path = None):
+    # data_path = "./data_phase_nodes/"
+    # images_path = data_path + "images_cropped_resized/"
+    # coco_eval_path = "./coco_eval_phase_cropped/"
+    num_images = 0
+    tp_sum = 0
+    fp_sum = 0
+    fn_sum = 0
+    mse_image_counter =0
+    iou_image_counter =0
+    ratio_gt_p = 0
+    ratio_p_gt = 0
+    mse = 0
+    precision = 0
+    recall = 0
+    iou = 0
+    all_img_mse = []
+    for filename in sorted(os.listdir(images_path)):  
+        # print(filename)
+        image_id = filename.split(".tif")[0]
+        #print(image_id)
+        train_file = json.load(open(data_path + train_file_name,'r'))
+        train_file_bbox_list = []
+        #print(train_file[int(image_id)]["file_name"])
+        if train_file[int(image_id)]["file_name"] == filename:
+            for i, annotation in enumerate(train_file[int(image_id)]["annotations"]):
+                train_file_bbox_list.append(train_file[int(image_id)]["annotations"][i]["bbox"])
+            # print("Number of Ground Truth boxes: ",len(train_file_bbox_list)) 
+            # print(train_file_bbox_list[0])
+
+        detectron_train_file = json.load(open(test_json_path + detectron_train_file_name,'r'))
+        detectron_train_file_bbox_list = []
+        
+        
+        try:
+            for i,annotation in enumerate(detectron_train_file["annotations"]):
+                if detectron_train_file["annotations"][i]["image_id"] == int(image_id):
+                    detectron_train_file_bbox_list.append(detectron_train_file["annotations"][i]["bbox"])
+                #annotation_image_id = detectron_train_file["annotations"][i]["image_id"]
+                #print(detectron_train_file["annotations"][i]["bbox"])
+        # print(len(detectron_train_file_bbox_list)) 
+    #     print(detectron_train_file_bbox_list[0])
+            num_ground_truth_boxes = len(detectron_train_file_bbox_list)
+        except:
+            num_ground_truth_boxes = 0
+      
+
+        coco_results_file = json.load(open(coco_eval_path + coco_file,'r'))
+        coco_results_file_bbox_list = []
+        coco_results_file_bbox_scores = []
+
+        for i,annotation in enumerate(coco_results_file):
+            #print(coco_results_file[i])
+            #print(image_id)
+            if coco_results_file[i]["image_id"] == int(image_id):
+                #print(true)
+                coco_results_file_bbox_list.append(coco_results_file[i]["bbox"])
+                coco_results_file_bbox_scores.append(coco_results_file[i]["score"])
+    # print(coco_results_file[i]["bbox"])
+    # print(len(coco_results_file_bbox_list))
+    # print((coco_results_file_bbox_list))  
+                  #print(detectron_train_file["annotations"][i]["bbox"])
+        # print("Number of Predicted bounding boxes: ", len(coco_results_file_bbox_list)) 
+        # print(coco_results_file_bbox_list[0])
+
+        #coco_results_file_bbox_list = nms_predicted_bounding_boxes(coco_results_file_bbox_list, coco_results_file_bbox_scores ,iou_thresh_nms=0.001)
+        
+        # Stat 1 - Ratio of boxes Predicted to Groud truth boxes
+        num_predicted_boxes = len(coco_results_file_bbox_list) 
+        
+
+        num_images = num_images + 1
+        if num_ground_truth_boxes>0:
+            
+            #Ratio of GT to P per image
+            ratio_gt_p = ratio_gt_p + (num_ground_truth_boxes/ max(num_predicted_boxes,1))
+            ratio_p_gt = ratio_p_gt + (num_predicted_boxes / max(num_ground_truth_boxes,1))
+            
+
+            # Stat 2 - MSE of fiber intersections from the matching boxes
+            iou_mat = get_iou_mat_image(detectron_train_file_bbox_list, coco_results_file_bbox_list)
+            matching_boxes_iou = get_matching_boxes_iou(detectron_train_file_bbox_list, iou_mat, iou_thresh = iou_thresh)
+            mse_image, iou_image, mse_list = compute_mse_matching_boxes(matching_boxes_iou, detectron_train_file_bbox_list, coco_results_file_bbox_list)
+            if mse_image>=0: #if no predicted boxes or no ground truth boxes then iou is nan
+                mse = mse + mse_image
+                mse_image_counter += 1
+            if iou_image>=0:#if no predicted boxes or no ground truth boxes then mse is nan
+                iou += iou_image
+                iou_image_counter += 1
+        
+            true_positive = len(matching_boxes_iou)
+            # print("num image: ", num_images)
+            # print("Num pred: ", num_predicted_boxes)
+            # print("num matched: ", len(matching_boxes_iou))
+            os.makedirs(bacteria_tp_path, exist_ok=True)
+            coord_file = open(bacteria_tp_path + image_id +".txt",'w')
+            coord_file.write(image_id)
+            coord_file.write(" ")
+            coord_file.write(str(true_positive))
+            coord_file.write("\n")
+            coord_file.close()
+    
+            # OLD METHOD OF PRECISION RECALL CALCULATION
+            # assumes 0 for no tp
+            # precision is average of averages
+#             if true_positive>0:
+#                 false_positive = num_predicted_boxes - true_positive
+#                 false_negative = num_ground_truth_boxes - true_positive
+
+#                 precision += true_positive/(true_positive + false_positive)
+
+#                 recall += true_positive/(true_positive + false_negative)
+            all_img_mse.extend(mse_list)
+
+            # UPDATED METHOD
+            # precision is total tp / total tp+fp in all images not average of preciosn for all images
+        
+            false_positive = abs(num_predicted_boxes - true_positive)
+            false_negative = abs(num_ground_truth_boxes - true_positive)
+            
+            tp_sum = tp_sum + true_positive
+            fp_sum = fp_sum + false_positive
+            fn_sum = fn_sum + false_negative
+        
+        if num_ground_truth_boxes==0:
+            #rint("no gt")
+            fp_sum = fp_sum + num_predicted_boxes         
+    
+    # Stat 1 - Ratio of boxes Predicted to Groud truth boxes
+    avg_ratio_gt_p = ratio_gt_p/num_images 
+    avg_ratio_p_gt = ratio_p_gt/num_images 
+    
+
+    # Stat 2 - MSE of fiber intersections from the matching boxes
+    try:
+        avg_mse = mse / mse_image_counter
+    except:
+        avg_mse = 0
+    
+    
+    try:
+        avg_iou = iou / iou_image_counter
+    except:
+        avg_iou = 0
+    try:
+        avg_prec = tp_sum / (tp_sum + fp_sum)
+        avg_recall = tp_sum / (tp_sum + fn_sum)
+    except:
+        avg_prec = 0
+        avg_recall = 0
+
+    if store_results_path:
+        result_file = open(store_results_path,'a+')
+        result_file.write(str(tp_sum))
+        result_file.write(",")
+        result_file.write(str(fp_sum))
+        result_file.write(",")
+        result_file.write(str(fn_sum))
+        result_file.write("\n")
+        result_file.close()
+        
+    if print_details:
+        # ap = precision/num_images
+        # ar = recall/num_images
+        print("Average Ground Truth to Predicted Ratio: ", avg_ratio_gt_p)
+        print("Average Predicted to Ground Truth Ratio: ", avg_ratio_p_gt)
+        print("Average Mean Squared Error of fiber intersections: ", avg_mse)
+        print("Average IoU of TP boxes: ", avg_iou)  
+        print("TP sum: ", tp_sum)
+        print("FP sum: ", fp_sum)
+        print("FN sum: ", fn_sum)
+        
+        print("Average Precision: ", avg_prec)    
+        print("Average Recall: ", avg_recall)  
+
+        pyplot.figure(figsize=(10, 3))
+        sns.distplot(all_img_mse)
+        pyplot.xlabel("Nodal Errors (in px sq)")
+        # plt.xlabel("")
+        pyplot.grid("on", alpha = 0.3)
+        pyplot.show()
+    return avg_prec, avg_recall
+
+video_map_path = args.video_map_path #"/alldata/medha/CleanCodeData/Data/videomap.txt"
+video_map = open(video_map_path,'r',)
+header = [x.strip() for x in (video_map.readline().strip()).split(",")]
+video_num_id = header.index("video_num")
+strain_id = header.index("strain")
+strain_map = {}
+for line in video_map.readlines():
+    line_details = [x.strip() for x in line.split(",")]
+    video_num = line_details[video_num_id]
+    strain = line_details[strain_id]
+    strain_map[video_num] = strain
+strain_map[""] = "all"
+
+#53 54 58 59 60 64 65 69 70 71 75 76 80 81 82 collagen
+#85 24 30 35 83 84 agar0.2
+#updated gar 19 22 29 30 83 84 
+#53 60 64 69 75 82 
+video_num = args.video
+src_path  = args.data_path #"/alldata/medha/CleanCodeData/DataFeatures/exp_collagen_train_veryhard/"
+print(src_path)
+video_path = f"data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/"
+data_path = os.path.join(src_path,video_path, "test/")
+print(data_path)
+images_path = data_path + "images/"
+coco_eval_path = data_path
+test_json_path = data_path
+precision = {}
+recall = {}
+store_file_path_main = src_path + "test_set_results"
+
+# ### Step 1: Detection Individual Models
+
+# In[15]:
+
+
+try:
+    difficulty_level = "Motility-low"
+    coco_file=f"coco_instances_results_{difficulty_level}.json"
+    detectron_train_file_name = f"boardetect_test_coco_format_{difficulty_level}.json"
+    train_file_name=f"test_{difficulty_level}.json"
+    
+    coco_file = "coco_instances_results_wiggle_diffusivity.json"
+    train_file_name="test_Motility-wiggle.json"
+    detectron_train_file_name = "boardetect_test_coco_format_Motility-wiggle.json"
+
+
+    store_file_path = store_file_path_main + f"_{difficulty_level}.txt" if store_file_path_main else None
+    if store_file_path:
+        result_file = open(store_file_path,'a+')
+        result_file.write(str(video_num))
+        result_file.write(",")
+        result_file.close()
+    #bacteria_tp_path = ./data_video22_feature_optical_flow_median_back_2pyr_18win_background_img/test/1/predicted/"
+    # visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+    #                   train_file_name=train_file_name,detectron_train_file_name=detectron_train_file_name, 
+    #                   coco_file=coco_file)
+
+    pr, rc = get_statistics(data_path, images_path, coco_eval_path,test_json_path, 
+                   train_file_name=train_file_name,detectron_train_file_name = detectron_train_file_name, 
+                   iou_thresh=0.1, coco_file=coco_file, print_details=False, store_results_path=store_file_path)
+    print("Precision: ", pr)
+    print("Recall: ", rc)
+    precision[ "Motility-low"] = pr
+    recall[ "Motility-low"] = rc
+except:
+    precision[ "Motility-low"] = 0.0
+    recall[ "Motility-low"] = 0.0
+    with open(store_file_path, "r+") as f:
+        current_position = previous_position = f.tell()
+        while f.readline():
+            previous_position = current_position
+            current_position = f.tell()
+        f.truncate(previous_position)
+    f.close()
+
+
+# In[16]:
+
+
+try:
+    difficulty_level = "Motility-wiggle"
+    coco_file=f"coco_instances_results_{difficulty_level}.json"
+    detectron_train_file_name = f"boardetect_test_coco_format_{difficulty_level}.json"
+    train_file_name=f"test_{difficulty_level}.json"
+    store_file_path = store_file_path_main + f"_{difficulty_level}.txt" if store_file_path_main else None
+    if store_file_path:
+            result_file = open(store_file_path,'a+')
+            result_file.write(str(video_num))
+            result_file.write(",")
+            result_file.close()
+    #bacteria_tp_path = ./data_video22_feature_optical_flow_median_back_2pyr_18win_background_img/test/1/predicted/"
+    # visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+    #                   train_file_name=train_file_name,detectron_train_file_name=detectron_train_file_name, 
+    #                   coco_file=coco_file)
+
+    pr, rc = get_statistics(data_path, images_path, coco_eval_path,test_json_path, 
+                   train_file_name=train_file_name,detectron_train_file_name = detectron_train_file_name, 
+                   iou_thresh=0.1, coco_file=coco_file, print_details=False, store_results_path=store_file_path)
+    print("Precision: ", pr)
+    print("Recall: ", rc)
+    precision["Motility-wiggle"] = pr
+    recall["Motility-wiggle"] = rc
+except:
+    precision["Motility-wiggle"] = 0.0
+    recall["Motility-wiggle"] = 0.0
+    with open(store_file_path, "r+") as f:
+        current_position = previous_position = f.tell()
+        while f.readline():
+            previous_position = current_position
+            current_position = f.tell()
+        f.truncate(previous_position)
+    f.close()
+
+
+# In[17]:
+
+
+try:
+    difficulty_level = "Motility-mid"
+    coco_file=f"coco_instances_results_{difficulty_level}.json"
+    detectron_train_file_name = f"boardetect_test_coco_format_{difficulty_level}.json"
+    train_file_name=f"test_{difficulty_level}.json"
+    store_file_path = store_file_path_main + f"_{difficulty_level}.txt" if store_file_path_main else None
+    if store_file_path:
+        result_file = open(store_file_path,'a+')
+        result_file.write(str(video_num))
+        result_file.write(",")
+        result_file.close()
+    #bacteria_tp_path = ./data_video22_feature_optical_flow_median_back_2pyr_18win_background_img/test/1/predicted/"
+    # visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+    #                   train_file_name=train_file_name,detectron_train_file_name=detectron_train_file_name, 
+    #                   coco_file=coco_file)
+
+    pr, rc = get_statistics(data_path, images_path, coco_eval_path,test_json_path, 
+                   train_file_name=train_file_name,detectron_train_file_name = detectron_train_file_name, 
+                   iou_thresh=0.1, coco_file=coco_file, print_details=False, store_results_path=store_file_path)
+    print("Precision: ", pr)
+    print("Recall: ", rc)
+    precision["Motility-mid"] = pr
+    recall["Motility-mid"] = rc
+except:
+    precision["Motility-mid"] = 0.0
+    recall["Motility-mid"] = 0.0
+    with open(store_file_path, "r+") as f:
+        current_position = previous_position = f.tell()
+        while f.readline():
+            previous_position = current_position
+            current_position = f.tell()
+        f.truncate(previous_position)
+    f.close()
+
+
+try:
+    difficulty_level = "Motility-high"
+    coco_file=f"coco_instances_results_{difficulty_level}.json"
+    detectron_train_file_name = f"boardetect_test_coco_format_{difficulty_level}.json"
+    train_file_name=f"test_{difficulty_level}.json"
+    store_file_path = store_file_path_main + f"_{difficulty_level}.txt" if store_file_path_main else None
+    if store_file_path:
+        result_file = open(store_file_path,'a+')
+        result_file.write(str(video_num))
+        result_file.write(",")
+        result_file.close()
+    #bacteria_tp_path = ./data_video22_feature_optical_flow_median_back_2pyr_18win_background_img/test/1/predicted/"
+    # visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+    #                   train_file_name=train_file_name,detectron_train_file_name=detectron_train_file_name, 
+    #                   coco_file=coco_file)
+
+    pr, rc = get_statistics(data_path, images_path, coco_eval_path,test_json_path, 
+                   train_file_name=train_file_name,detectron_train_file_name = detectron_train_file_name, 
+                   iou_thresh=0.1, coco_file=coco_file, print_details=False, store_results_path=store_file_path)
+    print("Precision: ", pr)
+    print("Recall: ", rc)
+    precision["Motility-high"] = pr
+    recall["Motility-high"] = rc
+except:
+    precision["Motility-high"] = 0.0
+    recall["Motility-high"] = 0.0
+    with open(store_file_path, "r+") as f:
+        current_position = previous_position = f.tell()
+        while f.readline():
+            previous_position = current_position
+            current_position = f.tell()
+        f.truncate(previous_position)
+    f.close()
+
+
+
+    
+# ### Step 2: Detection Combination Model
+
+# In[18]:
+
+coco_file = "coco_instances_results_combined.json"
+
+
+train_file_name=f"test_{args.all_labels_file}.json"
+detectron_train_file_name = f"boardetect_test_coco_format_{args.all_labels_file}.json"
+store_file_path = store_file_path_main + f"_combined.txt" if store_file_path_main else None
+if store_file_path:
+        result_file = open(store_file_path,'a+')
+        result_file.write(str(video_num))
+        result_file.write(",")
+        result_file.close()
+# visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+#                   train_file_name=train_file_name,
+#                   detectron_train_file_name = detectron_train_file_name, coco_file=coco_file)
+pr, rc = get_statistics(data_path, images_path, coco_eval_path, test_json_path, 
+              train_file_name=train_file_name, detectron_train_file_name = detectron_train_file_name,
+              iou_thresh=0.1,coco_file=coco_file,print_details=False, store_results_path=store_file_path)
+print("Precision: ", pr)
+print("Recall: ", rc)
+precision["Combination Model Detection"] = pr
+recall["Combination Model Detection"] = rc
+
+
+
+
+# ### Step 3: Filter on Predicted Bacteria Bounding Box Size
+
+# In[19]:
+
+
+coco_file = "coco_instances_results_combined_filter_box_size.json"
+train_file_name="test_All.json"
+detectron_train_file_name = "boardetect_test_coco_format_All.json"
+
+train_file_name=f"test_{args.all_labels_file}.json"
+detectron_train_file_name = f"boardetect_test_coco_format_{args.all_labels_file}.json"
+store_file_path = store_file_path_main + f"_filter_bbox.txt" if store_file_path_main else None
+if store_file_path:
+        result_file = open(store_file_path,'a+')
+        result_file.write(str(video_num))
+        result_file.write(",")
+        result_file.close()
+# visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+#                   train_file_name=train_file_name,
+#                   detectron_train_file_name = detectron_train_file_name, coco_file=coco_file)
+pr, rc = get_statistics(data_path, images_path, coco_eval_path, test_json_path, 
+              train_file_name=train_file_name, detectron_train_file_name = detectron_train_file_name,
+              iou_thresh=0.1,coco_file=coco_file,print_details=False, store_results_path=store_file_path)
+print("Precision: ", pr)
+print("Recall: ", rc)
+precision["Filter Box Size"] = pr
+recall["Filter Box Size"] = rc
+
+
+# ### Step 4: Filter on Predicted Bacteria Confidence Score
+
+# In[20]:
+
+
+coco_file = "coco_instances_results_combined_filter_conf_score.json"
+train_file_name="test_All.json"
+detectron_train_file_name = "boardetect_test_coco_format_All.json"
+
+train_file_name=f"test_{args.all_labels_file}.json"
+detectron_train_file_name = f"boardetect_test_coco_format_{args.all_labels_file}.json"
+store_file_path = store_file_path_main + f"_filter_conf_score.txt" if store_file_path_main else None
+if store_file_path:
+        result_file = open(store_file_path,'a+')
+        result_file.write(str(video_num))
+        result_file.write(",")
+        result_file.close()
+# visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+#                   train_file_name=train_file_name,
+#                   detectron_train_file_name = detectron_train_file_name, coco_file=coco_file)
+pr, rc = get_statistics(data_path, images_path, coco_eval_path, test_json_path, 
+              train_file_name=train_file_name, detectron_train_file_name = detectron_train_file_name,
+              iou_thresh=0.1,coco_file=coco_file,print_details=False, store_results_path=store_file_path)
+print("Precision: ", pr)
+print("Recall: ", rc)
+precision["Filter Conf Score"] = pr
+recall["Filter Conf Score"] = rc
+
+
+# ### Step 4: Filter using NMS
+
+# In[21]:
+
+
+coco_file = "coco_instances_results_combined_filter_nms.json"
+train_file_name="test_All.json"
+detectron_train_file_name = "boardetect_test_coco_format_All.json"
+
+train_file_name=f"test_{args.all_labels_file}.json"
+detectron_train_file_name = f"boardetect_test_coco_format_{args.all_labels_file}.json"
+
+store_file_path = store_file_path_main + f"_filter_nms.txt" if store_file_path_main else None
+if store_file_path:
+        result_file = open(store_file_path,'a+')
+        result_file.write(str(video_num))
+        result_file.write(",")
+        result_file.close()
+# visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+#                   train_file_name=train_file_name,
+#                   detectron_train_file_name = detectron_train_file_name, coco_file=coco_file)
+pr, rc = get_statistics(data_path, images_path, coco_eval_path, test_json_path, 
+              train_file_name=train_file_name, detectron_train_file_name = detectron_train_file_name,
+              iou_thresh=0.1,coco_file=coco_file,print_details=False, store_results_path=store_file_path)
+print("Precision: ", pr)
+print("Recall: ", rc)
+precision["Filter NMS"] = pr
+recall["Filter NMS"] = rc
+
+
+
+# ## Step 5: Tracking
+
+# In[22]:
+
+
+coco_file = f"./video{video_num}_tracking_predictions.json"
+train_file_name="test_All.json"
+detectron_train_file_name = "boardetect_test_coco_format_All.json"
+
+train_file_name=f"test_{args.all_labels_file}.json"
+detectron_train_file_name = f"boardetect_test_coco_format_{args.all_labels_file}.json"
+store_file_path = store_file_path_main + f"_tracking.txt" if store_file_path_main else None
+if store_file_path:
+        result_file = open(store_file_path,'a+')
+        result_file.write(str(video_num))
+        result_file.write(",")
+        result_file.close()
+# visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+#                   train_file_name=train_file_name,
+#                   detectron_train_file_name = detectron_train_file_name, coco_file=coco_file)
+pr, rc = get_statistics(data_path, images_path, coco_eval_path, test_json_path, 
+              train_file_name=train_file_name, detectron_train_file_name = detectron_train_file_name,
+              iou_thresh=0.1,coco_file=coco_file,print_details=False, store_results_path=store_file_path)
+print("Precision: ", pr)
+print("Recall: ", rc)
+precision["Tracking"] = pr
+recall["Tracking"] = rc
+
+
+
+
+# ### Step 6: Filter on Track length
+
+# In[23]:
+
+coco_file = "coco_instances_results_final.json"
+train_file_name="test_All.json"
+detectron_train_file_name = "boardetect_test_coco_format_All.json"
+
+train_file_name=f"test_{args.all_labels_file}.json"
+detectron_train_file_name = f"boardetect_test_coco_format_{args.all_labels_file}.json"
+
+
+
+store_file_path = store_file_path_main + f"_filter_track_length.txt" if store_file_path_main else None
+if store_file_path:
+        result_file = open(store_file_path,'a+')
+        result_file.write(str(video_num))
+        result_file.write(",")
+        result_file.close()
+# visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+#                   train_file_name=train_file_name,
+#                   detectron_train_file_name = detectron_train_file_name, coco_file=coco_file)
+pr, rc = get_statistics(data_path, images_path, coco_eval_path, test_json_path, 
+              train_file_name=train_file_name, detectron_train_file_name = detectron_train_file_name,
+              iou_thresh=0.1,coco_file=coco_file,print_details=False, store_results_path=store_file_path)
+print("Precision: ", pr)
+print("Recall: ", rc)
+precision["Filter Track Length"] = pr
+recall["Filter Track Length"] = rc
+
+
+
+# #motility analysis
+
+# coco_file = "coco_pred_motile.json"
+# train_file_name="test_All.json"
+# detectron_train_file_name = "coco_gt_motile.json"
+# store_file_path = store_file_path_main + f"_motile.txt" if store_file_path_main else None
+# if store_file_path:
+#         result_file = open(store_file_path,'a+')
+#         result_file.write(str(video_num))
+#         result_file.write(",")
+#         result_file.close()
+# # visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+# #                   train_file_name=train_file_name,
+# #                   detectron_train_file_name = detectron_train_file_name, coco_file=coco_file)
+# pr, rc = get_statistics(data_path, images_path, coco_eval_path, test_json_path, 
+#               train_file_name=train_file_name, detectron_train_file_name = detectron_train_file_name,
+#               iou_thresh=0.1,coco_file=coco_file,print_details=False, store_results_path=store_file_path)
+# print("Precision: ", pr)
+# print("Recall: ", rc)
+
+# coco_file = "coco_pred_non_motile.json"
+# train_file_name="test_All.json"
+# detectron_train_file_name = "coco_gt_non_motile.json"
+# store_file_path = store_file_path_main + f"_non_motile.txt" if store_file_path_main else None
+# if store_file_path:
+#         result_file = open(store_file_path,'a+')
+#         result_file.write(str(video_num))
+#         result_file.write(",")
+#         result_file.close()
+# # visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+# #                   train_file_name=train_file_name,
+# #                   detectron_train_file_name = detectron_train_file_name, coco_file=coco_file)
+# pr, rc = get_statistics(data_path, images_path, coco_eval_path, test_json_path, 
+#               train_file_name=train_file_name, detectron_train_file_name = detectron_train_file_name,
+#               iou_thresh=0.1,coco_file=coco_file,print_details=False, store_results_path=store_file_path)
+# print("Precision: ", pr)
+# print("Recall: ", rc)
+
+# In[24]:
+
+
+precision
+
+
+# In[25]:
+
+
+recall

MEMTrack/src/evaluation_step_wise_trackmate.py

@@ -0,0 +1,401 @@
+import random
+import os
+import cv2
+import json
+import matplotlib
+import argparse
+from matplotlib import image, patches
+import matplotlib.pyplot as pyplot
+import seaborn as sns
+import numpy as np
+import seaborn as sns
+import pandas as pd
+import matplotlib.pyplot as plt
+
+ap = argparse.ArgumentParser(description='Inference')
+ap.add_argument('--video_map_path', type=str, metavar='PATH')
+ap.add_argument('--data_path',  type=str, metavar='PATH')
+ap.add_argument('--video', type=str, metavar='PATH')
+ap.add_argument('--filename',type=str, metavar='PATH')
+args = ap.parse_args()
+
+def bb_intersection_over_union(boxA, boxB):
+    # determine the (x, y)-coordinates of the intersection rectangle
+    xA = max(boxA[0], boxB[0])
+    yA = max(boxA[1], boxB[1])
+    xB = min(boxA[2], boxB[2])
+    yB = min(boxA[3], boxB[3])
+
+     # compute the area of intersection rectangle
+    interArea = abs(max((xB - xA, 0)) * max((yB - yA), 0))
+    if interArea == 0:
+        return 0
+
+    # compute the area of both the prediction and ground-truth
+    # rectangles
+    boxAArea = abs((boxA[2] - boxA[0]) * (boxA[3] - boxA[1]))
+    boxBArea = abs((boxB[2] - boxB[0]) * (boxB[3] - boxB[1]))
+
+    # compute the intersection over union by taking the intersection
+    # area and dividing it by the sum of prediction + ground-truth
+    # areas - the interesection area
+    iou_box = interArea / float(boxAArea + boxBArea - interArea)
+
+    # return the intersection over union value
+    return iou_box
+
+def get_iou_mat_image(detectron_train_file_bbox_list, coco_results_file_bbox_list ):
+    # iou_mat = []
+    # for gt_box in detectron_train_file_bbox_list:
+    #     iou_row = []
+    #     for pred_box in coco_results_file_bbox_list:
+    #         gt_box_xyxy = [gt_box[0], gt_box[1], gt_box[0] + gt_box[2], gt_box[1] + gt_box[3]]
+    #         pred_box_xyxy = [pred_box[0], pred_box[1], pred_box[0] + pred_box[2], pred_box[1] + pred_box[3]]
+    #         # print(pred_box_xyxy)
+    #         iou_boxes = bb_intersection_over_union(gt_box_xyxy, pred_box_xyxy)
+    #         iou_row.append(iou_boxes)
+    #     iou_mat.append(iou_row)
+    iou_mat = np.zeros((len(detectron_train_file_bbox_list), len(coco_results_file_bbox_list)))
+    for i, gt_box in enumerate(detectron_train_file_bbox_list):
+        for j, pred_box in enumerate(coco_results_file_bbox_list):
+            gt_box_xyxy = [gt_box[0], gt_box[1], gt_box[0] + gt_box[2], gt_box[1] + gt_box[3]]
+            pred_box_xyxy = [pred_box[0], pred_box[1], pred_box[0] + pred_box[2], pred_box[1] + pred_box[3]]
+            iou_mat[i, j] = bb_intersection_over_union(gt_box_xyxy, pred_box_xyxy)
+    return iou_mat
+
+def get_matching_boxes_iou(detectron_train_file_bbox_list, iou_mat, iou_thresh = 0.75):
+    matching_iou_boxes = []
+    iou_mat = iou_mat.tolist()
+    for i in range(0, len(detectron_train_file_bbox_list)):
+        if iou_mat[i]:
+            iou_row_max = max(iou_mat[i])
+            iou_row_max_pred_id = iou_mat[i].index(iou_row_max)
+           # print("iou_mat:", iou_mat)           
+            # print((iou_row_max))
+            # print(iou_row_max_pred_id)
+            #print(iou_row_max)
+            if iou_row_max>iou_thresh:
+                matching_iou_boxes.append([i, iou_row_max_pred_id, iou_row_max])
+    # print(matching_iou_boxes)  
+    #print("Number of matching IOU Ground truth and Predicted boxes: " , len(matching_iou_boxes))
+    return matching_iou_boxes
+
+def compute_mse_matching_boxes(matching_iou_boxes, detectron_train_file_bbox_list, coco_results_file_bbox_list):
+    # Compute MSE of fiber intersection based on matching iou-boxes
+    mse = []
+    iou_sum = 0
+    for matching_bbox in matching_iou_boxes:
+           # ground truth
+            x = detectron_train_file_bbox_list[matching_bbox[0]][0]
+            y = detectron_train_file_bbox_list[matching_bbox[0]][1]
+            w = detectron_train_file_bbox_list[matching_bbox[0]][2]
+            h = detectron_train_file_bbox_list[matching_bbox[0]][3]
+            center_x_gt = x + (w/2)
+            center_y_gt = y + (h/2)
+
+            # predicted
+            x = coco_results_file_bbox_list[matching_bbox[1]][0]
+            y = coco_results_file_bbox_list[matching_bbox[1]][1]
+            w = coco_results_file_bbox_list[matching_bbox[1]][2]
+            h = coco_results_file_bbox_list[matching_bbox[1]][3]
+            center_x_pred = x + (w / 2)
+            center_y_pred = y + (h / 2)  
+
+            mse.append(((center_x_gt - center_x_pred)**2 + (center_y_gt - center_y_pred)** 2)**0.5)
+            iou_sum += matching_bbox[2]
+    
+   # np.array(matching_iou_boxes)[2].mean()
+    if len(matching_iou_boxes)>0:
+        return np.array(mse).mean(), iou_sum/len(matching_iou_boxes), mse
+    else:
+        return np.array(mse).mean(), -1, mse
+
+def nms_predicted_bounding_boxes(pred_bbox_list, pred_bbox_scores, iou_thresh_nms=0.95):
+    #NMS of predicted boxes
+    #print(len(pred_bbox_list))
+    final_tracking_bbox_list =[]
+    iou_mat = get_iou_mat_image(pred_bbox_list, pred_bbox_list)
+    
+    matching_pred_boxes_iou = get_matching_boxes_iou(pred_bbox_list, iou_mat, iou_thresh = iou_thresh_nms)
+    while(len(matching_pred_boxes_iou)>0):
+        sorted_bbox_list = sorted(zip(pred_bbox_scores, pred_bbox_list))
+        iou_mat = get_iou_mat_image(pred_bbox_list, [pred_bbox_list[-1]])
+        #print(iou_mat)
+        pred_bbox_list_temp = []
+        pred_bbox_scores_temp = []
+        for index, iou in enumerate(iou_mat):
+            if iou[0]<iou_thresh_nms:
+                pred_bbox_list_temp.append(pred_bbox_list[index])
+                pred_bbox_scores_temp.append(pred_bbox_scores[index])
+        # matching_pred_boxes_iou = get_matching_boxes_iou(pred_bbox_list, iou_mat, iou_thresh = iou_thresh_nms)
+        # print(matching_pred_boxes_iou)
+        final_tracking_bbox_list.append(pred_bbox_list[-1]) #add highest scored bbox to final list
+        # matching_pred_boxes_index=[]
+        # for bbox in matching_pred_boxes_iou:
+        #     matching_pred_boxes_index.append(bbox[0])
+        # print("hello",matching_pred_boxes_index)
+      
+        # for index, bbox in enumerate(pred_bbox_list):
+        #     if index not in matching_pred_boxes_index:
+        #         pred_bbox_list_temp.append(pred_bbox_list[index])
+        #         pred_bbox_scores_temp.append(pred_bbox_scores[index])
+        pred_bbox_list = pred_bbox_list_temp
+        pred_bbox_scores = pred_bbox_scores_temp
+        iou_mat = get_iou_mat_image(pred_bbox_list, pred_bbox_list)
+        matching_pred_boxes_iou = get_matching_boxes_iou(pred_bbox_list, iou_mat, iou_thresh = iou_thresh_nms)
+        
+      
+    #print(final_tracking_bbox_list + (pred_bbox_list))
+    # print(len(final_tracking_bbox_list))
+    # print(len(pred_bbox_list))
+    #print(len(final_tracking_bbox_list + (pred_bbox_list)))
+    return final_tracking_bbox_list + (pred_bbox_list)
+
+    
+def get_statistics(data_path, images_path, coco_eval_path, test_json_path, train_file_name="train.json", detectron_train_file_name = "boardetect_train_coco_format.json",iou_thresh=0.75, bacteria_tp_path="./bacteria-detections/", coco_file = "coco_instances_results.json",print_details=True, store_results_path = None):
+    # data_path = "./data_phase_nodes/"
+    # images_path = data_path + "images_cropped_resized/"
+    # coco_eval_path = "./coco_eval_phase_cropped/"
+    num_images = 0
+    tp_sum = 0
+    fp_sum = 0
+    fn_sum = 0
+    mse_image_counter =0
+    iou_image_counter =0
+    ratio_gt_p = 0
+    ratio_p_gt = 0
+    mse = 0
+    precision = 0
+    recall = 0
+    iou = 0
+    all_img_mse = []
+    for filename in sorted(os.listdir(images_path)):  
+        # print(filename)
+        image_id = filename.split(".tif")[0]
+        #print(image_id)
+        train_file = json.load(open(data_path + train_file_name,'r'))
+        train_file_bbox_list = []
+        #print(train_file[int(image_id)]["file_name"])
+        if train_file[int(image_id)]["file_name"] == filename:
+            for i, annotation in enumerate(train_file[int(image_id)]["annotations"]):
+                train_file_bbox_list.append(train_file[int(image_id)]["annotations"][i]["bbox"])
+            # print("Number of Ground Truth boxes: ",len(train_file_bbox_list)) 
+            # print(train_file_bbox_list[0])
+
+        detectron_train_file = json.load(open(test_json_path + detectron_train_file_name,'r'))
+        detectron_train_file_bbox_list = []
+
+        for i,annotation in enumerate(detectron_train_file["annotations"]):
+            #print(detectron_train_file["annotations"][i]["image_id"])
+            #print(image_id)
+            #if detectron_train_file["images"]["annotations"][i]["image_id"] == int(image_id):
+            if detectron_train_file["annotations"][i]["image_id"] == int(image_id):
+                #print(true)
+                detectron_train_file_bbox_list.append(detectron_train_file["annotations"][i]["bbox"])
+                #annotation_image_id = detectron_train_file["annotations"][i]["image_id"]
+                #print(detectron_train_file["annotations"][i]["bbox"])
+        # print(len(detectron_train_file_bbox_list)) 
+    #     print(detectron_train_file_bbox_list[0])
+
+        coco_results_file = json.load(open(coco_eval_path + coco_file,'r'))
+        coco_results_file_bbox_list = []
+        coco_results_file_bbox_scores = []
+
+        for i,annotation in enumerate(coco_results_file):
+            #print(coco_results_file[i])
+            #print(image_id)
+            if coco_results_file[i]["image_id"] == int(image_id):
+                #print(true)
+                coco_results_file_bbox_list.append(coco_results_file[i]["bbox"])
+                coco_results_file_bbox_scores.append(coco_results_file[i]["score"])
+    # print(coco_results_file[i]["bbox"])
+    # print(len(coco_results_file_bbox_list))
+    # print((coco_results_file_bbox_list))  
+                  #print(detectron_train_file["annotations"][i]["bbox"])
+        # print("Number of Predicted bounding boxes: ", len(coco_results_file_bbox_list)) 
+        # print(coco_results_file_bbox_list[0])
+
+        #coco_results_file_bbox_list = nms_predicted_bounding_boxes(coco_results_file_bbox_list, coco_results_file_bbox_scores ,iou_thresh_nms=0.001)
+        
+        # Stat 1 - Ratio of boxes Predicted to Groud truth boxes
+        num_predicted_boxes = len(coco_results_file_bbox_list) 
+        num_ground_truth_boxes = len(detectron_train_file_bbox_list)
+
+        if num_ground_truth_boxes>0:
+            
+            #Ratio of GT to P per image
+            ratio_gt_p = ratio_gt_p + (num_ground_truth_boxes/ max(num_predicted_boxes,1))
+            ratio_p_gt = ratio_p_gt + (num_predicted_boxes / max(num_ground_truth_boxes,1))
+            num_images = num_images + 1
+
+            # Stat 2 - MSE of fiber intersections from the matching boxes
+            iou_mat = get_iou_mat_image(detectron_train_file_bbox_list, coco_results_file_bbox_list)
+            matching_boxes_iou = get_matching_boxes_iou(detectron_train_file_bbox_list, iou_mat, iou_thresh = iou_thresh)
+            mse_image, iou_image, mse_list = compute_mse_matching_boxes(matching_boxes_iou, detectron_train_file_bbox_list, coco_results_file_bbox_list)
+            if mse_image>=0: #if no predicted boxes or no ground truth boxes then iou is nan
+                mse = mse + mse_image
+                mse_image_counter += 1
+            if iou_image>=0:#if no predicted boxes or no ground truth boxes then mse is nan
+                iou += iou_image
+                iou_image_counter += 1
+        
+            true_positive = len(matching_boxes_iou)
+            # print("num image: ", num_images)
+            # print("Num pred: ", num_predicted_boxes)
+            # print("num matched: ", len(matching_boxes_iou))
+            os.makedirs(bacteria_tp_path, exist_ok=True)
+            coord_file = open(bacteria_tp_path + image_id +".txt",'w')
+            coord_file.write(image_id)
+            coord_file.write(" ")
+            coord_file.write(str(true_positive))
+            coord_file.write("\n")
+            coord_file.close()
+    
+            # OLD METHOD OF PRECISION RECALL CALCULATION
+            # assumes 0 for no tp
+            # precision is average of averages
+#             if true_positive>0:
+#                 false_positive = num_predicted_boxes - true_positive
+#                 false_negative = num_ground_truth_boxes - true_positive
+
+#                 precision += true_positive/(true_positive + false_positive)
+
+#                 recall += true_positive/(true_positive + false_negative)
+            all_img_mse.extend(mse_list)
+
+            # UPDATED METHOD
+            # precision is total tp / total tp+fp in all images not average of preciosn for all images
+        
+            false_positive = num_predicted_boxes - true_positive
+            false_negative = num_ground_truth_boxes - true_positive
+            
+            tp_sum = tp_sum + true_positive
+            fp_sum = fp_sum + false_positive
+            fn_sum = fn_sum + false_negative
+            
+
+           
+    
+    # Stat 1 - Ratio of boxes Predicted to Groud truth boxes
+    avg_ratio_gt_p = ratio_gt_p/num_images 
+    avg_ratio_p_gt = ratio_p_gt/num_images 
+    
+
+    # Stat 2 - MSE of fiber intersections from the matching boxes
+    try:
+        avg_mse = mse / mse_image_counter
+    except:
+        avg_mse = 0
+    
+    
+    try:
+        avg_iou = iou / iou_image_counter
+    except:
+        avg_iou = 0
+    
+    try:
+        
+        avg_prec = tp_sum / (tp_sum + fp_sum)
+        avg_recall = tp_sum / (tp_sum + fn_sum)
+    except:
+        avg_prec = 0
+        avg_recall = 0
+    
+    if store_results_path:
+        result_file = open(store_results_path,'a+')
+        result_file.write(str(tp_sum))
+        result_file.write(",")
+        result_file.write(str(fp_sum))
+        result_file.write(",")
+        result_file.write(str(fn_sum))
+        result_file.write("\n")
+        result_file.close()
+        
+    if print_details:
+        # ap = precision/num_images
+        # ar = recall/num_images
+        print("Average Ground Truth to Predicted Ratio: ", avg_ratio_gt_p)
+        print("Average Predicted to Ground Truth Ratio: ", avg_ratio_p_gt)
+        print("Average Mean Squared Error of fiber intersections: ", avg_mse)
+        print("Average IoU of TP boxes: ", avg_iou)  
+        print("TP sum: ", tp_sum)
+        print("FP sum: ", fp_sum)
+        print("FN sum: ", fn_sum)
+        
+        print("Average Precision: ", avg_prec)    
+        print("Average Recall: ", avg_recall)  
+
+        pyplot.figure(figsize=(10, 3))
+        sns.distplot(all_img_mse)
+        pyplot.xlabel("Nodal Errors (in px sq)")
+        # plt.xlabel("")
+        pyplot.grid("on", alpha = 0.3)
+        pyplot.show()
+    return avg_prec, avg_recall
+
+video_map_path = args.video_map_path #"/alldata/medha/CleanCodeData/Data/videomap.txt"
+video_map = open(video_map_path,'r',)
+header = [x.strip() for x in (video_map.readline().strip()).split(",")]
+video_num_id = header.index("video_num")
+strain_id = header.index("strain")
+strain_map = {}
+for line in video_map.readlines():
+    line_details = [x.strip() for x in line.split(",")]
+    video_num = line_details[video_num_id]
+    strain = line_details[strain_id]
+    strain_map[video_num] = strain
+strain_map[""] = "all"
+
+#53 54 58 59 60 64 65 69 70 71 75 76 80 81 82 collagen
+#85 24 30 35 83 84 agar0.2
+#updated gar 19 22 29 30 83 84 
+#53 60 64 69 75 82 
+video_num = args.video
+src_path  = args.data_path #"/alldata/medha/CleanCodeData/DataFeatures/exp_collagen_train_veryhard/"
+print(src_path)
+video_path = f"data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/"
+data_path = os.path.join(src_path,video_path, "test/")
+print(data_path)
+images_path = data_path + "images/"
+coco_eval_path = data_path
+test_json_path = data_path
+precision = {}
+recall = {}
+store_file_path_main = src_path + "test_set_results"
+
+
+
+
+coco_file=args.filename #f"mosaic_coco.json"
+train_file_name="test_All.json"
+detectron_train_file_name = "boardetect_test_coco_format_All.json"
+store_file_path = store_file_path_main + f"_trackmate.txt" if store_file_path_main else None
+if store_file_path:
+        result_file = open(store_file_path,'a+')
+        result_file.write(str(video_num))
+        result_file.write(",")
+        result_file.close()
+# visualize_gt_pred(data_path, images_path, coco_eval_path, test_json_path, 
+#                   train_file_name=train_file_name,
+#                   detectron_train_file_name = detectron_train_file_name, coco_file=coco_file)
+pr, rc = get_statistics(data_path, images_path, coco_eval_path, test_json_path, 
+              train_file_name=train_file_name, detectron_train_file_name = detectron_train_file_name,
+              iou_thresh=0.1,coco_file=coco_file,print_details=False, store_results_path=store_file_path)
+print("Precision: ", pr)
+print("Recall: ", rc)
+precision["Filter Track Length"] = pr
+recall["Filter Track Length"] = rc
+
+
+
+
+# In[24]:
+
+
+precision
+
+
+# In[25]:
+
+
+recall

MEMTrack/src/inferenceBacteriaRetinanet_Motility.py

@@ -0,0 +1,338 @@
+# Some basic setup:
+# Setup detectron2 logger
+import detectron2
+from detectron2.utils.logger import setup_logger
+setup_logger()
+
+# import some common libraries
+import numpy as np
+import os, json, cv2, random
+import cv2
+import pandas as pd
+import shutil
+import numpy as np
+
+# import some common detectron2 utilities
+from detectron2 import model_zoo
+from detectron2.engine import DefaultPredictor
+from detectron2.config import get_cfg
+from detectron2.utils.visualizer import Visualizer
+from detectron2.data import MetadataCatalog, DatasetCatalog
+import json
+from detectron2.structures import BoxMode
+from detectron2.data import DatasetCatalog, MetadataCatalog
+import argparse
+from natsort import natsorted
+import PIL
+from PIL import Image
+import matplotlib.pyplot as plt
+from detectron2.engine import DefaultTrainer
+from detectron2.evaluation import COCOEvaluator
+from detectron2.config import get_cfg
+import os
+
+ap = argparse.ArgumentParser(description='Inference')
+ap.add_argument('--video', default='29', type=str, metavar='PATH')
+ap.add_argument('--source_path', default="./CombinationModel/data_feature_optical_flow_median_back_2pyr_18win_background_img/", type=str, metavar='PATH')
+ap.add_argument('--output_dir', default="CombinationModel/easy-optical_flow_median_back_2pyr_18win_00125_34k", type=str, metavar='CELL PATH')
+ap.add_argument('--test_dir', default="CombinationModel/easy-optical_flow_median_back_2pyr_18win_00125_34k", type=str, metavar='CELL PATH')
+ap.add_argument('--annotations_train', default="Easy", type=str, metavar='TRAIN')
+ap.add_argument('--annotations_test', default="Easy", type=str, metavar='TEST')
+ap.add_argument('--epochs', default="90", type=str, metavar='TEST')
+ap.add_argument('--lr', default="0.00125", type=str, metavar='TEST')
+ap.add_argument('--custom_test_dir', type=str, metavar='CELL PATH')
+ap.add_argument('--feature', default="images_feature/", type=str, metavar='FEATURE')
+
+args = ap.parse_args()
+print("lr: ", float(args.lr))
+coco_format_train = pd.DataFrame(columns=["file_name","height","width","annotations"])
+coco_format_train["annotations"] = coco_format_train["annotations"].astype('object')
+
+video_num = args.video
+
+#source = "./try_median_issue_data_feature_optical_flow_median_back_2pyr_18win_background_img/"
+source = args.source_path
+dest_train = source + "train/"
+if args.custom_test_dir:
+    dest_test = args.custom_test_dir
+else:
+    dest_test = args.test_dir + f"/data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/test/"
+#dest_test = f"./video{video_num}_feature_optical_flow_median_back_2pyr_18win/test/"
+
+val = True 
+
+
+output_dir = ("try_median_issue_optical_flow_median_back_2pyr_18win_00125_34k_background_img")
+#output_dir = ("./optical_flow_median_back/")
+output_dir = args.output_dir
+train_images_path = "/train" + args.feature #"/images_feature/"
+val_image_path = "/val" + args.feature #"/images_feature/"
+
+from contextlib import suppress
+
+with suppress(OSError):
+    os.remove(output_dir +'/boardetect_test_coco_format.json.lock')
+    os.remove(output_dir +'/boardetect_test_coco_format.json')
+    os.remove(source + 'test.json')
+
+#annotations_source = source + "annotation_easy_hard/"
+images_source = source + args.feature #"/images_feature/"
+
+if args.annotations_train == "Easy":
+    annotations_train = dest_train + "annotation_easy/" 
+elif args.annotations_train == "Hard":
+    annotations_train = dest_train + "annotation_hard/" 
+elif args.annotations_train == "VeryHard":
+    annotations_train = dest_train + "annotation_veryhard/" 
+elif args.annotations_train == "Easy+Hard":
+    annotations_train = dest_test + "annotation_easy_hard/" 
+elif args.annotations_train == "All":
+    annotations_train = dest_test + "annotation_easy_hard_veryhard/" 
+elif args.annotations_train == "Motility-low":
+    annotations_train = dest_train + "annotation_motility_low/" 
+elif args.annotations_train == "Motility-high":
+    annotations_train = dest_train + "annotation_motility_high/" 
+elif args.annotations_train == "Motility-wiggle":
+    annotations_train = dest_train + "annotation_motility_wiggle/" 
+elif args.annotations_train == "Motility-mid":
+    annotations_train = dest_train + "annotation_motility_mid/"
+elif args.annotations_train == "Motility-motile":
+    annotations_train = dest_train + "annotation_motility_wiggle_mid_high/"
+elif args.annotations_train == "Sticking-stick":
+    annotations_train = dest_train + "annotation_sticking_stick/" 
+elif args.annotations_train == "Sticking-motile":
+    annotations_train = dest_train + "annotation_sticking_motile/" 
+elif args.annotations_train == "Sticking-non_motile":
+    annotations_train = dest_train + "annotation_sticking_non_motile/" 
+elif args.annotations_train == "Motility-low-wiggle":
+    annotations_train = dest_train + "annotation_motility_low_wiggle/" 
+elif args.annotations_train == "Motility-mid-high":
+    annotations_train = dest_train + "annotation_motility_mid_high/"
+
+if args.annotations_test == "Easy":
+    annotations_test = dest_test + "annotation_easy/" 
+elif args.annotations_test == "Hard":
+    annotations_test = dest_test + "annotation_hard/" 
+elif args.annotations_test == "VeryHard":
+    annotations_test = dest_test + "annotation_veryhard/" 
+elif args.annotations_test == "Easy+Hard":
+    annotations_test = dest_test + "annotation_easy_hard/" 
+elif args.annotations_test == "All":
+    annotations_test = dest_test + "annotation_easy_hard_veryhard/" 
+elif args.annotations_test == "Motility-low":
+    annotations_test = dest_test + "annotation_motility_low/" 
+elif args.annotations_test == "Motility-high":
+    annotations_test = dest_test + "annotation_motility_high/" 
+elif args.annotations_test == "Motility-wiggle":
+    annotations_test = dest_test + "annotation_motility_wiggle/" 
+elif args.annotations_test == "Motility-mid":
+    annotations_test = dest_test + "annotation_motility_mid/" 
+elif args.annotations_test == "Motility-motile":
+    annotations_test = dest_test + "annotation_motility_wiggle_mid_high/"
+elif args.annotations_test == "Sticking-stick":
+    annotations_test = dest_test + "annotation_sticking_stick/"
+elif args.annotations_test == "Sticking-motile":
+    annotations_test = dest_test + "annotation_sticking_motile/" 
+elif args.annotations_test == "Sticking-non_motile":
+    annotations_test = dest_test + "annotation_sticking_non_motile/" 
+elif args.annotations_test == "Motility-low-wiggle":
+    annotations_test = dest_test + "annotation_motility_low_wiggle/" 
+elif args.annotations_test == "Motility-mid-high":
+    annotations_test = dest_test + "annotation_motility_mid_high/" 
+
+
+
+#To test a particular bacteria
+#annotations_test = dest_test + "bacteria/4/xy_coord/"
+images_train = dest_train + args.feature #"/images_feature/"
+images_test = dest_test + args.feature #"/images_feature/"
+test_image_path = images_test
+
+factor_w = 1#1024/1388
+factor_h = 1#1024/1040
+
+#function to get background frame
+#function to get prev frame
+#function to create new image given image num
+
+for txt_file in natsorted(os.listdir(annotations_train)):
+        width = 31
+        text_file = open(annotations_train + txt_file, 'r')
+        xy_coords = text_file.readlines()
+        boxes = []
+        res=pd.DataFrame(columns=["file_name","height","width","annotations"])
+        image = PIL.Image.open(images_train + txt_file[:-4] + ".tif").convert('L')
+        image_feature = PIL.Image.open(images_train + txt_file[:-4] + ".tif")
+        image = image_feature
+        #print(image.size)
+        res.at[0,"height"] = image.height
+        res.at[0,"width"] = image.width
+        res.at[0,"file_name"] = txt_file[:-4]+".tif"
+        bbox_mode = 0
+        category_id = 0
+        # image2 = image.resize((1024,1024))
+        # image2.save(images_resized_train + txt_file[:-4] + ".jpg")
+        for xy in xy_coords:
+            box = []
+            x = float(xy.split(" ")[0])
+            y = float(xy.split(" ")[1])
+            x1 = int(x*factor_w - (width // 2))
+            y1 = int(y*factor_h - (width // 2))
+            x2 = int(x*factor_w + (width // 2))
+            y2 = int(y*factor_h + (width // 2))
+            w = h = 31
+            box = [x1, y1, x2, y2]
+            boxes.append(np.array(box))
+            #print(np.array(box))
+
+        res["annotations"]=res["annotations"].astype('object')
+        annotation_df = pd.DataFrame(columns=["bbox","bbox_mode","category_id"])
+        annotation_df["bbox"] = boxes
+        annotation_df["bbox_mode"] = bbox_mode
+        annotation_df["category_id"] = category_id
+        annotations = annotation_df.T.to_dict().values()
+        l = []
+        for j in annotations:
+            l.append(j)
+        res.at[0,"annotations"] = l
+        coco_format_train = coco_format_train.append(res)
+        coco_format_train.reset_index(drop=True,inplace=True)
+        
+coco_format_train.reset_index(inplace=True)
+coco_format_train.rename(columns={"index":"image_id"},inplace=True)
+coco_format_train.to_json(source + "train.json",orient="records")
+
+coco_format_test = pd.DataFrame(columns=["file_name","height","width","annotations"])
+coco_format_test["annotations"] = coco_format_test["annotations"].astype('object')
+
+for txt_file in natsorted(os.listdir(annotations_test)):
+        width = 31
+        text_file = open(annotations_test + txt_file, 'r')
+        xy_coords = text_file.readlines()
+        boxes = []
+        res=pd.DataFrame(columns=["file_name","height","width","annotations"])
+        image = PIL.Image.open(images_test + txt_file[:-4] + ".tif").convert('L')
+        image_feature = PIL.Image.open(images_test + txt_file[:-4] + ".tif")
+        image = image_feature
+        #print(image.size)
+        res.at[0,"height"] = image.height
+        res.at[0,"width"] = image.width
+        res.at[0,"file_name"] = txt_file[:-4]+".tif"
+        bbox_mode = 0
+        category_id = 0
+        # image2 = image.resize((1024,1024))
+        # image2.save(images_resized_train + txt_file[:-4] + ".jpg")
+        for xy in xy_coords:
+            box = []
+            x = float(xy.split(" ")[0])
+            y = float(xy.split(" ")[1])
+            x1 = int(x*factor_w - (width // 2))
+            y1 = int(y*factor_h - (width // 2))
+            x2 = int(x*factor_w + (width // 2))
+            y2 = int(y*factor_h + (width // 2))
+            w = h = 31
+            box = [x1, y1, x2, y2]
+            boxes.append(np.array(box))
+            #print(np.array(box))
+
+        res["annotations"]=res["annotations"].astype('object')
+        annotation_df = pd.DataFrame(columns=["bbox","bbox_mode","category_id"])
+        annotation_df["bbox"] = boxes
+        annotation_df["bbox_mode"] = bbox_mode
+        annotation_df["category_id"] = category_id
+        annotations = annotation_df.T.to_dict().values()
+        l = []
+        for j in annotations:
+            l.append(j)
+        res.at[0,"annotations"] = l
+        coco_format_test = coco_format_test.append(res)
+        coco_format_test.reset_index(drop=True,inplace=True)
+        
+coco_format_test.reset_index(inplace=True)
+coco_format_test.rename(columns={"index":"image_id"},inplace=True)
+coco_format_test.to_json(source + "test.json",orient="records")
+
+
+
+def get_board_dicts(imgdir, mode):
+    if mode is 'train':
+        json_file = imgdir+"/train.json" #Fetch the json file
+    if mode is 'test':
+         json_file = imgdir+"/test.json" #Fetch the json file
+    with open(json_file) as f:
+        dataset_dicts = json.load(f)
+    for i in dataset_dicts:
+        filename = i["file_name"] 
+        if mode is 'train':
+            i["file_name"] = imgdir + train_images_path + filename 
+        if mode is 'test':
+             i["file_name"] =  test_image_path + filename 
+        for j in i["annotations"]:
+            j["bbox_mode"] = BoxMode.XYXY_ABS #Setting the required Box Mode
+            j["category_id"] = int(j["category_id"])
+    return dataset_dicts
+
+#Registering the Dataset
+for d in ["train", "test"]:
+    DatasetCatalog.register("boardetect_" + d, lambda d=d: get_board_dicts(source, d))
+    MetadataCatalog.get("boardetect_" + d).set(thing_classes=["node"])
+board_metadata = MetadataCatalog.get("boardetect_train")
+val_metadata = MetadataCatalog.get("boardetect_test")
+
+train_data = ("boardetect_train",)
+test_data = ("boardetect_test",)
+if val ==True:
+    val_data = ("boardetect_test",)
+else:
+    val_data = ("boardetect_train",)
+
+
+
+class CocoTrainer(DefaultTrainer):
+       
+    @classmethod
+    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
+        if output_folder is None:
+            os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
+            output_folder = cfg.OUTPUT_DIR
+
+
+        return COCOEvaluator(dataset_name, cfg, False, output_folder)
+
+cfg = get_cfg()
+cfg.MODEL.DEVICE = 'cuda:0' 
+cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/retinanet_R_101_FPN_3x.yaml")) #Get the basic model configuration from the model zoo 
+#Passing the Train and Validation sets
+cfg.DATASETS.TRAIN = train_data #("boardetect_train",)
+cfg.DATASETS.TEST = test_data #("boardetect_train",)
+cfg.OUTPUT_DIR = output_dir #("comparison-optical-flow")
+# Number of data loading threads
+cfg.DATALOADER.NUM_WORKERS = 4
+#cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
+#cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") #uncommwnt during inference
+cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") #uncommwnt during inference
+#cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_0029999.pth") #uncommwnt during inference
+# Number of images per batch across all machines.
+cfg.SOLVER.IMS_PER_BATCH = 4
+cfg.SOLVER.BASE_LR = float(args.lr) #0.0125  # pick a good LearningRate
+cfg.SOLVER.MAX_ITER = int(args.epochs) #30000  #No. of iterations
+# cfg.SOLVER.STEPS = (300,600)
+cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 80
+cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
+cfg.TEST.EVAL_PERIOD = 500 # No. of iterations after which the Validation Set is evaluated. 
+cfg.TEST.DETECTIONS_PER_IMAGE = 60
+cfg.SOLVER.CHECKPOINT_PERIOD = 500
+cfg.VIS_PERIOD = 500
+cfg.DATASETS.PRECOMPUTED_PROPOSAL_TOPK_TRAIN = 100
+cfg.DATASETS.PRECOMPUTED_PROPOSAL_TOPK_TEST = 100
+cfg.MODEL.RPN.POST_NMS_TOPK_TRAIN = 80 #80
+cfg.MODEL.RPN.POST_NMS_TOPK_TEST = 80
+os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
+trainer = CocoTrainer(cfg) 
+trainer.resume_or_load(resume=True) #resume=True for inference
+trainer.train()
+print(cfg.MODEL.WEIGHTS)
+
+os.system(f'cp {output_dir}/coco_instances_results.json {dest_test}/coco_instances_results_{args.annotations_test}.json')  
+os.system(f'cp {output_dir}/boardetect_test_coco_format.json {dest_test}/boardetect_test_coco_format_{args.annotations_test}.json')  
+os.system(f'cp {source}/test.json {dest_test}/test_{args.annotations_test}.json')  

MEMTrack/src/inferenceBacteriaRetinanet_Motility_v2.py

@@ -0,0 +1,230 @@
+# Some basic setup:
+# Setup detectron2 logger
+import detectron2
+from detectron2.utils.logger import setup_logger
+setup_logger()
+
+# import some common libraries
+import numpy as np
+import os, json, cv2, random
+import cv2
+import pandas as pd
+import shutil
+import numpy as np
+
+# import some common detectron2 utilities
+from detectron2 import model_zoo
+from detectron2.engine import DefaultPredictor
+from detectron2.config import get_cfg
+from detectron2.utils.visualizer import Visualizer
+from detectron2.data import MetadataCatalog, DatasetCatalog
+import json
+from detectron2.structures import BoxMode
+from detectron2.data import DatasetCatalog, MetadataCatalog
+import argparse
+from natsort import natsorted
+import PIL
+from PIL import Image
+import matplotlib.pyplot as plt
+from detectron2.engine import DefaultTrainer
+from detectron2.evaluation import COCOEvaluator
+from detectron2.config import get_cfg
+from detectron2.evaluation import inference_on_dataset
+from detectron2.data import build_detection_test_loader
+import os
+from contextlib import suppress
+
+
+def register_dataset_(annotations_test, images_test, dest_test):
+    factor_w = 1#1024/1388
+    factor_h = 1#1024/1040
+
+    coco_format_test = pd.DataFrame(columns=["file_name","height","width","annotations"])
+    coco_format_test["annotations"] = coco_format_test["annotations"].astype('object')
+    
+
+    for txt_file in natsorted(os.listdir(annotations_test)):
+            width = 31
+            text_file = open(annotations_test + txt_file, 'r')
+            xy_coords = text_file.readlines()
+            boxes = []
+            res=pd.DataFrame(columns=["file_name","height","width","annotations"])
+            image = PIL.Image.open(images_test + txt_file[:-4] + ".tif").convert('L')
+            image_feature = PIL.Image.open(images_test + txt_file[:-4] + ".tif")
+            image = image_feature
+            #print(image.size)
+            res.at[0,"height"] = image.height
+            res.at[0,"width"] = image.width
+            res.at[0,"file_name"] = txt_file[:-4]+".tif"
+            bbox_mode = 0
+            category_id = 0
+            # image2 = image.resize((1024,1024))
+            # image2.save(images_resized_train + txt_file[:-4] + ".jpg")
+            for xy in xy_coords:
+                box = []
+                x = float(xy.split(" ")[0])
+                y = float(xy.split(" ")[1])
+                x1 = int(x*factor_w - (width // 2))
+                y1 = int(y*factor_h - (width // 2))
+                x2 = int(x*factor_w + (width // 2))
+                y2 = int(y*factor_h + (width // 2))
+                w = h = 31
+                box = [x1, y1, x2, y2]
+                boxes.append(np.array(box))
+                #print(np.array(box))
+
+            res["annotations"]=res["annotations"].astype('object')
+            annotation_df = pd.DataFrame(columns=["bbox","bbox_mode","category_id"])
+            annotation_df["bbox"] = boxes
+            annotation_df["bbox_mode"] = bbox_mode
+            annotation_df["category_id"] = category_id
+            annotations = annotation_df.T.to_dict().values()
+            l = []
+            for j in annotations:
+                l.append(j)
+            res.at[0,"annotations"] = l
+            coco_format_test = coco_format_test.append(res)
+            coco_format_test.reset_index(drop=True,inplace=True)
+
+    coco_format_test.reset_index(inplace=True)
+    coco_format_test.rename(columns={"index":"image_id"},inplace=True)
+    coco_format_test.to_json(dest_test + "test.json",orient="records")
+    
+      
+    def get_board_dicts(imgdir, mode):
+        if mode == 'test':
+             json_file = imgdir+"/test.json" #Fetch the json file
+        with open(json_file) as f:
+            dataset_dicts = json.load(f)
+        for i in dataset_dicts:
+            filename = i["file_name"] 
+            if mode == 'test':
+                 i["file_name"] =  images_test + filename 
+            for j in i["annotations"]:
+                j["bbox_mode"] = BoxMode.XYXY_ABS #Setting the required Box Mode
+                j["category_id"] = int(j["category_id"])
+        return dataset_dicts
+
+    #Registering the Dataset
+    for d in ["test"]:
+        DatasetCatalog.register("boardetect_" + d, lambda d=d: get_board_dicts(dest_test, d))
+        MetadataCatalog.get("boardetect_" + d).set(thing_classes=["node"])
+
+    test_metadata = MetadataCatalog.get("boardetect_test")
+
+    test_data = ("boardetect_test",)
+    return 
+    
+def run_inference(video_num, output_dir, feature="images_feature/", annotations_test="Easy", test_dir=None, custom_test_dir=None, register_dataset=True):
+    annotations_test_args = annotations_test
+    if custom_test_dir:
+        dest_test = custom_test_dir
+    else:
+        dest_test = test_dir + f"/data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/test/"
+   
+    val_image_path = "/val" + feature 
+
+    if os.path.exists(output_dir +'/boardetect_test_coco_format.json.lock'):
+        os.remove(output_dir +'/boardetect_test_coco_format.json.lock')
+    if os.path.exists(output_dir +'/boardetect_test_coco_format.json'):
+        os.remove(output_dir +'/boardetect_test_coco_format.json')
+    if os.path.exists(output_dir +'/coco_instances_results.json'):
+        os.remove(output_dir +'/coco_instances_results.json')
+#     if os.path.exists(dest_test + 'test.json'):
+#         os.remove(dest_test + 'test.json')
+    
+    print("output_dir: ", output_dir)
+    print("dest_test: ", dest_test)
+    
+
+
+    if annotations_test_args == "Easy":
+        annotations_test = dest_test + "annotation_easy/" 
+    elif annotations_test_args == "Hard":
+        annotations_test = dest_test + "annotation_hard/" 
+    elif annotations_test_args == "VeryHard":
+        annotations_test = dest_test + "annotation_veryhard/" 
+    elif annotations_test_args == "Easy+Hard":
+        annotations_test = dest_test + "annotation_easy_hard/" 
+    elif annotations_test_args == "All":
+        annotations_test = dest_test + "annotation_easy_hard_veryhard/" 
+    elif annotations_test_args == "Motility-low":
+        annotations_test = dest_test + "annotation_motility_low/" 
+    elif annotations_test_args == "Motility-high":
+        annotations_test = dest_test + "annotation_motility_high/" 
+    elif annotations_test_args == "Motility-wiggle":
+        annotations_test = dest_test + "annotation_motility_wiggle/" 
+    elif annotations_test_args == "Motility-mid":
+        annotations_test = dest_test + "annotation_motility_mid/" 
+    elif annotations_test_args == "Motility-motile":
+        annotations_test = dest_test + "annotation_motility_wiggle_mid_high/"
+    elif annotations_test_args == "Sticking-stick":
+        annotations_test = dest_test + "annotation_sticking_stick/"
+    elif annotations_test_args == "Sticking-motile":
+        annotations_test = dest_test + "annotation_sticking_motile/" 
+    elif annotations_test_args == "Sticking-non_motile":
+        annotations_test = dest_test + "annotation_sticking_non_motile/" 
+    elif annotations_test_args == "Motility-low-wiggle":
+        annotations_test = dest_test + "annotation_motility_low_wiggle/" 
+    elif annotations_test_args == "Motility-mid-high":
+        annotations_test = dest_test + "annotation_motility_mid_high/" 
+
+
+    #To test a particular bacteria
+    images_test = dest_test + feature #"/images_feature/"
+    test_image_path = images_test
+
+    if register_dataset:
+        register_dataset_(annotations_test, images_test, dest_test)
+    
+    cfg = get_cfg()
+    cfg.MODEL.DEVICE = 'cpu' 
+    cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/retinanet_R_101_FPN_3x.yaml")) #Get the basic model configuration from the model zoo 
+    cfg.DATASETS.TEST = ("boardetect_test",) #("boardetect_train",)
+    cfg.OUTPUT_DIR = output_dir #("comparison-optical-flow")
+    # Number of data loading threads
+    cfg.DATALOADER.NUM_WORKERS = 4
+    cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_0089999.pth") #uncommwnt during inference
+    #cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_0029999.pth") #uncommwnt during inference
+    # Number of images per batch across all machines.
+    cfg.SOLVER.IMS_PER_BATCH = 8
+    # cfg.SOLVER.STEPS = (300,600)
+    cfg.TEST.DETECTIONS_PER_IMAGE = 60
+    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
+    print(cfg.MODEL.WEIGHTS)
+
+    predictor = DefaultPredictor(cfg)
+    evaluator = COCOEvaluator("boardetect_test", cfg, False, output_dir)
+    val_loader = build_detection_test_loader(cfg, "boardetect_test")
+    print(inference_on_dataset(predictor.model, val_loader, evaluator))
+
+    if os.path.exists(os.path.join(output_dir, "coco_instances_results.json")):
+        shutil.copy(os.path.join(output_dir, "coco_instances_results.json"), os.path.join(dest_test, f"coco_instances_results_{annotations_test_args}.json"))
+        
+    if os.path.exists(os.path.join(output_dir, "boardetect_test_coco_format.json")):
+        shutil.copy(os.path.join(output_dir, "boardetect_test_coco_format.json"), os.path.join(dest_test, f"boardetect_test_coco_format_{annotations_test_args}.json"))
+        
+    if os.path.exists(os.path.join(dest_test, "test.json")):
+        shutil.copy(os.path.join(dest_test, "test.json"), os.path.join(dest_test, f"test_{annotations_test_args}.json"))
+    return ("Inference done")
+
+if __name__ == "__main__":
+
+    ap = argparse.ArgumentParser(description='Inference')
+    ap.add_argument('--video', default='1', type=str, metavar='PATH')
+    ap.add_argument('--output_dir', type=str, metavar='CELL PATH')
+    ap.add_argument('--test_dir', type=str, metavar='CELL PATH')
+    ap.add_argument('--annotations_test', default="Easy", type=str, metavar='TEST')
+    ap.add_argument('--custom_test_dir', type=str, metavar='CELL PATH')
+    ap.add_argument('--feature', default="images_feature/", type=str, metavar='FEATURE')
+
+    args = ap.parse_args()
+    video_num = args.video 
+    test_dir = args.test_dir
+    custom_test_dir = args.custom_test_dir
+    output_dir = args.output_dir
+    feature = args.feature
+    args.annotations_test = annotations_test
+    
+    run_inference(video_num, output_dir, feature, annotations_test, test_dir, custom_test_dir)
+

MEMTrack/src/sort.py

@@ -0,0 +1,332 @@
+"""
+    SORT: A Simple, Online and Realtime Tracker
+    Copyright (C) 2016-2020 Alex Bewley alex@bewley.ai
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+"""
+from __future__ import print_function
+
+import os
+import numpy as np
+import matplotlib
+#matplotlib.use('TkAgg')
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from skimage import io
+
+import glob
+import time
+import argparse
+from filterpy.kalman import KalmanFilter
+
+np.random.seed(0)
+
+
+def linear_assignment(cost_matrix):
+  try:
+    import lap
+    _, x, y = lap.lapjv(cost_matrix, extend_cost=True)
+    return np.array([[y[i],i] for i in x if i >= 0]) #
+  except ImportError:
+    from scipy.optimize import linear_sum_assignment
+    x, y = linear_sum_assignment(cost_matrix)
+    return np.array(list(zip(x, y)))
+
+
+def iou_batch(bb_test, bb_gt):
+  """
+  From SORT: Computes IOU between two bboxes in the form [x1,y1,x2,y2]
+  """
+  bb_gt = np.expand_dims(bb_gt, 0)
+  bb_test = np.expand_dims(bb_test, 1)
+  
+  xx1 = np.maximum(bb_test[..., 0], bb_gt[..., 0])
+  yy1 = np.maximum(bb_test[..., 1], bb_gt[..., 1])
+  xx2 = np.minimum(bb_test[..., 2], bb_gt[..., 2])
+  yy2 = np.minimum(bb_test[..., 3], bb_gt[..., 3])
+  w = np.maximum(0., xx2 - xx1)
+  h = np.maximum(0., yy2 - yy1)
+  wh = w * h
+  o = wh / ((bb_test[..., 2] - bb_test[..., 0]) * (bb_test[..., 3] - bb_test[..., 1])                                      
+    + (bb_gt[..., 2] - bb_gt[..., 0]) * (bb_gt[..., 3] - bb_gt[..., 1]) - wh)                                              
+  return(o)  
+
+
+def convert_bbox_to_z(bbox):
+  """
+  Takes a bounding box in the form [x1,y1,x2,y2] and returns z in the form
+    [x,y,s,r] where x,y is the centre of the box and s is the scale/area and r is
+    the aspect ratio
+  """
+  w = bbox[2] - bbox[0]
+  h = bbox[3] - bbox[1]
+  x = bbox[0] + w/2.
+  y = bbox[1] + h/2.
+  s = w * h    #scale is just area
+  r = w / float(h)
+  return np.array([x, y, s, r]).reshape((4, 1))
+
+
+def convert_x_to_bbox(x,score=None):
+  """
+  Takes a bounding box in the centre form [x,y,s,r] and returns it in the form
+    [x1,y1,x2,y2] where x1,y1 is the top left and x2,y2 is the bottom right
+  """
+  w = np.sqrt(x[2] * x[3])
+  h = x[2] / w
+  if(score==None):
+    return np.array([x[0]-w/2.,x[1]-h/2.,x[0]+w/2.,x[1]+h/2.]).reshape((1,4))
+  else:
+    return np.array([x[0]-w/2.,x[1]-h/2.,x[0]+w/2.,x[1]+h/2.,score]).reshape((1,5))
+
+
+class KalmanBoxTracker(object):
+  """
+  This class represents the internal state of individual tracked objects observed as bbox.
+  """
+  count = 0
+  def __init__(self,bbox):
+    """
+    Initialises a tracker using initial bounding box.
+    """
+    #define constant velocity model
+    self.kf = KalmanFilter(dim_x=7, dim_z=4) 
+    self.kf.F = np.array([[1,0,0,0,1,0,0],[0,1,0,0,0,1,0],[0,0,1,0,0,0,1],[0,0,0,1,0,0,0],  [0,0,0,0,1,0,0],[0,0,0,0,0,1,0],[0,0,0,0,0,0,1]])
+    self.kf.H = np.array([[1,0,0,0,0,0,0],[0,1,0,0,0,0,0],[0,0,1,0,0,0,0],[0,0,0,1,0,0,0]])
+
+    self.kf.R[2:,2:] *= 10.
+    self.kf.P[4:,4:] *= 1000. #give high uncertainty to the unobservable initial velocities
+    self.kf.P *= 10.
+    self.kf.Q[-1,-1] *= 0.01
+    self.kf.Q[4:,4:] *= 0.01
+
+    self.kf.x[:4] = convert_bbox_to_z(bbox)
+    self.time_since_update = 0
+    self.id = KalmanBoxTracker.count
+    KalmanBoxTracker.count += 1
+    self.history = []
+    self.hits = 0
+    self.hit_streak = 0
+    self.age = 0
+
+  def update(self,bbox):
+    """
+    Updates the state vector with observed bbox.
+    """
+    self.time_since_update = 0
+    self.history = []
+    self.hits += 1
+    self.hit_streak += 1
+    self.kf.update(convert_bbox_to_z(bbox))
+
+  def predict(self):
+    """
+    Advances the state vector and returns the predicted bounding box estimate.
+    """
+    if((self.kf.x[6]+self.kf.x[2])<=0):
+      self.kf.x[6] *= 0.0
+    self.kf.predict()
+    self.age += 1
+    if(self.time_since_update>0):
+      self.hit_streak = 0
+    self.time_since_update += 1
+    self.history.append(convert_x_to_bbox(self.kf.x))
+    return self.history[-1]
+
+  def get_state(self):
+    """
+    Returns the current bounding box estimate.
+    """
+    return convert_x_to_bbox(self.kf.x)
+
+
+def associate_detections_to_trackers(detections,trackers,iou_threshold = 0.3):
+  """
+  Assigns detections to tracked object (both represented as bounding boxes)
+
+  Returns 3 lists of matches, unmatched_detections and unmatched_trackers
+  """
+  if(len(trackers)==0):
+    return np.empty((0,2),dtype=int), np.arange(len(detections)), np.empty((0,5),dtype=int)
+
+  iou_matrix = iou_batch(detections, trackers)
+
+  if min(iou_matrix.shape) > 0:
+    a = (iou_matrix > iou_threshold).astype(np.int32)
+    if a.sum(1).max() == 1 and a.sum(0).max() == 1:
+        matched_indices = np.stack(np.where(a), axis=1)
+    else:
+      matched_indices = linear_assignment(-iou_matrix)
+  else:
+    matched_indices = np.empty(shape=(0,2))
+
+  unmatched_detections = []
+  for d, det in enumerate(detections):
+    if(d not in matched_indices[:,0]):
+      unmatched_detections.append(d)
+  unmatched_trackers = []
+  for t, trk in enumerate(trackers):
+    if(t not in matched_indices[:,1]):
+      unmatched_trackers.append(t)
+
+  #filter out matched with low IOU
+  matches = []
+  for m in matched_indices:
+    if(iou_matrix[m[0], m[1]]<iou_threshold):
+      unmatched_detections.append(m[0])
+      unmatched_trackers.append(m[1])
+    else:
+      matches.append(m.reshape(1,2))
+  if(len(matches)==0):
+    matches = np.empty((0,2),dtype=int)
+  else:
+    matches = np.concatenate(matches,axis=0)
+
+  return matches, np.array(unmatched_detections), np.array(unmatched_trackers)
+
+
+class Sort(object):
+  def __init__(self, max_age=1, min_hits=3, iou_threshold=0.3, max_interpolation=1):
+    """
+    Sets key parameters for SORT
+    """
+    self.max_age = max_age
+    self.min_hits = min_hits
+    self.iou_threshold = iou_threshold
+    self.trackers = []
+    self.frame_count = 0
+    self.max_interpolation = max_interpolation
+
+  def update(self, dets=np.empty((0, 5))):
+    """
+    Params:
+      dets - a numpy array of detections in the format [[x1,y1,x2,y2,score],[x1,y1,x2,y2,score],...]
+    Requires: this method must be called once for each frame even with empty detections (use np.empty((0, 5)) for frames without detections).
+    Returns the a similar array, where the last column is the object ID.
+
+    NOTE: The number of objects returned may differ from the number of detections provided.
+    """
+    self.frame_count += 1
+    # get predicted locations from existing trackers.
+    trks = np.zeros((len(self.trackers), 5))
+    to_del = []
+    ret = []
+    for t, trk in enumerate(trks):
+      pos = self.trackers[t].predict()[0]
+      trk[:] = [pos[0], pos[1], pos[2], pos[3], 0]
+      if np.any(np.isnan(pos)):
+        to_del.append(t)
+    trks = np.ma.compress_rows(np.ma.masked_invalid(trks))
+    for t in reversed(to_del):
+      self.trackers.pop(t)
+    matched, unmatched_dets, unmatched_trks = associate_detections_to_trackers(dets,trks, self.iou_threshold)
+
+    # update matched trackers with assigned detections
+    for m in matched:
+      self.trackers[m[1]].update(dets[m[0], :])
+
+    # create and initialise new trackers for unmatched detections
+    for i in unmatched_dets:
+        trk = KalmanBoxTracker(dets[i,:])
+        self.trackers.append(trk)
+    i = len(self.trackers)
+    for trk in reversed(self.trackers):
+        d = trk.get_state()[0]
+        # updated (trk.time_since_update < 1) to (trk.time_since_update < self.max_age) --->Medha
+        if (trk.time_since_update < self.max_age) and (trk.hit_streak >= self.min_hits or self.frame_count <= self.min_hits):
+          ret.append(np.concatenate((d,[trk.id+1])).reshape(1,-1)) # +1 as MOT benchmark requires positive
+        i -= 1
+        # remove dead tracklet
+        if(trk.time_since_update > self.max_age):
+          self.trackers.pop(i)
+    if(len(ret)>0):
+      return np.concatenate(ret)
+    return np.empty((0,5))
+
+def parse_args():
+    """Parse input arguments."""
+    parser = argparse.ArgumentParser(description='SORT demo')
+    parser.add_argument('--display', dest='display', help='Display online tracker output (slow) [False]',action='store_true')
+    parser.add_argument("--seq_path", help="Path to detections.", type=str, default='data')
+    parser.add_argument("--phase", help="Subdirectory in seq_path.", type=str, default='train')
+    parser.add_argument("--max_age", 
+                        help="Maximum number of frames to keep alive a track without associated detections.", 
+                        type=int, default=1)
+    parser.add_argument("--min_hits", 
+                        help="Minimum number of associated detections before track is initialised.", 
+                        type=int, default=3)
+    parser.add_argument("--iou_threshold", help="Minimum IOU for match.", type=float, default=0.3)
+    args = parser.parse_args()
+    return args
+
+if __name__ == '__main__':
+  # all train
+  args = parse_args()
+  display = args.display
+  phase = args.phase
+  total_time = 0.0
+  total_frames = 0
+  colours = np.random.rand(32, 3) #used only for display
+  if(display):
+    if not os.path.exists('mot_benchmark'):
+      print('\n\tERROR: mot_benchmark link not found!\n\n    Create a symbolic link to the MOT benchmark\n    (https://motchallenge.net/data/2D_MOT_2015/#download). E.g.:\n\n    $ ln -s /path/to/MOT2015_challenge/2DMOT2015 mot_benchmark\n\n')
+      exit()
+    plt.ion()
+    fig = plt.figure()
+    ax1 = fig.add_subplot(111, aspect='equal')
+
+  if not os.path.exists('output'):
+    os.makedirs('output')
+  pattern = os.path.join(args.seq_path, phase, '*', 'det', 'det.txt')
+  for seq_dets_fn in glob.glob(pattern):
+    mot_tracker = Sort(max_age=args.max_age, 
+                       min_hits=args.min_hits,
+                       iou_threshold=args.iou_threshold) #create instance of the SORT tracker
+    seq_dets = np.loadtxt(seq_dets_fn, delimiter=',')
+    seq = seq_dets_fn[pattern.find('*'):].split(os.path.sep)[0]
+    
+    with open(os.path.join('output', '%s.txt'%(seq)),'w') as out_file:
+      print("Processing %s."%(seq))
+      for frame in range(int(seq_dets[:,0].max())):
+        frame += 1 #detection and frame numbers begin at 1
+        dets = seq_dets[seq_dets[:, 0]==frame, 2:7]
+        dets[:, 2:4] += dets[:, 0:2] #convert to [x1,y1,w,h] to [x1,y1,x2,y2]
+        total_frames += 1
+
+        if(display):
+          fn = os.path.join('mot_benchmark', phase, seq, 'img1', '%06d.jpg'%(frame))
+          im =io.imread(fn)
+          ax1.imshow(im)
+          plt.title(seq + ' Tracked Targets')
+
+        start_time = time.time()
+        trackers = mot_tracker.update(dets)
+        cycle_time = time.time() - start_time
+        total_time += cycle_time
+
+        for d in trackers:
+          print('%d,%d,%.2f,%.2f,%.2f,%.2f,1,-1,-1,-1'%(frame,d[4],d[0],d[1],d[2]-d[0],d[3]-d[1]),file=out_file)
+          if(display):
+            d = d.astype(np.int32)
+            ax1.add_patch(patches.Rectangle((d[0],d[1]),d[2]-d[0],d[3]-d[1],fill=False,lw=3,ec=colours[d[4]%32,:]))
+
+        if(display):
+          fig.canvas.flush_events()
+          plt.draw()
+          ax1.cla()
+
+  print("Total Tracking took: %.3f seconds for %d frames or %.1f FPS" % (total_time, total_frames, total_frames / total_time))
+
+  if(display):
+    print("Note: to get real runtime results run without the option: --display")

MEMTrack/src/trainBacteriaRetinanetMotionData_Motility.py

@@ -0,0 +1,318 @@
+# Some basic setup:
+# Setup detectron2 logger
+import detectron2
+from detectron2.utils.logger import setup_logger
+setup_logger()
+
+# import some common libraries
+import numpy as np
+import os, json, cv2, random
+import cv2
+import pandas as pd
+import shutil
+import numpy as np
+
+# import some common detectron2 utilities
+from detectron2 import model_zoo
+from detectron2.engine import DefaultPredictor
+from detectron2.config import get_cfg
+from detectron2.utils.visualizer import Visualizer
+from detectron2.data import MetadataCatalog, DatasetCatalog
+import json
+from detectron2.structures import BoxMode
+from detectron2.data import DatasetCatalog, MetadataCatalog
+import argparse
+from natsort import natsorted
+import PIL
+from PIL import Image
+import matplotlib.pyplot as plt
+from detectron2.engine import DefaultTrainer
+from detectron2.evaluation import COCOEvaluator
+from detectron2.config import get_cfg
+from contextlib import suppress
+import os
+
+ap = argparse.ArgumentParser(description='Training')
+ap.add_argument('--source_path', default="./CombinationModel/data_feature_optical_flow_median_back_2pyr_18win_background_img/", type=str, metavar='PATH')
+
+ap.add_argument('--output_dir', default="CombinationModel/easy-optical_flow_median_back_2pyr_18win_00125_34k", type=str, metavar='CELL PATH')
+
+ap.add_argument('--annotations_train', default="Easy", type=str, metavar='TRAIN')
+ap.add_argument('--annotations_test', default="Easy", type=str, metavar='TEST')
+ap.add_argument('--bbox_size', default=31, type=int)
+ap.add_argument('--epochs', default="90", type=str, metavar='TEST')
+ap.add_argument('--lr', default="0.00125", type=str, metavar='TEST')
+ap.add_argument('--feature', default="images_feature/", type=str, metavar='FEATURE')
+ap.add_argument("--cuda", default=0, type=int)
+
+args = ap.parse_args()
+
+width = args.bbox_size
+coco_format_train = pd.DataFrame(columns=["file_name","height","width","annotations"])
+coco_format_train["annotations"] = coco_format_train["annotations"].astype('object')
+
+#source = "./CombinationModel/data_feature_optical_flow_median_back_2pyr_18win_background_img/" #update
+source = args.source_path
+dest_train = source + "train/"
+dest_test = source + "val/" #update
+
+with suppress(OSError):
+    os.remove(source + 'test.json')
+    
+    
+val = True  #update
+#output_dir = ("CombinationModel/easy-optical_flow_median_back_2pyr_18win_00125_34k") #update
+output_dir = args.output_dir
+train_images_path = "/train/" + args.feature #"/images_feature/"
+val_image_path = "/val/" + args.feature #"/images_feature/"
+
+#annotations_source = source + "annotation_easy_hard/"
+images_source = source + args.feature #"images_feature/"
+
+if args.annotations_train == "Easy":
+    annotations_train = dest_train + "annotation_easy/" 
+elif args.annotations_train == "Hard":
+    annotations_train = dest_train + "annotation_hard/" 
+elif args.annotations_train == "VeryHard":
+    annotations_train = dest_train + "annotation_veryhard/" 
+elif args.annotations_train == "Easy+Hard":
+    annotations_train = dest_train + "annotation_easy_hard/" 
+elif args.annotations_train == "All":
+    annotations_train = dest_train + "annotation_easy_hard_veryhard/"
+elif args.annotations_train == "Motility-low":
+    annotations_train = dest_train + "annotation_motility_low/" 
+elif args.annotations_train == "Motility-high":
+    annotations_train = dest_train + "annotation_motility_high/" 
+elif args.annotations_train == "Motility-wiggle":
+    annotations_train = dest_train + "annotation_motility_wiggle/" 
+elif args.annotations_train == "Motility-mid":
+    annotations_train = dest_train + "annotation_motility_mid/" 
+elif args.annotations_train == "Sticking-stick":
+    annotations_train = dest_train + "annotation_sticking_stick/" 
+elif args.annotations_train == "Sticking-motile":
+    annotations_train = dest_train + "annotation_sticking_motile/" 
+elif args.annotations_train == "Sticking-non_motile":
+    annotations_train = dest_train + "annotation_sticking_non_motile/" 
+elif args.annotations_train == "Motility-low-wiggle":
+    annotations_train = dest_train + "annotation_motility_low_wiggle/" 
+elif args.annotations_train == "Motility-mid-high":
+    annotations_train = dest_train + "annotation_motility_mid_high/" 
+
+
+if args.annotations_test == "Easy":
+    annotations_test = dest_test + "annotation_easy/" 
+elif args.annotations_test == "Hard":
+    annotations_test = dest_test + "annotation_hard/" 
+elif args.annotations_test == "VeryHard":
+    annotations_test = dest_test + "annotation_veryhard/" 
+elif args.annotations_test == "Easy+Hard":
+    annotations_test = dest_test + "annotation_easy_hard/" 
+elif args.annotations_test == "All":
+    annotations_test = dest_test + "annotation_easy_hard_veryhard/" 
+elif args.annotations_test == "Motility-low":
+    annotations_test = dest_test + "annotation_motility_low/" 
+elif args.annotations_test == "Motility-high":
+    annotations_test = dest_test + "annotation_motility_high/" 
+elif args.annotations_test == "Motility-wiggle":
+    annotations_test = dest_test + "annotation_motility_wiggle/" 
+elif args.annotations_test == "Motility-mid":
+    annotations_test = dest_test + "annotation_motility_mid/" 
+elif args.annotations_test == "Sticking-stick":
+    annotations_test = dest_test + "annotation_sticking_stick/" 
+elif args.annotations_test == "Sticking-motile":
+    annotations_test = dest_test + "annotation_sticking_motile/" 
+elif args.annotations_test == "Sticking-non_motile":
+    annotations_test = dest_test + "annotation_sticking_non_motile/" 
+elif args.annotations_test == "Motility-low-wiggle":
+    annotations_test = dest_test + "annotation_motility_low_wiggle/" 
+elif args.annotations_test == "Motility-mid-high":
+    annotations_test = dest_test + "annotation_motility_mid_high/" 
+
+
+images_train = dest_train + args.feature #"images_feature/"
+images_test = dest_test + args.feature #"images_feature/"
+
+factor_w = 1#1024/1388
+factor_h = 1#1024/1040
+
+#function to get background frame
+#function to get prev frame
+#function to create new image given image num
+
+for txt_file in natsorted(os.listdir(annotations_train)):
+        #width = 31
+        text_file = open(annotations_train + txt_file, 'r')
+        xy_coords = text_file.readlines()
+        boxes = []
+        res=pd.DataFrame(columns=["file_name","height","width","annotations"])
+        image = PIL.Image.open(images_train + txt_file[:-4] + ".tif").convert('L')
+        image_feature = PIL.Image.open(images_train + txt_file[:-4] + ".tif")
+        image = image_feature
+        #print(image.size)
+        res.at[0,"height"] = image.height
+        res.at[0,"width"] = image.width
+        res.at[0,"file_name"] = txt_file[:-4]+".tif"
+        bbox_mode = 0
+        category_id = 0
+        # image2 = image.resize((1024,1024))
+        # image2.save(images_resized_train + txt_file[:-4] + ".jpg")
+        for xy in xy_coords:
+            box = []
+            x = float(xy.split(" ")[0])
+            y = float(xy.split(" ")[1])
+            x1 = int(x*factor_w - (width // 2))
+            y1 = int(y*factor_h - (width // 2))
+            x2 = int(x*factor_w + (width // 2))
+            y2 = int(y*factor_h + (width // 2))
+            w = h = 31
+            box = [x1, y1, x2, y2]
+            boxes.append(np.array(box))
+            #print(np.array(box))
+
+        res["annotations"]=res["annotations"].astype('object')
+        annotation_df = pd.DataFrame(columns=["bbox","bbox_mode","category_id"])
+        annotation_df["bbox"] = boxes
+        annotation_df["bbox_mode"] = bbox_mode
+        annotation_df["category_id"] = category_id
+        annotations = annotation_df.T.to_dict().values()
+        l = []
+        for j in annotations:
+            l.append(j)
+        res.at[0,"annotations"] = l
+        coco_format_train = coco_format_train.append(res)
+        coco_format_train.reset_index(drop=True,inplace=True)
+        
+coco_format_train.reset_index(inplace=True)
+coco_format_train.rename(columns={"index":"image_id"},inplace=True)
+coco_format_train.to_json(source + "train.json",orient="records")
+
+coco_format_test = pd.DataFrame(columns=["file_name","height","width","annotations"])
+coco_format_test["annotations"] = coco_format_test["annotations"].astype('object')
+
+for txt_file in natsorted(os.listdir(annotations_test)):
+        #width = 31
+        text_file = open(annotations_test + txt_file, 'r')
+        xy_coords = text_file.readlines()
+        boxes = []
+        res=pd.DataFrame(columns=["file_name","height","width","annotations"])
+        image = PIL.Image.open(images_test + txt_file[:-4] + ".tif").convert('L')
+        image_feature = PIL.Image.open(images_test + txt_file[:-4] + ".tif")
+        image = image_feature
+        #print(image.size)
+        res.at[0,"height"] = image.height
+        res.at[0,"width"] = image.width
+        res.at[0,"file_name"] = txt_file[:-4]+".tif"
+        bbox_mode = 0
+        category_id = 0
+        # image2 = image.resize((1024,1024))
+        # image2.save(images_resized_train + txt_file[:-4] + ".jpg")
+        for xy in xy_coords:
+            box = []
+            x = float(xy.split(" ")[0])
+            y = float(xy.split(" ")[1])
+            x1 = int(x*factor_w - (width // 2))
+            y1 = int(y*factor_h - (width // 2))
+            x2 = int(x*factor_w + (width // 2))
+            y2 = int(y*factor_h + (width // 2))
+            w = h = 31
+            box = [x1, y1, x2, y2]
+            boxes.append(np.array(box))
+            #print(np.array(box))
+
+        res["annotations"]=res["annotations"].astype('object')
+        annotation_df = pd.DataFrame(columns=["bbox","bbox_mode","category_id"])
+        annotation_df["bbox"] = boxes
+        annotation_df["bbox_mode"] = bbox_mode
+        annotation_df["category_id"] = category_id
+        annotations = annotation_df.T.to_dict().values()
+        l = []
+        for j in annotations:
+            l.append(j)
+        res.at[0,"annotations"] = l
+        coco_format_test = coco_format_test.append(res)
+        coco_format_test.reset_index(drop=True,inplace=True)
+        
+coco_format_test.reset_index(inplace=True)
+coco_format_test.rename(columns={"index":"image_id"},inplace=True)
+coco_format_test.to_json(source + "val.json",orient="records")
+
+
+
+def get_board_dicts(imgdir, mode):
+    if mode is 'train':
+        json_file = imgdir+"/train.json" #Fetch the json file
+    if mode is 'val':
+         json_file = imgdir+"/val.json" #Fetch the json file
+    with open(json_file) as f:
+        dataset_dicts = json.load(f)
+    for i in dataset_dicts:
+        filename = i["file_name"] 
+        if mode is 'train':
+            i["file_name"] = imgdir + train_images_path + filename 
+        if mode is 'val':
+             i["file_name"] = imgdir + val_image_path + filename 
+        for j in i["annotations"]:
+            j["bbox_mode"] = BoxMode.XYXY_ABS #Setting the required Box Mode
+            j["category_id"] = int(j["category_id"])
+    return dataset_dicts
+
+#Registering the Dataset
+for d in ["train", "val"]:
+    DatasetCatalog.register("boardetect_" + d, lambda d=d: get_board_dicts(source, d))
+    MetadataCatalog.get("boardetect_" + d).set(thing_classes=["node"])
+board_metadata = MetadataCatalog.get("boardetect_train")
+val_metadata = MetadataCatalog.get("boardetect_val")
+
+train_data = ("boardetect_train",)
+
+if val ==True:
+    val_data = ("boardetect_val",)
+else:
+    val_data = ("boardetect_train",)
+
+
+
+class CocoTrainer(DefaultTrainer):
+       
+    @classmethod
+    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
+        if output_folder is None:
+            os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
+            output_folder = cfg.OUTPUT_DIR
+
+
+        return COCOEvaluator(dataset_name, cfg, False, output_folder)
+
+cfg = get_cfg()
+cfg.MODEL.DEVICE = f"cuda:{args.cuda}"
+cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/retinanet_R_101_FPN_3x.yaml")) #Get the basic model configuration from the model zoo 
+#Passing the Train and Validation sets
+cfg.DATASETS.TRAIN = train_data #("boardetect_train",)
+cfg.DATASETS.TEST = val_data #("boardetect_train",)
+cfg.OUTPUT_DIR = output_dir #("comparison-optical-flow")
+# Number of data loading threads
+cfg.DATALOADER.NUM_WORKERS = 4
+#cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
+#cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") #uncommwnt during inference
+#cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") #uncommwnt during inference
+#cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_00044999.pth") #uncommwnt during inference
+# Number of images per batch across all machines.
+cfg.SOLVER.IMS_PER_BATCH = 4
+cfg.SOLVER.BASE_LR = float(args.lr) #0.00125  # pick a good LearningRate
+cfg.SOLVER.MAX_ITER = int(args.epochs) #90000  #No. of iterations
+# cfg.SOLVER.STEPS = (300,600)
+cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 80
+cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
+cfg.TEST.EVAL_PERIOD = 500 # No. of iterations after which the Validation Set is evaluated. 
+cfg.TEST.DETECTIONS_PER_IMAGE = 60
+cfg.SOLVER.CHECKPOINT_PERIOD = 5000
+cfg.VIS_PERIOD = 500
+cfg.DATASETS.PRECOMPUTED_PROPOSAL_TOPK_TRAIN = 100
+cfg.DATASETS.PRECOMPUTED_PROPOSAL_TOPK_TEST = 100
+cfg.MODEL.RPN.POST_NMS_TOPK_TRAIN = 80 #
+cfg.MODEL.RPN.POST_NMS_TOPK_TEST = 80
+os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
+trainer = CocoTrainer(cfg) 
+trainer.resume_or_load(resume=True) #resume=True for inference
+trainer.train()
+print(cfg.MODEL.WEIGHTS)

MEMTrack/src/trainBacteriaRetinanetMotionData_Motility_Val_loss.py

@@ -0,0 +1,333 @@
+# Some basic setup:
+# Setup detectron2 logger
+import detectron2
+from detectron2.utils.logger import setup_logger
+setup_logger()
+
+# import some common libraries
+import numpy as np
+import os, json, cv2, random
+import cv2
+import pandas as pd
+import shutil
+import numpy as np
+
+# import some common detectron2 utilities
+from detectron2 import model_zoo
+from detectron2.engine import DefaultPredictor
+from detectron2.config import get_cfg
+from detectron2.utils.visualizer import Visualizer
+from detectron2.data import MetadataCatalog, DatasetCatalog
+import json
+from detectron2.structures import BoxMode
+from detectron2.data import DatasetCatalog, MetadataCatalog
+import argparse
+from natsort import natsorted
+import PIL
+from PIL import Image
+import matplotlib.pyplot as plt
+from detectron2.engine import DefaultTrainer
+from detectron2.evaluation import COCOEvaluator
+from detectron2.config import get_cfg
+from contextlib import suppress
+import os
+from LossEvalHook import LossEvalHook 
+from detectron2.data import DatasetMapper, build_detection_test_loader
+
+ap = argparse.ArgumentParser(description='Training')
+ap.add_argument('--source_path', default="./CombinationModel/data_feature_optical_flow_median_back_2pyr_18win_background_img/", type=str, metavar='PATH')
+
+ap.add_argument('--output_dir', default="CombinationModel/easy-optical_flow_median_back_2pyr_18win_00125_34k", type=str, metavar='CELL PATH')
+
+ap.add_argument('--annotations_train', default="Easy", type=str, metavar='TRAIN')
+ap.add_argument('--annotations_test', default="Easy", type=str, metavar='TEST')
+ap.add_argument('--bbox_size', default=31, type=int)
+ap.add_argument('--epochs', default="90", type=str, metavar='TEST')
+ap.add_argument('--lr', default="0.00125", type=str, metavar='TEST')
+ap.add_argument('--feature', default="images_feature/", type=str, metavar='FEATURE')
+ap.add_argument("--cuda", default=0, type=int)
+ap.add_argument("--exp",  type=str)
+
+args = ap.parse_args()
+
+width = args.bbox_size
+coco_format_train = pd.DataFrame(columns=["file_name","height","width","annotations"])
+coco_format_train["annotations"] = coco_format_train["annotations"].astype('object')
+
+#source = "./CombinationModel/data_feature_optical_flow_median_back_2pyr_18win_background_img/" #update
+source = args.source_path
+dest_train = source + "train/"
+dest_test = source + "val/" #update
+
+with suppress(OSError):
+    os.remove(source + 'test.json')
+    
+    
+val = True  #update
+#output_dir = ("CombinationModel/easy-optical_flow_median_back_2pyr_18win_00125_34k") #update
+output_dir = args.output_dir
+train_images_path = "/train/" + args.feature #"/images_feature/"
+val_image_path = "/val/" + args.feature #"/images_feature/"
+
+#annotations_source = source + "annotation_easy_hard/"
+images_source = source + args.feature #"images_feature/"
+
+if args.annotations_train == "Easy":
+    annotations_train = dest_train + "annotation_easy/" 
+elif args.annotations_train == "Hard":
+    annotations_train = dest_train + "annotation_hard/" 
+elif args.annotations_train == "VeryHard":
+    annotations_train = dest_train + "annotation_veryhard/" 
+elif args.annotations_train == "Easy+Hard":
+    annotations_train = dest_train + "annotation_easy_hard/" 
+elif args.annotations_train == "All":
+    annotations_train = dest_train + "annotation_easy_hard_veryhard/"
+elif args.annotations_train == "Motility-low":
+    annotations_train = dest_train + "annotation_motility_low/" 
+elif args.annotations_train == "Motility-high":
+    annotations_train = dest_train + "annotation_motility_high/" 
+elif args.annotations_train == "Motility-wiggle":
+    annotations_train = dest_train + "annotation_motility_wiggle/" 
+elif args.annotations_train == "Motility-mid":
+    annotations_train = dest_train + "annotation_motility_mid/" 
+elif args.annotations_train == "Sticking-stick":
+    annotations_train = dest_train + "annotation_sticking_stick/" 
+elif args.annotations_train == "Sticking-motile":
+    annotations_train = dest_train + "annotation_sticking_motile/" 
+elif args.annotations_train == "Sticking-non_motile":
+    annotations_train = dest_train + "annotation_sticking_non_motile/" 
+elif args.annotations_train == "Motility-low-wiggle":
+    annotations_train = dest_train + "annotation_motility_low_wiggle/" 
+elif args.annotations_train == "Motility-mid-high":
+    annotations_train = dest_train + "annotation_motility_mid_high/" 
+
+
+if args.annotations_test == "Easy":
+    annotations_test = dest_test + "annotation_easy/" 
+elif args.annotations_test == "Hard":
+    annotations_test = dest_test + "annotation_hard/" 
+elif args.annotations_test == "VeryHard":
+    annotations_test = dest_test + "annotation_veryhard/" 
+elif args.annotations_test == "Easy+Hard":
+    annotations_test = dest_test + "annotation_easy_hard/" 
+elif args.annotations_test == "All":
+    annotations_test = dest_test + "annotation_easy_hard_veryhard/" 
+elif args.annotations_test == "Motility-low":
+    annotations_test = dest_test + "annotation_motility_low/" 
+elif args.annotations_test == "Motility-high":
+    annotations_test = dest_test + "annotation_motility_high/" 
+elif args.annotations_test == "Motility-wiggle":
+    annotations_test = dest_test + "annotation_motility_wiggle/" 
+elif args.annotations_test == "Motility-mid":
+    annotations_test = dest_test + "annotation_motility_mid/" 
+elif args.annotations_test == "Sticking-stick":
+    annotations_test = dest_test + "annotation_sticking_stick/" 
+elif args.annotations_test == "Sticking-motile":
+    annotations_test = dest_test + "annotation_sticking_motile/" 
+elif args.annotations_test == "Sticking-non_motile":
+    annotations_test = dest_test + "annotation_sticking_non_motile/" 
+elif args.annotations_test == "Motility-low-wiggle":
+    annotations_test = dest_test + "annotation_motility_low_wiggle/" 
+elif args.annotations_test == "Motility-mid-high":
+    annotations_test = dest_test + "annotation_motility_mid_high/" 
+
+
+images_train = dest_train + args.feature #"images_feature/"
+images_test = dest_test + args.feature #"images_feature/"
+
+factor_w = 1#1024/1388
+factor_h = 1#1024/1040
+
+#function to get background frame
+#function to get prev frame
+#function to create new image given image num
+
+for txt_file in natsorted(os.listdir(annotations_train)):
+        #width = 31
+        text_file = open(annotations_train + txt_file, 'r')
+        xy_coords = text_file.readlines()
+        boxes = []
+        res=pd.DataFrame(columns=["file_name","height","width","annotations"])
+        image = PIL.Image.open(images_train + txt_file[:-4] + ".tif").convert('L')
+        image_feature = PIL.Image.open(images_train + txt_file[:-4] + ".tif")
+        image = image_feature
+        #print(image.size)
+        res.at[0,"height"] = image.height
+        res.at[0,"width"] = image.width
+        res.at[0,"file_name"] = txt_file[:-4]+".tif"
+        bbox_mode = 0
+        category_id = 0
+        # image2 = image.resize((1024,1024))
+        # image2.save(images_resized_train + txt_file[:-4] + ".jpg")
+        for xy in xy_coords:
+            box = []
+            x = float(xy.split(" ")[0])
+            y = float(xy.split(" ")[1])
+            x1 = int(x*factor_w - (width // 2))
+            y1 = int(y*factor_h - (width // 2))
+            x2 = int(x*factor_w + (width // 2))
+            y2 = int(y*factor_h + (width // 2))
+            w = h = 31
+            box = [x1, y1, x2, y2]
+            boxes.append(np.array(box))
+            #print(np.array(box))
+
+        res["annotations"]=res["annotations"].astype('object')
+        annotation_df = pd.DataFrame(columns=["bbox","bbox_mode","category_id"])
+        annotation_df["bbox"] = boxes
+        annotation_df["bbox_mode"] = bbox_mode
+        annotation_df["category_id"] = category_id
+        annotations = annotation_df.T.to_dict().values()
+        l = []
+        for j in annotations:
+            l.append(j)
+        res.at[0,"annotations"] = l
+        coco_format_train = coco_format_train.append(res)
+        coco_format_train.reset_index(drop=True,inplace=True)
+        
+coco_format_train.reset_index(inplace=True)
+coco_format_train.rename(columns={"index":"image_id"},inplace=True)
+coco_format_train.to_json(source + f"train_{args.exp}.json",orient="records")
+
+coco_format_test = pd.DataFrame(columns=["file_name","height","width","annotations"])
+coco_format_test["annotations"] = coco_format_test["annotations"].astype('object')
+
+for txt_file in natsorted(os.listdir(annotations_test)):
+        #width = 31
+        text_file = open(annotations_test + txt_file, 'r')
+        xy_coords = text_file.readlines()
+        boxes = []
+        res=pd.DataFrame(columns=["file_name","height","width","annotations"])
+        image = PIL.Image.open(images_test + txt_file[:-4] + ".tif").convert('L')
+        image_feature = PIL.Image.open(images_test + txt_file[:-4] + ".tif")
+        image = image_feature
+        #print(image.size)
+        res.at[0,"height"] = image.height
+        res.at[0,"width"] = image.width
+        res.at[0,"file_name"] = txt_file[:-4]+".tif"
+        bbox_mode = 0
+        category_id = 0
+        # image2 = image.resize((1024,1024))
+        # image2.save(images_resized_train + txt_file[:-4] + ".jpg")
+        for xy in xy_coords:
+            box = []
+            x = float(xy.split(" ")[0])
+            y = float(xy.split(" ")[1])
+            x1 = int(x*factor_w - (width // 2))
+            y1 = int(y*factor_h - (width // 2))
+            x2 = int(x*factor_w + (width // 2))
+            y2 = int(y*factor_h + (width // 2))
+            w = h = 31
+            box = [x1, y1, x2, y2]
+            boxes.append(np.array(box))
+            #print(np.array(box))
+
+        res["annotations"]=res["annotations"].astype('object')
+        annotation_df = pd.DataFrame(columns=["bbox","bbox_mode","category_id"])
+        annotation_df["bbox"] = boxes
+        annotation_df["bbox_mode"] = bbox_mode
+        annotation_df["category_id"] = category_id
+        annotations = annotation_df.T.to_dict().values()
+        l = []
+        for j in annotations:
+            l.append(j)
+        res.at[0,"annotations"] = l
+        coco_format_test = coco_format_test.append(res)
+        coco_format_test.reset_index(drop=True,inplace=True)
+        
+coco_format_test.reset_index(inplace=True)
+coco_format_test.rename(columns={"index":"image_id"},inplace=True)
+coco_format_test.to_json(source + f"val_{args.exp}.json",orient="records")
+
+
+
+def get_board_dicts(imgdir, mode):
+    if mode is 'train':
+        json_file = imgdir+f"/train_{args.exp}.json" #Fetch the json file
+    if mode is 'val':
+         json_file = imgdir+f"/val_{args.exp}.json" #Fetch the json file
+    with open(json_file) as f:
+        dataset_dicts = json.load(f)
+    for i in dataset_dicts:
+        filename = i["file_name"] 
+        if mode is 'train':
+            i["file_name"] = imgdir + train_images_path + filename 
+        if mode is 'val':
+             i["file_name"] = imgdir + val_image_path + filename 
+        for j in i["annotations"]:
+            j["bbox_mode"] = BoxMode.XYXY_ABS #Setting the required Box Mode
+            j["category_id"] = int(j["category_id"])
+    return dataset_dicts
+
+#Registering the Dataset
+for d in ["train", "val"]:
+    DatasetCatalog.register("boardetect_" + d, lambda d=d: get_board_dicts(source, d))
+    MetadataCatalog.get("boardetect_" + d).set(thing_classes=["node"])
+board_metadata = MetadataCatalog.get("boardetect_train")
+val_metadata = MetadataCatalog.get("boardetect_val")
+
+train_data = ("boardetect_train",)
+
+if val ==True:
+    val_data = ("boardetect_val",)
+else:
+    val_data = ("boardetect_train",)
+
+
+
+class CocoTrainer(DefaultTrainer):
+       
+    @classmethod
+    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
+        if output_folder is None:
+            os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
+            output_folder = cfg.OUTPUT_DIR
+        return COCOEvaluator(dataset_name, cfg, False, output_folder)
+                             
+                     
+    def build_hooks(self):
+        hooks = super().build_hooks()
+        hooks.insert(-1,LossEvalHook(
+            cfg.TEST.EVAL_PERIOD,
+            self.model,
+            build_detection_test_loader(
+                self.cfg,
+                self.cfg.DATASETS.TEST[0],
+                DatasetMapper(self.cfg,True)
+            )
+        ))
+        return hooks
+
+cfg = get_cfg()
+cfg.MODEL.DEVICE = f"cuda:{args.cuda}"
+cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/retinanet_R_101_FPN_3x.yaml")) #Get the basic model configuration from the model zoo 
+#Passing the Train and Validation sets
+cfg.DATASETS.TRAIN = train_data #("boardetect_train",)
+cfg.DATASETS.TEST = val_data #("boardetect_train",)
+cfg.OUTPUT_DIR = output_dir #("comparison-optical-flow")
+# Number of data loading threads
+cfg.DATALOADER.NUM_WORKERS = 4
+#cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
+#cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") #uncommwnt during inference
+#cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") #uncommwnt during inference
+#cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_00044999.pth") #uncommwnt during inference
+# Number of images per batch across all machines.
+cfg.SOLVER.IMS_PER_BATCH = 4
+cfg.SOLVER.BASE_LR = float(args.lr) #0.00125  # pick a good LearningRate
+cfg.SOLVER.MAX_ITER = int(args.epochs) #90000  #No. of iterations
+# cfg.SOLVER.STEPS = (300,600)
+cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 80
+cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
+cfg.TEST.EVAL_PERIOD = 500 # No. of iterations after which the Validation Set is evaluated. 
+cfg.TEST.DETECTIONS_PER_IMAGE = 60
+cfg.SOLVER.CHECKPOINT_PERIOD = 500
+cfg.VIS_PERIOD = 500
+cfg.DATASETS.PRECOMPUTED_PROPOSAL_TOPK_TRAIN = 100
+cfg.DATASETS.PRECOMPUTED_PROPOSAL_TOPK_TEST = 100
+cfg.MODEL.RPN.POST_NMS_TOPK_TRAIN = 80 #
+cfg.MODEL.RPN.POST_NMS_TOPK_TEST = 80
+os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
+trainer = CocoTrainer(cfg) 
+trainer.resume_or_load(resume=True) #resume=True for inference
+trainer.train()
+print(cfg.MODEL.WEIGHTS)

MEMTrack/todo.txt

@@ -0,0 +1,7 @@
+1. Update low to no motility and wiggle to low
+2. add instructions for feature genration code
+3. update feature generation code
+4. update training scripts, add instructions for that
+5. add inference code, instructions for that
+6. inference on other data
+7. 

app.py

@@ -1,166 +1,43 @@
+import os
+import sys
 import cv2
-import gradio as gr
-import numpy as np
 import os
+import glob
 import shutil
-from pathlib import Path
-
-# def process_video(input_video):
-#     # Create a directory to store frames
-#     frames_dir = Path("frames")
-#     frames_dir.mkdir(exist_ok=True)
-    
-#     # Read the video
-#     cap = cv2.VideoCapture(input_video)
-#     count = 0
-#     frames = []
-    
-#     while True:
-#         ret, frame = cap.read()
-#         if not ret:
-#             break
-#         frame_path = str(frames_dir / f"frame_{count}.jpg")
-#         cv2.imwrite(frame_path, frame)
-#         frames.append(frame)
-#         count += 1
-    
-#     # Close video file
-#     cap.release()
-    
-#     # Process the first and the last frame (simple example: convert to grayscale)
-#     if frames:
-#         frames[0] = cv2.cvtColor(frames[0], cv2.COLOR_BGR2GRAY)
-#         frames[-1] = cv2.cvtColor(frames[-1], cv2.COLOR_BGR2GRAY)
-    
-#     # Create a processed video
-#     output_video_path = "processed_video.mp4"
-#     height, width = frames[0].shape[:2]
-#     fourcc = cv2.VideoWriter_fourcc(*'mp4v') # or use 'XVID' if mp4 does not work
-#     out = cv2.VideoWriter(output_video_path, fourcc, 60.0, (width, height), isColor=len(frames[0].shape) > 2)
-    
-#     for frame in frames:
-#         out.write(frame)
-    
-#     out.release()
-#     return output_video_path
-
-
-def process_video(input_video):
-    # Create a directory to store difference frames
-    diff_frames_dir = Path("diff_frames")
-    if os.path.exists(diff_frames_dir):
-        shutil.rmtree(diff_frames_dir)
-    diff_frames_dir.mkdir(exist_ok=True)
-    
-    frames_dir = Path("frames")
-    if os.path.exists(frames_dir):
-        shutil.rmtree(frames_dir)
-    frames_dir.mkdir(exist_ok=True)
-    
-    # Read the video
-    cap = cv2.VideoCapture(input_video)
-    count = 0
-    diff_frames = []
-    prev_frame = None
-    
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-            
-        # Convert frame to grayscale
-        gray_frame = frame #cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-        
-        frame_path = str(frames_dir / f"frame_{count}.jpg")
-        cv2.imwrite(frame_path, frame)
-        
-        if prev_frame is not None:
-            # Calculate difference with the previous frame
-            diff = cv2.absdiff(prev_frame, gray_frame)
-            diff_frames.append(gray_frame)
-            
-            # Save difference frame
-            diff_frame_path = str(diff_frames_dir / f"diff_frame_{count}.jpg")
-            cv2.imwrite(diff_frame_path, diff)
-        else:
-            # For the first frame, there is no previous frame to compare
-            diff_frames.append(gray_frame)
-            
-        prev_frame = gray_frame
-        count += 1
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    print(fps)
-    size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
-    print(size)
-    # Close video file
-    cap.release()
-    
-    # Create a video from the difference frames
-    output_video_path = "diff_video.mp4"
-    if diff_frames:
-        print("Here")
-        height, width = diff_frames[0].shape[:2]
-        fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Use 'mp4v' for MP4 or 'XVID' for AVI
-        out = cv2.VideoWriter(output_video_path, fourcc, fps, size, isColor=True)
-         
-        for diff_frame in diff_frames:
-            # Need to convert single channel image to three channels to write with VideoWriter
-#             three_channel_frame = cv2.cvtColor(diff_frame, cv2.COLOR_GRAY2BGR)
-            out.write(diff_frame)
-        
-        out.release()
-    
-    return output_video_path
-
-def show_preds_video(input_video):
-    cap = cv2.VideoCapture(input_video)
-    while(cap.isOpened()):
-        ret, frame = cap.read()
-        if ret:
-            frame_copy = frame.copy()
-            
-            yield frame_copy
-
-
-# # Define Gradio Interface
-# iface = gr.Interface(
-#     fn=process_video,
-#     inputs=gr.Video(label="Upload a Video"),
-#     outputs=gr.Video(label="Processed Video"),
-#     title="Video Frame Processor",
-#     description="Upload a video to split it into frames, process the first and last frame, and return the video."
-# )
-
-# # Define Gradio Interface
-# iface = gr.Interface(
-#     fn=show_preds_video,
-#     inputs=gr.Video(label="Upload a Video"),
-#     outputs=[gr.components.Image(type="numpy", label="Output Image")],
-#     cache_examples=False,
-#     title="Video Frame Processor",
-#     description="Upload a video to split it into frames, process the first and last frame, and return the video."
-# )
-
-
-# if __name__ == "__main__":
-#     iface.launch(share=True)
-
-
-# # import spaces
-import gradio as gr
-import cv2
-import numpy as np
+import gdown
+import zipfile
+# import spaces
 import time
 import random
+import gradio as gr
+import numpy as np
 from PIL import Image
-# from transparent_background import Remover
-
-# @spaces.GPU()
-def doo(video, mode, progress=gr.Progress()):
+from pathlib import Path
+sys.path.insert(1, "MEMTrack/src")
+from data_prep_utils import process_data
+from data_feature_gen import create_train_data, create_test_data
+from inferenceBacteriaRetinanet_Motility_v2 import run_inference
+from GenerateTrackingData import gen_tracking_data
+from Tracking import track_bacteria
+from TrackingAnalysis import analyse_tracking
+from GenerateVideo import gen_tracking_video
+
+def find_and_return_csv_files(folder_path, search_pattern):
+    search_pattern = f"{folder_path}/{search_pattern}*.csv"
+    csv_files = list(glob.glob(search_pattern))
+    return csv_files
+
+def read_video(video, raw_frame_dir, progress=gr.Progress()):
+    # read video and save frames
+    video_dir  = str(random.randint(111111111, 999999999))
+    images_dir = "Images without Labels"
+    frames_dir = os.path.join(raw_frame_dir, video_dir, images_dir)
+    os.makedirs(frames_dir, exist_ok=True)
+    count = 0
+    frames = []
+    
     cap = cv2.VideoCapture(video)
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # Get total frames
-    writer = None
-    tmpname = random.randint(111111111, 999999999)
     processed_frames = 0
 
     while cap.isOpened():
@@ -169,28 +46,133 @@ def doo(video, mode, progress=gr.Progress()):
         if ret is False:
             break
             
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        img = Image.fromarray(frame).convert('RGB')
-
-        if writer is None:
-            writer = cv2.VideoWriter(str(tmpname) + '.mp4', cv2.VideoWriter_fourcc(*'mp4v'), cap.get(cv2.CAP_PROP_FPS), img.size)
-
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        frame_path = os.path.join(frames_dir, f"{count}.jpg")
+        cv2.imwrite(frame_path, frame)
+        frames.append(frame)
+        count += 1
         processed_frames += 1
         print(f"Processing frame {processed_frames}")
-        progress(processed_frames / total_frames, desc=f"Processing frame {processed_frames}/{total_frames}")
-        out = img 
-        writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))
+        progress(processed_frames / total_frames, desc=f"Reading frame {processed_frames}/{total_frames}")
+        
 
     cap.release()
-    writer.release()
-    return str(tmpname) + '.mp4'
+    return video_dir
+    
+def download_and_unzip_google_drive_file(file_id, output_path, unzip_path):
+    url = f'https://drive.google.com/uc?id={file_id}'
+    url="https://drive.usercontent.google.com/download?id=1agsLD5HV_VmDNpDhjHXTCAVmGUm2IQ6p&export=download&&confirm=t"
+    gdown.download(url, output_path, quiet=False, )
 
-title = "🎞️ Video Tool 🎥"
+    with zipfile.ZipFile(output_path, 'r') as zip_ref:
+        zip_ref.extractall(unzip_path)    
+    
+
+# @spaces.GPU()
+def doo(video, progress=gr.Progress()):
+    # download and unzip models
+    file_id = '1agsLD5HV_VmDNpDhjHXTCAVmGUm2IQ6p'
+    output_path = 'models.zip'
+    unzip_path = './'
+
+    download_and_unzip_google_drive_file(file_id, output_path, unzip_path)
+
+    # Initialize paths and variables
+    raw_frame_dir = "raw_data/" # Path to raw videos vefore processing (same format as sample data)
+    final_data_dir = "data" # root directory to store processed videos
+    out_sub_dir = "bacteria" # sub directory to store processed videos
+    target_data_sub_dir = os.path.join(final_data_dir, out_sub_dir)
+    feature_dir = "DataFeatures" # directory to store processed videos
+    test_video_list = ["video1"] # list of videos to generate features for
+    exp_name = "collagen_motility_inference" # name of experiment
+    feature_data_path = os.path.join(feature_dir, exp_name)
+    
+    
+#     #path to saved models
+#     no_motility_model_path = "models/motility/no/collagen_optical_flow_median_bkg_more_data_90k/"
+#     low_motility_model_path = "models/motility/low/collagen_optical_flow_median_bkg_more_data_90k/"
+#     mid_motility_model_path = "models/motility/mid/collagen_optical_flow_median_bkg_more_data_90k/"
+#     high_motility_model_path = "models/motility/high/collagen_optical_flow_median_bkg_more_data_90k/"
+    
+    
+#     # Clear previous results and data
+#     if os.path.exists(final_data_dir):
+#         shutil.rmtree(final_data_dir)
+#     if os.path.exists(raw_frame_dir):
+#         shutil.rmtree(raw_frame_dir)
+#     if os.path.exists(feature_dir):
+#         shutil.rmtree(feature_dir)
+    
+#     # Read video and store frames separately for object detection model
+#     video_dir = read_video(video, raw_frame_dir, progress=gr.Progress())   
+#     # Process raw frames and store in acceptable format
+#     progress(1 / 3, desc=f"Processing Frames {1}/{3}")
+#     video_num = process_data(video_dir, raw_frame_dir, final_data_dir, out_sub_dir)    
+#     progress(3 / 3, desc=f"Processing Frames {3}/{3}")
+#     # generate features for raw frames for the object detector model
+#     progress(1 / 3, desc=f"Generating Features {1}/{3}")
+#     create_test_data(target_data_sub_dir, feature_dir, exp_name, test_video_list)
+#     progress(3 / 3, desc=f"Features Generated {3}/{3}")
+    
+
+#     progress(1 / 3, desc=f"Loading Models {1}/{3}")
+#     # Run Object Detection Code
+#     for video_num in [1]:
+#         #To genearate testing files for all motilities
+#         run_inference(video_num=video_num, output_dir=no_motility_model_path, 
+#                       annotations_test="All", test_dir=feature_data_path, register_dataset=True)
+#         progress(3 / 3, desc=f"Models Loaded{3}/{3}")
+#         run_inference(video_num=video_num, output_dir=mid_motility_model_path, 
+#                       annotations_test="Motility-mid", test_dir=feature_data_path, register_dataset=False)
+#         progress(1 / 3, desc=f"Running Bacteria Detection {1}/{3}")
+        
+#         run_inference(video_num=video_num, output_dir=high_motility_model_path, 
+#                       annotations_test="Motility-high", test_dir=feature_data_path, register_dataset=False)
+#         progress(2 / 3, desc=f"Running Bacteria Detection {2}/{3}")
+        
+#         run_inference(video_num=video_num, output_dir=low_motility_model_path, 
+#                       annotations_test="Motility-low", test_dir=feature_data_path, register_dataset=False)
+#         progress(3 / 3, desc=f"Running Bacteria Detection {3}/{3}")
+        
+#     # Tracking where GT is present
+#     progress(0 / 3, desc=f"Tracking {0}/{3}")
+    for video_num in [1]:
+#         gen_tracking_data(video_num=video_num, data_path=feature_data_path, filter_thresh=0.3)
+#         progress(1 / 3, desc=f"Tracking {1}/{3}")
+#         track_bacteria(video_num=video_num, max_age=35, max_interpolation=35, data_path=feature_data_path)
+#         progress(2 / 3, desc=f"Tracking {2}/{3}")
+        folder_path = analyse_tracking(video_num=video_num, data_feature_path=feature_data_path, data_root_path=final_data_dir, plot=True)
+        progress(3 / 3, desc=f"Tracking {3}/{3}")
+        output_video = gen_tracking_video(video_num=video_num, fps=60, data_path=feature_data_path)
+        final_video = os.path.basename(output_video)
+        shutil.copy(output_video, final_video)
+        print(output_video)
+        print(final_video)
+        
+        search_pattern = "TrackedRawData"
+        tracking_preds = find_and_return_csv_files(folder_path, search_pattern)
+         
+ 
+    return final_video, tracking_preds #str(tmpname) + '.mp4'
+
+
+examples = [['./sample_videos/control_4_h264.mp4']]
+
+title = "🎞️ MEMTrack Bacteria Tracking Video Tool"
+description = "Upload a video or selct from example to track. <br><br> If the input video does not play on browser, ensure its in a browser accetable format. Output will be generated iirespective of playback on browser. Refer: https://colab.research.google.com/drive/1U5pX_9iaR_T8knVV7o4ftKdDoGndCdEM?usp=sharing"
 
 iface = gr.Interface(
     fn=doo,
-    inputs="video",
-    outputs="video",
-    title=title
+    inputs=gr.Video(label="Input Video"),
+    outputs=[
+        gr.Video(label="Tracked Video"),
+        gr.File(label="CSV Data")
+    ],
+    examples=examples,
+    title=title,
+    description=description
 )
-iface.launch()
+
+
+if __name__ == "__main__":
+    iface.launch(share=True)

requirements.txt

@@ -2,4 +2,38 @@ gradio
 torch
 torchvision
 opencv-python
-tqdm
+tqdm
+pip==22.0.3
+numpy
+filterpy==1.4.5
+json5==0.9.6
+jsonlines==3.0.0
+jsonpatch==1.32
+jsonpointer==2.1
+jsonschema==3.2.0
+jupyter-client==7.0.3
+jupyter-core==4.8.1
+jupyter-server==1.11.0
+jupyterlab==3.1.13
+jupyterlab-pygments==0.1.2
+jupyterlab-server==2.8.1
+jupyterlab-widgets==1.0.2
+lap==0.4.0
+matplotlib==3.3.4
+motmetrics==1.2.5
+natsort==7.1.1
+numpy==1.19.5
+opencv-python==4.5.5.64
+opencv-python-headless==4.5.4.60
+openpyxl
+pandas==1.1.5
+plotly==5.11.0
+scikit-image==0.17.2
+scikit-learn==0.24.2
+scipy==1.5.4
+seaborn==0.11.2
+torch==1.9.1
+torchfile==0.1.0
+torchmetrics==0.5.1
+torchvision==0.10.1
+tqdm==4.62.3