181 lines
6.4 KiB
Python
181 lines
6.4 KiB
Python
import numpy as np
|
|
import torch
|
|
|
|
from body_model import OP_EDGE_LIST
|
|
from geometry.rotation import batch_rodrigues
|
|
from geometry.plane import (
|
|
parse_floor_plane,
|
|
compute_plane_intersection,
|
|
fit_plane,
|
|
)
|
|
|
|
|
|
def compute_world2prior(floor_plane, trans, root_orient, origin):
|
|
"""
|
|
Computes rotation and translation from the camera frame to the canonical coordinate system
|
|
used by the motion and initial state priors.
|
|
- floor_plane : B x 3
|
|
- trans : B x 3
|
|
- root_orient : B x 3
|
|
- origin: B x 3 desired origin (first joint of human model)
|
|
- returns R (B, 3, 3), t (B, 3), root_height (B)
|
|
"""
|
|
B = trans.size(0)
|
|
if floor_plane.size(1) == 3:
|
|
floor_plane_4d = parse_floor_plane(floor_plane)
|
|
else:
|
|
floor_plane_4d = floor_plane
|
|
floor_normal = floor_plane_4d[:, :3]
|
|
floor_trans, _ = compute_plane_intersection(trans, -floor_normal, floor_plane_4d)
|
|
|
|
# compute prior frame axes within the camera frame
|
|
# up is the floor_plane normal
|
|
up_axis = floor_normal
|
|
# right is body -x direction projected to floor plane
|
|
root_orient_mat = batch_rodrigues(root_orient)
|
|
body_right = -root_orient_mat[:, :, 0]
|
|
floor_body_right, s = compute_plane_intersection(trans, body_right, floor_plane_4d)
|
|
right_axis = floor_body_right - floor_trans
|
|
# body right may not actually intersect - in this case must negate axis because we have the -x
|
|
right_axis = torch.where(s.reshape((B, 1)) < 0, -right_axis, right_axis)
|
|
right_axis = right_axis / torch.norm(right_axis, dim=1, keepdim=True)
|
|
# forward is their cross product
|
|
fwd_axis = torch.linalg.cross(up_axis, right_axis)
|
|
fwd_axis = fwd_axis / torch.norm(fwd_axis, dim=1, keepdim=True)
|
|
|
|
# prior frame is right, fwd, up
|
|
prior_R = torch.stack([right_axis, fwd_axis, up_axis], dim=2)
|
|
world2prior_R = prior_R.transpose(-1, -2)
|
|
|
|
# translation takes translation to origin plus offset to the floor
|
|
world2prior_t = -trans
|
|
|
|
# compute the distance from root to ground plane
|
|
_, s_root = compute_plane_intersection(origin, -floor_normal, floor_plane_4d)
|
|
root_height = s_root.reshape(B, 1)
|
|
|
|
if False:
|
|
# apply transform to origin
|
|
origin_prior = torch.matmul(world2prior_R, (origin - trans)[..., None])[..., 0]
|
|
root_height = s_root.reshape(B, 1) - origin_prior[:, 2:3]
|
|
|
|
return world2prior_R, world2prior_t, root_height
|
|
|
|
|
|
def find_cliques(edges):
|
|
cliques = []
|
|
N = edges.shape[0]
|
|
empty = np.ones(N, dtype=bool)
|
|
while empty.sum() > 0:
|
|
idcs = np.where(empty)[0]
|
|
i = idcs[0]
|
|
cur_clique = [i]
|
|
for j in idcs[1:]:
|
|
if all(edges[c, j] for c in cur_clique):
|
|
cur_clique.append(j)
|
|
empty[cur_clique] = False
|
|
cliques.append(cur_clique)
|
|
return cliques
|
|
|
|
|
|
def estimate_floor_planes(
|
|
smpl_joints, valid_mask, thresh=0.5, group=False, flatten=False
|
|
):
|
|
"""
|
|
:param smpl_joints (B, T, J, 3)
|
|
:param valid_mask (B, T)
|
|
:param thresh (optional, default 0.5) threshold for separating
|
|
:param group (optional, default False) whether to group the planes
|
|
:param flatten (optional, default False) whether to estimate one plane for everyone
|
|
returns floor_plane (C, 3), floor_idcs (B,)
|
|
"""
|
|
B, T = valid_mask.shape
|
|
device = valid_mask.device
|
|
feet_joints = smpl_joints[..., 10:12, :] # (B, T, 2, 3)
|
|
points = [feet_joints[b, valid_mask[b]].reshape(-1, 3) for b in range(B)]
|
|
|
|
if flatten:
|
|
planes = fit_plane(torch.cat(points, dim=0))[None]
|
|
labels = torch.zeros(B, device=device, dtype=torch.long)
|
|
return planes[..., :3] * planes[..., 3:], labels
|
|
|
|
# estimate floors independently
|
|
planes = torch.stack([fit_plane(pts) for pts in points], dim=0) # (B, 4)
|
|
normals, offsets = planes[..., :3], planes[..., 3:] # (B, 3), (B, 1)
|
|
sgn = torch.sign(offsets)
|
|
normals, offsets = sgn * normals, sgn * offsets
|
|
planes = normals * offsets # (B, 3)
|
|
labels = torch.arange(B, device=device)
|
|
|
|
if not group:
|
|
return planes, labels # (B, 3), (B,)
|
|
|
|
# asymetric, project each plane vector onto the other normals
|
|
diffs = offsets - (planes[None] * normals[:, None]).sum(dim=-1)
|
|
print(diffs)
|
|
edges = diffs < thresh
|
|
edges = edges.T & edges
|
|
clusters = find_cliques(edges) # list of indices of clusters
|
|
planes = []
|
|
labels = torch.zeros(B, dtype=torch.long, device=device)
|
|
for c, idcs in enumerate(clusters):
|
|
plane = fit_plane(torch.cat([points[i] for i in idcs], dim=0))
|
|
planes.append(plane)
|
|
labels[idcs] = c
|
|
planes = torch.stack(planes, dim=0) # (C, 4)
|
|
return planes[..., :3] * planes[..., 3:], labels
|
|
|
|
|
|
def estimate_initial_trans(body_pose, joints3d_op, joints2d_op, focal):
|
|
"""
|
|
use focal length and bone lengths to approximate distance from camera
|
|
(based on PROX https://github.com/mohamedhassanmus/prox/blob/master/prox/fitting.py)
|
|
:param body_pose
|
|
:param joints3d_op OP keypoints in 3d
|
|
:param joints2d_op OP keypoints in 2d
|
|
"""
|
|
B, T, N, _ = joints2d_op.shape
|
|
joints2d_obs = joints2d_op[:, :, :, :2] # (B, T, N, 3)
|
|
joints2d_conf = joints2d_op[:, :, :, 2]
|
|
|
|
# find least-occluded 2d frame
|
|
num_2d_vis = torch.sum(joints2d_conf > 0.0, dim=2)
|
|
best_2d_idx = torch.max(num_2d_vis, dim=1)[1]
|
|
|
|
# calculate bone lengths and confidence in each bone length
|
|
bone3d = []
|
|
bone2d = []
|
|
conf2d = []
|
|
for pair in OP_EDGE_LIST:
|
|
diff3d = torch.norm(
|
|
joints3d_op[:, 0, pair[0], :] - joints3d_op[:, 0, pair[1], :],
|
|
dim=1,
|
|
) # does not change over time
|
|
diff2d = torch.norm(
|
|
joints2d_obs[:, :, pair[0], :] - joints2d_obs[:, :, pair[1], :],
|
|
dim=2,
|
|
)
|
|
minconf2d = torch.min(
|
|
joints2d_conf[:, :, pair[0]], joints2d_conf[:, :, pair[1]]
|
|
)
|
|
bone3d.append(diff3d)
|
|
bone2d.append(diff2d)
|
|
conf2d.append(minconf2d)
|
|
|
|
bone3d = torch.stack(bone3d, dim=1)
|
|
bone2d = torch.stack(bone2d, dim=2)
|
|
bone2d = bone2d[np.arange(B), best_2d_idx, :]
|
|
conf2d = torch.stack(conf2d, dim=2)
|
|
conf2d = conf2d[np.arange(B), best_2d_idx, :]
|
|
|
|
# mean over all bones
|
|
mean_bone3d = torch.mean(bone3d, dim=1)
|
|
mean_bone2d = torch.mean(bone2d * (conf2d > 0.0), dim=1)
|
|
|
|
# approx z based on ratio
|
|
init_z = focal * (mean_bone3d / mean_bone2d)
|
|
init_trans = torch.zeros(B, T, 3)
|
|
init_trans[:, :, 2] = init_z.unsqueeze(1).expand(B, T).detach()
|
|
|
|
return init_trans
|