update track_demo and readme

Author: JackWoo0831

README.md

@@ -2,16 +2,19 @@
 
 ## ❗❗Important Notes
 
-Compared to the previous version, this is an ***entirely new version (branch v2)***!!!
+There has been a major update recently, and I **re-organized** all the codes for accuracy and readability. More importantly, **two new sota trackers are added** (ImproAssoc and TrackTrack), and **TensorRT** engine is supported!
 
-**Please use this version directly, as I have almost rewritten all the code to ensure better readability and improved results, as well as to correct some errors in the past code.**
+The new version is on **branch v2.1**:
 
 ```bash 
 git clone https://github.com/JackWoo0831/Yolov7-tracker.git
-git checkout v2  # change to v2 branch !!
+git checkout v2.1  # change to v2.1 branch !!
 ```
 
-🙌 ***If you have any suggestions for adding trackers***, please leave a comment in the Issues section with the paper title or link! Everyone is welcome to contribute to making this repo better.
+🙌 ***The QQ Group is established and welcome to join!*** You can raise bugs, suggestions, or work together on interesting CV/AI projects in the QQ group!
+However, bugs or issues should still be prioritized in the **Issue section in Github** for others to see.
+
+<img src="figure/GroupQRcode.jpg" alt="group" style="width:40%;">
 
 <div align="center">
 
@@ -21,9 +24,24 @@ git checkout v2  # change to v2 branch !!
 
 ## 🗺️ Latest News
 
-- ***2025.4.14*** Fix some minor bugs [issue#144](https://github.com/JackWoo0831/Yolov7-tracker/issues/144), and fix the lost tracklet bugs in sort.
-- ***2025.4.7*** Add more Re-ID modules (ShuffleNet, VehicleNet, MobileNet), fix some bugs (such as abandon bbox aspect ratio updating if the tracklet is not activated), and add some functions (customized low filter threshold, fuse detection score, etc.)
-- ***2025.4.3*** Support the newest ultralytics version (YOLO v3 ~ v12) and fix some bugs of hybrid sort.
+- ***2025.7.8*** New version 2.1 released. Add ImproAssoc, TrackTrack and support TensorRT. The other details are as follows:
+
+<details>
+<summary>Update details</summary>
+
+
+1. Re annotate and organize all functions in `matching.py`
+2. For camera motion compensation, custom feature extraction algorithms (SIFT, ORB, ECC) can be used, and the `--cmc_method parameter` can be specified when running `track.py` (or `track_demo.py`).
+3. For methods such as BoT SORT and ByteTrack, the original low confidence screening threshold is fixed at 0.1 You can now manually set the `--conf_thresh_low` parameter when running `track.py`.
+4. Add the `init_thresh` parameter as the initialization target threshold, abandoning the original `args.conf + 0.1` setting. Specify the `--init_thresh` parameter when running `track.py`.
+5. In ReID feature extraction, the original crop size was a fixed value of `(h, w) = (128, 64)`, which can now be manually set. When running `track.py`, specify the `--reid_crop_size` parameter, for example, `--reid_crop_size 32 64`.
+6. Inherit all Trackers from the BaseTracker class to achieve good code reuse
+7. Fix the reid similarity calculation bug in Strongsort
+8. Abandon cython.bbox for better compatibility with numpy versions
+9. Abandon np.float, etc. for better compatibility with numpy versions
+10. Reorganize requirements.txt
+</details>
+
 
 ## ❤️ Introduction
 
@@ -45,7 +63,9 @@ and the tracker supports:
 - Strong SORT ([IEEE TMM 2023](https://arxiv.org/pdf/2202.13514))
 - Sparse Track ([arxiv 2306](https://arxiv.org/pdf/2306.05238))
 - UCMC Track ([AAAI 2024](http://arxiv.org/abs/2312.08952))
-- Hybrid SORT([AAAI 2024](https://ojs.aaai.org/index.php/AAAI/article/view/28471))
+- Hybrid SORT ([AAAI 2024](https://ojs.aaai.org/index.php/AAAI/article/view/28471))
+- ImproAssoc ([CVPRW 2023](https://openaccess.thecvf.com/content/CVPR2023W/E2EAD/papers/Stadler_An_Improved_Association_Pipeline_for_Multi-Person_Tracking_CVPRW_2023_paper.pdf))
+- TrackTrack ([CVPR 2025](https://openaccess.thecvf.com/content/CVPR2025/html/Shim_Focusing_on_Tracks_for_Online_Multi-Object_Tracking_CVPR_2025_paper.html))
 
 and the reid model supports:
 
@@ -188,6 +208,12 @@ For example:
 python tracker/track_demo.py --obj M0203.mp4 --detector yolo_ultra_v8 --tracker deepsort --kalman_format byte --detector_model_path weights/yolov8l_UAVDT_60epochs_20230509.pt --save_images
 ```
 
+or
+
+```bash
+python tracker/track_demo.py --obj /root/datasets/visdrone/images/val/seq/ --detector yolox --tracker bytetrack --kalman_format byte --detector_model_path weights/yolox_m_VisDrone_55epochs_20230509.pth.tar --yolox_exp_file ./tracker/yolox_utils/yolox_m.py --save_images
+```
+
 **If you want to run trackers on dataset**:
 
 ```bash
@@ -240,6 +266,19 @@ In addition, you can also specify
 > 
 > 2. The code does not contain the camera motion compensation part between every two frame, please refer to [https://github.com/corfyi/UCMCTrack/issues/12](https://github.com/corfyi/UCMCTrack/issues/12). From my perspective, since the algorithm name is 'uniform', the update of compensation between every two frames is not necessary.
 
+### ✨ TensorRT Convert and Inference
+
+This code supports **fully automatic** generation and reasoning of Tensor RT engine, **which can be used for both detection model and ReID model**. If you have not converted Tensor RT engine, just add `--trt` parameter when running, for example:
+
+```bash
+python tracker/track.py --dataset mot17 --detector yolox --tracker ocsort --kalman_format ocsort --detector_model_path weights/bytetrack_m_mot17.pth.tar --reid --reid_model shufflenet_v2_x1_0 --reid_model_path shufflenetv2_x1-5666bf0f80.pth --trt
+```
+
+If you already have engine, just write the relevant path as engine, and parameter `--trt` can be omitted:
+
+```bash
+python tracker/track.py --dataset visdrone_part --detector b8_ultra --tracker deepsort --kalman_format byte --detector_model_path weights/yolov8l_VisDroneDet_35epochs_20230605.engine --reid deepsort --reid_model_path weights/ckpt.engine
+```
 
 ### ✅ Evaluation 
 

README_CN.md

@@ -2,22 +2,39 @@
 
 ## ❗❗重要提示
 
-与之前的版本相比，这是一个***全新的版本（分支v2）***！！！
+最近有一个重大更新, 为了准确性和可读性, 我**重新组织了**所有代码. 更重要的是, **添加了两个新的sota跟踪器**(ImpraAssoc和TrackTrack), 并支持**TensorRT engine**.
 
-**请直接使用这个版本，因为我几乎重写了所有代码，以确保更好的可读性和改进的结果，并修正了以往代码中的一些错误。**
+新版本发布于**branch v2.1**:
 
 ```bash 
 git clone https://github.com/JackWoo0831/Yolov7-tracker.git
-git checkout v2  # change to v2 branch !!
+git checkout v2.1  # change to v2.1 branch !!
 ```
 
-🙌 ***如果您有任何关于添加跟踪器的建议***，请在Issues部分留言并附上论文标题或链接！欢迎大家一起来让这个repo变得更好
+🙌 ***QQ交流群已建立, 欢迎加入!***，您可以在QQ群中提出bug、意见建议或者一起来做有趣的CV/AI项目！
+然而，**Bug或问题还是优先在issue区提出，以便让更多人看到.**
+
+<img src="figure/GroupQRcode.jpg" alt="group" style="width:40%;">
 
 ## 🗺️ 最近更新
 
-- ***2025.4.14*** 修复[issue#144](https://github.com/JackWoo0831/Yolov7-tracker/issues/144)中提到的一些bug，修复sort对丢失轨迹处理的bug.
-- ***2025.4.7*** 增加更多Re-ID模型 (ShuffleNet, VehicleNet, MobileNet), 修复一些bug (例如在轨迹为非活动状态时停止更新边界框长宽), 增加一些小功能 (例如可以修改两阶段关联策略的最低阈值，原来是固定的0.1; 增加将IoU和检测置信度融合的选项)
-- ***2025.4.3*** 增加了ultralytics库最新版本的支持，修复了hybrid sort中的一些bug.
+- ***2025.7.8*** 新版本2.1发布. 添加ImproAssoc, TrackTrack并支持TensorRT. 其他细节如下:
+
+<details>
+<summary>更新细节</summary>
+
+
+1. 重新注释整理matching.py中所有函数
+2. 对于相机运动补偿, 可自定义特征提取子的算法(SIFT, ORB, ECC), 运行`track.py`时指定`--cmc_method`参数.
+3. 对于BoT-SORT, ByteTrack等方法, 原先的低置信度筛选阈值被固定设置为`0.1`. 现在可以手动设置, 运行`track.py`(或`track_demo.py`)时指定`--conf_thresh_low`参数.
+4. 加入`init_thresh`参数作为初始化目标阈值, 弃用原本的`args.conf + 0.1`定值. 运行`track.py`时指定`--init_thresh`参数.
+5. 在ReID特征提取中, 原本的裁剪-resize大小为定值`(h, w) = (128, 64)`, 现在可以手动设置, 运行`track.py`时指定`--reid_crop_size`参数, 例如`--reid_crop_size 32 64`.
+6. 将所有Tracker继承BaseTracker类以实现良好的代码复用
+7. 修复strongsort的reid相似度计算bug
+8. 弃用cython_bbox以获得更好的numpy版本兼容
+9. 弃用np.float等以获得更好的numpy版本兼容
+10. 重新整理requirements.txt
+</details>
 
 
 ## ❤️ 介绍
@@ -41,6 +58,8 @@ git checkout v2  # change to v2 branch !!
 - Sparse Track ([arxiv 2306](https://arxiv.org/pdf/2306.05238))
 - UCMC Track ([AAAI 2024](http://arxiv.org/abs/2312.08952))
 - Hybrid SORT([AAAI 2024](https://ojs.aaai.org/index.php/AAAI/article/view/28471))
+- ImproAssoc ([CVPRW 2023](https://openaccess.thecvf.com/content/CVPR2023W/E2EAD/papers/Stadler_An_Improved_Association_Pipeline_for_Multi-Person_Tracking_CVPRW_2023_paper.pdf))
+- TrackTrack ([CVPR 2025](https://openaccess.thecvf.com/content/CVPR2025/html/Shim_Focusing_on_Tracks_for_Online_Multi-Object_Tracking_CVPR_2025_paper.html))
 
 REID模型支持：
 
@@ -231,12 +250,30 @@ python tracker/track.py --dataset ${dataset name, related with the yaml file} --
 
 - Hybrid SORT: `python tracker/track.py --dataset visdrone_part --detector yolo_ultra --tracker hybridsort --kalman_format hybridsort --detector_model_path weights/yolov8l_VisDrone_35epochs_20230509.pt --save_images`
 
+- ImproAssoc: `python tracker/track.py --dataset visdrone_part --detector yolo_ultra --tracker improassoc --kalman_format bot --detector_model_path weights/yolov8l_VisDrone_35epochs_20230509.pt --save_images`
+
+- TrackTrack: `python tracker/track.py --dataset visdrone_part --detector yolo_ultra --tracker tracktrack --kalman_format bot --detector_model_path weights/yolov8l_VisDrone_35epochs_20230509.pt --save_images --nms_thresh 0.95 --reid`
+
 >**UCMC Track的重要提示：**
 > 
 > 1. 相机参数. UCMC Track需要相机的内参和外参. 请按照`tracker/cam_ram_files/uavdt/M0101.txt`的格式组织. 一个视频序列对应一个txt文件. 如果您没有标记的参数, 可以参考原始仓库中的估算工具箱([https://github.com/corfyi/UCMCTrack](https://github.com/corfyi/UCMCTrack)).
 > 
 > 2. 该代码不包含每两帧之间的相机运动补偿部分, 请参阅[https://github.com/corfyi/UCMCTrack/issues/12](https://github.com/corfyi/UCMCTrack/issues/12). 在我看来, 既然算法叫"统一相机运动补偿", 因此不需要每两帧之间再更新补偿. 
 
+### ✨ TensorRT的转换与推理
+
+该代码支持**全自动**的Tensor RT engine的生成与推理, **既可以用于检测模型, 也可以用于ReID模型**. 如果您还没有转换Tensor RT engine, 只需在运行时加上`--trt`参数, 例如:
+
+```bash
+python tracker/track.py --dataset mot17 --detector yolox --tracker ocsort --kalman_format ocsort --detector_model_path weights/bytetrack_m_mot17.pth.tar --reid --reid_model shufflenet_v2_x1_0 --reid_model_path shufflenetv2_x1-5666bf0f80.pth --trt
+```
+
+如果已有engine, 则直接将相关路径写成engine, 参数`--trt`可以省略:
+
+```bash
+python tracker/track.py --dataset visdrone_part --detector b8_ultra --tracker deepsort --kalman_format byte --detector_model_path weights/yolov8l_VisDroneDet_35epochs_20230605.engine --reid deepsort --reid_model_path weights/ckpt.engine
+```
+
 ### ✅ 评估
 
 马上推出！作为备选项，你可以使用这个repo： [Easier to use TrackEval repo](https://github.com/JackWoo0831/Easier_To_Use_TrackEval).

figure/GroupQRcode.jpg

@@ -120,7 +120,6 @@ def get_args():
     parser.add_argument('--trace', type=bool, default=False, help='traced model of YOLO v7')
     # other model path
     parser.add_argument('--reid_model_path', type=str, default='./weights/osnet_x0_25.pth', help='path for reid model path')
-    parser.add_argument('--dhn_path', type=str, default='./weights/DHN.pth', help='path of DHN path for DeepMOT')
 
 
     """other options"""

tracker/track.py

@@ -65,6 +65,13 @@
     logger.warning('Load yolov8 fail. If you want to use yolov8, please check the installation.')
     pass
 
+# TensorRT
+try:
+    from accelerations.tensorrt_tools import TensorRTConverter, TensorRTInference
+except Exception as e:
+    logger.warning(e)
+    logger.warning('Load TensorRT fail. If you want to convert model to TensorRT, please install the packages.')
+
 TRACKER_DICT = {
     'sort': SortTracker, 
     'bytetrack': ByteTracker, 
@@ -94,6 +101,7 @@ def get_args():
 
     parser.add_argument('--kalman_format', type=str, default='default', help='use what kind of Kalman, sort, deepsort, byte, etc.')
     parser.add_argument('--img_size', type=int, default=1280, help='image size, [h, w]')
+    parser.add_argument('--reid_crop_size', type=int, default=[128, 64], nargs='+', help='crop size in reid model, [h, w]')
 
     parser.add_argument('--conf_thresh', type=float, default=0.2, help='filter tracks')
     parser.add_argument('--conf_thresh_low', type=float, default=0.1, help='filter low conf detections, used in two-stage association')
@@ -111,7 +119,6 @@ def get_args():
     parser.add_argument('--trace', type=bool, default=False, help='traced model of YOLO v7')
     # other model path
     parser.add_argument('--reid_model_path', type=str, default='./weights/osnet_x0_25.pth', help='path for reid model path')
-    parser.add_argument('--dhn_path', type=str, default='./weights/DHN.pth', help='path of DHN path for DeepMOT')
 
 
     """other options"""
@@ -130,6 +137,10 @@ def get_args():
 
     """camera parameter"""
     parser.add_argument('--camera_parameter_folder', type=str, default='./tracker/cam_param_files', help='folder path of camera parameter files')
+    
+    """tensorrt options"""
+    parser.add_argument('--trt', action='store_true', help='use tensorrt engine to detect and reid')
+    
     return parser.parse_args()
 
 def main(args):
@@ -156,48 +167,81 @@ def main(args):
         model.to(device)
         model.eval()
 
-        logger.info(f"loading detector {args.detector} checkpoint {args.detector_model_path}")
-        ckpt = torch.load(args.detector_model_path, map_location=device)
-        model.load_state_dict(ckpt['model'])
-        logger.info("loaded checkpoint done")
-        model = fuse_model(model)
-
-        stride = None  # match with yolo v7
-
-        logger.info(f'Now detector is on device {next(model.parameters()).device}')
+        if args.trt:  # convert trt
+            # check if need to convert
+            if not args.detector_model_path.endswith('.engine'):
+                trt_converter = TensorRTConverter(model, input_shape=[3, *model_img_size], ckpt_path=args.detector_model_path, 
+                                                min_opt_max_batch=[1, 1, 1], device=device, load_ckpt=True, ckpt_key='model')
+                trt_converter.export()
+                model = TensorRTInference(engine_path=trt_converter.trt_model, min_opt_max_batch=[1, 1, 1], device=device)
+            else:   
+                model = TensorRTInference(engine_path=args.detector_model_path, min_opt_max_batch=[1, 1, 1], device=device)
+
+        else:  # normal load
+            logger.info(f"loading detector {args.detector} checkpoint {args.detector_model_path}")
+            ckpt = torch.load(args.detector_model_path, map_location=device)
+            model.load_state_dict(ckpt['model'])
+            logger.info("loaded checkpoint done")
+            model = fuse_model(model)
+            logger.info(f'Now detector is on device {next(model.parameters()).device}')
+
+        stride = None  # match with yolo v7        
 
     elif args.detector == 'yolov7':
 
-        logger.info(f"loading detector {args.detector} checkpoint {args.detector_model_path}")
-        model = attempt_load(args.detector_model_path, map_location=device)
+        if args.trt:
+            # check if need to convert
+            stride = 32
+            model_img_size = check_img_size(args.img_size, s=32)
+            if not args.detector_model_path.endswith('.engine'):
+                model = attempt_load(args.detector_model_path, map_location=device)
+                trt_converter = TensorRTConverter(model, input_shape=[3, *model_img_size], ckpt_path=args.detector_model_path, 
+                                                min_opt_max_batch=[1, 1, 1], device=device, load_ckpt=False)
+                trt_converter.export()
+                model = TensorRTInference(engine_path=trt_converter.trt_model, min_opt_max_batch=[1, 1, 1], device=device)
+            else:
+                model = TensorRTInference(engine_path=args.detector_model_path, min_opt_max_batch=[1, 1, 1], device=device)
 
-        # get inference img size
-        stride = int(model.stride.max())  # model stride
-        model_img_size = check_img_size(args.img_size, s=stride)  # check img_size
+        else:
+            logger.info(f"loading detector {args.detector} checkpoint {args.detector_model_path}")
+            model = attempt_load(args.detector_model_path, map_location=device)
 
-        # Traced model
-        model = TracedModel(model, device=device, img_size=args.img_size)
-        # model.half()
+            # Traced model
+            model = TracedModel(model, device=device, img_size=args.img_size)
+            # model.half()
 
-        logger.info("loaded checkpoint done")
+            logger.info("loaded checkpoint done")
+            logger.info(f'Now detector is on device {next(model.parameters()).device}')
 
-        logger.info(f'Now detector is on device {next(model.parameters()).device}')
+            # get inference img size
+            stride = int(model.stride.max())  # model stride
+            model_img_size = check_img_size(args.img_size, s=stride)  # check img_size
 
     elif 'ultra' in args.detector:
 
-        logger.info(f"loading detector {args.detector} checkpoint {args.detector_model_path}")
-        model = YOLO(args.detector_model_path)
+        if args.trt:
+            # for ultralytics, we use the api provided by official ultralytics
+            # check if need to convert
+            if not args.detector_model_path.endswith('.engine'):
+                model = YOLO(args.detector_model_path)
+                model = YOLO(model.export(format="engine"))
+            else:
+                model = YOLO(args.detector_model_path)
+
+        else:
+            logger.info(f"loading detector {args.detector} checkpoint {args.detector_model_path}")
+            model = YOLO(args.detector_model_path)
+
+            logger.info("loaded checkpoint done")
 
         model_img_size = [None, None]  
         stride = None 
 
-        logger.info("loaded checkpoint done")
-
     else:
         logger.error(f"detector {args.detector} is not supprted")
         logger.error("If you want to use the yolo v8 by ultralytics, please specify the `--detector` \
-                as the string including the substring `ultra`, \
-                such as `yolo_ultra_v8` or `yolo11_ultralytics`")
+                     as the string including the substring `ultra`, \
+                     such as `yolo_ultra_v8` or `yolo11_ultralytics`")
         exit(0)
 
     """3. load sequences"""