add timer and save video and modify readme

Author: JackWoo0831

.gitignore

@@ -10,7 +10,7 @@ runs/*
 wandb/*
 track_result.txt
 .idea/
-tracker/results/*
+track_results/*
 *.mp4
 *.mkv
 temp.py

README.md

@@ -8,11 +8,16 @@ Compared to the previous version, this is an ***entirely new version (branch v2)
 
 ```bash 
 git clone https://github.com/JackWoo0831/Yolov7-tracker.git
+git checkout v2  # change to v2 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.
 
+<div align="center">
 
+**Language**: English | [简体中文](README_CN.md)
+
+</div>
 
 ## ❤️ Introduction
 
@@ -47,8 +52,8 @@ The highlights are:
 ## 🗺️ Roadmap
 
 - [ ] Add StrongSort
-- [ ] Add save video function
-- [ ] Add timer function to calculate fps
+- [ x ] Add save video function
+- [ x ] Add timer function to calculate fps
 
 ##  🔨 Installation
 

README_CN.md

@@ -8,6 +8,7 @@
 
 ```bash 
 git clone https://github.com/JackWoo0831/Yolov7-tracker.git
+git checkout v2  # change to v2 branch !!
 ```
 
 🙌 ***如果您有任何关于添加跟踪器的建议***，请在Issues部分留言并附上论文标题或链接！欢迎大家一起来让这个repo变得更好
@@ -47,8 +48,8 @@ REID模型支持：
 ## 🗺️ 路线图
 
 - [ ] Add StrongSort
-- [ ] Add save video function
-- [ ] Add timer function to calculate fps
+- [ x ] Add save video function
+- [ x ] Add timer function to calculate fps
 
 ##  🔨 安装
 

tracker/timer.py renamed to tracker/my_timer.py

@@ -14,7 +14,8 @@
 
 from tracking_utils.envs import select_device
 from tracking_utils.tools import *
-from tracking_utils.visualization import plot_img
+from tracking_utils.visualization import plot_img, save_video
+from my_timer import Timer
 
 from tracker_dataloader import TestDataset
 
@@ -187,6 +188,11 @@ def main(args, dataset_cfgs):
 
 
     """4. Tracking"""
+
+    # set timer 
+    timer = Timer()
+    seq_fps = []
+
     for seq in seqs:
         logger.info(f'--------------tracking seq {seq}--------------')
 
@@ -202,6 +208,10 @@ def main(args, dataset_cfgs):
         results = []
 
         for frame_idx, (ori_img, img) in process_bar:
+
+            # start timing this frame
+            timer.tic()
+
             if args.detector == 'yolov8':
                 img = img.squeeze(0).cpu().numpy()
 
@@ -257,13 +267,28 @@ def main(args, dataset_cfgs):
 
             results.append((frame_idx + 1, cur_id, cur_tlwh, cur_cls, cur_score))
 
+            timer.toc()
+
             if args.save_images:
                 plot_img(img=ori_img, frame_id=frame_idx, results=[cur_tlwh, cur_id, cur_cls], 
                          save_dir=os.path.join(save_dir, 'vis_results'))
 
-        save_results(folder_name=os.path.join(args.datasets, SPLIT), 
+        save_results(folder_name=os.path.join(args.dataset, SPLIT), 
                      seq_name=seq, 
                      results=results)
+        
+        # show the fps
+        seq_fps.append(frame_idx / timer.total_time)
+        logger.info(f'fps of seq {seq}: {seq_fps[-1]}')
+        timer.clear()
+        
+        if args.save_videos:
+            save_video(images_path=os.path.join(save_dir, 'vis_results'))
+            logger.info(f'save video of {seq} done')
+
+    # show the average fps
+    logger.info(f'average fps: {np.mean(seq_fps)}')
+
 
 if __name__ == '__main__':
 

tracker/track.py

@@ -14,7 +14,7 @@
 
 from tracking_utils.envs import select_device
 from tracking_utils.tools import *
-from tracking_utils.visualization import plot_img
+from tracking_utils.visualization import plot_img, save_video
 
 from tracker_dataloader import TestDataset, DemoDataset
 
@@ -254,6 +254,10 @@ def main(args):
     save_results(folder_name=os.path.join(save_dir, 'txt_results'), 
                     seq_name='demo', 
                     results=results)
+    
+    if args.save_videos:
+        save_video(images_path=os.path.join(save_dir, 'vis_results'))
+        logger.info(f'save video done')
 
 if __name__ == '__main__':
 

tracker/track_demo.py

@@ -1,6 +1,7 @@
 import cv2 
 import os 
 import numpy as np 
+from PIL import Image
 
 def plot_img(img, frame_id, results, save_dir):
     """
@@ -43,4 +44,21 @@ def get_color(idx):
     idx = idx * 3
     color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)
 
-    return color
+    return color
+
+def save_video(images_path):
+    """
+    save images (frames) to a video
+    """
+
+    images_list = sorted(os.listdir(images_path))
+    save_video_path = os.path.join(images_path, images_path.split('/')[-1] + '.mp4')
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+
+    img0 = Image.open(os.path.join(images_path, images_list[0]))
+    vw = cv2.VideoWriter(save_video_path, fourcc, 15, img0.size)
+
+    for image_name in images_list:
+        image = cv2.imread(filename=os.path.join(images_path, image_name))
+        vw.write(image)