334 lines
11 KiB
Python
334 lines
11 KiB
Python
import copy
|
|
import math
|
|
import time
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from skelda import utils_pose
|
|
|
|
# ==================================================================================================
|
|
|
|
core_joints = [
|
|
"shoulder_left",
|
|
"shoulder_right",
|
|
"hip_left",
|
|
"hip_right",
|
|
"elbow_left",
|
|
"elbow_right",
|
|
"knee_left",
|
|
"knee_right",
|
|
"wrist_left",
|
|
"wrist_right",
|
|
"ankle_left",
|
|
"ankle_right",
|
|
]
|
|
|
|
# ==================================================================================================
|
|
|
|
|
|
def undistort_points(points: np.ndarray, caminfo: dict):
|
|
K = np.array(caminfo["K"], dtype=np.float32)
|
|
DC = np.array(caminfo["DC"][0:5], dtype=np.float32)
|
|
w = caminfo["width"]
|
|
h = caminfo["height"]
|
|
|
|
# Undistort camera matrix
|
|
newK, _ = cv2.getOptimalNewCameraMatrix(K, DC, (w, h), 1, (w, h))
|
|
caminfo = copy.deepcopy(caminfo)
|
|
caminfo["K"] = newK.tolist()
|
|
caminfo["DC"] = [0.0, 0.0, 0.0, 0.0, 0.0]
|
|
|
|
# Undistort points
|
|
pshape = points.shape
|
|
points = np.reshape(points, (-1, 1, 2))
|
|
points = cv2.undistortPoints(points, K, DC, P=newK)
|
|
points = points.reshape(pshape)
|
|
|
|
return points, caminfo
|
|
|
|
|
|
# ==================================================================================================
|
|
|
|
|
|
def get_camera_P(cam):
|
|
"""Calculate opencv-style projection matrix"""
|
|
|
|
R = np.array(cam["R"])
|
|
T = np.array(cam["T"])
|
|
Tr = R @ (T * -1)
|
|
P = cam["K"] @ np.hstack([R, Tr])
|
|
return P
|
|
|
|
|
|
# ==================================================================================================
|
|
|
|
|
|
def calc_pair_score(pair, poses_2d, camparams, roomparams, joint_names_2d, use_joints):
|
|
"""Triangulates a pair of persons and scores them based on the reprojection error"""
|
|
|
|
cam1 = camparams[pair[0][0]]
|
|
cam2 = camparams[pair[0][1]]
|
|
pose1 = np.array(poses_2d[pair[0][0]][pair[0][2]])
|
|
pose2 = np.array(poses_2d[pair[0][1]][pair[0][3]])
|
|
|
|
# Select core joints
|
|
jids = [joint_names_2d.index(j) for j in use_joints]
|
|
pose1 = pose1[jids]
|
|
pose2 = pose2[jids]
|
|
|
|
poses_3d, score = calc_pose_scored(pose1, pose2, cam1, cam2, roomparams)
|
|
return poses_3d, score
|
|
|
|
|
|
# ==================================================================================================
|
|
|
|
|
|
def calc_pose_scored(pose1, pose2, cam1, cam2, roomparams):
|
|
"""Triangulates a pair of persons and scores them based on the reprojection error"""
|
|
|
|
# Mask out invisible joints
|
|
mask1a = pose1[:, 2] >= 0.1
|
|
mask2a = pose2[:, 2] >= 0.1
|
|
mask = mask1a & mask2a
|
|
|
|
# If no joints are visible return a low score
|
|
if np.sum(mask) == 0:
|
|
pose3d = np.zeros([len(pose1), 4])
|
|
score = 0.0
|
|
return pose3d, score
|
|
|
|
# Triangulate points
|
|
points1 = pose1[mask, 0:2].T
|
|
points2 = pose2[mask, 0:2].T
|
|
P1 = get_camera_P(cam1)
|
|
P2 = get_camera_P(cam2)
|
|
points3d = cv2.triangulatePoints(P1, P2, points1, points2)
|
|
points3d = points3d / points3d[3, :]
|
|
points3d = points3d[0:3, :].T
|
|
pose3d = np.zeros([len(pose1), 4])
|
|
pose3d[mask] = np.concatenate([points3d, np.ones([points3d.shape[0], 1])], axis=-1)
|
|
|
|
# If the triangulated points are outside the room drop it
|
|
mean = np.mean(pose3d[mask][:, 0:3], axis=0)
|
|
mins = np.min(pose3d[mask][:, 0:3], axis=0)
|
|
maxs = np.max(pose3d[mask][:, 0:3], axis=0)
|
|
rsize = np.array(roomparams["room_size"]) / 2
|
|
rcent = np.array(roomparams["room_center"])
|
|
center_outside = np.any((mean > rsize + rcent) | (mean < -rsize + rcent))
|
|
limb_outside = np.any((maxs > rsize + rcent + 0.1) | (mins < -rsize + rcent - 0.1))
|
|
if center_outside or limb_outside:
|
|
pose3d[:, 3] = 0.001
|
|
score = 0.001
|
|
return pose3d, score
|
|
|
|
# Calculate reprojection error
|
|
poses_3d = np.expand_dims(pose3d, axis=0)
|
|
repro1, _ = utils_pose.project_poses(poses_3d, cam1)
|
|
repro2, _ = utils_pose.project_poses(poses_3d, cam2)
|
|
repro1 = repro1[0]
|
|
repro2 = repro2[0]
|
|
error1 = np.linalg.norm(pose1[mask, 0:2] - repro1[mask, 0:2], axis=1)
|
|
error2 = np.linalg.norm(pose2[mask, 0:2] - repro2[mask, 0:2], axis=1)
|
|
|
|
# Set errors of invisible reprojections to a high value
|
|
penalty = (cam1["width"] + cam1["height"]) / 2
|
|
mask1b = (repro1[:, 2] < 0.1)[mask]
|
|
mask2b = (repro2[:, 2] < 0.1)[mask]
|
|
error1[mask1b] = penalty
|
|
error2[mask2b] = penalty
|
|
|
|
# Convert errors to a score
|
|
scale = (cam1["width"] + cam1["height"]) / 2
|
|
error1 = error1.clip(0, scale / 4)
|
|
error2 = error2.clip(0, scale / 4)
|
|
error1 = error1 / scale
|
|
error2 = error2 / scale
|
|
error = (error1 + error2) / 2
|
|
scores = 1.0 / (1.0 + error * 10)
|
|
score = np.mean(scores)
|
|
|
|
# Add score to 3D pose
|
|
full_scores = np.zeros([poses_3d.shape[1], 1])
|
|
full_scores[mask] = np.expand_dims(scores, axis=-1)
|
|
pose3d[:, 3] = full_scores[:, 0]
|
|
|
|
return pose3d, score
|
|
|
|
|
|
# ==================================================================================================
|
|
|
|
|
|
def calc_grouping(all_pairs):
|
|
"""Groups pairs that share a person"""
|
|
|
|
# Calculate the pose center for each pair
|
|
min_score = 0.9
|
|
for i in range(len(all_pairs)):
|
|
pair = all_pairs[i]
|
|
pose_3d = pair[2][0]
|
|
mask = pose_3d[:, 3] > min_score
|
|
center = np.mean(pose_3d[mask, 0:3], axis=0)
|
|
all_pairs[i] = all_pairs[i] + [center]
|
|
|
|
groups = []
|
|
for i in range(len(all_pairs)):
|
|
pair = all_pairs[i]
|
|
|
|
# Create new group if non exists
|
|
if len(groups) == 0:
|
|
groups.append([pair[4], pair[2][0], [pair]])
|
|
continue
|
|
|
|
# Check if the pair matches to an existing group
|
|
max_center_dist = 0.9
|
|
max_joint_avg_dist = 0.3
|
|
best_dist = math.inf
|
|
best_group = -1
|
|
for j in range(len(groups)):
|
|
g0 = groups[j]
|
|
center0 = g0[0]
|
|
center1 = pair[4]
|
|
if np.linalg.norm(center0 - center1) < max_center_dist:
|
|
pose0 = g0[1]
|
|
pose1 = pair[2][0]
|
|
|
|
# Calculate the distance between the two poses
|
|
mask0 = pose0[:, 3] > min_score
|
|
mask1 = pose1[:, 3] > min_score
|
|
mask = mask0 & mask1
|
|
dists = np.linalg.norm(pose0[mask, 0:3] - pose1[mask, 0:3], axis=1)
|
|
dist = np.mean(dists)
|
|
if dist < max_joint_avg_dist:
|
|
if dist < best_dist:
|
|
best_dist = dist
|
|
best_group = j
|
|
if best_group >= 0:
|
|
# Add pair to existing group and update the mean positions
|
|
group = groups[best_group]
|
|
new_center = (group[0] * len(group[1]) + pair[4]) / (len(group[1]) + 1)
|
|
new_pose = (group[1] * len(group[1]) + pair[2][0]) / (len(group[1]) + 1)
|
|
group[2].append(pair)
|
|
group[0] = new_center
|
|
group[1] = new_pose
|
|
else:
|
|
# Create new group if no match was found
|
|
groups.append([pair[4], pair[2][0], [pair]])
|
|
|
|
return groups
|
|
|
|
|
|
# ==================================================================================================
|
|
|
|
|
|
def merge_group(poses_3d: np.ndarray):
|
|
"""Merges a group of poses into a single pose"""
|
|
|
|
# Merge poses to create initial pose
|
|
# Use only those triangulations with a high score
|
|
min_score = 0.9
|
|
mask = poses_3d[:, :, 3:4] > min_score
|
|
sum_poses = np.sum(poses_3d * mask, axis=0)
|
|
sum_mask = np.sum(mask, axis=0)
|
|
initial_pose_3d = np.divide(
|
|
sum_poses, sum_mask, where=(sum_mask > 0), out=np.zeros_like(sum_poses)
|
|
)
|
|
|
|
# Drop outliers that are far away from the other proposals
|
|
max_dist = 0.3
|
|
distances = np.linalg.norm(
|
|
poses_3d[:, :, :3] - initial_pose_3d[np.newaxis, :, :3], axis=2
|
|
)
|
|
dist_mask = distances <= max_dist
|
|
mask = mask & np.expand_dims(dist_mask, axis=-1)
|
|
|
|
# Select the best-k proposals for each joint that are closest to the initial pose
|
|
keep_best = 3
|
|
sorted_indices = np.argsort(distances, axis=0)
|
|
best_k_mask = np.zeros_like(mask, dtype=bool)
|
|
num_joints = poses_3d.shape[1]
|
|
for i in range(num_joints):
|
|
valid_indices = sorted_indices[:, i][mask[sorted_indices[:, i], i, 0]]
|
|
best_k_mask[valid_indices[:keep_best], i, 0] = True
|
|
mask = mask & best_k_mask
|
|
|
|
# Final pose computation with combined masks
|
|
sum_poses = np.sum(poses_3d * mask, axis=0)
|
|
sum_mask = np.sum(mask, axis=0)
|
|
final_pose_3d = np.divide(
|
|
sum_poses, sum_mask, where=(sum_mask > 0), out=np.zeros_like(sum_poses)
|
|
)
|
|
|
|
return final_pose_3d
|
|
|
|
|
|
# ==================================================================================================
|
|
|
|
|
|
def get_3d_pose(poses_2d, camparams, roomparams, joint_names_2d):
|
|
|
|
# Undistort 2D points
|
|
for i, cam in enumerate(camparams):
|
|
poses = np.array(poses_2d[i])
|
|
poses[:, :, 0:2], camparams[i] = undistort_points(poses[:, :, 0:2], cam)
|
|
poses_2d[i] = poses.tolist()
|
|
|
|
# Create pairs of persons
|
|
num_persons = [len(p) for p in poses_2d]
|
|
all_pairs = []
|
|
for i in range(len(poses_2d)):
|
|
poses = poses_2d[i]
|
|
for j in range(i + 1, len(poses_2d)):
|
|
poses2 = poses_2d[j]
|
|
for k in range(len(poses)):
|
|
for l in range(len(poses2)):
|
|
pid1 = sum(num_persons[:i]) + k
|
|
pid2 = sum(num_persons[:j]) + l
|
|
all_pairs.append([(i, j, k, l), (pid1, pid2)])
|
|
|
|
# Calculate pair scores
|
|
for i in range(len(all_pairs)):
|
|
pair = all_pairs[i]
|
|
pose_3d, score = calc_pair_score(
|
|
pair, poses_2d, camparams, roomparams, joint_names_2d, core_joints
|
|
)
|
|
all_pairs[i].append((pose_3d, score))
|
|
|
|
# import draw_utils
|
|
# poses3D = np.array([pose_3d])
|
|
# _ = draw_utils.utils_view.show_poses3d(
|
|
# poses3D, core_joints, {}, camparams
|
|
# )
|
|
# draw_utils.utils_view.show_plots()
|
|
|
|
# Drop pairs with low scores
|
|
min_score = 0.9
|
|
all_pairs = [p for p in all_pairs if p[2][1] > min_score]
|
|
|
|
# Calculate full 3D poses
|
|
poses_3d = []
|
|
for pair in all_pairs:
|
|
cam1 = camparams[pair[0][0]]
|
|
cam2 = camparams[pair[0][1]]
|
|
pose1 = np.array(poses_2d[pair[0][0]][pair[0][2]])
|
|
pose2 = np.array(poses_2d[pair[0][1]][pair[0][3]])
|
|
|
|
pose_3d, _ = calc_pose_scored(pose1, pose2, cam1, cam2, roomparams)
|
|
pair.append(pose_3d)
|
|
|
|
# Group pairs that share a person
|
|
groups = calc_grouping(all_pairs)
|
|
|
|
# Merge groups
|
|
poses_3d = []
|
|
for group in groups:
|
|
poses = np.array([p[3] for p in group[2]])
|
|
pose_3d = merge_group(poses)
|
|
poses_3d.append(pose_3d)
|
|
|
|
if len(poses_3d) > 0:
|
|
poses3D = np.array(poses_3d)
|
|
else:
|
|
poses3D = np.zeros([1, len(joint_names_2d), 4])
|
|
return poses3D
|