Upload pose.py

src/pose.py

@@ -0,0 +1,1482 @@
+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
+import utils
+
+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_all(
+    smpl,
+    smpl_output,
+    save_path,
+    fname,
+):
+    utils.save_mesh_with_colors(
+        smpl_output.vertices[0].cpu().numpy(),
+        smpl.faces,
+        save_path / f"{fname}.ply",
+    )
+
+
+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_all(
+        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()