Merge pull request #40 from kadirnar/tracker

The code architecture has been updated.

README.md

@@ -18,31 +18,59 @@ The TorchYolo library aims to support all YOLO models(like YOLOv5, YOLOv6, YOLOv
 pip install torchyolo
 ```
 ### Prediction
+First download the [default_config.yaml](https://github.com/kadirnar/torchyolo/blob/tracker/torchyolo/default_config.yaml) file.
+
 ```python
 from torchyolo import YoloHub
-predictor = YoloHub(
-  model_type="yolov5", 
-  model_path="yolov5s.pt", 
-  device='cpu', 
-  image_size=640
-)
-predictor.conf_thres = 0.25
-predictor.iou_thres = 0.45
-predictor.save = True
-predictor.show = False
-image = "data/highway.jpg"
-result = predictor.predict(image)
+
+model = YoloHub(config_path="torchyolo/default_config.yaml")
+result = model.predict(tracker=True)
+```
+
+### Configuration
+```yaml
+TRACKER_CONFIG:
+    # The name of the tracker
+    TRACKER_TYPE: NORFAIR_TRACK
+    # The path of the config file
+    CONFIG_PATH: torchyolo/configs/tracker/norfair_track.yaml
+    # The path of the model file
+    WEIGHT_PATH: osnet_x1_0_msmt17.pt
+
+
+DETECTOR_CONFIG:
+  # The name of the detector
+  DETECTOR_TYPE: yolov8 # yolov7
+  # The threshold for the IOU score
+  IOU_TH: 0.45
+  # The threshold for the confidence score
+  CONF_TH: 0.25
+  # The size of the image
+  IMAGE_SIZE: 640
+  # The path of the weight file
+  MODEL_PATH: yolov8s.pt
+  # The device to run the detector
+  DEVICE: cuda:0
+  # F16 precision
+  HALF: False
+
+
+DATA_CONFIG:
+  # The path of the input video
+  INPUT_PATH: ../test.mp4
+  # The path of the output video
+  OUTPUT_PATH: Results
+  # Save the video
+  SHOW: False 
+  # Show the video
+  SAVE: True
 ```
 
 ## Model Architecture
 ```python
 from torchyolo import YoloHub
 
-model = YoloHub(
-  model_type="yolov8", 
-  model_path="yolov8n.pt", 
-  device="cuda:0", 
-  image_size=640)
+model = YoloHub(config_path="torchyolo/default_config.yaml")
 result = model.view_model(file_format="pdf")
 ```
 
@@ -62,29 +90,13 @@ A part of the code is borrowed from [SAHI](https://github.com/obss/sahi). Many t
 
 ### Citation
 ```bibtex
-@article{li2022yolov6,
-  title={YOLOv6: A single-stage object detection framework for industrial applications},
-  author={Li, Chuyi and Li, Lulu and Jiang, Hongliang and Weng, Kaiheng and Geng, Yifei and Li, Liang and Ke, Zaidan and Li, Qingyuan and Cheng, Meng and Nie, Weiqiang and others},
-  journal={arXiv preprint arXiv:2209.02976},
-  year={2022}
-}
-```
-```bibtex
 @article{wang2022yolov7,
   title={{YOLOv7}: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors},
   author={Wang, Chien-Yao and Bochkovskiy, Alexey and Liao, Hong-Yuan Mark},
   journal={arXiv preprint arXiv:2207.02696},
   year={2022}
 }
 ```
-```bibtex
- @article{yolox2021,
-  title={YOLOX: Exceeding YOLO Series in 2021},
-  author={Ge, Zheng and Liu, Songtao and Wang, Feng and Li, Zeming and Sun, Jian},
-  journal={arXiv preprint arXiv:2107.08430},
-  year={2021}
-}
-```
 ```bibtex
 @software{glenn_jocher_2020_4154370,
   author       = {Glenn Jocher and,Alex Stoken and,Jirka Borovec and,NanoCode012 and,ChristopherSTAN and,Liu Changyu and,Laughing and,tkianai and,Adam Hogan and,lorenzomammana and,yxNONG and,AlexWang1900 and,Laurentiu Diaconu and,Marc and,wanghaoyang0106 and,ml5ah and,Doug and,Francisco Ingham and,Frederik and,Guilhen and,Hatovix and,Jake Poznanski and,Jiacong Fang and,Lijun Yu 于力军 and,changyu98 and,Mingyu Wang and,Naman Gupta and,Osama Akhtar and,PetrDvoracek and,Prashant Rai},
@@ -98,3 +110,38 @@ A part of the code is borrowed from [SAHI](https://github.com/obss/sahi). Many t
   url= {https://doi.org/10.5281/zenodo.4154370}
 }
 ```
+```bibtex
+@article{cao2022observation,
+  title={Observation-Centric SORT: Rethinking SORT for Robust Multi-Object Tracking},
+  author={Cao, Jinkun and Weng, Xinshuo and Khirodkar, Rawal and Pang, Jiangmiao and Kitani, Kris},
+  journal={arXiv preprint arXiv:2203.14360},
+  year={2022}
+}
+```
+```bibtex
+@article{zhang2022bytetrack,
+  title={ByteTrack: Multi-Object Tracking by Associating Every Detection Box},
+  author={Zhang, Yifu and Sun, Peize and Jiang, Yi and Yu, Dongdong and Weng, Fucheng and Yuan, Zehuan and Luo, Ping and Liu, Wenyu and Wang, Xinggang},
+  booktitle={Proceedings of the European Conference on Computer Vision (ECCV)},
+  year={2022}
+}
+```
+```bibtex
+@article{du2022strongsort,
+  title={Strongsort: Make deepsort great again},
+  author={Du, Yunhao and Song, Yang and Yang, Bo and Zhao, Yanyun},
+  journal={arXiv preprint arXiv:2202.13514},
+  year={2022}
+}
+```
+```bibtex
+@inproceedings{Bewley2016_sort,
+  author={Bewley, Alex and Ge, Zongyuan and Ott, Lionel and Ramos, Fabio and Upcroft, Ben},
+  booktitle={2016 IEEE International Conference on Image Processing (ICIP)},
+  title={Simple online and realtime tracking},
+  year={2016},
+  pages={3464-3468},
+  keywords={Benchmark testing;Complexity theory;Detectors;Kalman filters;Target tracking;Visualization;Computer Vision;Data Association;Detection;Multiple Object Tracking},
+  doi={10.1109/ICIP.2016.7533003}
+}
+```

torchyolo/__init__.py

@@ -1,3 +1,3 @@
 from torchyolo.predict import YoloHub
 
-__version__ = "0.2.3"
+__version__ = "0.3.0"

torchyolo/automodel.py

@@ -1,64 +1,24 @@
-from typing import Optional
+from torchyolo.utils.config_utils import get_config
 
 MODEL_TYPE_TO_MODEL_CLASS_NAME = {
     "yolov5": "Yolov5DetectionModel",
-    "yolov6": "Yolov6DetectionModel",
     "yolov7": "Yolov7DetectionModel",
     "yolov8": "Yolov8DetectionModel",
-    "yolox": "YoloxDetectionModel",
 }
 
 
 class AutoDetectionModel:
     def from_pretrained(
-        model_type: str,
-        model_path: Optional[str] = None,
-        config_path: Optional[str] = None,
-        device: Optional[str] = None,
-        confidence_threshold: float = 0.3,
-        iou_threshold: float = 0.5,
-        image_size: int = None,
-        **kwargs,
+        config_path: str,
     ):
-        """
-        Loads a DetectionModel from given path.
-        Args:
-            model_type: str
-                Name of the detection framework (example: "yolov5", "mmdet", "detectron2")
-            model_path: str
-                Path of the detection model (ex. 'model.pt')
-            config_path: str
-                Path of the config file (ex. 'mmdet/configs/cascade_rcnn_r50_fpn_1x.py')
-            device: str
-                Device, "cpu" or "cuda:0"
-            mask_threshold: float
-                Value to threshold mask pixels, should be between 0 and 1
-            confidence_threshold: float
-                All predictions with score < confidence_threshold will be discarded
-            category_mapping: dict: str to str
-                Mapping from category id (str) to category name (str) e.g. {"1": "pedestrian"}
-            category_remapping: dict: str to int
-                Remap category ids based on category names, after performing inference e.g. {"car": 3}
-            load_at_init: bool
-                If True, automatically loads the model at initalization
-            image_size: int
-                Inference input size.
-        Returns:
-            Returns an instance of a DetectionModel
-        Raises:
-            ImportError: If given {model_type} framework is not installed
-        """
 
+        config = get_config(config_path)
+        model_type = config.DETECTOR_CONFIG.DETECTOR_TYPE
         model_class_name = MODEL_TYPE_TO_MODEL_CLASS_NAME[model_type]
         DetectionModel = getattr(
             __import__(f"torchyolo.modelhub.{model_type}", fromlist=[model_class_name]), model_class_name
         )
+
         return DetectionModel(
-            model_path=model_path,
             config_path=config_path,
-            device=device,
-            confidence_threshold=confidence_threshold,
-            iou_threshold=iou_threshold,
-            image_size=image_size,
-            **kwargs,
         )

torchyolo/configs/tracker/byte_track.yaml

@@ -0,0 +1,7 @@
+BYTE_TRACK:
+  # The name of the sort
+  TRACK_BUFFER: 25
+  # The buffer for the track
+  FRAME_RATE: 30
+  # The frame rate of the video
+

torchyolo/configs/tracker/norfair_track.yaml

@@ -0,0 +1,10 @@
+NORFAIR_TRACK:
+  DISTANCE_FUNCTION: "frobenius" # mean_manhattan, mean_euclidean, iou, iou_opt
+  DISTANCE_THRESHOLD: 500
+  HIT_COUNTER_MAX: 15
+  INITIALIZATION_DELAY: null
+  POINTWISE_HIT_COUNTER_MAX: 4
+  DETECTION_THRESHOLD: 0
+  PAST_DETECTIONS_LENGTH: 4
+  REID_DISTANCE_THRESHOLD: 0
+  REID_HIT_COUNTER_MAX: null

torchyolo/configs/tracker/oc_sort.yaml

@@ -0,0 +1,18 @@
+OC_SORT:
+  # The name of the sort
+  CONF_THRESHOLD: 0.05 
+  # The threshold for the confidence score
+  IOU_THRESHOLD: 0.3
+  # The threshold for the IOU score
+  MAX_AGE: 30
+  # The maximum age of the track
+  MIN_HITS: 3
+  # The minimum number of hits for the track
+  DELTA_T: 3
+  # The time interval between two frames
+  ASSO_FUNC: "iou"  # giou, ciou, diou, ct_dist
+  # The association function
+  INERTIA: 0.2
+  # The inertia of the track
+  USE_BYTE: False
+  # Whether to use byte as the unit of the bounding box

torchyolo/configs/tracker/sort_track.yaml

@@ -0,0 +1,6 @@
+SORT_TRACK:
+  # The name of the sort
+  MAX_AGE: 1
+  # The maximum number of frames to keep alive a track without associated detections
+  MIN_HITS: 3
+  # The minimum number of associated detections before track initialization

torchyolo/configs/tracker/strong_sort.yaml

@@ -0,0 +1,9 @@
+STRONG_SORT:
+  ECC: True              # activate camera motion compensation
+  MC_LAMBDA: 0.995       # matching with both appearance (1 - MC_LAMBDA) and motion cost
+  EMA_ALPHA: 0.9         # updates  appearance  state in  an exponential moving average manner
+  MAX_DIST: 0.2          # The matching threshold. Samples with larger distance are considered an invalid match
+  MAX_IOU_DISTANCE: 0.7  # Gating threshold. Associations with cost larger than this value are disregarded.
+  MAX_AGE: 30            # Maximum number of missed misses before a track is deleted
+  N_INIT: 3              # Number of frames that a track remains in initialization phase
+  NN_BUDGET: 100         # Maximum size of the appearance descriptors gallery