abc

tracking/utils/plotting.py

@@ -285,6 +285,38 @@ def boxing_img(det, img, line_width=3):
     
     return imgx
 
+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
+
+
+def get_subimgs(imgs, tracks, scale=2):
+    bboxes = []
+    if len(tracks):
+        bboxes = array2list(tracks)
+    
+    subimgs = []
+    for i, boxes in enumerate(bboxes):
+        fid = int(boxes[0, 7])
+        
+        for *xyxy, tid, conf, cls, fid, bid in boxes:
+            pt2 = [p/scale for p in xyxy]
+            x1, y1, x2, y2 = (int(pt2[0]), int(pt2[1])), (int(pt2[2]), int(pt2[3]))
+            subimgs.append((int(fid), int(bid), imgs[fid-1][y1:y2, x1:x2]))
+    
+    return subimgs
+
 def draw_tracking_boxes(imgs, tracks, scale=2):
     '''需要确保 imgs 覆盖tracks中的帧ID数
         tracks: [x1, y1, x2, y2, track_id, score, cls, frame_index, box_index]
@@ -295,27 +327,15 @@ def draw_tracking_boxes(imgs, tracks, scale=2):
        
     '''
     
-    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)
+
+    bboxes = []
+    if len(tracks):
+        bboxes = array2list(tracks)
     
     # if len(bboxes)!=len(imgs):
     #     return False, imgs
     
-    subimgs = []
+    annimgs = []
     for i, boxes in enumerate(bboxes):
         fid = int(boxes[0, 7])
         annotator = Annotator(imgs[fid-1].copy())
@@ -331,11 +351,11 @@ def draw_tracking_boxes(imgs, tracks, scale=2):
             
             pt2 = [p/scale for p in xyxy]
             annotator.box_label(pt2, label, color=color)
-        
+
         img = annotator.result()
-        subimgs.append((fid, img))
+        annimgs.append((int(fid), img))
         
-    return subimgs
+    return annimgs