lichen/detecttracking
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

modify dotrack module

Browse Source
This commit is contained in:
王庆刚
2024-06-03 15:25:39 +08:00
parent d1ea304491
commit f90ef72cbf
19 changed files with 502 additions and 420 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
track_reid.py
View File

@ -110,8 +110,6 @@ def inference_image(image, detections):
def init_trackers(tracker_yaml = None, bs=1):
"""
Initialize trackers for object tracking during prediction.
@ -177,7 +175,7 @@ def run(
save_dir = Path(project) / Path(source).stem
if save_dir.exists():
print(Path(source).stem)
return
# return
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
@ -215,6 +213,9 @@ def run(
track_boxes = np.empty((0, 9), dtype = np.float32)
det_boxes = np.empty((0, 9), dtype = np.float32)
DetBoxes = np.empty((0, 6), dtype = np.float32)
TrackerBoxes = np.empty((0, 9), dtype = np.float32)
TrackerFeats = np.empty((0, 256), dtype = np.float32)
features_dict = {}
TracksDict = {}
@ -267,38 +268,40 @@ def run(
det = det.cpu().numpy()
det = np.concatenate([det[:, :4], np.arange(nd).reshape(-1, 1), det[:, 4:]], axis=-1)
DetBoxes = np.concatenate([DetBoxes, det[:, :6]], axis=0)
#
def static_estimate(box1, box2, TH1=8, TH2=12):
dij_abs = max(np.abs(box1 - box2))
dij_euc = max([np.linalg.norm((box1[:2] - box2[:2])),
np.linalg.norm((box1[2:4] - box2[2:4]))
])
if dij_abs < TH1 and dij_euc < TH2:
return True
else:
return False
## ============================================================ 前后帧相同 boxes 的特征赋值
# def static_estimate(box1, box2, TH1=8, TH2=12):
# dij_abs = max(np.abs(box1 - box2))
# dij_euc = max([np.linalg.norm((box1[:2] - box2[:2])),
# np.linalg.norm((box1[2:4] - box2[2:4]))
# ])
# if dij_abs < TH1 and dij_euc < TH2:
# return True
# else:
# return False
nw = 3 # 向前递推检查的窗口大小
nf = len(BoxesFeats) # 已经检测+特征提取的帧数
feat_curr = [None] * nd # nd: 当前帧检测出的boxes数
for ii in range(nd):
box = det[ii, :4]
# nw = 3 # 向前递推检查的窗口大小
# nf = len(BoxesFeats) # 已经检测+特征提取的帧数
# feat_curr = [None] * nd # nd: 当前帧检测出的boxes数
# for ii in range(nd):
# box = det[ii, :4]
kk=1
feat = None
while kk <= nw and nf>=kk:
ki = -1 * kk
boxes_ = BoxesFeats[ki][0]
feats_ = BoxesFeats[ki][1]
# kk=1
# feat = None
# while kk <= nw and nf>=kk:
# ki = -1 * kk
# boxes_ = BoxesFeats[ki][0]
# feats_ = BoxesFeats[ki][1]
flag = [jj for jj in range(len(boxes_)) if static_estimate(box, boxes_[jj, :4])]
if len(flag) == 1:
feat = feats_[flag[0]]
break
kk += 1
if feat is not None:
feat_curr[ii] = feat
# flag = [jj for jj in range(len(boxes_)) if static_estimate(box, boxes_[jj, :4])]
# if len(flag) == 1:
# feat = feats_[flag[0]]
# break
# kk += 1
# if feat is not None:
# feat_curr[ii] = feat
@ -319,7 +322,11 @@ def run(
'''================== 1. 存储 dets/subimgs/features Dict ============='''
imgs, features = inference_image(im0, tracks)
BoxesFeats.append((tracks, features))
TrackerFeats = np.concatenate([TrackerFeats, features], axis=0)
imgdict = {}
@ -532,7 +539,10 @@ def main_loop(opt):
# r"D:\datasets\ym\广告板遮挡测试\8\2500441577966_20240508-175946_front_addGood_70f75407b7ae_155_17788571404.mp4"
# ]
# files = [r"D:\datasets\ym\videos\标记视频\test_20240402-173935_6920152400975_back_174037372.mp4"]
# files = [r"D:\datasets\ym\视频\20240529\110518062-090ac04c-0a8c-479f-bc18-cb3553c90683-0_seek0.017962635633665514.mp4"]
files = [r"D:\datasets\ym\视频\20240529\110518060-550b7c4d-9946-4aa4-9131-81008692cd65-1_seek0.7670042724609232.mp4"]
for file in files:
optdict["source"] = file
run(**optdict)