fix bug of yolo fromat data and generate videos

Author: JackWoo0831

tracker/track.py

@@ -36,10 +36,13 @@
 
 DATASET_ROOT = '/data/wujiapeng/datasets/VisDrone2019/VisDrone2019'  # your dataset root
 
-# CATEGORY_NAMES = ['car', 'van', 'truck', 'bus']
-CATEGORY_NAMES = ['pedestrain', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor']
+CATEGORY_NAMES = ['car', 'van', 'truck', 'bus']
+# CATEGORY_NAMES = ['pedestrain', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor']
 CATEGORY_DICT = {i: CATEGORY_NAMES[i] for i in range(len(CATEGORY_NAMES))}  # show class
 
+# IGNORE_SEQS = []
+IGNORE_SEQS = ['uav0000073_00600_v', 'uav0000088_00290_v']  # ignore seqs
+
 timer = Timer()
 seq_fps = []  # list to store time used for every seq
 def main(opts):
@@ -59,7 +62,10 @@ def main(opts):
     elif opts.tracker == 'strongsort':
         opts.kalman_format = 'strongsort'
 
-        
+    # NOTE: if save video, you must save image
+    if opts.save_videos:
+        opts.save_images = True
+
     """
     1. load model
     """
@@ -82,15 +88,17 @@ def main(opts):
         with open(f'./{opts.dataset}/test.txt', 'r') as f:
             lines = f.readlines()
             for line in lines:
-                if line[-2] not in seqs:
-                    seqs.append(line[-2])
+                elems = line.split('/')  # devide path by / in order to get sequence name(elems[-2])
+                if elems[-2] not in seqs:
+                    seqs.append(elems[-2])
 
     elif opts.data_format == 'origin':
         DATA_ROOT = os.path.join(DATASET_ROOT, 'VisDrone2019-MOT-test-dev/sequences')
         seqs = os.listdir(DATA_ROOT)
     else:
         raise NotImplementedError
     seqs = sorted(seqs)
+    seqs = [seq for seq in seqs if seq not in IGNORE_SEQS]
     print(f'Seqs will be evalueated, total{len(seqs)}:')
     print(seqs)
 
@@ -105,8 +113,9 @@ def main(opts):
     for seq in seqs:
         print(f'--------------tracking seq {seq}--------------')
 
-        path = os.path.join(DATA_ROOT, seq) if opts.data_format == 'origin' else seq
-        loader = tracker_dataloader.TrackerLoader(path, opts.img_size, opts.data_format)
+        path = os.path.join(DATA_ROOT, seq) if opts.data_format == 'origin' else os.path.join('./', f'{opts.dataset}', 'test.txt')
+
+        loader = tracker_dataloader.TrackerLoader(path, opts.img_size, opts.data_format, seq)
 
         data_loader = torch.utils.data.DataLoader(loader, batch_size=1)
 
@@ -120,27 +129,33 @@ def main(opts):
             pbar.update()
             timer.tic()  # start timing this img
 
-            out = model(img.to(device))  # model forward 
+            if not i % opts.detect_per_frame:  # if it's time to detect
+
+                out = model(img.to(device))  # model forward             
+                out = out[0]  # NOTE: for yolo v7
+            
+                if len(out.shape) == 3:  # case (bs, num_obj, ...)
+                    # out = out.squeeze()
+                    # NOTE: assert batch size == 1
+                    out = out.squeeze(0)
+                    img0 = img0.squeeze(0)
+                # remove some low conf detections
+                out = out[out[:, 4] > 0.001]
+                
 
-            out = out[0]  # NOTE: for yolo v7
-          
-            if len(out.shape) == 3:  # case (bs, num_obj, ...)
-                # out = out.squeeze()
-                # NOTE: assert batch size == 1
-                out = out.squeeze(0)
-                img0 = img0.squeeze(0)
-            # remove some low conf detections
-            out = out[out[:, 4] > 0.001]
+                # NOTE: yolo v7 origin out format: [xc, yc, w, h, conf, cls0_conf, cls1_conf, ..., clsn_conf]
+                if opts.det_output_format == 'yolo':
+                    cls_conf, cls_idx = torch.max(out[:, 5:], dim=1)
+                    # out[:, 4] *= cls_conf  # fuse object and cls conf
+                    out[:, 5] = cls_idx
+                    out = out[:, :6]
 
-          
-            # NOTE: yolo v7 origin out format: [xc, yc, w, h, conf, cls0_conf, cls1_conf, ..., clsn_conf]
-            if opts.det_output_format == 'yolo':
-                cls_conf, cls_idx = torch.max(out[:, 5:], dim=1)
-                # out[:, 4] *= cls_conf  # fuse object and cls conf
-                out[:, 5] = cls_idx
-                out = out[:, :6]
-           
-            current_tracks = tracker.update(out, img0)  # List[class(STracks)]
+                current_tracks = tracker.update(out, img0)  # List[class(STracks)]
+            else:  # otherwize
+                # make the img shape (bs, C, H, W) as (C, H, W)
+                if len(img0.shape) == 4:
+                    img0 = img0.squeeze(0)
+                current_tracks = tracker.update_without_detection(None, img0)
 
             # save results
             cur_tlwh, cur_id, cur_cls = [], [], []
@@ -171,19 +186,17 @@ def main(opts):
         # every time assign a different name
         save_results(folder_name, seq, results)
 
+        ## finally, save videos
+        if opts.save_images and opts.save_videos:
+            save_videos(seq_names=seq)
+            
     """
     3. evaluate results
     """
     print(f'average fps: {np.mean(seq_fps)}')
     evaluate(sorted(os.listdir(f'./tracker/results/{folder_name}')), 
                 sorted([seq + '.txt' for seq in seqs]), data_type='visdrone', result_folder=folder_name)  
 
-    """
-    4. save videos
-    """
-    if opts.save_videos:
-        save_videos(seq_names='uav0000119_02301_v')
-
 
 def save_results(folder_name, seq_name, results, data_type='default'):
     """
@@ -283,15 +296,14 @@ def get_color(idx):
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
 
-    parser.add_argument('--dataset', type=str, default='visdrone', help='visdrone or mot')
+    parser.add_argument('--dataset', type=str, default='visdrone', help='visdrone, or mot')
     parser.add_argument('--data_format', type=str, default='origin', help='format of reading dataset')
     parser.add_argument('--det_output_format', type=str, default='yolo', help='data format of output of detector, yolo or other')
 
     parser.add_argument('--tracker', type=str, default='bytetrack', help='sort, deepsort, etc')
 
     parser.add_argument('--model_path', type=str, default=None, help='model path')
 
-    parser.add_argument('--trace', action='store_true', help='trace model')
     parser.add_argument('--img_size', nargs='+', type=int, default=[1280, 1280], help='[train, test] image sizes')
 
     """For tracker"""
@@ -312,6 +324,9 @@ def get_color(idx):
 
     parser.add_argument('--save_images', action='store_true', help='save tracking results (image)')
     parser.add_argument('--save_videos', action='store_true', help='save tracking results (video)')
+
+    # detect per several frames
+    parser.add_argument('--detect_per_frame', type=int, default=1, help='choose how many frames per detect')
 
 
     opts = parser.parse_args()