import argparse import math from pathlib import Path import cv2 import numpy as np import PIL.Image as Image import selfcontact import selfcontact.losses import shapely.geometry import torch import torch.nn as nn import torch.optim as optim import torchgeometry import tqdm import trimesh from skimage import measure import fist_pose import hist_cub import losses import pose_estimation import spin PE_KSP_TO_SPIN = { "Head": "Head", "Neck": "Neck", "Right Shoulder": "Right ForeArm", "Right Arm": "Right Arm", "Right Hand": "Right Hand", "Left Shoulder": "Left ForeArm", "Left Arm": "Left Arm", "Left Hand": "Left Hand", "Spine": "Spine1", "Hips": "Hips", "Right Upper Leg": "Right Upper Leg", "Right Leg": "Right Leg", "Right Foot": "Right Foot", "Left Upper Leg": "Left Upper Leg", "Left Leg": "Left Leg", "Left Foot": "Left Foot", "Left Toe": "Left Toe", "Right Toe": "Right Toe", } MODELS_DIR = "models" def parse_args(): parser = argparse.ArgumentParser() parser.add_argument( "--pose-estimation-model-path", type=str, default=f"./{MODELS_DIR}/hrn_w48_384x288.onnx", help="Pose Estimation model", ) parser.add_argument( "--contact-model-path", type=str, default=f"./{MODELS_DIR}/contact_hrn_w32_256x192.onnx", help="Contact model", ) parser.add_argument( "--device", type=str, default="cuda", choices=["cpu", "cuda"], help="Torch device", ) parser.add_argument( "--spin-model-path", type=str, default=f"./{MODELS_DIR}/spin_model_smplx_eft_18.pt", help="SPIN model path", ) parser.add_argument( "--smpl-type", type=str, default="smplx", choices=["smplx"], help="SMPL model type", ) parser.add_argument( "--smpl-model-dir", type=str, default=f"./{MODELS_DIR}/models/smplx", help="SMPL model dir", ) parser.add_argument( "--smpl-mean-params-path", type=str, default=f"./{MODELS_DIR}/data/smpl_mean_params.npz", help="SMPL mean params", ) parser.add_argument( "--essentials-dir", type=str, default=f"./{MODELS_DIR}/smplify-xmc-essentials", help="SMPL Essentials folder for contacts", ) parser.add_argument( "--parametrization-path", type=str, default=f"./{MODELS_DIR}/smplx_parametrization/parametrization.npy", help="Parametrization path", ) parser.add_argument( "--bone-parametrization-path", type=str, default=f"./{MODELS_DIR}/smplx_parametrization/bone_to_param2.npy", help="Bone parametrization path", ) parser.add_argument( "--foot-inds-path", type=str, default=f"./{MODELS_DIR}/smplx_parametrization/foot_inds.npy", help="Foot indinces", ) parser.add_argument( "--save-path", type=str, required=True, help="Path to save the results", ) parser.add_argument( "--img-path", type=str, required=True, help="Path to img to test", ) parser.add_argument( "--use-contacts", action="store_true", help="Use contact model", ) parser.add_argument( "--use-msc", action="store_true", help="Use MSC loss", ) parser.add_argument( "--use-natural", action="store_true", help="Use regularity", ) parser.add_argument( "--use-cos", action="store_true", help="Use cos model", ) parser.add_argument( "--use-angle-transf", action="store_true", help="Use cube foreshortening transformation", ) parser.add_argument( "--c-mse", type=float, default=0, help="MSE weight", ) parser.add_argument( "--c-par", type=float, default=10, help="Parallel weight", ) parser.add_argument( "--c-f", type=float, default=1000, help="Cos coef", ) parser.add_argument( "--c-parallel", type=float, default=100, help="Parallel weight", ) parser.add_argument( "--c-reg", type=float, default=1000, help="Regularity weight", ) parser.add_argument( "--c-cont2d", type=float, default=1, help="Contact 2D weight", ) parser.add_argument( "--c-msc", type=float, default=17_500, help="MSC weight", ) parser.add_argument( "--fist", nargs="+", type=str, choices=list(fist_pose.INT_TO_FIST), ) args = parser.parse_args() return args def freeze_layers(model): for module in model.modules(): if type(module) is False: continue if isinstance(module, nn.modules.batchnorm._BatchNorm): module.eval() for m in module.parameters(): m.requires_grad = False if isinstance(module, nn.Dropout): module.eval() for m in module.parameters(): m.requires_grad = False def project_and_normalize_to_spin(vertices_3d, camera): vertices_2d = vertices_3d # [:, :2] scale, translate = camera[0], camera[1:] translate = scale.new_zeros(3) translate[:2] = camera[1:] vertices_2d = vertices_2d + translate vertices_2d = scale * vertices_2d + 1 vertices_2d = spin.constants.IMG_RES / 2 * vertices_2d return vertices_2d def project_and_normalize_to_spin_legs(vertices_3d, A, camera): A, J = A A = A[0] J = J[0] L = vertices_3d.new_tensor( [ [0.98619063, 0.16560926, 0.00127302], [-0.16560601, 0.98603675, 0.01749799], [0.00164258, -0.01746717, 0.99984609], ] ) R = vertices_3d.new_tensor( [ [0.9910211, -0.13368178, -0.0025208], [0.13367888, 0.99027076, 0.03864949], [-0.00267045, -0.03863944, 0.99924965], ] ) scale = camera[0] R = A[2, :3, :3] @ R # 2 - right L = A[1, :3, :3] @ L # 1 - left r = J[5] - J[2] l = J[4] - J[1] rleg = scale * spin.constants.IMG_RES / 2 * R @ r lleg = scale * spin.constants.IMG_RES / 2 * L @ l rleg = rleg[:2] lleg = lleg[:2] return rleg, lleg def rotation_matrix_to_angle_axis(rotmat): bs, n_joints, *_ = rotmat.size() rotmat = torch.cat( [ rotmat.view(-1, 3, 3), rotmat.new_tensor([0, 0, 1], dtype=torch.float32) .view(bs, 3, 1) .expand(n_joints, -1, -1), ], dim=-1, ) aa = torchgeometry.rotation_matrix_to_angle_axis(rotmat) aa = aa.reshape(bs, 3 * n_joints) return aa def get_smpl_output(smpl, rotmat, betas, use_betas=True, zero_hands=False): if smpl.name() == "SMPL": smpl_output = smpl( betas=betas if use_betas else None, body_pose=rotmat[:, 1:], global_orient=rotmat[:, 0].unsqueeze(1), pose2rot=False, ) elif smpl.name() == "SMPL-X": rotmat = rotation_matrix_to_angle_axis(rotmat) if zero_hands: for i in [20, 21]: rotmat[:, 3 * i : 3 * (i + 1)] = 0 for i in [12, 15]: # neck, head rotmat[:, 3 * i + 1] = 0 # y smpl_output = smpl( betas=betas if use_betas else None, body_pose=rotmat[:, 3:], global_orient=rotmat[:, :3], pose2rot=True, ) else: raise NotImplementedError return smpl_output, rotmat def get_predictions(model_hmr, smpl, input_img, use_betas=True, zero_hands=False): input_img = input_img.unsqueeze(0) rotmat, betas, camera = model_hmr(input_img) smpl_output, rotmat = get_smpl_output( smpl, rotmat, betas, use_betas=use_betas, zero_hands=zero_hands ) rotmat = rotmat.squeeze(0) betas = betas.squeeze(0) camera = camera.squeeze(0) z = smpl_output.joints z = z.squeeze(0) return rotmat, betas, camera, smpl_output, z def get_pred_and_data( model_hmr, smpl, selector, input_img, use_betas=True, zero_hands=False ): rotmat, betas, camera, smpl_output, zz = get_predictions( model_hmr, smpl, input_img, use_betas=use_betas, zero_hands=zero_hands ) joints = smpl_output.joints.squeeze(0) joints_2d = project_and_normalize_to_spin(joints, camera) rleg, lleg = project_and_normalize_to_spin_legs(joints, smpl_output.A, camera) joints_2d_orig = joints_2d joints_2d = joints_2d[selector] vertices = smpl_output.vertices.squeeze(0) vertices_2d = project_and_normalize_to_spin(vertices, camera) zz = zz[selector] return ( rotmat, betas, camera, joints_2d, zz, vertices_2d, smpl_output, (rleg, lleg), joints_2d_orig, ) def normalize_keypoints_to_spin(keypoints_2d, img_size): h, w = img_size if h > w: # vertically ax1 = 1 ax2 = 0 else: # horizontal ax1 = 0 ax2 = 1 shift = (img_size[ax1] - img_size[ax2]) / 2 scale = spin.constants.IMG_RES / img_size[ax2] keypoints_2d_normalized = np.copy(keypoints_2d) keypoints_2d_normalized[:, ax2] -= shift keypoints_2d_normalized *= scale return keypoints_2d_normalized, shift, scale, ax2 def unnormalize_keypoints_from_spin(keypoints_2d, shift, scale, ax2): keypoints_2d_normalized = np.copy(keypoints_2d) keypoints_2d_normalized /= scale keypoints_2d_normalized[:, ax2] += shift return keypoints_2d_normalized def get_vertices_in_heatmap(contact_heatmap): contact_heatmap_size = contact_heatmap.shape[:2] label = measure.label(contact_heatmap) y_data_conts = [] for i in range(1, label.max() + 1): predicted_kps_contact = np.vstack(np.nonzero(label == i)[::-1]).T.astype( "float" ) predicted_kps_contact_scaled, *_ = normalize_keypoints_to_spin( predicted_kps_contact, contact_heatmap_size ) y_data_cont = torch.from_numpy(predicted_kps_contact_scaled).int().tolist() y_data_cont = shapely.geometry.MultiPoint(y_data_cont).convex_hull y_data_conts.append(y_data_cont) return y_data_conts def get_contact_heatmap(model_contact, img_path, thresh=0.5): contact_heatmap = pose_estimation.infer_single_image( model_contact, img_path, input_img_size=(192, 256), return_kps=False, ) contact_heatmap = contact_heatmap.squeeze(0) contact_heatmap_orig = contact_heatmap.copy() mi = contact_heatmap.min() ma = contact_heatmap.max() contact_heatmap = (contact_heatmap - mi) / (ma - mi) contact_heatmap_ = ((contact_heatmap > thresh) * 255).astype("uint8") contact_heatmap = np.repeat(contact_heatmap[..., None], repeats=3, axis=-1) contact_heatmap = (contact_heatmap * 255).astype("uint8") return contact_heatmap_, contact_heatmap, contact_heatmap_orig def discretize(parametrization, n_bins=100): bins = np.linspace(0, 1, n_bins + 1) inds = np.digitize(parametrization, bins) disc_parametrization = bins[inds - 1] return disc_parametrization def get_mapping_from_params_to_verts(verts, params): mapping = {} for v, t in zip(verts, params): mapping.setdefault(t, []).append(v) return mapping def find_contacts(y_data_conts, keypoints_2d, bone_to_params, thresh=12, step=0.0072246375): n_bins = int(math.ceil(1 / step)) - 1 # mean face's circumradius contact = [] contact_2d = [] for_mask = [] for y_data_cont in y_data_conts: contact_loc = [] contact_2d_loc = [] buffer = y_data_cont.buffer(thresh) mask_add = False for i, j in pose_estimation.SKELETON: verts, t3d = bone_to_params[(i, j)] if len(verts) == 0: continue t3d = discretize(t3d, n_bins=n_bins) t3d_to_verts = get_mapping_from_params_to_verts(verts, t3d) t3d_to_verts_sorted = sorted(t3d_to_verts.items(), key=lambda x: x[0]) t3d_sorted_np = np.array([x for x, _ in t3d_to_verts_sorted]) line = shapely.geometry.LineString([keypoints_2d[i], keypoints_2d[j]]) lint = buffer.intersection(line) if len(lint.boundary.geoms) < 2: continue t2d_start = line.project(lint.boundary.geoms[0], normalized=True) t2d_end = line.project(lint.boundary.geoms[1], normalized=True) assert t2d_start <= t2d_end t2ds = discretize( np.linspace(t2d_start, t2d_end, n_bins + 1), n_bins=n_bins ) to_add = False for t2d in t2ds: if t2d < t3d_sorted_np[0] or t2d > t3d_sorted_np[-1]: continue t2d_ind = np.searchsorted(t3d_sorted_np, t2d) c = t3d_to_verts_sorted[t2d_ind][1] contact_loc.extend(c) to_add = True mask_add = True if t2d_ind + 1 < len(t3d_to_verts_sorted): c = t3d_to_verts_sorted[t2d_ind + 1][1] contact_loc.extend(c) if t2d_ind > 0: c = t3d_to_verts_sorted[t2d_ind - 1][1] contact_loc.extend(c) if to_add: contact_2d_loc.append((i, j, t2d_start + 0.5 * (t2d_end - t2d_start))) if mask_add: for_mask.append(buffer.exterior.coords.xy) contact_loc = sorted(set(contact_loc)) contact_loc = np.array(contact_loc, dtype="int") contact.append(contact_loc) contact_2d.append(contact_2d_loc) for_mask = [np.stack((x, y), axis=0).T[:, None].astype("int") for x, y in for_mask] return contact, contact_2d, for_mask def optimize( model_hmr, smpl, selector, input_img, keypoints_2d, optimizer, args, loss_mse=None, loss_parallel=None, c_mse=0.0, c_new_mse=1.0, c_beta=1e-3, sc_crit=None, msc_crit=None, contact=None, n_steps=60, i_ini=0, ): mean_zfoot_val = {} with tqdm.trange(n_steps) as pbar: for i in pbar: global_step = i + i_ini optimizer.zero_grad() ( rotmat_pred, betas_pred, camera_pred, keypoints_3d_pred, z, vertices_2d_pred, smpl_output, (rleg, lleg), joints_2d_orig, ) = get_pred_and_data( model_hmr, smpl, selector, input_img, ) keypoints_2d_pred = keypoints_3d_pred[:, :2] loss = l2 = 0.0 if c_mse > 0 and loss_mse is not None: l2 = loss_mse(keypoints_2d_pred, keypoints_2d) loss = loss + c_mse * l2 vertices_pred = smpl_output.vertices lpar = z_loss = loss_sh = 0.0 if c_new_mse > 0 and loss_parallel is not None: Ltan, Lcos, Lpar, Lspine, Lgr, Lstraight3d, Lcon2d = loss_parallel( keypoints_3d_pred, keypoints_2d, z, (rleg, lleg), global_step=global_step, ) lpar = ( Ltan + c_new_mse * (args.c_f * Lcos + args.c_parallel * Lpar) + Lspine + args.c_reg * Lgr + args.c_reg * Lstraight3d + args.c_cont2d * Lcon2d ) loss = loss + 300 * lpar for side in ["left", "right"]: attr = f"{side}_foot_inds" if hasattr(loss_parallel, attr): foot_inds = getattr(loss_parallel, attr) zind = 1 if attr not in mean_zfoot_val: with torch.no_grad(): mean_zfoot_val[attr] = torch.median( vertices_pred[0, foot_inds, zind], dim=0 ).values loss_foot = ( (vertices_pred[0, foot_inds, zind] - mean_zfoot_val[attr]) ** 2 ).sum() loss = loss + args.c_reg * loss_foot if hasattr(loss_parallel, "silhuette_vertices_inds"): inds = loss_parallel.silhuette_vertices_inds loss_sh = ( (vertices_pred[0, inds, 1] - loss_parallel.ground) ** 2 ).sum() loss = loss + args.c_reg * loss_sh lbeta = (betas_pred**2).mean() lcam = ((torch.exp(-camera_pred[0] * 10)) ** 2).mean() loss = loss + c_beta * lbeta + lcam lgsc_a = gsc_contact_loss = faces_angle_loss = 0.0 if sc_crit is not None: gsc_contact_loss, faces_angle_loss = sc_crit( vertices_pred, ) lgsc_a = 1000 * gsc_contact_loss + 0.1 * faces_angle_loss loss = loss + lgsc_a msc_loss = 0.0 if contact is not None and len(contact) > 0 and msc_crit is not None: if not isinstance(contact, list): contact = [contact] for cntct in contact: msc_loss = msc_crit( cntct, vertices_pred, ) loss = loss + args.c_msc * msc_loss loss.backward() optimizer.step() epoch_loss = loss.item() pbar.set_postfix( **{ "l": f"{epoch_loss:.3}", "l2": f"{l2:.3}", "par": f"{lpar:.3}", "beta": f"{lbeta:.3}", "cam": f"{lcam:.3}", "z": f"{z_loss:.3}", "gsc_contact": f"{float(gsc_contact_loss):.3}", "faces_angle": f"{float(faces_angle_loss):.3}", "msc": f"{float(msc_loss):.3}", } ) with torch.no_grad(): ( rotmat_pred, betas_pred, camera_pred, keypoints_3d_pred, z, vertices_2d_pred, smpl_output, (rleg, lleg), joints_2d_orig, ) = get_pred_and_data( model_hmr, smpl, selector, input_img, zero_hands=True, ) return ( rotmat_pred, betas_pred, camera_pred, keypoints_3d_pred, vertices_2d_pred, smpl_output, z, joints_2d_orig, ) def optimize_ft( theta, camera, smpl, selector, keypoints_2d, args, loss_mse=None, loss_parallel=None, c_mse=0.0, c_new_mse=1.0, sc_crit=None, msc_crit=None, contact=None, n_steps=60, i_ini=0, zero_hands=False, fist=None, ): mean_zfoot_val = {} theta = theta.detach().clone() camera = camera.detach().clone() rotmat_pred = nn.Parameter(theta) camera_pred = nn.Parameter(camera) optimizer = torch.optim.Adam( [ rotmat_pred, camera_pred, ], lr=1e-3, ) global_step = i_ini with tqdm.trange(n_steps) as pbar: for i in pbar: global_step = i + i_ini optimizer.zero_grad() global_orient = rotmat_pred[:3] body_pose = rotmat_pred[3:] smpl_output = smpl( global_orient=global_orient.unsqueeze(0), body_pose=body_pose.unsqueeze(0), pose2rot=True, ) z = smpl_output.joints z = z.squeeze(0) joints = smpl_output.joints.squeeze(0) joints_2d = project_and_normalize_to_spin(joints, camera_pred) rleg, lleg = project_and_normalize_to_spin_legs( joints, smpl_output.A, camera_pred ) joints_2d = joints_2d[selector] z = z[selector] keypoints_3d_pred = joints_2d keypoints_2d_pred = keypoints_3d_pred[:, :2] lprior = ((rotmat_pred - theta) ** 2).sum() + ( (camera_pred - camera) ** 2 ).sum() loss = lprior l2 = 0.0 if c_mse > 0 and loss_mse is not None: l2 = loss_mse(keypoints_2d_pred, keypoints_2d) loss = loss + c_mse * l2 vertices_pred = smpl_output.vertices lpar = z_loss = loss_sh = 0.0 if c_new_mse > 0 and loss_parallel is not None: Ltan, Lcos, Lpar, Lspine, Lgr, Lstraight3d, Lcon2d = loss_parallel( keypoints_3d_pred, keypoints_2d, z, (rleg, lleg), global_step=global_step, ) lpar = ( Ltan + c_new_mse * (args.c_f * Lcos + args.c_parallel * Lpar) + Lspine + args.c_reg * Lgr + args.c_reg * Lstraight3d + args.c_cont2d * Lcon2d ) loss = loss + 300 * lpar for side in ["left", "right"]: attr = f"{side}_foot_inds" if hasattr(loss_parallel, attr): foot_inds = getattr(loss_parallel, attr) zind = 1 if attr not in mean_zfoot_val: with torch.no_grad(): mean_zfoot_val[attr] = torch.median( vertices_pred[0, foot_inds, zind], dim=0 ).values loss_foot = ( (vertices_pred[0, foot_inds, zind] - mean_zfoot_val[attr]) ** 2 ).sum() loss = loss + args.c_reg * loss_foot if hasattr(loss_parallel, "silhuette_vertices_inds"): inds = loss_parallel.silhuette_vertices_inds loss_sh = ( (vertices_pred[0, inds, 1] - loss_parallel.ground) ** 2 ).sum() loss = loss + args.c_reg * loss_sh lgsc_a = gsc_contact_loss = faces_angle_loss = 0.0 if sc_crit is not None: gsc_contact_loss, faces_angle_loss = sc_crit(vertices_pred) lgsc_a = 1000 * gsc_contact_loss + 0.1 * faces_angle_loss loss = loss + lgsc_a msc_loss = 0.0 if contact is not None and len(contact) > 0 and msc_crit is not None: if not isinstance(contact, list): contact = [contact] for cntct in contact: msc_loss = msc_crit( cntct, vertices_pred, ) loss = loss + args.c_msc * msc_loss loss.backward() optimizer.step() epoch_loss = loss.item() pbar.set_postfix( **{ "l": f"{epoch_loss:.3}", "l2": f"{l2:.3}", "par": f"{lpar:.3}", "z": f"{z_loss:.3}", "gsc_contact": f"{float(gsc_contact_loss):.3}", "faces_angle": f"{float(faces_angle_loss):.3}", "msc": f"{float(msc_loss):.3}", } ) rotmat_pred = rotmat_pred.detach() if zero_hands: for i in [20, 21]: rotmat_pred[3 * i : 3 * (i + 1)] = 0 for i in [12, 15]: # neck, head rotmat_pred[3 * i + 1] = 0 # y global_orient = rotmat_pred[:3] body_pose = rotmat_pred[3:] left_hand_pose = None right_hand_pose = None if fist is not None: left_hand_pose = rotmat_pred.new_tensor(fist_pose.LEFT_RELAXED).unsqueeze(0) right_hand_pose = rotmat_pred.new_tensor(fist_pose.RIGHT_RELAXED).unsqueeze(0) for f in fist: pp = fist_pose.INT_TO_FIST[f] if pp is not None: pp = rotmat_pred.new_tensor(pp).unsqueeze(0) if f.startswith("lf"): left_hand_pose = pp elif f.startswith("rf"): right_hand_pose = pp elif f.startswith("l"): body_pose[19 * 3 : 19 * 3 + 3] = pp left_hand_pose = None elif f.startswith("r"): body_pose[20 * 3 : 20 * 3 + 3] = pp right_hand_pose = None else: raise RuntimeError(f"No such hand pose: {f}") with torch.no_grad(): smpl_output = smpl( global_orient=global_orient.unsqueeze(0), body_pose=body_pose.unsqueeze(0), left_hand_pose=left_hand_pose, right_hand_pose=right_hand_pose, pose2rot=True, ) return rotmat_pred, smpl_output def create_bone(i, j, keypoints_2d): a = keypoints_2d[i] b = keypoints_2d[j] ab = b - a ab = torch.nn.functional.normalize(ab, dim=0) return ab def is_parallel_to_plane(bone, thresh=21): return abs(bone[0]) > math.cos(math.radians(thresh)) def is_close_to_plane(bone, plane, thresh): dist = abs(bone[0] - plane) return dist < thresh def get_selector(): selector = [] for kp in pose_estimation.KPS: tmp = spin.JOINT_NAMES.index(PE_KSP_TO_SPIN[kp]) selector.append(tmp) return selector def calc_cos(joints_2d, joints_3d): cos = [] for i, j in pose_estimation.SKELETON: a = joints_2d[i] - joints_2d[j] a = nn.functional.normalize(a, dim=0) b = joints_3d[i] - joints_3d[j] b = nn.functional.normalize(b, dim=0)[:2] c = (a * b).sum() cos.append(c) cos = torch.stack(cos, dim=0) return cos def get_natural(keypoints_2d, vertices, right_foot_inds, left_foot_inds, loss_parallel, smpl): height_2d = ( keypoints_2d.max(dim=0).values[0] - keypoints_2d.min(dim=0).values[0] ).item() plane_2d = keypoints_2d.max(dim=0).values[0].item() ground_parallel = [] parallel_in_3d = [] parallel3d_bones = set() # parallel chains for i, j, k in [ ("Right Upper Leg", "Right Leg", "Right Foot"), ("Right Leg", "Right Foot", "Right Toe"), # to remove? ("Left Upper Leg", "Left Leg", "Left Foot"), ("Left Leg", "Left Foot", "Left Toe"), # to remove? ("Right Shoulder", "Right Arm", "Right Hand"), ("Left Shoulder", "Left Arm", "Left Hand"), # ("Hips", "Spine", "Neck"), # ("Spine", "Neck", "Head"), ]: i = pose_estimation.KPS.index(i) j = pose_estimation.KPS.index(j) k = pose_estimation.KPS.index(k) upleg_leg = create_bone(i, j, keypoints_2d) leg_foot = create_bone(j, k, keypoints_2d) if is_parallel_to_plane(upleg_leg) and is_parallel_to_plane(leg_foot): if is_close_to_plane( upleg_leg, plane_2d, thresh=0.1 * height_2d ) or is_close_to_plane(leg_foot, plane_2d, thresh=0.1 * height_2d): ground_parallel.append(((i, j), 1)) ground_parallel.append(((j, k), 1)) if (upleg_leg * leg_foot).sum() > math.cos(math.radians(21)): parallel_in_3d.append(((i, j), (j, k))) parallel3d_bones.add((i, j)) parallel3d_bones.add((j, k)) # parallel feets for i, j in [ ("Right Foot", "Right Toe"), ("Left Foot", "Left Toe"), ]: i = pose_estimation.KPS.index(i) j = pose_estimation.KPS.index(j) if (i, j) in parallel3d_bones: continue foot_toe = create_bone(i, j, keypoints_2d) if is_parallel_to_plane(foot_toe, thresh=25): if "Right" in pose_estimation.KPS[i]: loss_parallel.right_foot_inds = right_foot_inds else: loss_parallel.left_foot_inds = left_foot_inds loss_parallel.ground_parallel = ground_parallel loss_parallel.parallel_in_3d = parallel_in_3d vertices_np = vertices[0].cpu().numpy() if len(ground_parallel) > 0: # Silhuette veritices mesh = trimesh.Trimesh(vertices=vertices_np, faces=smpl.faces, process=False) silhuette_vertices_mask_1 = np.abs(mesh.vertex_normals[..., 2]) < 2e-1 height_3d = vertices_np[:, 1].max() - vertices_np[:, 1].min() plane_3d = vertices_np[:, 1].max() silhuette_vertices_mask_2 = ( np.abs(vertices_np[:, 1] - plane_3d) < 0.15 * height_3d ) silhuette_vertices_mask = np.logical_and( silhuette_vertices_mask_1, silhuette_vertices_mask_2 ) (silhuette_vertices_inds,) = np.where(silhuette_vertices_mask) if len(silhuette_vertices_inds) > 0: loss_parallel.silhuette_vertices_inds = silhuette_vertices_inds loss_parallel.ground = plane_3d def get_cos(keypoints_3d_pred, use_angle_transf, loss_parallel): keypoints_2d_pred = keypoints_3d_pred[:, :2] with torch.no_grad(): cos_r = calc_cos(keypoints_2d_pred, keypoints_3d_pred) alpha = torch.acos(cos_r) if use_angle_transf: leg_inds = [ 5, 6, # right leg 7, 8, # left leg ] foot_inds = [15, 16] nleg_inds = sorted( set(range(len(pose_estimation.SKELETON))) - set(leg_inds) - set(foot_inds) ) alpha[nleg_inds] = alpha[nleg_inds] - alpha[nleg_inds].min() amli = alpha[leg_inds].min() leg_inds.extend(foot_inds) alpha[leg_inds] = alpha[leg_inds] - amli angles = alpha.detach().cpu().numpy() angles = hist_cub.cub( angles / (math.pi / 2), a=1.2121212121212122, b=-1.105527638190953, c=0.787878787878789, ) * (math.pi / 2) alpha = alpha.new_tensor(angles) loss_parallel.cos = torch.cos(alpha) return cos_r def get_contacts( args, sc_module, y_data_conts, keypoints_2d, vertices, bone_to_params, loss_parallel, ): use_contacts = args.use_contacts use_msc = args.use_msc c_mse = args.c_mse if use_contacts: assert c_mse == 0 contact, contact_2d, _ = find_contacts( y_data_conts, keypoints_2d, bone_to_params ) if len(contact_2d) > 0: loss_parallel.contact_2d = contact_2d if len(contact) == 0: _, contact = sc_module.verts_in_contact(vertices, return_idx=True) contact = contact.cpu().numpy().ravel() elif use_msc: _, contact = sc_module.verts_in_contact(vertices, return_idx=True) contact = contact.cpu().numpy().ravel() else: contact = np.array([]) return contact def save_mesh( smpl, smpl_output, save_path, fname, ): mesh = trimesh.Trimesh( vertices=smpl_output.vertices[0].cpu().numpy(), faces=smpl.faces, process=False, ) rot = trimesh.transformations.rotation_matrix(np.pi, [1, 0, 0]) mesh.apply_transform(rot) rot = trimesh.transformations.rotation_matrix(np.pi, [0, 1, 0]) mesh.apply_transform(rot) mesh.export(save_path / f"{fname}.glb") def eft_step( model_hmr, smpl, selector, input_img, keypoints_2d, optimizer, args, loss_mse, loss_parallel, c_beta, sc_module, y_data_conts, bone_to_params, ): ( _, _, _, keypoints_3d_pred, _, smpl_output, _, _, ) = optimize( model_hmr, smpl, selector, input_img, keypoints_2d, optimizer, args, loss_mse=loss_mse, loss_parallel=loss_parallel, c_mse=1, c_new_mse=0, c_beta=c_beta, sc_crit=None, msc_crit=None, contact=None, n_steps=60 + 90, ) # find contacts vertices = smpl_output.vertices.detach() contact = get_contacts( args, sc_module, y_data_conts, keypoints_2d, vertices, bone_to_params, loss_parallel, ) return vertices, keypoints_3d_pred, contact def dc_step( model_hmr, smpl, selector, input_img, keypoints_2d, optimizer, args, loss_mse, loss_parallel, c_mse, c_new_mse, c_beta, sc_crit, msc_crit, contact, use_contacts, use_msc, ): rotmat_pred, *_ = optimize( model_hmr, smpl, selector, input_img, keypoints_2d, optimizer, args, loss_mse=loss_mse, loss_parallel=loss_parallel, c_mse=c_mse, c_new_mse=c_new_mse, c_beta=c_beta, sc_crit=sc_crit, msc_crit=msc_crit if use_contacts or use_msc else None, contact=contact if use_contacts or use_msc else None, n_steps=60 if c_new_mse > 0 or use_contacts or use_msc else 0, # + 60,, i_ini=60 + 90, ) return rotmat_pred def us_step( model_hmr, smpl, selector, input_img, rotmat_pred, keypoints_2d, args, loss_mse, loss_parallel, c_mse, c_new_mse, sc_crit, msc_crit, contact, use_contacts, use_msc, save_path, ): (_, _, camera_pred_us, _, _, _, smpl_output_us, _, _,) = get_pred_and_data( model_hmr, smpl, selector, input_img, use_betas=False, zero_hands=True, ) _, smpl_output_us = optimize_ft( rotmat_pred, camera_pred_us, smpl, selector, keypoints_2d, args, loss_mse=loss_mse, loss_parallel=loss_parallel, c_mse=c_mse, c_new_mse=c_new_mse, sc_crit=sc_crit, msc_crit=msc_crit if use_contacts or use_msc else None, contact=contact if use_contacts or use_msc else None, n_steps=60 if use_contacts or use_msc else 0, # + 60, i_ini=60 + 90 + 60, zero_hands=True, fist=args.fist, ) save_mesh( smpl, smpl_output_us, save_path, "us", ) def main(): args = parse_args() print(args) # models model_pose = cv2.dnn.readNetFromONNX( args.pose_estimation_model_path ) # "hrn_w48_384x288.onnx" model_contact = cv2.dnn.readNetFromONNX( args.contact_model_path ) # "contact_hrn_w32_256x192.onnx" device = ( torch.device(args.device) if torch.cuda.is_available() else torch.device("cpu") ) model_hmr = spin.hmr(args.smpl_mean_params_path) # "smpl_mean_params.npz" model_hmr.to(device) checkpoint = torch.load( args.spin_model_path, # "spin_model_smplx_eft_18.pt" map_location="cpu" ) smpl = spin.SMPLX( args.smpl_model_dir, # "models/smplx" batch_size=1, create_transl=False, use_pca=False, flat_hand_mean=args.fist is not None, ) smpl.to(device) selector = get_selector() use_contacts = args.use_contacts use_msc = args.use_msc bone_to_params = np.load(args.bone_parametrization_path, allow_pickle=True).item() foot_inds = np.load(args.foot_inds_path, allow_pickle=True).item() left_foot_inds = foot_inds["left_foot_inds"] right_foot_inds = foot_inds["right_foot_inds"] if use_contacts: model_type = args.smpl_type sc_module = selfcontact.SelfContact( essentials_folder=args.essentials_dir, # "smplify-xmc-essentials" geothres=0.3, euclthres=0.02, test_segments=True, compute_hd=True, model_type=model_type, device=device, ) sc_module.to(device) sc_crit = selfcontact.losses.SelfContactLoss( contact_module=sc_module, inside_loss_weight=0.5, outside_loss_weight=0.0, contact_loss_weight=0.5, align_faces=True, use_hd=True, test_segments=True, device=device, model_type=model_type, ) sc_crit.to(device) msc_crit = losses.MimickedSelfContactLoss(geodesics_mask=sc_module.geomask) msc_crit.to(device) else: sc_module = None sc_crit = None msc_crit = None loss_mse = losses.MSE([1, 10, 13]) # Neck + Right Upper Leg + Left Upper Leg ignore = ( (1, 2), # Neck + Right Shoulder (1, 5), # Neck + Left Shoulder (9, 10), # Hips + Right Upper Leg (9, 13), # Hips + Left Upper Leg ) loss_parallel = losses.Parallel( skeleton=pose_estimation.SKELETON, ignore=ignore, ) c_mse = args.c_mse c_new_mse = args.c_par c_beta = 1e-3 if c_mse > 0: assert c_new_mse == 0 elif c_mse == 0: assert c_new_mse > 0 root_path = Path(args.save_path) root_path.mkdir(exist_ok=True, parents=True) path_to_imgs = Path(args.img_path) if path_to_imgs.is_dir(): path_to_imgs = path_to_imgs.iterdir() else: path_to_imgs = [path_to_imgs] for img_path in path_to_imgs: if not any( img_path.name.lower().endswith(ext) for ext in [".jpg", ".png", ".jpeg"] ): continue img_name = img_path.stem # use 2d keypoints detection ( img_original, predicted_keypoints_2d, _, _, ) = pose_estimation.infer_single_image( model_pose, img_path, input_img_size=pose_estimation.IMG_SIZE, return_kps=True, ) save_path = root_path / img_name save_path.mkdir(exist_ok=True, parents=True) img_original = cv2.cvtColor(img_original, cv2.COLOR_BGR2RGB) img_size_original = img_original.shape[:2] keypoints_2d, *_ = normalize_keypoints_to_spin( predicted_keypoints_2d, img_size_original ) keypoints_2d = torch.from_numpy(keypoints_2d) keypoints_2d = keypoints_2d.to(device) ( predicted_contact_heatmap, predicted_contact_heatmap_raw, very_hm_raw, ) = get_contact_heatmap(model_contact, img_path) predicted_contact_heatmap_raw = Image.fromarray( predicted_contact_heatmap_raw ).resize(img_size_original[::-1]) predicted_contact_heatmap_raw = cv2.resize(very_hm_raw, img_size_original[::-1]) if c_new_mse == 0: predicted_contact_heatmap_raw = None y_data_conts = get_vertices_in_heatmap(predicted_contact_heatmap) model_hmr.load_state_dict(checkpoint["model"], strict=True) model_hmr.train() freeze_layers(model_hmr) _, input_img = spin.process_image(img_path, input_res=spin.constants.IMG_RES) input_img = input_img.to(device) optimizer = optim.Adam( filter(lambda p: p.requires_grad, model_hmr.parameters()), lr=1e-6, ) vertices, keypoints_3d_pred, contact = eft_step( model_hmr, smpl, selector, input_img, keypoints_2d, optimizer, args, loss_mse, loss_parallel, c_beta, sc_module, y_data_conts, bone_to_params, ) if args.use_natural: get_natural( keypoints_2d, vertices, right_foot_inds, left_foot_inds, loss_parallel, smpl, ) if args.use_cos: get_cos(keypoints_3d_pred, args.use_angle_transf, loss_parallel) rotmat_pred = dc_step( model_hmr, smpl, selector, input_img, keypoints_2d, optimizer, args, loss_mse, loss_parallel, c_mse, c_new_mse, c_beta, sc_crit, msc_crit, contact, use_contacts, use_msc, ) us_step( model_hmr, smpl, selector, input_img, rotmat_pred, keypoints_2d, args, loss_mse, loss_parallel, c_mse, c_new_mse, sc_crit, msc_crit, contact, use_contacts, use_msc, save_path, ) if __name__ == "__main__": main()