1075 lines
49 KiB
Python
1075 lines
49 KiB
Python
import numpy as np
|
||
import cv2
|
||
from easymocap.mytools.camera_utils import Undistort
|
||
from easymocap.mytools.debug_utils import mywarn
|
||
from .triangulate import batch_triangulate, project_wo_dist
|
||
from collections import defaultdict
|
||
LOG_FILE = 'log.txt'
|
||
LOG_LEVEL = 0 #2
|
||
FULL_LOG = (lambda x: print(x, file=open(LOG_FILE, 'a'))) if LOG_LEVEL > 1 else (lambda x: None)
|
||
LOG = (lambda x: print(x, file=open(LOG_FILE, 'a'))) if LOG_LEVEL > 0 else (lambda x: None)
|
||
|
||
def LOG_ARRAY(array2d, format='{:>8.2f} '):
|
||
res = ''
|
||
for i in range(array2d.shape[0]):
|
||
for j in range(array2d.shape[1]):
|
||
res += format.format(array2d[i, j])
|
||
res += '\n'
|
||
return res
|
||
|
||
class MatchBase:
|
||
def __init__(self, mode, cfg) -> None:
|
||
self.mode = mode
|
||
self.cfg = cfg
|
||
print('[{}]'.format(self.__class__.__name__))
|
||
print(self.cfg)
|
||
self.max_id = 0
|
||
|
||
def make_grids(self, grids, grids_step):
|
||
grid_x = np.arange(grids[0][0], grids[1][0], grids_step)
|
||
grid_y = np.arange(grids[0][1], grids[1][1], grids_step)
|
||
grid_z = np.arange(grids[0][2], grids[1][2], grids_step)
|
||
grid_xyz = np.meshgrid(grid_x, grid_y, grid_z)
|
||
grid_xyz = np.stack(grid_xyz, axis=-1)
|
||
grids = grid_xyz.reshape(-1, 3)
|
||
print('[{}] Generate {} => {} grids'.format(self.__class__.__name__, grid_xyz.shape, grids.shape[0]))
|
||
return grids
|
||
|
||
@staticmethod
|
||
def stack_array(arrays):
|
||
dimGroups = [0]
|
||
results = []
|
||
views_all = []
|
||
for nv, array in enumerate(arrays):
|
||
dimGroups.append(dimGroups[-1] + array.shape[0])
|
||
views_all.extend([nv for _ in range(array.shape[0])])
|
||
results.append(array)
|
||
results = np.concatenate(results, axis=0)
|
||
return results, np.array(views_all), np.array(dimGroups)
|
||
|
||
@staticmethod
|
||
def undistort(points, cameras):
|
||
nViews = len(points)
|
||
pelvis_undis = []
|
||
for nv in range(nViews):
|
||
camera = {key:cameras[key][nv] for key in ['R', 'T', 'K', 'dist']}
|
||
if points[nv].shape[0] > 0:
|
||
pelvis = Undistort.points(points[nv], camera['K'], camera['dist'])
|
||
else:
|
||
pelvis = points[nv].copy()
|
||
pelvis_undis.append(pelvis)
|
||
return pelvis_undis
|
||
|
||
@staticmethod
|
||
def distance_by_triangulate(p_src, p_dst, camera_src, camera_dst, ranges):
|
||
dist = np.zeros((p_src.shape[0], p_dst.shape[0]), dtype=np.float32)
|
||
# generate (m, n) points and distance
|
||
idx_src = np.arange(p_src.shape[0])
|
||
idx_dst = np.arange(p_dst.shape[0])
|
||
idx_src, idx_dst = np.meshgrid(idx_src, idx_dst)
|
||
idx_src = idx_src.reshape(-1)
|
||
idx_dst = idx_dst.reshape(-1)
|
||
p_src = p_src[idx_src]
|
||
p_dst = p_dst[idx_dst]
|
||
keypoints = np.stack([p_src, p_dst], axis=0)
|
||
keypoints_flat = keypoints.reshape(keypoints.shape[0], -1, keypoints.shape[-1])
|
||
P = np.stack([camera_src['P'], camera_dst['P']], axis=0)
|
||
k3d = batch_triangulate(keypoints_flat, P, min_view=2)
|
||
repro, depth = project_wo_dist(k3d, P)
|
||
dist_repro = np.linalg.norm(repro[..., :2] - keypoints_flat[..., :2], axis=-1).mean(axis=0)
|
||
valid = (k3d[:, 0] > ranges[0][0]) & (k3d[:, 0] < ranges[1][0]) & \
|
||
(k3d[:, 1] > ranges[0][1]) & (k3d[:, 1] < ranges[1][1]) & \
|
||
(k3d[:, 2] > ranges[0][2]) & (k3d[:, 2] < ranges[1][2])
|
||
dist_repro[~valid] = 1e5
|
||
dist[idx_src, idx_dst] = dist_repro
|
||
return dist
|
||
|
||
def calculate_distance(self, pelvis_undis, cameras, dimGroups):
|
||
DIST_MAX = 10000.
|
||
distance = np.zeros((dimGroups[-1], dimGroups[-1]), dtype=np.float32) + DIST_MAX
|
||
nViews = len(dimGroups) - 1
|
||
ray0 = np.array([0, 0, 1], dtype=np.float32).reshape(1, 3, 1)
|
||
ray_cam = cameras['R'].transpose(0, 2, 1) @ ray0
|
||
ray_cam = ray_cam[..., 0]
|
||
cos_theta = np.sum(ray_cam[:, None] * ray_cam[None], axis=-1)
|
||
theta = np.rad2deg(np.arccos(np.clip(cos_theta, -1., 1.)))
|
||
valid_theta = np.logical_and(theta > self.cfg.valid_angle[0], theta < self.cfg.valid_angle[1])
|
||
for src in range(nViews - 1):
|
||
for dst in range(src + 1, nViews):
|
||
# TODO: 计算两个射线的夹角
|
||
# 这里对于不相邻或者对角的视角，我们直接跳过距离的计算
|
||
# 这样后面在进行初始化的时候就无法挑到两个比较接近的视角了
|
||
# if not valid_theta[src, dst]:
|
||
# continue
|
||
p_src = pelvis_undis[src][:, None] #(m, 2)
|
||
p_dst = pelvis_undis[dst][:, None] #(n, 2)
|
||
if p_src.shape[0] == 0 or p_dst.shape[0] == 0:
|
||
continue
|
||
camera_src = {key:cameras[key][src] for key in ['R', 'T', 'K', 'dist', 'P']}
|
||
camera_dst = {key:cameras[key][dst] for key in ['R', 'T', 'K', 'dist', 'P']}
|
||
dist = self.distance_by_triangulate(p_src, p_dst, camera_src, camera_dst, self.cfg.valid_ranges)
|
||
distance[dimGroups[src]:dimGroups[src+1], dimGroups[dst]:dimGroups[dst+1]] = dist
|
||
distance[dimGroups[dst]:dimGroups[dst+1], dimGroups[src]:dimGroups[src+1]] = dist.T
|
||
cameras['valid_theta'] = valid_theta
|
||
return distance
|
||
|
||
def calculate_repro(self, results, pelvis_undis, cameras, views_all):
|
||
nViews = len(cameras['P'])
|
||
n3D = len(results)
|
||
distance = np.zeros((pelvis_undis.shape[0], n3D), dtype=np.float32)
|
||
if n3D == 0:
|
||
return distance
|
||
keypoints3d = np.stack([d['pelvis'] for d in results], axis=0)
|
||
Pall = np.stack([cameras['P'][nv] for nv in range(nViews)])
|
||
# k2d: (nViews, nPerson, nPoints, 3)
|
||
k2d, depth = project_wo_dist(keypoints3d, Pall, einsum='vab,pkb->vpka')
|
||
repro_select = k2d[views_all]
|
||
# dist: (nPoints, n3D)
|
||
dist = np.linalg.norm(repro_select[..., :2] - pelvis_undis[:, None, None, :2], axis=-1).mean(axis=2)
|
||
# for nv in range(nViews):
|
||
return dist
|
||
|
||
def triangulate_and_repro(self, cameras, views, proposals):
|
||
Pall = np.stack([cameras['P'][v] for v in views])
|
||
kpts = np.stack(proposals)
|
||
kpts = kpts[:, None]
|
||
k3d = batch_triangulate(kpts, Pall)
|
||
k2d, depth = project_wo_dist(k3d, Pall)
|
||
dist_repro = np.linalg.norm(k2d[..., :2] - kpts[..., :2], axis=-1).mean(axis=-1)
|
||
return k3d, dist_repro, depth
|
||
|
||
|
||
@staticmethod
|
||
def check_is_best_3d_of_2d(distance, idx3d, idx2d, visited3d):
|
||
isbest3d = True
|
||
distance_2d = distance[idx2d]
|
||
for i3d in distance_2d.argsort():
|
||
if i3d != idx3d and i3d not in visited3d:
|
||
isbest3d = False
|
||
break
|
||
elif i3d == idx3d:
|
||
break
|
||
return isbest3d
|
||
|
||
@staticmethod
|
||
def sort_with_affinity(distance, dimGroups, INLIER_REPRO):
|
||
nViews = len(dimGroups) - 1
|
||
# 排序计算affinity
|
||
count_rows = np.zeros((dimGroups[-1]), dtype=int)
|
||
distance_rows = np.zeros((dimGroups[-1]))
|
||
for nv in range(nViews):
|
||
if dimGroups[nv] == dimGroups[nv+1]:continue
|
||
valid_view = np.clip((distance[:, dimGroups[nv]:dimGroups[nv+1]] < INLIER_REPRO).sum(axis=-1), 0, 1)
|
||
count_rows += valid_view # 最多也只累计一个
|
||
distance_rows += valid_view * (distance[:, dimGroups[nv]:dimGroups[nv+1]].min(axis=-1))
|
||
index = list(range(dimGroups[-1]))
|
||
# index.sort(key=lambda x: (-count_rows[x], distance_rows[x]))
|
||
# sort with 2D confidence
|
||
# index.sort(key=lambda x: -pelvis_all[x, 2])
|
||
# sort with valid matches
|
||
# 选择2D的依据改为：根据有效的2D重投影距离的数量
|
||
valid_count = (distance < INLIER_REPRO * 2).sum(axis=0)
|
||
index = (-valid_count).argsort()
|
||
return index
|
||
|
||
def assign_by_3D(self, used_index, distance, pelvis_all, views_all, dimGroups, cameras):
|
||
INLIER_TRACK = self.cfg.track_pixel
|
||
INLIER_REPRO = self.cfg.max_pixel
|
||
# 使用前一帧的可见性来进行排序
|
||
index_3d = list(range(len(self.results)))
|
||
index_3d.sort(key=lambda x:-len(self.results[x]['views']))
|
||
results = []
|
||
visited3d = set()
|
||
for idx3d in index_3d:
|
||
visited3d.add(idx3d)
|
||
self.results[idx3d]['tracked'] = False
|
||
pid = self.results[idx3d]['id']
|
||
dist = distance[:, idx3d]
|
||
FULL_LOG('[Assign 3D] Check 3D {}'.format(pid))
|
||
FULL_LOG('[Assign 3D] Distance {}'.format(LOG_ARRAY(dist[None])))
|
||
current = []
|
||
views = []
|
||
proposal = dist.argsort()
|
||
# 初始化一下：
|
||
for idx2d in proposal:
|
||
# 不满足视角关系
|
||
# if not valid_theta[views_all[row], views_all[idx2d]]:
|
||
# continue
|
||
# 不满足距离关系
|
||
if dist[idx2d] > INLIER_TRACK:
|
||
break
|
||
if used_index[idx2d] > -1:
|
||
continue
|
||
if views_all[idx2d] in views:
|
||
continue
|
||
if not self.check_is_best_3d_of_2d(distance, idx3d, idx2d, visited3d):
|
||
continue
|
||
if len(current) == 1: # 已经有一个了，如果还要再添加，那么需要判断一下三角化出来的距离关系
|
||
k3d, dist_repro, depth = self.triangulate_and_repro(cameras, views + [views_all[idx2d]], current + [pelvis_all[idx2d]])
|
||
_dist = np.linalg.norm(k3d[:, :3] - self.results[idx3d]['pelvis'][:, :3], axis=-1).mean()
|
||
if _dist > self.cfg.max_movement:
|
||
continue
|
||
# 找到了合理的pair，作为一个良好的初始化
|
||
current.append(pelvis_all[idx2d])
|
||
views.append(views_all[idx2d])
|
||
used_index[idx2d] = pid
|
||
FULL_LOG(f'[Assign 3D] First track 3D {pid} with {idx2d}, view ({views_all[idx2d]})')
|
||
if len(current) == 2:
|
||
break
|
||
if len(current) < 2:
|
||
# 没有找到良好的初始化
|
||
continue
|
||
for idx2d in proposal:
|
||
# 这个视角已经有了 ｜ 这个2D已经被使用过了
|
||
if views_all[idx2d] in views:
|
||
continue
|
||
if used_index[idx2d] > -1:
|
||
continue
|
||
if not self.check_is_best_3d_of_2d(distance, idx3d, idx2d, visited3d):
|
||
continue
|
||
# 尝试添加
|
||
FULL_LOG('[Assign 3D] 3D {} add {}, distance={:.2f}'.format(pid, idx2d, dist[idx2d]))
|
||
new = current + [pelvis_all[idx2d]]
|
||
views_new = views + [views_all[idx2d]]
|
||
k3d, dist_repro, depth = self.triangulate_and_repro(cameras, views_new, new)
|
||
_dist = np.linalg.norm(k3d[:, :3] - self.results[idx3d]['pelvis'][:, :3], axis=-1).mean()
|
||
flag_movement = _dist < self.cfg.max_movement
|
||
flag_depth = (depth > 0.5).all()
|
||
flag_repro = dist_repro.mean() < INLIER_REPRO
|
||
flag = flag_repro & flag_depth
|
||
FULL_LOG('[Assign 3D] repro: \n{}, \ndepth: \n{}'.format(LOG_ARRAY(dist_repro[None]), LOG_ARRAY(depth.T)))
|
||
if flag:
|
||
# 添加
|
||
current = new
|
||
views = views_new
|
||
used_index[idx2d] = pid
|
||
FULL_LOG('[Assign 3D] {} => {}'.format(idx2d, np.where(used_index == pid)[0]))
|
||
else:
|
||
FULL_LOG('[Assign 3D] Failed')
|
||
# check the results
|
||
if len(views) < self.cfg.min_views: #不足以添加
|
||
continue
|
||
k3d, dist_repro, depth = self.triangulate_and_repro(cameras, views, current)
|
||
select = np.where(used_index == pid)[0]
|
||
results.append({
|
||
'id': pid,
|
||
'pelvis': k3d,
|
||
'keypoints3d': k3d, # 这里保存两个，这样即使后面覆盖掉了keypoints3d还能取出pelvis来
|
||
'views': views_all[select],
|
||
'select': select,
|
||
'indices': select - dimGroups[views_all[select]],
|
||
'frames': self.results[idx3d]['frames'] + [self.frames]
|
||
})
|
||
self.results[idx3d]['tracked'] = True
|
||
for res in results:
|
||
text = f''' - Track {res['id']} with {len(res['views'])} views
|
||
views: {' '.join(list(map(lambda x:'{:2d}'.format(x), res['views'])))}
|
||
id : {' '.join(list(map(lambda x:'{:2d}'.format(x), res['select'])))}'''
|
||
LOG(text)
|
||
print(text)
|
||
for res in self.results:
|
||
if not res['tracked']:
|
||
mywarn('- 3D {} not tracked'.format(res['id']))
|
||
# 对于没有被跟踪到的：检查是否有两个距离很小的视角
|
||
# 如果有，并且被其他人占用了，那么把这个2D也给他；在极端情况下，有的视角下会有人恰好被另一个人挡住
|
||
print(res)
|
||
if len(res['frames']) < 3:
|
||
mywarn('- 3D {} not tracked, but only {} frames'.format(res['id'], len(res['frames'])))
|
||
else:
|
||
pass
|
||
# import ipdb; ipdb.set_trace()
|
||
return results
|
||
|
||
def find_initial_3_pair(self, distance, pelvis_all, views_all, dimGroups):
|
||
# 生成所有可能的候选的3个pair
|
||
index_0 = np.arange(pelvis_all.shape[0])
|
||
index_0 = np.stack(np.meshgrid(index_0, index_0, index_0), axis=-1).reshape(-1, 3)
|
||
flag_order = (index_0[:, 0] < index_0[:, 1]) & (index_0[:, 1] < index_0[:, 2])
|
||
# flag_views = (views_all[index_0[:, 0]] != views_all[index_0[:, 1]]) & \
|
||
# (views_all[index_0[:, 1]] != views_all[index_0[:, 2]]) & \
|
||
# (views_all[index_0[:, 0]] != views_all[index_0[:, 2]])
|
||
valid_index = index_0[flag_order]
|
||
distance_circle = distance[valid_index[:, 0], valid_index[:, 1]] + \
|
||
distance[valid_index[:, 1], valid_index[:, 2]] + \
|
||
distance[valid_index[:, 2], valid_index[:, 0]]
|
||
distance_circle = distance_circle / 3
|
||
valid_dist = distance_circle < self.cfg.max_pixel
|
||
valid_ = valid_index[valid_dist]
|
||
dist_sum = distance_circle[valid_dist]
|
||
arg_idx = dist_sum.argsort()
|
||
FULL_LOG('[Assign 2D] find {} 3 pair: '.format(len(arg_idx)))
|
||
return valid_[arg_idx], dist_sum[arg_idx]
|
||
|
||
def try_to_add_index(self, dist_row, cameras, pelvis_all, views_all, dimGroups,
|
||
used_index, views, current, pid):
|
||
INLIER_REPRO = self.cfg.max_pixel
|
||
proposal = dist_row.argsort()
|
||
indices = []
|
||
for idx2d in proposal:
|
||
if dist_row[idx2d] > INLIER_REPRO:
|
||
break
|
||
# 这个视角已经有了 ｜ 这个2D已经被使用过了
|
||
if views_all[idx2d] in views:
|
||
continue
|
||
if used_index[idx2d] > -1:
|
||
continue
|
||
FULL_LOG('[Assign 2D] Try to add {}, distance={:.2f}'.format(idx2d, dist_row[idx2d]))
|
||
# 尝试三角化并进行重投影
|
||
new = current + [pelvis_all[idx2d]]
|
||
views_new = views + [views_all[idx2d]]
|
||
k3d, dist_repro, depth = self.triangulate_and_repro(cameras, views_new, new)
|
||
flag_depth = (depth > 0.5).all()
|
||
flag_repro = dist_repro.mean() < INLIER_REPRO
|
||
flag = flag_repro & flag_depth
|
||
FULL_LOG('[Assign 2D] repro: \n{}, \ndepth: \n{}'.format(LOG_ARRAY(dist_repro[None]), LOG_ARRAY(depth.T)))
|
||
if flag:
|
||
# 添加
|
||
current.append(pelvis_all[idx2d])
|
||
views.append(views_all[idx2d])
|
||
indices.append(idx2d)
|
||
FULL_LOG('[Assign 2D] Add {}'.format(idx2d ))
|
||
else:
|
||
FULL_LOG('[Assign 2D] Failed')
|
||
return indices
|
||
|
||
def assign_by_2D_3pair(self, results, distance, dimGroups, used_index, valid_3pairs, views_all, pelvis_all, cameras):
|
||
INLIER_REPRO = self.cfg.max_pixel
|
||
for ipair, valid_3pair in enumerate(valid_3pairs):
|
||
# 先检查是否被使用过了
|
||
if (used_index[valid_3pair] > -1).any():
|
||
continue
|
||
# 先检查是否是合理的
|
||
FULL_LOG('[Assign 2D] Check 3 pair {}'.format(valid_3pair))
|
||
k3d, dist_repro, depth = self.triangulate_and_repro(cameras, views_all[valid_3pair], pelvis_all[valid_3pair])
|
||
flag_depth = (depth > 0.5).all()
|
||
flag_repro = dist_repro.mean() < INLIER_REPRO
|
||
# TODO: flag range
|
||
flag = flag_repro & flag_depth
|
||
if not flag: continue
|
||
# 添加其余的点
|
||
pid = self.max_id
|
||
self.max_id += 1
|
||
dist_pair = distance[valid_3pair].mean(axis=0)
|
||
views = views_all[valid_3pair].tolist()
|
||
current = [pelvis_all[i] for i in valid_3pair]
|
||
indices = self.try_to_add_index(dist_pair, cameras, pelvis_all, views_all, dimGroups,
|
||
used_index, views, current, pid)
|
||
select = np.array(valid_3pair.tolist() + indices)
|
||
k3d, dist_repro, depth = self.triangulate_and_repro(cameras, views, current)
|
||
|
||
used_index[select] = pid
|
||
results.append({
|
||
'id': pid,
|
||
'pelvis': k3d,
|
||
'keypoints3d': k3d, # 这里保存两个，这样即使后面覆盖掉了keypoints3d还能取出pelvis来
|
||
'views': views_all[select],
|
||
'select': select,
|
||
'indices': select - dimGroups[views_all[select]],
|
||
'frames': [self.frames],
|
||
})
|
||
return results
|
||
|
||
def assign_by_2D(self, used_index, distance, pelvis_all, views_all, dimGroups, cameras):
|
||
def log_index_2d(index2d):
|
||
return '({}|{}-{})'.format(index2d, views_all[index2d], index2d-dimGroups[views_all[index2d]])
|
||
def log_indexes_2d(index2d_):
|
||
return ', '.join(['({}|{}-{})'.format(index2d, views_all[index2d], index2d-dimGroups[views_all[index2d]]) for index2d in index2d_])
|
||
INLIER_REPRO = self.cfg.max_pixel
|
||
new_id_start = 10000
|
||
new_max_id = new_id_start
|
||
valid_3pairs, dist_3pair = self.find_initial_3_pair(distance, pelvis_all, views_all, dimGroups=dimGroups)
|
||
results = []
|
||
if valid_3pairs.sum() > 0:
|
||
results = self.assign_by_2D_3pair(results, distance, dimGroups, used_index, valid_3pairs, views_all, pelvis_all, cameras)
|
||
valid_theta = cameras['valid_theta']
|
||
nViews = len(dimGroups)-1
|
||
# 排序计算affinity
|
||
count_rows = np.zeros((dimGroups[-1]), dtype=int)
|
||
distance_rows = np.zeros((dimGroups[-1]))
|
||
for nv in range(nViews):
|
||
if dimGroups[nv] == dimGroups[nv+1]:continue
|
||
valid_view = np.clip((distance[:, dimGroups[nv]:dimGroups[nv+1]] < INLIER_REPRO).sum(axis=-1), 0, 1)
|
||
count_rows += valid_view # 最多也只累计一个
|
||
distance_rows += valid_view * (distance[:, dimGroups[nv]:dimGroups[nv+1]].min(axis=-1))
|
||
index = list(range(dimGroups[-1]))
|
||
# index.sort(key=lambda x: (-count_rows[x], distance_rows[x]))
|
||
# sort with 2D confidence
|
||
# index.sort(key=lambda x: -pelvis_all[x, 2])
|
||
# sort with valid matches
|
||
# 选择2D的依据改为：根据有效的2D重投影距离的数量
|
||
valid_count = (distance < INLIER_REPRO * 2).sum(axis=0)
|
||
index = (-valid_count).argsort()
|
||
|
||
visited2d = set()
|
||
for row in index:
|
||
visited2d.add(row)
|
||
if used_index[row] > -1:continue
|
||
FULL_LOG('[Assign 2D] Check 2D {}'.format(log_index_2d(row)))
|
||
pid = new_max_id
|
||
new_max_id += 1
|
||
dist_row = distance[row]
|
||
proposal = dist_row.argsort()
|
||
current = [pelvis_all[row]]
|
||
views = [views_all[row]]
|
||
used_index[row] = pid
|
||
# 初始化一下：
|
||
for idx2d in proposal:
|
||
# 不满足视角关系
|
||
if not valid_theta[views_all[row], views_all[idx2d]]:
|
||
continue
|
||
# 不满足距离关系
|
||
if dist_row[idx2d] > INLIER_REPRO:
|
||
break
|
||
if used_index[idx2d] > -1:
|
||
continue
|
||
if views_all[idx2d] in views:
|
||
continue
|
||
# self.triangulate_and_repro(cameras, [views_all[18], views_all[34]], [pelvis_all[18], pelvis_all[34]])
|
||
# 2D的时候不能选择是最好的，因为2D可能还有其他视角的在
|
||
# 顶多判断一下，是对于这个视角来说最好的
|
||
# if not self.check_is_best_3d_of_2d(distance, row, idx2d, visited2d):
|
||
# continue
|
||
# 找到了合理的pair，作为一个良好的初始化
|
||
current.append(pelvis_all[idx2d])
|
||
views.append(views_all[idx2d])
|
||
used_index[idx2d] = pid
|
||
FULL_LOG(f'[Assign 2D] Init with {log_index_2d(idx2d)}')
|
||
break
|
||
if len(current) < 2:
|
||
# 没有找到良好的初始化
|
||
continue
|
||
for idx2d in proposal:
|
||
if dist_row[idx2d] > INLIER_REPRO:
|
||
break
|
||
# 这个视角已经有了 ｜ 这个2D已经被使用过了
|
||
if views_all[idx2d] in views:
|
||
continue
|
||
if used_index[idx2d] > -1:
|
||
continue
|
||
# if not self.check_is_best_3d_of_2d(distance, row, idx2d, visited2d):
|
||
# continue
|
||
# 尝试三角化并进行重投影
|
||
new = current + [pelvis_all[idx2d]]
|
||
views_new = views + [views_all[idx2d]]
|
||
k3d, dist_repro, depth = self.triangulate_and_repro(cameras, views_new, new)
|
||
flag_depth = (depth > 0.5).all()
|
||
flag_repro = dist_repro.mean() < INLIER_REPRO
|
||
flag = flag_repro & flag_depth
|
||
FULL_LOG('[Assign 2D] repro: \n{}, \ndepth: \n{}'.format(LOG_ARRAY(dist_repro[None]), LOG_ARRAY(depth.T)))
|
||
if flag:
|
||
# 添加
|
||
current = new
|
||
views = views_new
|
||
used_index[idx2d] = pid
|
||
_current_id = np.where(used_index == pid)[0]
|
||
FULL_LOG('[Assign 2D] {} => {}'.format(idx2d, log_indexes_2d(_current_id)))
|
||
else:
|
||
FULL_LOG('[Assign 2D] Failed')
|
||
if len(views) < self.cfg.min_views_init: #不足以添加
|
||
continue
|
||
k3d, dist_repro, depth = self.triangulate_and_repro(cameras, views, current)
|
||
select = np.where(used_index == pid)[0]
|
||
final_id = self.max_id
|
||
self.max_id += 1
|
||
used_index[select] = final_id
|
||
results.append({
|
||
'id': final_id,
|
||
'pelvis': k3d,
|
||
'keypoints3d': k3d, # 这里保存两个，这样即使后面覆盖掉了keypoints3d还能取出pelvis来
|
||
'views': views_all[select],
|
||
'select': select,
|
||
'indices': select - dimGroups[views_all[select]],
|
||
'frames': [self.frames],
|
||
})
|
||
for res in results:
|
||
text = f''' - Init {res['id']} with {len(res['views'])} views
|
||
views: {res['views']}
|
||
id : {res['select']}'''
|
||
LOG(text)
|
||
print(text)
|
||
return results
|
||
|
||
class MatchRoot(MatchBase):
|
||
def __init__(self, mode, cfg):
|
||
super().__init__(mode, cfg)
|
||
self.results = []
|
||
self.frames = -1
|
||
|
||
def __call__(self, pelvis, cameras, self_results=None):
|
||
"""
|
||
cameras: {K, R, T, dist, P}
|
||
"""
|
||
self.frames += 1
|
||
LOG('>>> Current frames: {}'.format(self.frames))
|
||
if self_results is None:
|
||
self_results = self.results
|
||
nViews = len(pelvis)
|
||
pelvis_all, views_all, dimGroups = self.stack_array(pelvis)
|
||
# Undistort
|
||
pelvis_undis = self.undistort(pelvis, cameras)
|
||
pelvis_undis_all, _, _ = self.stack_array(pelvis_undis)
|
||
# distance3D => 2D
|
||
distance3d_2d = self.calculate_repro(self_results, pelvis_all, cameras, views_all)
|
||
# FULL_LOG('distance3d_2d: {}'.format(LOG_ARRAY(distance3d_2d)))
|
||
# distance: triangulate and project
|
||
distance2d_2d = self.calculate_distance(pelvis_undis, cameras, dimGroups)
|
||
# FULL_LOG('distance2d_2d: {}'.format(LOG_ARRAY(distance2d_2d)))
|
||
# set assign index
|
||
used_index = np.zeros((dimGroups[-1]), dtype=int) - 1
|
||
results = []
|
||
# assign by 3D => 2D
|
||
results3d = self_results
|
||
if len(results3d) > 0:
|
||
results3d = self.assign_by_3D(used_index, distance3d_2d, pelvis_undis_all, views_all, dimGroups, cameras)
|
||
# assign by 2D + 2D
|
||
results2d = self.assign_by_2D(used_index, distance2d_2d, pelvis_undis_all, views_all, dimGroups, cameras)
|
||
results = results3d + results2d
|
||
# distance = np.linalg.norm(keypoints3d[:, None, ..., :3] - keypoints3d[None, ..., :3], axis=-1).mean(axis=-1)
|
||
# print(LOG_ARRAY(distance, format='{:6.2f}'))
|
||
results.sort(key=lambda x: -len(x['views']))
|
||
results = results[:self.cfg.max_person]
|
||
|
||
if self.mode == 'track':
|
||
self.results = results
|
||
results.sort(key=lambda x:x['id'])
|
||
# TODO: 增加结果的NMS检查和合并
|
||
if len(results) == 0:
|
||
keypoints3d = np.zeros((0, 25, 3))
|
||
else:
|
||
keypoints3d = np.stack([d['keypoints3d'] for d in results])
|
||
return {'keypoints3d': keypoints3d, 'results': results}
|
||
|
||
class MatchTwoRoot(MatchRoot):
|
||
def __init__(self, mode, cfg):
|
||
keys = ['pelvis', 'neck']
|
||
self._max_id_add = -1
|
||
self._max_id = {key: 0 for key in keys}
|
||
self.current = 'pelvis'
|
||
self._results = {key: [] for key in keys}
|
||
super().__init__(mode, cfg)
|
||
self.results_limb = []
|
||
self.mapping = {key: {} for key in keys}
|
||
|
||
@property
|
||
def max_id_add(self):
|
||
self._max_id_add += 1
|
||
return self._max_id_add
|
||
|
||
@property
|
||
def max_id(self):
|
||
return self._max_id[self.current]
|
||
|
||
@max_id.setter
|
||
def max_id(self, index):
|
||
self._max_id[self.current] = index
|
||
|
||
@property
|
||
def results(self):
|
||
return self._results[self.current]
|
||
|
||
@results.setter
|
||
def results(self, val):
|
||
self._results[self.current] = val
|
||
|
||
@staticmethod
|
||
def check_tracked(key, record_pelvis, current_3d, mapping):
|
||
for ires, res in enumerate(record_pelvis):
|
||
pid = res['id']
|
||
res['limb_id'] = -1
|
||
if pid in mapping[key]:
|
||
p3d = mapping[key][pid]
|
||
res['limb_id'] = p3d
|
||
current_3d[p3d][key] = ires
|
||
|
||
def __call__(self, cameras, openpose):
|
||
pelvis_id = 8
|
||
neck_id = 1
|
||
pelvis = [openpose[v][pelvis_id] for v in range(len(openpose))]
|
||
neck = [openpose[v][neck_id] for v in range(len(openpose))]
|
||
self.current = 'pelvis'
|
||
record_pelvis = super().__call__(pelvis, cameras)['results']
|
||
self.current = 'neck'
|
||
record_neck = super().__call__(neck, cameras)['results']
|
||
current_3d = {p['id']: {'pelvis': -1, 'neck': -1} for p in self.results_limb}
|
||
# 先检查是否已经track过了
|
||
self.check_tracked('pelvis', record_pelvis, current_3d, self.mapping)
|
||
self.check_tracked('neck', record_neck, current_3d, self.mapping)
|
||
# 先整体记录一下ID；然后如果某一帧有丢掉的；就更新
|
||
for p in self.results_limb:
|
||
# 检查一下当前帧
|
||
current_a, current_b = current_3d[p['id']]['pelvis'], current_3d[p['id']]['neck']
|
||
if current_a != -1 and current_b != -1:
|
||
assert current_a < len(record_pelvis) and current_b < len(record_neck), 'Index Error {}/{}, {}/{}'.format(current_a, current_b, len(record_pelvis), len(record_neck))
|
||
p['pelvis'] = record_pelvis[current_a]['pelvis']
|
||
p['neck'] = record_neck[current_b]['pelvis']
|
||
elif current_a == -1 and current_b != -1:
|
||
# a没有检测到，但b检测到了
|
||
# 保持相对值
|
||
mywarn('Missing Pelvis')
|
||
p['neck'] = record_neck[current_b]['pelvis']
|
||
pre_direc = p['pelvis'][:, :3] - p['neck'][:, :3]
|
||
p['pelvis'][:, :3] = p['neck'][:, :3] + pre_direc
|
||
# 得把补全的这个点设置回去
|
||
self._results['pelvis'].append({
|
||
'id': p['pelvis_id'],
|
||
'pelvis': p['pelvis'],
|
||
'views': [],
|
||
'frames': [],
|
||
'indices': [],
|
||
'limb_id': p['id'],
|
||
})
|
||
elif current_a != -1 and current_b == -1:
|
||
mywarn('Missing Neck')
|
||
pre_direc = p['neck'][:, :3] - p['pelvis'][:, :3]
|
||
p['pelvis'] = record_pelvis[current_a]['pelvis']
|
||
p['neck'][:, :3] = p['pelvis'][:, :3] + pre_direc
|
||
# 得把补全的这个点设置回去
|
||
self._results['neck'].append({
|
||
'id': p['neck_id'],
|
||
'pelvis': p['neck'],
|
||
'views': [],
|
||
'frames': [],
|
||
'indices': [],
|
||
'limb_id': p['id'],
|
||
})
|
||
else:
|
||
import ipdb; ipdb.set_trace()
|
||
raise NotImplementedError
|
||
# 遍历所有没有跟踪上的组合
|
||
n_pelvis = len(record_pelvis)
|
||
n_neck = len(record_neck)
|
||
dist = np.zeros((n_pelvis, n_neck))
|
||
# TODO: 用2D PAF来关联
|
||
for i in range(n_pelvis):
|
||
if record_pelvis[i]['limb_id'] > -1:
|
||
continue
|
||
for j in range(n_neck):
|
||
if record_neck[j]['limb_id'] > -1:
|
||
continue
|
||
pa = record_pelvis[i]['pelvis']
|
||
pb = record_neck[j]['pelvis']
|
||
length = np.linalg.norm(pa[:, :3] - pb[:, :3])
|
||
dist[i, j] = length
|
||
LIMB_MEAN = 0.489
|
||
dist_to_mean = np.exp(-(dist - LIMB_MEAN)**2/(2*(LIMB_MEAN/3)**2))
|
||
for i in range(n_pelvis):
|
||
if record_pelvis[i]['limb_id'] > -1:
|
||
continue
|
||
for j in range(n_neck):
|
||
if record_neck[j]['limb_id'] > -1:
|
||
continue
|
||
pa = record_pelvis[i]['pelvis']
|
||
pb = record_neck[j]['pelvis']
|
||
if dist_to_mean[i, j] > 0.8:
|
||
# 可以接受
|
||
limb = {
|
||
'id': self.max_id_add,
|
||
'pelvis_id': record_pelvis[i]['id'],
|
||
'neck_id': record_neck[j]['id'],
|
||
'pelvis': pa,
|
||
'neck': pb,
|
||
'frames': [self.frames],
|
||
}
|
||
self.mapping['pelvis'][limb['pelvis_id']] = limb['id']
|
||
self.mapping['neck'][limb['neck_id']] = limb['id']
|
||
self.results_limb.append(limb)
|
||
# 丢掉没有跟踪上的
|
||
results = []
|
||
for limb in self.results_limb:
|
||
k3d = np.vstack([limb['pelvis'], limb['neck']])
|
||
results.append({
|
||
'id': limb['id'],
|
||
'keypoints3d': k3d,
|
||
})
|
||
return {'results': results}
|
||
|
||
|
||
class MatchTorso(MatchBase):
|
||
def __init__(self, mode, cfg):
|
||
super().__init__(mode, cfg)
|
||
self.results = []
|
||
self.frames = -1
|
||
|
||
@staticmethod
|
||
def stack_pafs(pafs):
|
||
dimGroups = [0]
|
||
results = defaultdict(list)
|
||
views_all = []
|
||
for nv, paf in enumerate(pafs):
|
||
src = paf['src']
|
||
dimGroups.append(dimGroups[-1] + src.shape[0])
|
||
views_all.extend([nv for _ in range(src.shape[0])])
|
||
results['src'].append(src)
|
||
results['dst'].append(paf['dst'])
|
||
results['value'].append(paf['value'])
|
||
results = {key: np.concatenate(val, axis=0) for key, val in results.items()}
|
||
return results, np.array(views_all), np.array(dimGroups)
|
||
|
||
def check_used_index(self, info_limb, index, info_joints):
|
||
idx_src = info_limb['src'][index]
|
||
idx_dst = info_limb['dst'][index]
|
||
if info_joints['src']['used_index'][idx_src] > -1:
|
||
return True
|
||
if info_joints['dst']['used_index'][idx_dst] > -1:
|
||
return True
|
||
return False
|
||
|
||
def set_used_index(self, info_limb, index, info_joints, pid):
|
||
idx_src = info_limb['src'][index]
|
||
idx_dst = info_limb['dst'][index]
|
||
info_joints['src']['used_index'][idx_src] = pid
|
||
info_joints['dst']['used_index'][idx_dst] = pid
|
||
return True
|
||
|
||
def triangulate_limb(self, info_limb, info_joints, index, views, cameras):
|
||
flag = True
|
||
k3d_all = []
|
||
dist_all = []
|
||
for key in ['src', 'dst']:
|
||
proposals = []
|
||
for idx in index:
|
||
idx_ = info_limb[key][idx]
|
||
proposals.append(info_joints[key]['detect_undis'][idx_])
|
||
k3d, dist_repro, depth = self.triangulate_and_repro(cameras, views, proposals)
|
||
dist_all.append(dist_repro)
|
||
k3d_all.append(k3d)
|
||
k3d_all = np.vstack(k3d_all)
|
||
limb_length = np.linalg.norm(k3d_all[1, ..., :3] - k3d_all[0, ..., :3])
|
||
if limb_length < 0.3 or limb_length > 0.7:
|
||
flag = False
|
||
dist_all = np.stack(dist_all)
|
||
dist_all = np.max(dist_all, axis=0)
|
||
return flag, k3d_all, dist_all
|
||
|
||
# def assign_limb_by_2D(self, used_index, distance, pelvis_all, views_all, dimGroups, cameras):
|
||
def assign_limb_by_2D(self, info_limb, info_joints, distance, views_all, dimGroups, cameras):
|
||
def log_index_2d(index2d):
|
||
src = info_limb['src'][index2d]
|
||
dst = info_limb['dst'][index2d]
|
||
src = src - info_joints['src']['dimGroups'][views_all[index2d]]
|
||
dst = dst - info_joints['dst']['dimGroups'][views_all[index2d]]
|
||
return '({}|{}-({},{}))'.format(index2d, views_all[index2d], src, dst)
|
||
# def log_indexes_2d(index2d_):
|
||
# return ', '.join(['({}|{}-{})'.format(index2d, views_all[index2d], index2d-dimGroups[views_all[index2d]]) for index2d in index2d_])
|
||
|
||
INLIER_REPRO = self.cfg.max_pixel
|
||
valid_theta = cameras['valid_theta']
|
||
index = self.sort_with_affinity(distance, dimGroups, INLIER_REPRO)
|
||
visited2d = set()
|
||
results = []
|
||
new_id_start = 10000
|
||
new_max_id = new_id_start
|
||
for row in index:
|
||
visited2d.add(row)
|
||
if self.check_used_index(info_limb, row, info_joints):
|
||
continue
|
||
pid = new_max_id
|
||
new_max_id += 1
|
||
FULL_LOG('[Assign 2D] Check 2D {}'.format(log_index_2d(row)))
|
||
dist_row = distance[row]
|
||
proposal = dist_row.argsort()
|
||
# 尝试初始化
|
||
views = [views_all[row]]
|
||
current = [row]
|
||
for idx2d in proposal:
|
||
# 不满足视角关系
|
||
if not valid_theta[views_all[row], views_all[idx2d]]:
|
||
continue
|
||
# 不满足距离关系
|
||
if dist_row[idx2d] > INLIER_REPRO:
|
||
break
|
||
if self.check_used_index(info_limb, idx2d, info_joints):
|
||
continue
|
||
if views_all[idx2d] in views:
|
||
continue
|
||
# 检查骨长
|
||
flag, k3d, repro_error = self.triangulate_limb(info_limb, info_joints, [row, idx2d], [views_all[row], views_all[idx2d]], cameras)
|
||
length = np.linalg.norm(k3d[1, ..., :3] - k3d[0, ..., :3])
|
||
if flag:
|
||
views.append(views_all[idx2d])
|
||
current.append(idx2d)
|
||
FULL_LOG(f'[Assign 2D] Init with {log_index_2d(idx2d)}, length={length:.4f}')
|
||
break
|
||
else:
|
||
FULL_LOG(f'[Assign 2D] Init failed with {log_index_2d(idx2d)}, length = {length:.4f}')
|
||
if len(current) < 2:
|
||
# 没有找到良好的初始化
|
||
FULL_LOG(f'[Assign 2D] Cannot find a good initialization pair {log_index_2d(row)}')
|
||
continue
|
||
for idx2d in proposal:
|
||
if dist_row[idx2d] > INLIER_REPRO:break
|
||
# 这个视角已经有了 ｜ 这个2D已经被使用过了
|
||
if views_all[idx2d] in views:
|
||
continue
|
||
if self.check_used_index(info_limb, idx2d, info_joints):
|
||
continue
|
||
FULL_LOG('[Assign 2D] Try to add 2D {} => {}'.format(idx2d, log_index_2d(idx2d)))
|
||
# 尝试三角化并进行重投影
|
||
new = current + [idx2d]
|
||
views_new = views + [views_all[idx2d]]
|
||
flag_limb, k3d, dist_repro = self.triangulate_limb(info_limb, info_joints, new, views_new, cameras)
|
||
# flag_depth = (depth > 0.5).all()
|
||
flag_depth = True
|
||
flag_repro = dist_repro.mean() < INLIER_REPRO
|
||
flag = flag_repro & flag_depth & flag_limb
|
||
FULL_LOG('[Assign 2D] repro: \n{}'.format(LOG_ARRAY(dist_repro[None])))
|
||
if flag:
|
||
# 添加
|
||
current = new
|
||
views = views_new
|
||
self.set_used_index(info_limb, idx2d, info_joints, pid)
|
||
FULL_LOG('[Assign 2D] {} => {}'.format(idx2d, current))
|
||
else:
|
||
FULL_LOG('[Assign 2D] Failed')
|
||
new = None
|
||
views_new = None
|
||
if len(views) < self.cfg.min_views: #不足以添加
|
||
continue
|
||
flag_limb, k3d, dist_repro = self.triangulate_limb(info_limb, info_joints, current, views, cameras)
|
||
final_id = self.max_id
|
||
self.max_id += 1
|
||
results.append({
|
||
'id': final_id,
|
||
'torso': k3d,
|
||
'keypoints3d': k3d, # 这里保存两个，这样即使后面覆盖掉了keypoints3d还能取出pelvis来
|
||
'views': views,
|
||
'select': current,
|
||
# 'indices': select - dimGroups[views_all[select]],
|
||
'frames': [self.frames],
|
||
})
|
||
for res in results:
|
||
text = f''' - Init {res['id']} with {len(res['views'])} views
|
||
views: {res['views']}
|
||
id : {res['select']}'''
|
||
LOG(text)
|
||
print(text)
|
||
return results
|
||
|
||
def calculte_distance_src_dst(self, src, dst, cameras):
|
||
info = {}
|
||
for name, detect in zip(['src', 'dst'], [src, dst]):
|
||
detect_all, views_all, dimGroups = self.stack_array(detect)
|
||
# Undistort
|
||
detect_undis = self.undistort(detect, cameras)
|
||
detect_undis_all, _, _ = self.stack_array(detect_undis)
|
||
# # distance3D => 2D
|
||
# distance3d_2d = self.calculate_repro(self.results, pelvis_all, cameras, views_all)
|
||
# FULL_LOG('distance3d_2d: {}'.format(LOG_ARRAY(distance3d_2d)))
|
||
# distance: triangulate and project
|
||
distance2d_2d = self.calculate_distance(detect_undis, cameras, dimGroups)
|
||
info[name] = {
|
||
'detect_all': detect_all,
|
||
'views_all': views_all,
|
||
'dimGroups': dimGroups,
|
||
'distance2d_2d': distance2d_2d,
|
||
'detect_undis': detect_undis_all,
|
||
'used_index': np.zeros((dimGroups[-1]), dtype=int) - 1
|
||
}
|
||
return info
|
||
|
||
def get_valid_limbs(self, pafs, info_joint):
|
||
nViews = len(pafs)
|
||
valid_paf = []
|
||
for nv in range(nViews):
|
||
paf = pafs[nv]
|
||
src, dst = np.where(paf > 0.3)
|
||
value = paf[src, dst]
|
||
valid_paf.append({
|
||
'src': src + info_joint['src']['dimGroups'][nv],
|
||
'dst': dst + info_joint['dst']['dimGroups'][nv],
|
||
'value': value,
|
||
'view': nv,
|
||
})
|
||
results, views_all, dimGroups = self.stack_pafs(valid_paf)
|
||
return results, views_all, dimGroups
|
||
|
||
def calculate_distance_limb(self, results, cameras, dimGroups, distance_src, distance_dst):
|
||
src_idx, dst_idx = results['src'], results['dst']
|
||
src_idx0, src_idx1 = np.meshgrid(src_idx, src_idx)
|
||
dist_src_src = distance_src[src_idx0, src_idx1]
|
||
dst_idx0, dst_idx1 = np.meshgrid(dst_idx, dst_idx)
|
||
dist_dst_dst = distance_dst[dst_idx0, dst_idx1]
|
||
# TODO: 考虑每个视角的 limb的置信度，joint的置信度
|
||
dist_spatial = np.maximum(dist_src_src, dist_dst_dst)
|
||
return dist_spatial
|
||
|
||
def __call__(self, cameras, openpose, openpose_paf):
|
||
"""
|
||
cameras: {K, R, T, dist, P}
|
||
"""
|
||
self.frames += 1
|
||
pelvis_id = 8
|
||
neck_id = 1
|
||
nViews = len(openpose)
|
||
LOG('>>> Current frames: {}'.format(self.frames))
|
||
pelvis = [openpose[v][pelvis_id] for v in range(len(openpose))]
|
||
neck = [openpose[v][neck_id] for v in range(len(openpose))]
|
||
info_joint = self.calculte_distance_src_dst(pelvis, neck, cameras)
|
||
pafs = [openpose_paf[v][(pelvis_id, neck_id)] for v in range(len(openpose_paf))]
|
||
info_limb, views_all, dimGroups = self.get_valid_limbs(pafs, info_joint)
|
||
distance2d_2d = self.calculate_distance_limb(info_limb, cameras, dimGroups,
|
||
info_joint['src']['distance2d_2d'], info_joint['dst']['distance2d_2d'])
|
||
results = self.assign_limb_by_2D(info_limb, info_joint, distance2d_2d, views_all, dimGroups, cameras)
|
||
results.sort(key=lambda x: -len(x['views']))
|
||
results = results[:self.cfg.max_person]
|
||
# if self.mode == 'track':
|
||
# self.results = results
|
||
results.sort(key=lambda x:x['id'])
|
||
|
||
# TODO: 增加结果的NMS检查和合并
|
||
if len(results) == 0:
|
||
keypoints3d = np.zeros((0, 2, 3))
|
||
else:
|
||
keypoints3d = np.stack([d['keypoints3d'] for d in results])
|
||
return {'keypoints3d': keypoints3d, 'results': results}
|
||
|
||
class TriangulateAll:
|
||
def __init__(self, mode) -> None:
|
||
self.mode = mode
|
||
|
||
def __call__(self, bbox, keypoints, cameras, results):
|
||
for res in results:
|
||
bbox_, k2d, Pall = [], [], []
|
||
for i in range(len(res['views'])):
|
||
v = res['views'][i]
|
||
bbox_.append(bbox[v][res['indices'][i]])
|
||
k2d.append(keypoints[v][res['indices'][i]])
|
||
Pall.append(cameras['P'][v])
|
||
k2d = np.stack(k2d)
|
||
Pall = np.stack(Pall)
|
||
bbox_ = np.stack(bbox_)
|
||
if self.mode == 'naive':
|
||
k3d = batch_triangulate(k2d, Pall)
|
||
elif self.mode == 'robust':
|
||
from easymocap.mytools.triangulator import iterative_triangulate
|
||
k3d, k2d = iterative_triangulate(k2d, Pall,
|
||
dist_max=25)
|
||
res['keypoints3d'] = k3d
|
||
res['keypoints2d'] = k2d
|
||
res['bbox'] = bbox_
|
||
return {'keypoints3d': np.stack([d['keypoints3d'] for d in results]), 'results': results}
|
||
|
||
class MatchHandLR:
|
||
def __init__(self, mode, cfg):
|
||
self.model_l = MatchRoot(mode,cfg)
|
||
self.model_r = MatchRoot(mode,cfg)
|
||
def __call__(self, pelvis_l, pelvis_r, cameras):
|
||
ret = {}
|
||
outl = self.model_l(pelvis_l, cameras)
|
||
outr = self.model_r(pelvis_r, cameras)
|
||
for k in outl.keys():
|
||
ret[k+'_l'] = outl[k]
|
||
for k in outr.keys():
|
||
ret[k+'_r'] = outr[k]
|
||
return ret
|
||
|
||
class MatchBodyHand:
|
||
def __init__(self, mode) -> None:
|
||
pass
|
||
|
||
def projectPoints(self, X, K, R, t, Kd):
|
||
x = R @ X + t
|
||
x[0:2,:] = x[0:2,:]/x[2,:]#到归一化平面
|
||
r = x[0,:]*x[0,:] + x[1,:]*x[1,:]
|
||
|
||
x[0,:] = x[0,:]*(1 + Kd[0]*r + Kd[1]*r*r + Kd[4]*r*r*r) + 2*Kd[2]*x[0,:]*x[1,:] + Kd[3]*(r + 2*x[0,:]*x[0,:])
|
||
x[1,:] = x[1,:]*(1 + Kd[0]*r + Kd[1]*r*r + Kd[4]*r*r*r) + 2*Kd[3]*x[0,:]*x[1,:] + Kd[2]*(r + 2*x[1,:]*x[1,:])
|
||
x[0,:] = K[0,0]*x[0,:] + K[0,1]*x[1,:] + K[0,2]
|
||
x[1,:] = K[1,0]*x[0,:] + K[1,1]*x[1,:] + K[1,2]
|
||
return x
|
||
def match3d_step(self, results, keypoints3d,wristid):
|
||
match_results=(np.zeros((len(keypoints3d)),dtype=np.int)-1).tolist()
|
||
vis = (np.zeros((len(keypoints3d)))-1).tolist()
|
||
dis = []
|
||
for i in range(len(keypoints3d)):
|
||
for j in range(len(results)):
|
||
dis.append([i,j,((keypoints3d[i][wristid][:3]-results[j]['pelvis'][0,:3].reshape(-1))**2).sum()])
|
||
if(len(dis)>0):
|
||
dis = np.array(dis)
|
||
dis = dis[np.argsort(dis[:,-1])]
|
||
for i in range(len(dis)):
|
||
bid =int(dis[i][0])
|
||
hid =int(dis[i][1])
|
||
if vis[bid]>=0 or hid in vis:
|
||
continue
|
||
if dis[i][2]>0.5:
|
||
continue
|
||
tmp_results = results[hid].copy()
|
||
# tmp_results['dis_bh'] = dis[i][2]
|
||
match_results[bid]=tmp_results
|
||
vis[bid]=hid
|
||
return match_results
|
||
def match2d_step(self, bbox_hand, keypoints3d, wristid, results_match_l, cameras):
|
||
lack_body_id=[]
|
||
mv_use_hand=[]
|
||
for i in range(cameras['R'].shape[0]):
|
||
mv_use_hand.append([])
|
||
|
||
for i in range(len(results_match_l)):
|
||
if isinstance(results_match_l[i],int) and results_match_l[i]==-1:
|
||
lack_body_id.append(i)
|
||
else:
|
||
mv = results_match_l[i]['views']#[cid]
|
||
indices = results_match_l[i]['indices']#[cid]
|
||
for j in range(len(mv)):
|
||
mv_use_hand[mv[j]].append(indices[j])
|
||
|
||
wrist3dkpts = keypoints3d[lack_body_id,wristid,:3] #(nperson,3)每个人呢的wrist关键点
|
||
dis = []
|
||
for nv in range(len(bbox_hand)):
|
||
for hid in range(len(bbox_hand[nv])):
|
||
if hid in mv_use_hand[nv]:
|
||
continue
|
||
if bbox_hand[nv][hid][-1]==0:
|
||
continue
|
||
bx_ = bbox_hand[nv][hid]
|
||
k2d = np.array([(bx_[0]+bx_[2])/2,(bx_[1]+bx_[3])/2,bx_[-1]])
|
||
K = cameras['K'][nv]
|
||
Kd = cameras['dist'][nv].reshape(5)
|
||
R = cameras['R'][nv]
|
||
t = cameras['T'][nv]
|
||
wristkpts2d = self.projectPoints(wrist3dkpts.T[0:3,:], K, R, t, Kd).T
|
||
for bid in range(len(lack_body_id)):
|
||
D = ((wristkpts2d[bid][:2]-k2d[:2])**2).sum()
|
||
dis.append([D,lack_body_id[bid],nv,hid]) # 误差，3d身体id ,视角编号 ,2d图像上手box id
|
||
|
||
if(len(dis)>0):
|
||
vis = (np.zeros((len(keypoints3d)))-1).tolist()
|
||
dis = np.array(dis)
|
||
dis = dis[np.argsort(dis[:,0])]
|
||
# TODO 判断dis大小，将dis过大的删除掉
|
||
for i in range(len(dis)):
|
||
bid = int(dis[i][1])
|
||
nv = int(dis[i][2])
|
||
hid = int(dis[i][3])
|
||
if vis[bid]>=0 or hid in vis or results_match_l[bid]!=-1:
|
||
continue
|
||
if dis[i][0]>50: #人和手的在2D中距离
|
||
continue
|
||
results_match_l[bid]={
|
||
'views': np.array([nv]),
|
||
'indices': np.array([hid]), # ?indices是在对应的视角下第几个Box
|
||
# 'dis_bh': dis[i][0]
|
||
}
|
||
|
||
vis[bid]=hid
|
||
return results_match_l
|
||
|
||
|
||
def __call__(self, results_l, results_r, keypoints3d, cameras, bbox_handl, bbox_handr):
|
||
'''
|
||
results: list nhand
|
||
keypoints3d: (nperson,25,3)
|
||
'''
|
||
results_match_l = self.match3d_step(results_l, keypoints3d, 7)
|
||
results_match_r = self.match3d_step(results_r, keypoints3d, 4)
|
||
|
||
if(-1 in results_match_l):
|
||
# TODO: dis为空，则表示没有身体，或者所有视角都未检测到手，尝试启动单视角检测
|
||
# TODO: dis不为空，也有可能有的身体缺少与手的匹配，可以尝试单视角检测，或者之后尝试补全。
|
||
# 单视角匹配，从匹配列表中找出-1的部分，将其投影到多视角中，在多视角找出未被选择的box，然后匹配，记录在a
|
||
results_match_l = self.match2d_step(bbox_handl, keypoints3d, 7, results_match_l, cameras)
|
||
if(-1 in results_match_r):
|
||
results_match_r = self.match2d_step(bbox_handr, keypoints3d, 4, results_match_r, cameras)
|
||
return {'match3d_l':results_match_l ,'match3d_r':results_match_r}
|