96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
# -*- 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)
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|