174 lines
4.3 KiB
Python
174 lines
4.3 KiB
Python
# -*- 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)
|
|
# =============================================================================
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|