box select in a track and feat simi modify in tracker

This commit is contained in:
王庆刚
2025-01-14 19:00:59 +08:00
parent 744fb7b7b2
commit bfe7bc0fd5
11 changed files with 157 additions and 22 deletions

View File

@ -38,6 +38,7 @@ import glob
import numpy as np
import pickle
import torch
from scipy.spatial.distance import cdist
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
@ -222,8 +223,20 @@ def yolo_resnet_tracker(
这里frame_index 也可以用视频的 帧ID 代替, box_index 保持不变
'''
det_tracking = Boxes(det, im0.shape).cpu().numpy()
tracks = tracker.update(det_tracking, im0)
tracks, outfeats = tracker.update(det_tracking, im0)
simdict, simdict1 = {}, {}
for fid, bid, mfeat, cfeat, features in outfeats:
if mfeat is not None and cfeat is not None:
simi = 1 - np.maximum(0.0, cdist(mfeat[None, :], cfeat[None, :], "cosine"))[0][0]
simdict.update({f"{int(frameId)}_{int(bid)}":simi})
if cfeat is not None and len(features)>=2:
mfeat = features[-2]
simi = 1 - np.maximum(0.0, cdist(mfeat[None, :], cfeat[None, :], "cosine"))[0][0]
simdict1.update({f"{int(frameId)}_{int(bid)}":simi})
if len(tracks) > 0:
tracks[:, 7] = frameId
# trackerBoxes = np.concatenate([trackerBoxes, tracks], axis=0)
@ -239,7 +252,10 @@ def yolo_resnet_tracker(
"bboxes": det,
"tboxes": tracks,
"imgs": imgdict,
"feats": featdict}
"feats": featdict,
"featsimi": simdict, # 当前 box 特征和该轨迹 smooth_feat 特征的相似度
"featsimi1": simdict1 # 当前 box 特征和该轨迹前一个 box 特征的相似度
}
yoloResnetTracker.append(frameDict)
# imgs, features = inference_image(im0, tracks)
@ -247,8 +263,15 @@ def yolo_resnet_tracker(
'''================== 2. 提取手势位置 ==================='''
for *xyxy, id, conf, cls, fid, bid in reversed(tracks):
name = ('' if id==-1 else f'id:{int(id)} ') + names[int(cls)]
label = None if hide_labels else (name if hide_conf else f'{name} {conf:.2f}')
name = ('' if id==-1 else f'id:{int(id)} ') + names[int(cls)]
if f"{int(frameId)}_{int(bid)}" in simdict.keys():
sim = simdict[f"{int(frameId)}_{int(bid)}"]
label = f"{name} {sim:.2f}"
else:
label = None if hide_labels else name
# label = None if hide_labels else (name if hide_conf else f'{name} {conf:.1f}')
if id >=0 and cls==0:
color = colors(int(cls), True)
@ -489,7 +512,7 @@ def run(
'''
det_tracking = Boxes(det, im0.shape).cpu().numpy()
tracks = tracker.update(det_tracking, im0)
tracks, outfeats = tracker.update(det_tracking, im0)
if len(tracks) == 0:
continue