update
This commit is contained in:
173
ytracking/tracking/test_merge.py
Normal file
173
ytracking/tracking/test_merge.py
Normal file
@ -0,0 +1,173 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Created on Fri Feb 23 11:04:48 2024
|
||||
|
||||
@author: ym
|
||||
"""
|
||||
import numpy as np
|
||||
import cv2
|
||||
from scipy.spatial.distance import cdist
|
||||
# from trackers.utils import matching
|
||||
|
||||
def readDict(boxes, feat_dicts):
|
||||
feat = []
|
||||
for i in range(boxes.shape[0]):
|
||||
tid, fid, bid = int(boxes[i, 4]), int(boxes[i, 7]), int(boxes[i, 8])
|
||||
feat.append(feat_dicts[fid][bid])
|
||||
|
||||
# img = feat_dicts[fid][f'{bid}_img']
|
||||
# cv2.imwrite(f'./result/imgs/{tid}_{fid}_{bid}.png', img)
|
||||
|
||||
return np.asarray(feat, dtype=np.float32)
|
||||
|
||||
|
||||
|
||||
def track_equal_track(atrack, btrack, feat_dicts):
|
||||
# boxes: [x, y, w, h, track_id, score, cls, frame_index, box_index]
|
||||
aboxes = atrack.boxes
|
||||
bboxes = btrack.boxes
|
||||
|
||||
''' 1. 判断轨迹在时序上是否有交集 '''
|
||||
afids = aboxes[:, 7].astype(np.int_)
|
||||
bfids = bboxes[:, 7].astype(np.int_)
|
||||
# 帧索引交集
|
||||
interfid = set(afids).intersection(set(bfids))
|
||||
|
||||
# 或者直接判断帧索引是否有交集,返回 Ture or False
|
||||
# interfid = set(afids).isdisjoint(set(bfids))
|
||||
|
||||
|
||||
''' 2. 轨迹空间iou'''
|
||||
alabel = np.array([0] * afids.size, dtype=np.int_)
|
||||
blabel = np.array([1] * bfids.size, dtype=np.int_)
|
||||
|
||||
label = np.concatenate((alabel, blabel), axis=0)
|
||||
fids = np.concatenate((afids, bfids), axis=0)
|
||||
indices = np.argsort(fids)
|
||||
idx_pair = []
|
||||
for i in range(len(indices)-1):
|
||||
idx1, idx2 = indices[i], indices[i+1]
|
||||
if label[idx1] != label[idx2] and fids[idx2] - fids[idx1] == 1:
|
||||
if label[idx1] == 0:
|
||||
a_idx = idx1
|
||||
b_idx = idx2-alabel.size
|
||||
else:
|
||||
a_idx = idx2
|
||||
b_idx = idx1-alabel.size
|
||||
|
||||
idx_pair.append((a_idx, b_idx))
|
||||
|
||||
ious = []
|
||||
for a, b in idx_pair:
|
||||
abox, bbox = aboxes[a, :], bboxes[b, :]
|
||||
|
||||
xa1, ya1 = abox[0] - abox[2]/2, abox[1] - abox[3]/2
|
||||
xa2, ya2 = abox[0] + abox[2]/2, abox[1] + abox[3]/2
|
||||
|
||||
xb1, yb1 = bbox[0] - bbox[2]/2, bbox[1] - bbox[3]/2
|
||||
xb2, yb2 = bbox[0] + bbox[2]/2, bbox[1] + bbox[3]/2
|
||||
|
||||
|
||||
inter = (np.minimum(xb2, xa2) - np.maximum(xb1, xa1)).clip(0) * \
|
||||
(np.minimum(yb2, ya2) - np.maximum(yb1, ya1)).clip(0)
|
||||
|
||||
# Union Area
|
||||
box1_area = abox[2] * abox[3]
|
||||
box2_area = bbox[2] * bbox[3]
|
||||
union = box1_area + box2_area - inter + 1e-6
|
||||
|
||||
ious.append(inter/union)
|
||||
|
||||
''' 3. 轨迹特征相似度判断'''
|
||||
afeat = readDict(aboxes, feat_dicts)
|
||||
bfeat = readDict(bboxes, feat_dicts)
|
||||
feat = np.concatenate((afeat, bfeat), axis=0)
|
||||
|
||||
emb_simil = 1-np.maximum(0.0, cdist(feat, feat, 'cosine'))
|
||||
emb_ = 1-cdist(np.mean(afeat, axis=0)[None, :], np.mean(bfeat, axis=0)[None, :], 'cosine')
|
||||
|
||||
cont1 = False if len(interfid) else True
|
||||
cont2 = all(iou>0.5 for iou in ious)
|
||||
cont3 = emb_[0, 0]>0.75
|
||||
|
||||
cont = cont1 and cont2 and cont3
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
return cont
|
||||
|
||||
|
||||
|
||||
def track_equal_str(atrack, btrack):
|
||||
if atrack == btrack:
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
def merge_track(Residual):
|
||||
out_list = []
|
||||
alist = [t for t in Residual]
|
||||
while alist:
|
||||
atrack = alist[0]
|
||||
cur_list = []
|
||||
cur_list.append(atrack)
|
||||
alist.pop(0)
|
||||
|
||||
blist = [b for b in alist]
|
||||
alist = []
|
||||
for btrack in blist:
|
||||
if track_equal_str(atrack, btrack):
|
||||
cur_list.append(btrack)
|
||||
else:
|
||||
alist.append(btrack)
|
||||
|
||||
out_list.append(cur_list)
|
||||
return out_list
|
||||
|
||||
def main():
|
||||
Residual = ['a', 'b', 'c', 'd', 'a', 'b', 'c', 'b', 'c', 'd']
|
||||
out_list = merge_track(Residual)
|
||||
|
||||
print(Residual)
|
||||
print(out_list)
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
main()
|
||||
|
||||
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# for i, atrack in enumerate(input_list):
|
||||
# cur_list = []
|
||||
# cur_list.append(atrack)
|
||||
# del input_list[i]
|
||||
#
|
||||
# for j, btrack in enumerate(input_list):
|
||||
# if track_equal(atrack, btrack):
|
||||
# cur_list.append(btrack)
|
||||
# del input_list[j]
|
||||
#
|
||||
# out_list.append(cur_list)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
Reference in New Issue
Block a user