Speed up pair creation and matching.

This commit is contained in:
Daniel
2024-08-06 10:55:53 +02:00
parent 6002aeddde
commit 36c236b8d8
3 changed files with 605 additions and 520 deletions

File diff suppressed because it is too large Load Diff

View File

@ -319,6 +319,7 @@ def main():
all_ids = []
all_paths = []
times = []
last_poses_3d = np.array([])
for label in tqdm.tqdm(labels):
images_2d = []
@ -377,7 +378,7 @@ def main():
poses2D = np.zeros([len(images_2d), 1, len(joint_names_3d), 3])
else:
poses3D = triangulate_poses.get_3d_pose(
poses_2d, label["cameras"], roomparams, joint_names_2d, minscore
poses_2d, label["cameras"], roomparams, joint_names_2d, last_poses_3d, minscore
)
poses2D = []
for cam in label["cameras"]:
@ -392,6 +393,7 @@ def main():
drop_few_limbs=(dataset_use != "mvor"),
)
poses3D = add_missing_joints(poses3D, joint_names_3d)
last_poses_3d = poses3D
time_3d = time.time() - start
print("3D time:", time_3d)

View File

@ -1,6 +1,5 @@
import copy
import math
import time
import cv2
import numpy as np
@ -66,6 +65,29 @@ def get_camera_P(cam):
# ==================================================================================================
def calc_pose_score(pose1, pose2, dist1, cam1, joint_names, use_joints):
"""Calculates the score between two poses"""
# Select core joints
jids = [joint_names.index(j) for j in use_joints]
pose1 = pose1[jids]
pose2 = pose2[jids]
dist1 = dist1[jids]
mask = (pose1[:, 2] > 0.1) & (pose2[:, 2] > 0.1)
if np.sum(mask) < 3:
return 0.0
iscale = (cam1["width"] + cam1["height"]) / 2
scores = score_projection(pose1, pose2, dist1, mask, iscale)
score = np.mean(scores)
return score
# ==================================================================================================
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"""
@ -79,14 +101,38 @@ def calc_pair_score(pair, poses_2d, camparams, roomparams, joint_names_2d, use_j
pose1 = pose1[jids]
pose2 = pose2[jids]
poses_3d, score = calc_pose_scored(pose1, pose2, cam1, cam2, roomparams)
poses_3d, score = triangulate_and_score(pose1, pose2, cam1, cam2, roomparams)
return poses_3d, score
# ==================================================================================================
def calc_pose_scored(pose1, pose2, cam1, cam2, roomparams):
def score_projection(pose1, repro1, dists1, mask, iscale):
min_score = 0.1
error1 = np.linalg.norm(pose1[mask, 0:2] - repro1[mask, 0:2], axis=1)
# Set errors of invisible reprojections to a high value
penalty = iscale
mask1b = (repro1[:, 2] < min_score)[mask]
error1[mask1b] = penalty
# Scale error by image size and distance to the camera
error1 = error1.clip(0, iscale / 4) / iscale
dscale1 = np.sqrt(np.mean(dists1[mask]) / 3.5)
error1 = error1 * dscale1
# Convert errors to a score
score1 = 1.0 / (1.0 + error1 * 10)
return score1
# ==================================================================================================
def triangulate_and_score(pose1, pose2, cam1, cam2, roomparams):
"""Triangulates a pair of persons and scores them based on the reprojection error"""
# Mask out invisible joints
@ -129,27 +175,16 @@ def calc_pose_scored(pose1, pose2, cam1, cam2, roomparams):
poses_3d = np.expand_dims(pose3d, axis=0)
repro1, dists1 = utils_pose.project_poses(poses_3d, cam1, calc_dists=True)
repro2, dists2 = utils_pose.project_poses(poses_3d, cam2, calc_dists=True)
error1 = np.linalg.norm(pose1[mask, 0:2] - repro1[0, mask, 0:2], axis=1)
error2 = np.linalg.norm(pose2[mask, 0:2] - repro2[0, mask, 0:2], axis=1)
repro1, dists1 = repro1[0], dists1[0]
repro2, dists2 = repro2[0], dists2[0]
# Set errors of invisible reprojections to a high value
penalty = (cam1["width"] + cam1["height"]) / 2
mask1b = (repro1[0, :, 2] < min_score)[mask]
mask2b = (repro2[0, :, 2] < min_score)[mask]
error1[mask1b] = penalty
error2[mask2b] = penalty
# Convert errors to a score
# Scale by image size and distance to the camera
# Calculate scores for each view
iscale = (cam1["width"] + cam1["height"]) / 2
error1 = error1.clip(0, iscale / 4) / iscale
error2 = error2.clip(0, iscale / 4) / iscale
dscale1 = np.sqrt(np.mean(dists1[0, mask]) / 3.5)
dscale2 = np.sqrt(np.mean(dists2[0, mask]) / 3.5)
error1 = error1 * dscale1
error2 = error2 * dscale2
error = (error1 + error2) / 2
scores = 1.0 / (1.0 + error * 10)
score1 = score_projection(pose1, repro1, dists1, mask, iscale)
score2 = score_projection(pose2, repro2, dists2, mask, iscale)
# Combine scores
scores = (score1 + score2) / 2
# Drop lowest scores
drop_k = math.floor(len(pose1) * 0.2)
@ -283,7 +318,14 @@ def merge_group(poses_3d: np.ndarray, min_score: float):
# ==================================================================================================
def get_3d_pose(poses_2d, camparams, roomparams, joint_names_2d, min_score=0.95):
def get_3d_pose(
poses_2d,
camparams,
roomparams,
joint_names_2d,
last_poses_3d=np.array([]),
min_score=0.95,
):
"""Triangulates 3D poses from 2D poses of multiple views"""
# Convert poses and camparams to numpy arrays
@ -314,10 +356,42 @@ def get_3d_pose(poses_2d, camparams, roomparams, joint_names_2d, min_score=0.95)
cam["P"] = get_camera_P(cam)
camparams[i] = cam
# Project last 3D poses to 2D
last_poses_2d = []
last_poses_3d = np.asarray(last_poses_3d)
if last_poses_3d.size > 0:
for i in range(len(camparams)):
poses2d, dists = utils_pose.project_poses(last_poses_3d, camparams[i])
last_poses_2d.append((poses2d, dists))
# Check matches to old poses
threshold = min_score - 0.2
scored_pasts = {}
if last_poses_3d.size > 0:
for i in range(len(camparams)):
scored_pasts[i] = {}
poses = poses_2d[i]
last_poses, dists = last_poses_2d[i]
for j in range(len(last_poses)):
scored_pasts[i][j] = []
for k in range(len(poses)):
score = calc_pose_score(
poses[k],
last_poses[j],
dists[j],
camparams[i],
joint_names_2d,
core_joints,
)
if score > threshold:
scored_pasts[i][j].append(k)
# Create pairs of persons
# Checks if the person was already matched to the last frame and if so only creates pairs with those
# Else it creates all possible pairs
num_persons = [len(p) for p in poses_2d]
all_pairs = []
for i in range(len(poses_2d)):
for i in range(len(camparams)):
poses = poses_2d[i]
for j in range(i + 1, len(poses_2d)):
poses2 = poses_2d[j]
@ -325,6 +399,15 @@ def get_3d_pose(poses_2d, camparams, roomparams, joint_names_2d, min_score=0.95)
for l in range(len(poses2)):
pid1 = sum(num_persons[:i]) + k
pid2 = sum(num_persons[:j]) + l
match = False
if last_poses_3d.size > 0:
for m in range(len(last_poses_3d)):
if k in scored_pasts[i][m] and l in scored_pasts[j][m]:
match = True
all_pairs.append([(i, j, k, l), (pid1, pid2)])
elif k in scored_pasts[i][m] or l in scored_pasts[j][m]:
match = True
if not match:
all_pairs.append([(i, j, k, l), (pid1, pid2)])
# Calculate pair scores
@ -356,7 +439,7 @@ def get_3d_pose(poses_2d, camparams, roomparams, joint_names_2d, min_score=0.95)
pose1 = poses_2d[pair[0][0]][pair[0][2]]
pose2 = poses_2d[pair[0][1]][pair[0][3]]
pose_3d, _ = calc_pose_scored(pose1, pose2, cam1, cam2, roomparams)
pose_3d, _ = triangulate_and_score(pose1, pose2, cam1, cam2, roomparams)
pair.append(pose_3d)
# Merge groups