更新 detacttracking

detecttracking/tracking/utils/annotator.py

@@ -0,0 +1,96 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Sep 20 14:21:13 2023
+
+@author: ym
+"""
+import cv2
+
+# import sys
+# sys.path.append(r"D:\DeepLearning\yolov5")
+# from ultralytics.utils.plotting import Annotator, colors
+
+from .plotting import Annotator, colors
+
+class TrackAnnotator(Annotator): 
+    
+    def plotting_track(self, track, names='abc'):
+        """
+        track：[x, y, w, h, track_id, score, cls, frame_index]
+        
+        boxes: [x1, y1, x2, y2, track_id, score, cls, frame_index, box_index]
+                0   1   2   3      4        5     6       7           8
+        id：跟踪id，从 1 开始计数，
+        frame_index: 帧索引，从 1 开始计数
+        cls：类别编号，从 0 开始计数，用作 names 的 key 值
+        """
+        
+        if track.size==0: return
+    
+        id, cls = track[0, 4],  track[0, 6]
+        if id >=0 and cls==0:
+            color = colors(int(cls), True)
+        elif id >=0 and cls!=0:
+            color = colors(int(id), True)   # 不存在 id = 0，不会和上面产生冲突
+        else:
+            color = colors(19, True)     # 19为调色板的最后一个元素        
+        
+        nb = track.shape[0]
+        for i in range(nb):
+            if i == 0: 
+                # label = f'{int(track[i, 4])}:({int(track[i, 7])})'
+                label = f'ID_{int(track[i, 4])}'
+            elif i == nb-1:
+                label = ''
+                # label = f'{int(track[i, 4])}:({int(track[i, 7])})&{int(nb)}'
+            else:
+                label = ''
+                
+            self.circle_label(track[i, :], label, color=color)
+    
+    
+    def circle_label(self, track,  label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
+        """
+        绘制选定 track 的轨迹
+        """
+        if track.size==0: return
+        
+        x, y = int((track[0]+track[2])/2), int((track[1]+track[3])/2)
+        
+        cv2.circle(self.im, (x, y), 6, color, 2)
+        # txt_color = (0,0,0)
+        
+        if label:
+            tf = max(self.lw - 1, 1)  # font thickness
+            w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0]  # text width, height
+            outside = x + w <= self.im.shape[1]-3
+            # p2 = x + w, y - h - 3 if outside else y + h + 3
+            # cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA)  # filled
+            
+            cv2.putText(self.im,
+                        label, (x-10 if outside else x-w+2,  y-20),
+                        0,
+                        # self.lw / 3,
+                        self.lw/2,
+                        txt_color,
+                        thickness=tf,
+                        lineType=cv2.LINE_AA)
+            
+            
+
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+            
+