modified for site test

tracking/utils/plotting.py

@@ -1,4 +1,4 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics YOLO 🚀, AGPL-3.0 license
 
 import contextlib
 import math
@@ -284,5 +284,59 @@ def boxing_img(det, img, line_width=3):
     imgx = annotator.result()
     
     return imgx
+
+def draw_tracking_boxes(imgs, tracks, scale=2):
+    '''tracks: [x1, y1, x2, y2, track_id, score, cls, frame_index, box_index]
+                0   1   2   3      4        5     6       7           8
+       关键：
+         (1) imgs中的次序和 track 中的 fid 对应
+         (2) img 尺度小对于xyxy减半
+       
+    '''
+    
+    def array2list(bboxes):
+        track_fids = np.unique(bboxes[:, 7].astype(int))
+        track_fids.sort()
+        
+        lboxes = []
+        for f_id in track_fids:
+            # print(f"The ID is: {t_id}")
+            idx = np.where(bboxes[:, 7] == f_id)[0]
+            box = bboxes[idx, :]  
+            lboxes.append(box)
+            
+            assert len(set(box[:, 4])) == len(box), "Please check!!!"
+            
+        return lboxes
+    
+    bboxes = array2list(tracks)
+    
+    if len(bboxes)!=len(imgs):
+        return []
+    
+    subimgs = []
+    for i, boxes in enumerate(bboxes):
+        annotator = Annotator(imgs[i].copy())
+        for *xyxy, tid, conf, cls, fid, bid in boxes:
+            label = f'id:{int(tid)}_{int(cls)}_{conf:.2f}'
+            
+            if cls==0:
+                color = colors(int(cls), True)
+            elif tid>0 and cls!=0:
+                color = colors(int(tid), True)
+            else:
+                color = colors(19, True)     # 19为调色板的最后一个元素
+            
+            pt2 = [p/scale for p in xyxy]
+            annotator.box_label(pt2, label, color=color)
+        
+        img = annotator.result()
+        subimgs.append(img)
+        
+    return subimgs
+
+
+
+