83 lines
3.0 KiB
Python
83 lines
3.0 KiB
Python
|
import numpy as np
|
||
|
|
import cv2
|
||
|
|
|
||
|
|
MIN_PIXEL = 50
|
||
|
|
def callback_select_bbox_corner(start, end, annots, select, **kwargs):
|
||
|
|
if start is None or end is None:
|
||
|
|
select['corner'] = -1
|
||
|
|
return 0
|
||
|
|
if start[0] == end[0] and start[1] == end[1]:
|
||
|
|
return 0
|
||
|
|
# 判断选择了哪个角点
|
||
|
|
annots = annots['annots']
|
||
|
|
start = np.array(start)[None, :]
|
||
|
|
if select['bbox'] == -1 and select['corner'] == -1:
|
||
|
|
for i in range(len(annots)):
|
||
|
|
l, t, r, b = annots[i]['bbox'][:4]
|
||
|
|
corners = np.array([(l, t), (l, b), (r, t), (r, b)])
|
||
|
|
dist = np.linalg.norm(corners - start, axis=1)
|
||
|
|
mindist = dist.min()
|
||
|
|
if mindist < MIN_PIXEL:
|
||
|
|
mincor = dist.argmin()
|
||
|
|
select['bbox'] = i
|
||
|
|
select['corner'] = mincor
|
||
|
|
break
|
||
|
|
else:
|
||
|
|
select['corner'] = -1
|
||
|
|
elif select['bbox'] != -1 and select['corner'] == -1:
|
||
|
|
i = select['bbox']
|
||
|
|
l, t, r, b = annots[i]['bbox'][:4]
|
||
|
|
corners = np.array([(l, t), (l, b), (r, t), (r, b)])
|
||
|
|
dist = np.linalg.norm(corners - start, axis=1)
|
||
|
|
mindist = dist.min()
|
||
|
|
if mindist < MIN_PIXEL:
|
||
|
|
mincor = dist.argmin()
|
||
|
|
select['corner'] = mincor
|
||
|
|
elif select['bbox'] != -1 and select['corner'] != -1:
|
||
|
|
# Move the corner
|
||
|
|
x, y = end
|
||
|
|
(i, j) = [(0, 1), (0, 3), (2, 1), (2, 3)][select['corner']]
|
||
|
|
data = annots[select['bbox']]
|
||
|
|
data['bbox'][i] = x
|
||
|
|
data['bbox'][j] = y
|
||
|
|
elif select['bbox'] == -1 and select['corner'] != -1:
|
||
|
|
select['corner'] = -1
|
||
|
|
|
||
|
|
def callback_select_bbox_center(click, annots, select, **kwargs):
|
||
|
|
if click is None:
|
||
|
|
return 0
|
||
|
|
annots = annots['annots']
|
||
|
|
bboxes = np.array([d['bbox'] for d in annots])
|
||
|
|
center = (bboxes[:, [2, 3]] + bboxes[:, [0, 1]])/2
|
||
|
|
click = np.array(click)[None, :]
|
||
|
|
dist = np.linalg.norm(click - center, axis=1)
|
||
|
|
mindist, minid = dist.min(), dist.argmin()
|
||
|
|
if mindist < MIN_PIXEL:
|
||
|
|
select['bbox'] = minid
|
||
|
|
|
||
|
|
def auto_pose_track(self, param, **kwargs):
|
||
|
|
"auto tracking with poses"
|
||
|
|
MAX_SPEED = 100
|
||
|
|
if self.frame == 0:
|
||
|
|
return 0
|
||
|
|
previous = self.previous()
|
||
|
|
annots = param['annots']['annots']
|
||
|
|
keypoints_pre = np.array([d['keypoints'] for d in previous['annots']])
|
||
|
|
keypoints_now = np.array([d['keypoints'] for d in annots])
|
||
|
|
conf = np.sqrt(keypoints_now[:, None, :, -1] * keypoints_pre[None, :, :, -1])
|
||
|
|
diff = np.linalg.norm(keypoints_now[:, None, :, :] - keypoints_pre[None, :, :, :], axis=-1)
|
||
|
|
dist = np.sum(diff * conf, axis=-1)/np.sum(conf, axis=-1)
|
||
|
|
nows, pres = np.where(dist < MAX_SPEED)
|
||
|
|
edges = []
|
||
|
|
for n, p in zip(nows, pres):
|
||
|
|
edges.append((n, p, dist[n, p]))
|
||
|
|
edges.sort(key=lambda x:x[2])
|
||
|
|
used_n, used_p = [], []
|
||
|
|
for n, p, _ in edges:
|
||
|
|
if n in used_n or p in used_p:
|
||
|
|
continue
|
||
|
|
annots[n]['personID'] = previous['annots'][p]['personID']
|
||
|
|
used_n.append(n)
|
||
|
|
used_p.append(p)
|
||
|
|
# TODO:stop when missing
|
||
|
|
|