This is a code repository for training the IPMAN-R model and evaluating performance 
"""
This file contains definitions of useful data stuctures and the paths
for the datasets and data files necessary to run the code.
Things you need to change: *_ROOT that indicate the path to each dataset
"""
import os
import time
import yaml
import shutil
import argparse
import operator
import itertools
from os.path import join
from os.path import join
from loguru import logger
import numpy as np
from functools import reduce
from yacs.config import CfgNode as CN
from typing import Dict, List, Union, Any
from flatten_dict import flatten, unflatten

# Please update your own paths to dataset folders here:
H36M_ROOT = 'data/human36m'
RICH_ROOT = 'data/RICH/'
LSP_ROOT = 'data/LSP'
LSP_ORIGINAL_ROOT = 'data/LSP_ORIGINAL'
LSPET_ROOT = 'data/hr-lspet'
MPII_ROOT = 'data/MPII-pose'
COCO_ROOT = 'data/COCO/images'
MPI_INF_3DHP_ROOT = 'data/mpi_inf_3dhp'
PW3D_ROOT = 'data/3DPW___v1'
UPI_S1H_ROOT = ''
AGORA_ROOT = ''

# Output folder to save test/train npz files
DATASET_NPZ_PATH = 'data/dataset_extras'

# Output folder to store the openpose detections
# This is requires only in case you want to regenerate 
# the .npz files with the annotations.
OPENPOSE_PATH = ''

# Path to test/train npz files
DATASET_FILES = [ {'h36m-p1': join(DATASET_NPZ_PATH, 'h36m_valid_protocol1.npz'),
                   'h36m-p2': join(DATASET_NPZ_PATH, 'h36m_valid_protocol2.npz'),
                   'h36m-test-s1': join(DATASET_NPZ_PATH, 'h36m_only_S1_valid_protocol1.npz'),
                   'rich-val': join(DATASET_NPZ_PATH, 'rich_world_val_onlycam0_new.npz'),
                   'rich-test': join(DATASET_NPZ_PATH, 'rich_world_test_new.npz'),
                   'rich-test-onlycam0': join(DATASET_NPZ_PATH, 'rich_world_test_onlycam0_new.npz'),
                   'rich-test-last2seq-onlycam0': join(DATASET_NPZ_PATH, 'rich_world_test_last2seq_only_cam0_new.npz'),
                   'lsp': join(DATASET_NPZ_PATH, 'lsp_dataset_test.npz'),
                   'mpi-inf-3dhp': join(DATASET_NPZ_PATH, 'mpi_inf_3dhp_valid.npz'),
                   '3dpw': join(DATASET_NPZ_PATH, '3dpw_test.npz'),
                   'agora_val':  join(DATASET_NPZ_PATH, 'agora_val_0.npz'),
                  },

                  {
                   # 'h36m': join(DATASET_NPZ_PATH, 'h36m_train_world.npz'),
                   'rich': join(DATASET_NPZ_PATH, 'rich_world_val_new.npz'),
                   # 'agora': join(DATASET_NPZ_PATH, 'agora_train.npz'),
                   'agora1': join(DATASET_NPZ_PATH, 'agora_train_0.npz'),
                   'agora2': join(DATASET_NPZ_PATH, 'agora_train_1.npz'),
                   'agora3': join(DATASET_NPZ_PATH, 'agora_train_2.npz'),
                   'agora4': join(DATASET_NPZ_PATH, 'agora_train_3.npz'),
                   'agora5': join(DATASET_NPZ_PATH, 'agora_train_4.npz'),
                   'lsp-orig': join(DATASET_NPZ_PATH, 'lsp_dataset_original_train.npz'),
                   'mpii': join(DATASET_NPZ_PATH, 'mpii_train.npz'),
                   'coco': join(DATASET_NPZ_PATH, 'coco_2014_train.npz'),
                   'lspet': join(DATASET_NPZ_PATH, 'hr-lspet_train.npz'),
                   'mpii-eft': join(DATASET_NPZ_PATH, 'mpii_train_eft.npz'),
                   'coco-eft': join(DATASET_NPZ_PATH, 'coco_2014_train_eft.npz'),
                   'lspet-eft': join(DATASET_NPZ_PATH, 'hr-lspet_train_eft.npz'),
                   'mpi-inf-3dhp': join(DATASET_NPZ_PATH, 'mpi_inf_3dhp_train.npz')
                  }
                ]

DATASET_FOLDERS = {'h36m': H36M_ROOT,
                   'rich': RICH_ROOT,
                   # 'rich-val-onlycam0': RICH_ROOT,
                   # 'rich-test-onlycam0': RICH_ROOT,
                   'rich-val': RICH_ROOT,
                   'rich-test': RICH_ROOT,
                   'rich-test-onlycam0': RICH_ROOT,
                   'rich-val-onlycam1': RICH_ROOT,
                   'rich-val-onlycam5': RICH_ROOT,
                   'rich-test-last2seq-onlycam0': RICH_ROOT,
                   # 'agora': AGORA_ROOT,
                   'agora1': AGORA_ROOT,
                   'agora2': AGORA_ROOT,
                   'agora3': AGORA_ROOT,
                   'agora4': AGORA_ROOT,
                   'agora5': AGORA_ROOT,
                   'agora_val': AGORA_ROOT,
                   'h36m-p1': H36M_ROOT,
                   'h36m-p2': H36M_ROOT,
                   'h36m-test-s1': H36M_ROOT,
                   'h36m-train-small': H36M_ROOT,
                   'lsp-orig': LSP_ORIGINAL_ROOT,
                   'lsp': LSP_ROOT,
                   'lspet': LSPET_ROOT,
                   'lspet-eft': LSPET_ROOT,
                   'mpi-inf-3dhp': MPI_INF_3DHP_ROOT,
                   'mpii': MPII_ROOT,
                   'mpii-eft': MPII_ROOT,
                   'coco': COCO_ROOT,
                   'coco-eft': COCO_ROOT,
                   '3dpw': PW3D_ROOT,
                   'upi-s1h': UPI_S1H_ROOT,
                }

CUBE_PARTS_FILE = 'data/cube_parts.npy'
JOINT_REGRESSOR_TRAIN_EXTRA = 'data/J_regressor_extra.npy'
JOINT_REGRESSOR_H36M = 'data/J_regressor_h36m.npy'
VERTEX_TEXTURE_FILE = 'data/vertex_texture.npy'
STATIC_FITS_DIR = 'data/static_fits'
SMPL_MEAN_PARAMS = 'data/smpl_mean_params.npz'
SMPL_MODEL_DIR = 'data/smpl'

##### CONFIGS #####
hparams = CN()

hparams.name='default'
hparams.time_to_run=np.inf
hparams.resume=False
hparams.num_workers=8
hparams.pin_memory=True
hparams.log_dir='logs'
hparams.checkpoint=None
hparams.from_json=None
hparams.pretrained_checkpoint=None
hparams.num_epochs=200
hparams.lr=5e-5
hparams.cop_w=10.
hparams.cop_k=100.
hparams.in_alpha1=1.
hparams.in_alpha2=0.5
hparams.out_alpha1=1.
hparams.out_alpha2=0.15
hparams.contact_thresh=0.1
hparams.batch_size=64
hparams.summary_steps=100
hparams.test_steps=1000
hparams.run_eval_h36m=False
hparams.run_eval_rich_val=False
hparams.run_eval_rich_test=False
hparams.run_eval_3dpw=False
hparams.checkpoint_steps=10000
hparams.img_res=224
hparams.rot_factor=30
hparams.noise_factor=0.4
hparams.scale_factor=0.25
hparams.ignore_3d=False
hparams.is_agora=False
hparams.is_others=False
hparams.shape_loss_weight=0
hparams.keypoint_loss_weight=5.
hparams.pose_loss_weight=1.
hparams.beta_loss_weight=0.001
hparams.stability_loss_weight=1.
hparams.inside_push_loss_weight=1.
hparams.outside_pull_loss_weight=1.
hparams.openpose_train_weight=0.
hparams.gt_train_weight=1.
hparams.run_smplify=False
hparams.smplify_threshold=100.
hparams.num_smplify_iters=100
hparams.shuffle_train=True


def get_hparams_defaults():
    """Get a yacs hparamsNode object with default values for my_project."""
    # Return a clone so that the defaults will not be altered
    # This is for the "local variable" use pattern
    return hparams.clone()


def update_hparams(hparams_file):
    hparams = get_hparams_defaults()
    hparams.merge_from_file(hparams_file)
    return hparams.clone()


def update_hparams_from_dict(cfg_dict):
    hparams = get_hparams_defaults()
    cfg = hparams.load_cfg(str(cfg_dict))
    hparams.merge_from_other_cfg(cfg)
    return hparams.clone()


def get_grid_search_configs(config, excluded_keys=[]):
    """
    :param config: dictionary with the configurations
    :return: The different configurations
    """

    def bool_to_string(x: Union[List[bool], bool]) -> Union[List[str], str]:
        """
        boolean to string conversion
        :param x: list or bool to be converted
        :return: string converted thinghat
        """
        if isinstance(x, bool):
            return [str(x)]
        for i, j in enumerate(x):
            x[i] = str(j)
        return x

    # exclude from grid search

    flattened_config_dict = flatten(config, reducer='path')
    hyper_params = []

    for k,v in flattened_config_dict.items():
        if isinstance(v,list):
            if k in excluded_keys:
                flattened_config_dict[k] = ['+'.join(v)]
            elif len(v) > 1:
                hyper_params += [k]

        if isinstance(v, list) and isinstance(v[0], bool) :
            flattened_config_dict[k] = bool_to_string(v)

        if not isinstance(v,list):
            if isinstance(v, bool):
                flattened_config_dict[k] = bool_to_string(v)
            else:
                flattened_config_dict[k] = [v]

    keys, values = zip(*flattened_config_dict.items())
    experiments = [dict(zip(keys, v)) for v in itertools.product(*values)]

    for exp_id, exp in enumerate(experiments):
        for param in excluded_keys:
            exp[param] = exp[param].strip().split('+')
        for param_name, param_value in exp.items():
            # print(param_name,type(param_value))
            if isinstance(param_value, list) and (param_value[0] in ['True', 'False']):
                exp[param_name] = [True if x == 'True' else False for x in param_value]
            if param_value in ['True', 'False']:
                if param_value == 'True':
                    exp[param_name] = True
                else:
                    exp[param_name] = False


        experiments[exp_id] = unflatten(exp, splitter='path')

    return experiments, hyper_params


def run_grid_search_experiments(
        cfg_id,
        cfg_file,
        use_cluster,
        bid,
        memory,
        script='main.py',
        cmd_opts=[],
        gpu_min_mem=10000,
        gpu_arch=('tesla', 'quadro', 'rtx'),
):
    cfg = yaml.safe_load(open(cfg_file))
    # parse config file to get a list of configs and related hyperparameters
    different_configs, hyperparams = get_grid_search_configs(
        cfg,
        excluded_keys=[],
    )
    logger.info(f'Grid search hparams: \n {hyperparams}')

    different_configs = [update_hparams_from_dict(c) for c in different_configs]
    logger.info(f'======> Number of experiment configurations is {len(different_configs)}')

    config_to_run = CN(different_configs[cfg_id])

    if use_cluster:
        execute_task_on_cluster(
            script=script,
            output_dir=config_to_run.log_dir,
            exp_name=config_to_run.name,
            num_exp=len(different_configs),
            cfg_file=cfg_file,
            bid_amount=bid,
            num_workers=config_to_run.num_workers,
            memory=memory,
            exp_opts=cmd_opts,
            gpu_min_mem=gpu_min_mem,
            gpu_arch=gpu_arch,
        )
        exit()

    # ==== create logdir using hyperparam settings
    logtime = time.strftime('%d-%m-%Y_%H-%M-%S')
    logdir = f'{logtime}_{config_to_run.name}'

    def get_from_dict(dict, keys):
        return reduce(operator.getitem, keys, dict)

    for hp in hyperparams:
        v = get_from_dict(different_configs[cfg_id], hp.split('/'))
        logdir += f'_{hp.replace("/", ".").replace("_", "").lower()}-{v}'

    logdir = os.path.join(config_to_run.log_dir, logdir)

    os.makedirs(logdir, exist_ok=True)
    shutil.copy(src=cfg_file, dst=os.path.join(logdir, 'config.yaml'))

    # add extra paramters for checkpoints and tensorboard
    config_to_run['checkpoint_dir'] = os.path.join(logdir, 'checkpoints')
    os.makedirs(config_to_run['checkpoint_dir'], exist_ok=True)
    config_to_run['summary_dir'] = os.path.join(logdir, 'tensorboard')
    os.makedirs(config_to_run['summary_dir'], exist_ok=True)

    def save_dict_to_yaml(obj, filename, mode='w'):
        with open(filename, mode) as f:
            yaml.dump(obj, f, default_flow_style=False)

    # save config
    save_dict_to_yaml(
        unflatten(flatten(config_to_run)),
        os.path.join(logdir, 'config_to_run.yaml')
    )

    return config_to_run

config.py
sha2nkt/ipman-r
FOCAL_LENGTH = 5000.
IMG_RES = 224
CONTACT_THRESH = 0.1 # gmof_rho in overlay vis = contact thresh
#IMG_RES = 256

# Mean and standard deviation for normalizing input image
IMG_NORM_MEAN = [0.485, 0.456, 0.406]
IMG_NORM_STD = [0.229, 0.224, 0.225]

# Dictionary of ground plane offsets in datasets
GROUND_OFFSETS = {
    'h36m': -0.095,
}

"""
We create a superset of joints containing the OpenPose joints together with the ones that each dataset provides.
We keep a superset of 24 joints such that we include all joints from every dataset.
If a dataset doesn't provide annotations for a specific joint, we simply ignore it.
The joints used here are the following:
"""
JOINT_NAMES = [
# 25 OpenPose joints (in the order provided by OpenPose)
'OP Nose',
'OP Neck',
'OP RShoulder',
'OP RElbow',
'OP RWrist',
'OP LShoulder',
'OP LElbow',
'OP LWrist',
'OP MidHip',
'OP RHip',
'OP RKnee',
'OP RAnkle',
'OP LHip',
'OP LKnee',
'OP LAnkle',
'OP REye',
'OP LEye',
'OP REar',
'OP LEar',
'OP LBigToe',
'OP LSmallToe',
'OP LHeel',
'OP RBigToe',
'OP RSmallToe',
'OP RHeel',
# 24 Ground Truth joints (superset of joints from different datasets)
'Right Ankle',
'Right Knee',
'Right Hip',
'Left Hip',
'Left Knee',
'Left Ankle',
'Right Wrist',
'Right Elbow',
'Right Shoulder',
'Left Shoulder',
'Left Elbow',
'Left Wrist',
'Neck (LSP)',
'Top of Head (LSP)',
'Pelvis (MPII)',
'Thorax (MPII)',
'Spine (H36M)',
'Jaw (H36M)',
'Head (H36M)',
'Nose',
'Left Eye',
'Right Eye',
'Left Ear',
'Right Ear'
]

# Dict containing the joints in numerical order
JOINT_IDS = {JOINT_NAMES[i]: i for i in range(len(JOINT_NAMES))}

# Map joints to SMPL joints
JOINT_MAP = {
'OP Nose': 24, 'OP Neck': 12, 'OP RShoulder': 17,
'OP RElbow': 19, 'OP RWrist': 21, 'OP LShoulder': 16,
'OP LElbow': 18, 'OP LWrist': 20, 'OP MidHip': 0,
'OP RHip': 2, 'OP RKnee': 5, 'OP RAnkle': 8,
'OP LHip': 1, 'OP LKnee': 4, 'OP LAnkle': 7,
'OP REye': 25, 'OP LEye': 26, 'OP REar': 27,
'OP LEar': 28, 'OP LBigToe': 29, 'OP LSmallToe': 30,
'OP LHeel': 31, 'OP RBigToe': 32, 'OP RSmallToe': 33, 'OP RHeel': 34,
'Right Ankle': 8, 'Right Knee': 5, 'Right Hip': 45,
'Left Hip': 46, 'Left Knee': 4, 'Left Ankle': 7,
'Right Wrist': 21, 'Right Elbow': 19, 'Right Shoulder': 17,
'Left Shoulder': 16, 'Left Elbow': 18, 'Left Wrist': 20,
'Neck (LSP)': 47, 'Top of Head (LSP)': 48,
'Pelvis (MPII)': 49, 'Thorax (MPII)': 50,
'Spine (H36M)': 51, 'Jaw (H36M)': 52,
'Head (H36M)': 53, 'Nose': 24, 'Left Eye': 26,
'Right Eye': 25, 'Left Ear': 28, 'Right Ear': 27
}

# Joint selectors
# Indices to get the 14 LSP joints from the 17 H36M joints
H36M_TO_J17 = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9]
H36M_TO_J14 = H36M_TO_J17[:14]
# Indices to get the 14 LSP joints from the ground truth joints
J24_TO_J17 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 18, 14, 16, 17]
J24_TO_J14 = J24_TO_J17[:14]

# Permutation of SMPL pose parameters when flipping the shape
SMPL_JOINTS_FLIP_PERM = [0, 2, 1, 3, 5, 4, 6, 8, 7, 9, 11, 10, 12, 14, 13, 15, 17, 16, 19, 18, 21, 20, 23, 22]
SMPL_POSE_FLIP_PERM = []
for i in SMPL_JOINTS_FLIP_PERM:
    SMPL_POSE_FLIP_PERM.append(3*i)
    SMPL_POSE_FLIP_PERM.append(3*i+1)
    SMPL_POSE_FLIP_PERM.append(3*i+2)
# Permutation indices for the 24 ground truth joints
J24_FLIP_PERM = [5, 4, 3, 2, 1, 0, 11, 10, 9, 8, 7, 6, 12, 13, 14, 15, 16, 17, 18, 19, 21, 20, 23, 22]
# Permutation indices for the full set of 49 joints
J49_FLIP_PERM = [0, 1, 5, 6, 7, 2, 3, 4, 8, 12, 13, 14, 9, 10, 11, 16, 15, 18, 17, 22, 23, 24, 19, 20, 21]\
              + [25+i for i in J24_FLIP_PERM]

constants.py
sha2nkt/ipman-r
from .base_dataset import BaseDataset
from .base_dataset_eval import BaseDatasetEval
from .mixed_dataset import MixedDataset
from .mixed_dataset_eval import MixedDatasetEval

datasets/__init__.py
sha2nkt/ipman-r
from __future__ import division

import torch
from torch.utils.data import Dataset
from torchvision.transforms import Normalize
import numpy as np
import cv2
import ipdb
from os.path import join

import config
import constants
from utils.imutils import crop, flip_img, flip_pose, flip_kp, transform, rot_aa

class BaseDataset(Dataset):
    """
    Base Dataset Class - Handles data loading and augmentation.
    Able to handle heterogeneous datasets (different annotations available for different datasets).
    You need to update the path to each dataset in utils/config.py.
    """

    def __init__(self, options, dataset, ignore_3d=False, use_augmentation=True, is_train=True):
        super(BaseDataset, self).__init__()
        self.dataset = dataset
        self.is_train = is_train
        self.options = options
        self.img_dir = config.DATASET_FOLDERS[dataset]
        self.normalize_img = Normalize(mean=constants.IMG_NORM_MEAN, std=constants.IMG_NORM_STD)
        print(f'Loading dataset: {config.DATASET_FILES[is_train][dataset]}')
        self.data = np.load(config.DATASET_FILES[is_train][dataset])
        self.imgname = self.data['imgname']
        if 'agora' in self.dataset:
            self.imgpath = self.data['imgpath']
        
        # Get paths to gt masks, if available
        try:
            self.maskname = self.data['maskname']
        except KeyError:
            pass
        try:
            self.partname = self.data['partname']
        except KeyError:
            pass

        # Bounding boxes are assumed to be in the center and scale format
        self.scale = self.data['scale']
        self.center = self.data['center']
        
        # If False, do not do augmentation
        self.use_augmentation = use_augmentation

        # Get gt SMPL parameters, if available
        try:
            self.pose = self.data['pose'].astype(np.float)
            self.betas = self.data['shape'].astype(np.float)
            if 'has_smpl' in self.data:
                self.has_smpl = self.data['has_smpl']
            else:
                self.has_smpl = np.ones(len(self.imgname))
        except KeyError:
            self.has_smpl = np.zeros(len(self.imgname))
        # Get gt SMPL parameters in world coordinates, if available
        try:
            self.pose_world = self.data['pose_world'].astype(np.float)
            self.transl_world = self.data['transl_world'].astype(np.float)
            if 'has_smpl_world' in self.data:
                self.has_smpl_world = self.data['has_smpl_world']
            else:
                self.has_smpl_world = np.ones(len(self.imgname))
        except KeyError:
            self.has_smpl_world = np.zeros(len(self.imgname))

        if ignore_3d:
            self.has_smpl = np.zeros(len(self.imgname))
            self.has_smpl_world = np.zeros(len(self.imgname))

        # Get gt 3D pose, if available
        try:
            self.pose_3d = self.data['S']
            self.has_pose_3d = 1
        except KeyError:
            self.has_pose_3d = 0
        if ignore_3d:
            self.has_pose_3d = 0

        try:
            self.pose_3d_world = self.data['S_world']
            self.has_pose_3d_world = 1
        except KeyError:
            self.has_pose_3d_world = 0
        if ignore_3d:
            self.has_pose_3d_world = 0

        
        # Get 2D keypoints
        try:
            keypoints_gt = self.data['part']
        except KeyError:
            keypoints_gt = np.zeros((len(self.imgname), 24, 3))
        try:
            keypoints_openpose = self.data['openpose']
        except KeyError:
            keypoints_openpose = np.zeros((len(self.imgname), 25, 3))
        self.keypoints = np.concatenate([keypoints_openpose, keypoints_gt], axis=1)

        # Get gender data, if available
        try:
            gender = self.data['gender']
            if 'agora' in self.dataset:
                self.gender = np.array([0 if str(g) == 'male' else 1 for g in gender]).astype(np.int32)
            else:
                self.gender = np.array([0 if str(g) == 'm' else 1 for g in gender]).astype(np.int32)
        except KeyError:
            self.gender = -1 * np.ones(len(self.imgname)).astype(np.int32)

        # Get camera parameters
        try:
            self.cam_r = self.data['cam_r']
            self.cam_t = self.data['cam_t']
            if 'rich' in self.dataset:
                self.cam_t = self.cam_t.swapaxes(2,1) # rich cam_t is (b, 1, 3) whereas h36m cam_t is (b, 3, 1)
            self.cam_k = self.data['cam_k']
        except KeyError:
            self.cam_r = np.zeros((len(self.imgname), 3, 3))
            self.cam_t = np.zeros((len(self.imgname), 3, 1))
            self.cam_k = np.zeros((len(self.imgname), 3, 3))

        # get bos labels
        try:
            self.in_bos_label = self.data['in_bos_label']
            self.contact_label = self.data['contact_label']
        except KeyError:
            self.in_bos_label = np.zeros((len(self.imgname)))
            self.contact_label = np.zeros((len(self.imgname)))

        # get ground offset if rich, otherwise take constant h36m
        try:
            self.ground_offset = self.data['ground_offset']
        except KeyError:
            self.ground_offset = np.zeros((len(self.imgname)))

        self.length = self.scale.shape[0]

    def augm_params(self):
        """Get augmentation parameters."""
        flip = 0            # flipping
        pn = np.ones(3)  # per channel pixel-noise
        rot = 0            # rotation
        sc = 1            # scaling
        if self.is_train:
            # We flip with probability 1/2
            if np.random.uniform() <= 0.5:
                flip = 1
            
            # Each channel is multiplied with a number 
            # in the area [1-opt.noiseFactor,1+opt.noiseFactor]
            pn = np.random.uniform(1-self.options.noise_factor, 1+self.options.noise_factor, 3)
            
            # The rotation is a number in the area [-2*rotFactor, 2*rotFactor]
            rot = min(2*self.options.rot_factor,
                    max(-2*self.options.rot_factor, np.random.randn()*self.options.rot_factor))
            
            # The scale is multiplied with a number
            # in the area [1-scaleFactor,1+scaleFactor]
            sc = min(1+self.options.scale_factor,
                    max(1-self.options.scale_factor, np.random.randn()*self.options.scale_factor+1))
            # but it is zero with probability 3/5
            if np.random.uniform() <= 0.6:
                rot = 0
        
        return flip, pn, rot, sc

    def rgb_processing(self, rgb_img, center, scale, rot, flip, pn):
        """Process rgb image and do augmentation."""
        rgb_img = crop(rgb_img, center, scale,
                      [constants.IMG_RES, constants.IMG_RES], rot=rot)
        # flip the image 
        if flip:
            rgb_img = flip_img(rgb_img)
        # in the rgb image we add pixel noise in a channel-wise manner
        rgb_img[:,:,0] = np.minimum(255.0, np.maximum(0.0, rgb_img[:,:,0]*pn[0]))
        rgb_img[:,:,1] = np.minimum(255.0, np.maximum(0.0, rgb_img[:,:,1]*pn[1]))
        rgb_img[:,:,2] = np.minimum(255.0, np.maximum(0.0, rgb_img[:,:,2]*pn[2]))
        # (3,224,224),float,[0,1]
        rgb_img = np.transpose(rgb_img.astype('float32'),(2,0,1))/255.0
        return rgb_img

    def j2d_processing(self, kp, center, scale, r, f):
        """Process gt 2D keypoints and apply all augmentation transforms."""
        nparts = kp.shape[0]
        for i in range(nparts):
            kp[i,0:2] = transform(kp[i,0:2]+1, center, scale, 
                                  [constants.IMG_RES, constants.IMG_RES], rot=r)
        # convert to normalized coordinates
        kp[:,:-1] = 2.*kp[:,:-1]/constants.IMG_RES - 1.
        # flip the x coordinates
        if f:
             kp = flip_kp(kp)
        kp = kp.astype('float32')
        return kp

    def j3d_processing(self, S, r, f):
        """Process gt 3D keypoints and apply all augmentation transforms."""
        # in-plane rotation
        rot_mat = np.eye(3)
        if not r == 0:
            rot_rad = -r * np.pi / 180
            sn,cs = np.sin(rot_rad), np.cos(rot_rad)
            rot_mat[0,:2] = [cs, -sn]
            rot_mat[1,:2] = [sn, cs]
        S[:, :-1] = np.einsum('ij,kj->ki', rot_mat, S[:, :-1])
        # flip the x coordinates
        if f:
            S = flip_kp(S)
        S = S.astype('float32')
        return S

    def pose_processing(self, pose, r, f):
        """Process SMPL theta parameters  and apply all augmentation transforms."""
        # rotation or the pose parameters
        pose[:3] = rot_aa(pose[:3], r)
        # flip the pose parameters
        if f:
            pose = flip_pose(pose)
        # (72),float
        pose = pose.astype('float32')
        return pose

    def __getitem__(self, index):
        item = {}
        scale = self.scale[index].copy()
        center = self.center[index].copy()
        scale_original = scale.copy()
        center_original = center.copy()

        # Get augmentation parameters
        flip, pn, rot, sc = self.augm_params()

        # Load image
        if 'agora' in self.dataset:
            imgname = join(self.imgpath[index])
        elif 'mpi-inf-3dhp' in self.dataset:
            imgname = join(self.img_dir, self.imgname[index]).replace('frame_', '')
        elif 'eft' in self.dataset:
            imgname = self.imgname[index]
        else:
            imgname = join(self.img_dir, self.imgname[index])
        try:
            img = cv2.imread(imgname)[:,:,::-1].copy().astype(np.float32)
        except TypeError:
            print(imgname)
        orig_shape = np.array(img.shape)[:2]

        # Get SMPL parameters, if available
        if self.has_smpl[index]:
            pose = self.pose[index].copy()
            betas = self.betas[index].copy()
        else:
            pose = np.zeros(72)
            betas = np.zeros(10)

        # Get SMPL parameters in world coordinates, if available
        if self.has_smpl_world[index]:
            pose_world = self.pose_world[index].copy()
            transl_world = self.transl_world[index].copy()
        else:
            pose_world = np.zeros(72)
            transl_world = np.zeros(3)

        # Process image
        try:
            img = self.rgb_processing(img, center, sc*scale, rot, flip, pn)
        except:
            img = self.rgb_processing(img, center_original, sc * scale_original, rot, flip, pn)
        img = torch.from_numpy(img).float()
        # Store image before normalization to use it in visualization
        item['img'] = self.normalize_img(img)
        item['pose'] = torch.from_numpy(self.pose_processing(pose, rot, flip)).float()
        item['betas'] = torch.from_numpy(betas).float()
        item['pose_world'] = torch.from_numpy(pose_world).float() # Todo: should you do some rot, flip processing here?
        item['transl_world'] = torch.from_numpy(transl_world).float()
        item['imgname'] = imgname

        # Get 3D pose, if available
        if self.has_pose_3d:
            S = self.pose_3d[index].copy()
            item['pose_3d'] = torch.from_numpy(self.j3d_processing(S, rot, flip)).float()
        else:
            item['pose_3d'] = torch.zeros(24,4, dtype=torch.float32)

        if self.has_pose_3d_world:
            item['pose_3d_world'] = torch.from_numpy(self.pose_3d_world[index].copy()).float()
        else:
            item['pose_3d_world'] = torch.zeros(24,4, dtype=torch.float32)

        # Get 2D keypoints and apply augmentation transforms
        keypoints = self.keypoints[index].copy()
        item['keypoints'] = torch.from_numpy(self.j2d_processing(keypoints, center, sc*scale, rot, flip)).float()

        item['has_smpl'] = self.has_smpl[index]
        item['has_smpl_world'] = self.has_smpl_world[index]
        item['has_pose_3d'] = self.has_pose_3d
        item['has_pose_3d_world'] = self.has_pose_3d_world
        item['scale'] = float(sc * scale)
        item['center'] = center.astype(np.float32)
        item['orig_shape'] = orig_shape
        item['is_flipped'] = flip
        item['rot_angle'] = np.float32(rot)
        item['gender'] = self.gender[index]
        item['cam_r'] = self.cam_r[index]
        item['cam_t'] = self.cam_t[index]
        item['cam_k'] = self.cam_k[index]
        item['in_bos_label'] = self.in_bos_label[index]
        item['contact_label'] = self.contact_label[index]
        item['ground_offset'] = self.ground_offset[index]
        item['sample_index'] = index
        item['dataset_name'] = self.dataset

        try:
            item['maskname'] = self.maskname[index]
        except AttributeError:
            item['maskname'] = ''
        try:
            item['partname'] = self.partname[index]
        except AttributeError:
            item['partname'] = ''

        return item

    def __len__(self):
        return len(self.imgname)

datasets/base_dataset.py
sha2nkt/ipman-r
from __future__ import division

import torch
from torch.utils.data import Dataset
from torchvision.transforms import Normalize
import numpy as np
import cv2
from os.path import join

import config
import ipdb, os
import constants
from utils.imutils import crop, flip_img, flip_pose, flip_kp, transform, rot_aa


class BaseDatasetEval(Dataset):
    """
    Base Dataset Class - Handles data loading and augmentation.
    Able to handle heterogeneous datasets (different annotations available for different datasets).
    You need to update the path to each dataset in utils/config.py.
    """

    def __init__(self, options, dataset, ignore_3d=False, use_augmentation=True, is_train=True):
        super(BaseDatasetEval, self).__init__()
        self.dataset = dataset
        self.is_train = is_train
        self.options = options
        self.img_dir = config.DATASET_FOLDERS[dataset]
        self.normalize_img = Normalize(mean=constants.IMG_NORM_MEAN, std=constants.IMG_NORM_STD)
        print(f'Loading dataset: {config.DATASET_FILES[is_train][dataset]}')
        self.data = np.load(config.DATASET_FILES[is_train][dataset])
        self.imgname = self.data['imgname']
        if 'agora' in self.dataset:
            self.imgpath = self.data['imgpath']

        # Get paths to gt masks, if available
        try:
            self.maskname = self.data['maskname']
        except KeyError:
            pass
        try:
            self.partname = self.data['partname']
        except KeyError:
            pass

        # Bounding boxes are assumed to be in the center and scale format
        self.scale = self.data['scale']
        self.center = self.data['center']

        # If False, do not do augmentation
        self.use_augmentation = use_augmentation

        # Get gt SMPL parameters, if available
        try:
            self.pose = self.data['pose'].astype(np.float)
            self.betas = self.data['shape'].astype(np.float)
            if 'has_smpl' in self.data:
                self.has_smpl = self.data['has_smpl']
            else:
                self.has_smpl = np.ones(len(self.imgname))
        except KeyError:
            self.has_smpl = np.zeros(len(self.imgname))
        # Get gt SMPL parameters in world coordinates, if available
        try:
            self.pose_world = self.data['pose_world'].astype(np.float)
            self.transl_world = self.data['transl_world'].astype(np.float)
            if 'has_smpl_world' in self.data:
                self.has_smpl_world = self.data['has_smpl_world']
            else:
                self.has_smpl_world = np.ones(len(self.imgname))
        except KeyError:
            self.has_smpl_world = np.zeros(len(self.imgname))

        if ignore_3d:
            self.has_smpl = np.zeros(len(self.imgname))
            self.has_smpl_world = np.zeros(len(self.imgname))

        # Get gt 3D pose, if available
        try:
            self.pose_3d = self.data['S']
            self.has_pose_3d = 1
        except KeyError:
            self.has_pose_3d = 0
        if ignore_3d:
            self.has_pose_3d = 0

        try:
            self.pose_3d_world = self.data['S_world']
            self.has_pose_3d_world = 1
        except KeyError:
            self.has_pose_3d_world = 0
        if ignore_3d:
            self.has_pose_3d_world = 0


        # Get 2D keypoints
        try:
            keypoints_gt = self.data['part']
        except KeyError:
            keypoints_gt = np.zeros((len(self.imgname), 24, 3))
        try:
            keypoints_openpose = self.data['openpose']
        except KeyError:
            keypoints_openpose = np.zeros((len(self.imgname), 25, 3))
        self.keypoints = np.concatenate([keypoints_openpose, keypoints_gt], axis=1)

        # Get gender data, if available
        try:
            gender = self.data['gender']
            if 'agora' in self.dataset:
                self.gender = np.array([0 if str(g) == 'male' else 1 for g in gender]).astype(np.int32)
            else:
                self.gender = np.array([0 if str(g) == 'm' else 1 for g in gender]).astype(np.int32)
        except KeyError:
            self.gender = -1 * np.ones(len(self.imgname)).astype(np.int32)

        # Get camera parameters
        try:
            self.cam_r = self.data['cam_r']
            self.cam_t = self.data['cam_t']
            self.cam_k = self.data['cam_k']
        except KeyError:
            self.cam_r = np.zeros((len(self.imgname), 3, 3))
            self.cam_t = np.zeros((len(self.imgname), 3, 1))
            self.cam_k = np.zeros((len(self.imgname), 2, 1))

        # get bos labels
        try:
            self.in_bos_label = self.data['in_bos_label']
            self.contact_label = self.data['contact_label']
        except KeyError:
            self.in_bos_label = np.zeros((len(self.imgname)))
            self.contact_label = np.zeros((len(self.imgname)))

        # get ground offset if rich, otherwise take constant h36m
        try:
            self.ground_offset = self.data['ground_offset']
        except KeyError:
            self.ground_offset = np.zeros((len(self.imgname)))

        self.length = self.scale.shape[0]

    def augm_params(self):
        """Get augmentation parameters."""
        flip = 0  # flipping
        pn = np.ones(3)  # per channel pixel-noise
        rot = 0  # rotation
        sc = 1  # scaling
        if self.is_train:
            # We flip with probability 1/2
            if np.random.uniform() <= 0.5:
                flip = 1

            # Each channel is multiplied with a number
            # in the area [1-opt.noiseFactor,1+opt.noiseFactor]
            pn = np.random.uniform(1 - self.options.noise_factor, 1 + self.options.noise_factor, 3)

            # The rotation is a number in the area [-2*rotFactor, 2*rotFactor]
            rot = min(2 * self.options.rot_factor,
                      max(-2 * self.options.rot_factor, np.random.randn() * self.options.rot_factor))

            # The scale is multiplied with a number
            # in the area [1-scaleFactor,1+scaleFactor]
            sc = min(1 + self.options.scale_factor,
                     max(1 - self.options.scale_factor, np.random.randn() * self.options.scale_factor + 1))
            # but it is zero with probability 3/5
            if np.random.uniform() <= 0.6:
                rot = 0

        return flip, pn, rot, sc

    def rgb_processing(self, rgb_img, center, scale, rot, flip, pn):
        """Process rgb image and do augmentation."""
        rgb_img = crop(rgb_img, center, scale,
                       [constants.IMG_RES, constants.IMG_RES], rot=rot)
        # flip the image
        if flip:
            rgb_img = flip_img(rgb_img)
        # in the rgb image we add pixel noise in a channel-wise manner
        rgb_img[:, :, 0] = np.minimum(255.0, np.maximum(0.0, rgb_img[:, :, 0] * pn[0]))
        rgb_img[:, :, 1] = np.minimum(255.0, np.maximum(0.0, rgb_img[:, :, 1] * pn[1]))
        rgb_img[:, :, 2] = np.minimum(255.0, np.maximum(0.0, rgb_img[:, :, 2] * pn[2]))
        # (3,224,224),float,[0,1]
        rgb_img = np.transpose(rgb_img.astype('float32'), (2, 0, 1)) / 255.0
        return rgb_img

    def j2d_processing(self, kp, center, scale, r, f):
        """Process gt 2D keypoints and apply all augmentation transforms."""
        nparts = kp.shape[0]
        for i in range(nparts):
            kp[i, 0:2] = transform(kp[i, 0:2] + 1, center, scale,
                                   [constants.IMG_RES, constants.IMG_RES], rot=r)
        # convert to normalized coordinates
        kp[:, :-1] = 2. * kp[:, :-1] / constants.IMG_RES - 1.
        # flip the x coordinates
        if f:
            kp = flip_kp(kp)
        kp = kp.astype('float32')
        return kp

    def j3d_processing(self, S, r, f):
        """Process gt 3D keypoints and apply all augmentation transforms."""
        # in-plane rotation
        rot_mat = np.eye(3)
        if not r == 0:
            rot_rad = -r * np.pi / 180
            sn, cs = np.sin(rot_rad), np.cos(rot_rad)
            rot_mat[0, :2] = [cs, -sn]
            rot_mat[1, :2] = [sn, cs]
        S[:, :-1] = np.einsum('ij,kj->ki', rot_mat, S[:, :-1])
        # flip the x coordinates
        if f:
            S = flip_kp(S)
        S = S.astype('float32')
        return S

    def pose_processing(self, pose, r, f):
        """Process SMPL theta parameters  and apply all augmentation transforms."""
        # rotation or the pose parameters
        pose[:3] = rot_aa(pose[:3], r)
        # flip the pose parameters
        if f:
            pose = flip_pose(pose)
        # (72),float
        pose = pose.astype('float32')
        return pose

    def __getitem__(self, index):
        item = {}
        scale = self.scale[index].copy()
        center = self.center[index].copy()

        # Get augmentation parameters
        flip, pn, rot, sc = self.augm_params()

        # Load image
        if self.dataset == 'h36m-p1' or self.dataset == 'h36m-p2':
            imgname = self.imgname[index].replace('images/','Images/')
            # imgname_splits = imgname.split('_')
            # imgname = os.path.join(imgname_splits[0], '_'.join(imgname_splits[1:-1]), imgname_splits[-1])
            # imgname = join(self.img_dir, imgname)
        else:
            imgname = join(self.img_dir, self.imgname[index])
        try:
            img = cv2.imread(imgname)[:, :, ::-1].copy().astype(np.float32)
        except TypeError:
            print(imgname)
        orig_shape = np.array(img.shape)[:2]

        # Get SMPL parameters, if available
        if self.has_smpl[index]:
            pose = self.pose[index].copy()
            betas = self.betas[index].copy()
        else:
            pose = np.zeros(72)
            betas = np.zeros(10)

        # Get SMPL parameters in world coordinates, if available
        if self.has_smpl_world[index]:
            pose_world = self.pose_world[index].copy()
            transl_world = self.transl_world[index].copy()
        else:
            pose_world = np.zeros(72)
            transl_world = np.zeros(3)

        # Process image
        img = self.rgb_processing(img, center, sc * scale, rot, flip, pn)
        img = torch.from_numpy(img).float()
        # Store image before normalization to use it in visualization
        item['img'] = self.normalize_img(img)
        item['pose'] = torch.from_numpy(self.pose_processing(pose, rot, flip)).float()
        item['betas'] = torch.from_numpy(betas).float()
        item['pose_world'] = torch.from_numpy(pose_world).float() # Todo: should you do some rot, flip processing here?
        item['transl_world'] = torch.from_numpy(transl_world).float()
        item['imgname'] = imgname

        # Get 3D pose, if available
        if self.has_pose_3d:
            S = self.pose_3d[index].copy()
            item['pose_3d'] = torch.from_numpy(self.j3d_processing(S, rot, flip)).float()
        else:
            item['pose_3d'] = torch.zeros(24, 4, dtype=torch.float32)

        if self.has_pose_3d:
            item['pose_3d_world'] = self.pose_3d_world[index].copy()
        else:
            item['pose_3d_world'] = torch.zeros(24,4, dtype=torch.float32)

        # Get 2D keypoints and apply augmentation transforms
        keypoints = self.keypoints[index].copy()
        item['keypoints'] = torch.from_numpy(self.j2d_processing(keypoints, center, sc * scale, rot, flip)).float()

        item['has_smpl'] = self.has_smpl[index]
        item['has_smpl_world'] = self.has_smpl_world[index]
        item['has_pose_3d'] = self.has_pose_3d
        item['scale'] = float(sc * scale)
        item['center'] = center.astype(np.float32)
        item['orig_shape'] = orig_shape
        item['is_flipped'] = flip
        item['rot_angle'] = np.float32(rot)
        item['gender'] = self.gender[index]
        item['cam_r'] = self.cam_r[index]
        item['cam_t'] = self.cam_t[index]
        item['cam_k'] = self.cam_k[index]
        item['in_bos_label'] = self.in_bos_label[index]
        item['contact_label'] = self.contact_label[index]
        item['ground_offset'] = self.ground_offset[index]
        item['sample_index'] = index
        item['dataset_name'] = self.dataset

        try:
            item['maskname'] = self.maskname[index]
        except AttributeError:
            item['maskname'] = ''
        try:
            item['partname'] = self.partname[index]
        except AttributeError:
            item['partname'] = ''

        return item

    def __len__(self):
        return len(self.imgname)
datasets/base_dataset_eval.py
sha2nkt/ipman-r
import cv2
import time
import torch
import numpy as np
from os.path import join
from loguru import logger
import albumentations as A
from torch.utils.data import Dataset
from torchvision.transforms import Normalize
# from nvidia.dali.plugin.pytorch import DALIGenericIterator as PyTorchIterator

from ..models import SMPL
from ..core import constants, config
from ..utils.image_utils import random_crop
from ..core.config import DATASET_FILES, DATASET_FOLDERS
from ..utils.image_utils import crop, flip_img, flip_pose, flip_kp, transform, \
    rot_aa, random_crop, read_img
from .occlusion import load_occluders, occlude_with_objects


# from ..utils.image_utils import crop_cv2 as crop


class BaseDataset(Dataset):
    """
    Base Dataset Class - Handles data loading and augmentation.
    Able to handle heterogeneous datasets (different annotations available for different datasets).
    You need to update the path to each dataset in utils/config.py.
    """

    def __init__(self, options, dataset, ignore_3d=False,
                 use_augmentation=True, is_train=True, num_images=0, occluders=None):
        super(BaseDataset, self).__init__()
        self.dataset = dataset
        self.is_train = is_train
        self.options = options
        self.img_dir = DATASET_FOLDERS[dataset]
        self.normalize_img = Normalize(mean=constants.IMG_NORM_MEAN, std=constants.IMG_NORM_STD)
        self.data = np.load(DATASET_FILES[is_train][dataset])
        self.imgname = self.data['imgname']
        self.imgpath = self.data['imgpath']

        self.dataset_dict = {dataset: 0}

        if num_images > 0:
            # select a random subset of the dataset
            rand = np.random.randint(0, len(self.imgname), size=(num_images))
            logger.info(f'{rand.shape[0]} images are randomly sampled from {self.dataset}')
            self.imgname = self.imgname[rand]
            self.data_subset = {}
            for f in self.data.files:
                self.data_subset[f] = self.data[f][rand]
            self.data = self.data_subset

        # Get paths to gt masks, if available
        try:
            self.maskname = self.data['maskname']
        except KeyError:
            pass
        try:
            self.partname = self.data['partname']
        except KeyError:
            pass

        # Bounding boxes are assumed to be in the center and scale format
        self.scale = self.data['scale']
        self.center = self.data['center']

        # If False, do not do augmentation
        self.use_augmentation = use_augmentation

        # Get gt SMPL parameters, if available
        try:
            self.pose = self.data['pose'].astype(np.float)
            self.betas = self.data['shape'].astype(np.float)
            if 'has_smpl' in self.data:
                self.has_smpl = self.data['has_smpl']
            else:
                self.has_smpl = np.ones(len(self.imgname))
        except KeyError:
            self.has_smpl = np.zeros(len(self.imgname))
        if ignore_3d:
            self.has_smpl = np.zeros(len(self.imgname))

        # Get gt 3D pose, if available
        try:
            self.pose_3d = self.data['S']
            self.has_pose_3d = 1
        except KeyError:
            self.has_pose_3d = 0
        if ignore_3d:
            self.has_pose_3d = 0

        # Get 2D keypoints
        try:
            keypoints_gt = self.data['part']
        except KeyError:
            keypoints_gt = np.zeros((len(self.imgname), 24, 3))
        try:
            keypoints_openpose = self.data['openpose']
        except KeyError:
            keypoints_openpose = np.zeros((len(self.imgname), 25, 3))
        self.keypoints = np.concatenate([keypoints_openpose, keypoints_gt], axis=1)

        # Get gender data, if available
        try:
            gender = self.data['gender']
            self.gender = np.array([0 if str(g) == 'male' else 1 for g in gender]).astype(np.int32)
        except KeyError:
            self.gender = -1 * np.ones(len(self.imgname)).astype(np.int32)

        self.occluders = None
        if is_train and self.options.USE_SYNTHETIC_OCCLUSION:
            if occluders is None:
                logger.info(f'Loading synthetic occluders for {dataset}...')
                self.occluders = load_occluders(pascal_voc_root_path=config.PASCAL_ROOT)
                logger.info('Found {} suitable objects'.format(len(self.occluders)))
            else:
                logger.info('Using mixed/eft dataset occluders')
                self.occluders = occluders

        # evaluation variables
        # self.joint_mapper_h36m = constants.H36M_TO_J17 if dataset == 'mpi-inf-3dhp' \
        #     else constants.H36M_TO_J14
        # self.joint_mapper_gt = constants.J24_TO_J17 if dataset == 'mpi-inf-3dhp' \
        #     else constants.J24_TO_J14
        # self.J_regressor = torch.from_numpy(np.load(config.JOINT_REGRESSOR_H36M)).float()
        #
        # self.smpl = SMPL(
        #     config.SMPL_MODEL_DIR,
        #     batch_size=1,
        #     create_transl=False
        # )

        self.length = self.scale.shape[0]
        logger.info(f'Loaded {self.dataset} dataset, num samples {self.length}')

    def augm_params(self):
        """Get augmentation parameters."""
        flip = 0  # flipping
        pn = np.ones(3)  # per channel pixel-noise
        rot = 0  # rotation
        sc = 1  # scaling
        if self.is_train:
            # We flip with probability 1/2
            if np.random.uniform() <= self.options.FLIP_PROB:
                flip = 1

            # Each channel is multiplied with a number
            # in the area [1-opt.noiseFactor,1+opt.noiseFactor]
            pn = np.random.uniform(1 - self.options.NOISE_FACTOR, 1 + self.options.NOISE_FACTOR, 3)

            # The rotation is a number in the area [-2*rotFactor, 2*rotFactor]
            rot = min(2 * self.options.ROT_FACTOR,
                      max(-2 * self.options.ROT_FACTOR, np.random.randn() * self.options.ROT_FACTOR))

            # The scale is multiplied with a number
            # in the area [1-scaleFactor,1+scaleFactor]
            sc = min(1 + self.options.SCALE_FACTOR,
                     max(1 - self.options.SCALE_FACTOR, np.random.randn() * self.options.SCALE_FACTOR + 1))
            # but it is zero with probability 3/5
            if np.random.uniform() <= 0.6:
                rot = 0

        return flip, pn, rot, sc

    def rgb_processing(self, rgb_img, center, scale, rot, flip, pn, img_res):
        """Process rgb image and do augmentation."""
        rgb_img = crop(rgb_img, center, scale,
                       [img_res, img_res], rot=rot)
        # flip the image
        if flip:
            rgb_img = flip_img(rgb_img)

        if self.occluders is not None:
            rgb_img = occlude_with_objects(rgb_img, self.occluders)

        if self.is_train:
            albumentation_aug = A.Compose(transforms=[A.MotionBlur(p=0.5)])
            rgb_img = albumentation_aug(image=rgb_img)['image']

        # in the rgb image we add pixel noise in a channel-wise manner
        rgb_img[:, :, 0] = np.minimum(255.0, np.maximum(0.0, rgb_img[:, :, 0] * pn[0]))
        rgb_img[:, :, 1] = np.minimum(255.0, np.maximum(0.0, rgb_img[:, :, 1] * pn[1]))
        rgb_img[:, :, 2] = np.minimum(255.0, np.maximum(0.0, rgb_img[:, :, 2] * pn[2]))
        # (3,224,224),float,[0,1]
        rgb_img = np.transpose(rgb_img.astype('float32'), (2, 0, 1)) / 255.0
        return rgb_img

    def j2d_processing(self, kp, center, scale, r, f):
        """Process gt 2D keypoints and apply all augmentation transforms."""
        nparts = kp.shape[0]
        for i in range(nparts):
            kp[i, 0:2] = transform(kp[i, 0:2] + 1, center, scale,
                                   [self.options.IMG_RES, self.options.IMG_RES], rot=r)
        # convert to normalized coordinates
        kp[:, :-1] = 2. * kp[:, :-1] / self.options.IMG_RES - 1.
        # flip the x coordinates
        if f:
            kp = flip_kp(kp)
        kp = kp.astype('float32')
        return kp

    def j3d_processing(self, S, r, f):
        """Process gt 3D keypoints and apply all augmentation transforms."""
        # in-plane rotation
        rot_mat = np.eye(3)
        if not r == 0:
            rot_rad = -r * np.pi / 180
            sn, cs = np.sin(rot_rad), np.cos(rot_rad)
            rot_mat[0, :2] = [cs, -sn]
            rot_mat[1, :2] = [sn, cs]
        S[:, :-1] = np.einsum('ij,kj->ki', rot_mat, S[:, :-1])
        # flip the x coordinates
        if f:
            S = flip_kp(S)
        S = S.astype('float32')
        return S

    def pose_processing(self, pose, r, f):
        """Process SMPL theta parameters  and apply all augmentation transforms."""
        # rotation or the pose parameters
        pose[:3] = rot_aa(pose[:3], r)
        # flip the pose parameters
        if f:
            pose = flip_pose(pose)
        # (72),float
        pose = pose.astype('float32')
        return pose

    def __getitem__(self, index):
        item = {}
        scale = self.scale[index].copy()
        center = self.center[index].copy()
        keypoints = self.keypoints[index].copy()

        # EFT dataset bboxes are a bit large, make them smaller
        if self.dataset in ['mpii', 'coco', 'lspet']:
            scale /= 1.1

        # Get augmentation parameters
        flip, pn, rot, sc = self.augm_params()

        # apply crop augmentation
        if self.is_train and self.options.CROP_FACTOR > 0:
            if np.random.rand() < self.options.CROP_PROB:
                center, scale = random_crop(center, scale, crop_scale_factor=1 - self.options.CROP_FACTOR, axis='y')

        load_start = time.perf_counter()
        # Load image
        # imgname = join(self.img_dir, self.imgname[index])
        imgname = join(self.imgpath[index])
        try:
            # cv_img = np.load(imgname)['arr_0']
            cv_img = read_img(imgname)
            # cv_img = cv2.imread(imgname)[:,:,::-1].copy().astype(np.float32)
        except TypeError:
            logger.info(imgname)

        orig_shape = np.array(cv_img.shape)[:2]
        load_time = time.perf_counter() - load_start
        # print(f'loading: {time.perf_counter() - loading_start}s.')

        # Get SMPL parameters, if available
        if self.has_smpl[index]:
            pose = self.pose[index].copy()
            betas = self.betas[index].copy()
        else:
            pose = np.zeros(72)
            betas = np.zeros(10)

        proc_start = time.perf_counter()
        try:
            # Process image
            img = self.rgb_processing(cv_img, center, sc * scale, rot, flip, pn,
                                      img_res=self.options.IMG_RES)
        except:
            import IPython;
            IPython.embed()

        img = torch.from_numpy(img).float()
        # print(f'crop-augment: {time.perf_counter() - crop_start}s.')
        proc_time = time.perf_counter() - proc_start

        if not self.is_train:
            disp_img = self.rgb_processing(cv_img, center, sc * scale, rot, flip, pn, img_res=self.options.RENDER_RES)
            disp_img = torch.from_numpy(disp_img).float()
            item['disp_img'] = self.normalize_img(disp_img)

        # Store image before normalization to use it in visualization
        item['img'] = self.normalize_img(img)
        item['pose'] = torch.from_numpy(self.pose_processing(pose, rot, flip)).float()
        item['betas'] = torch.from_numpy(betas).float()
        item['imgname'] = imgname

        # Get 3D pose, if available
        if self.has_pose_3d:
            S = self.pose_3d[index].copy()
            item['pose_3d'] = torch.from_numpy(self.j3d_processing(S, rot, flip)).float()
        else:
            item['pose_3d'] = torch.zeros(24, 4, dtype=torch.float32)

        # prepare pose_3d for evaluation
        # For 3DPW get the 14 common joints from the rendered shape
        # if not self.is_train:
        #     if '3dpw' in self.dataset:
        #         J_regressor_batch = self.J_regressor[None, :].expand(1, -1, -1)
        #         gt_vertices = self.smpl(
        #             global_orient=item['pose'].unsqueeze(0)[:, :3],
        #             body_pose=item['pose'].unsqueeze(0)[:, 3:],
        #             betas=item['betas'].unsqueeze(0),
        #         ).vertices
        #         pose_3d = torch.matmul(J_regressor_batch, gt_vertices)
        #         pelvis = pose_3d[:, [0], :].clone()
        #         pose_3d = pose_3d[:, self.joint_mapper_h36m, :]
        #         pose_3d = pose_3d - pelvis
        #         item['pose_3d'] = pose_3d[0].float()
        #     else:
        #         item['pose_3d'] = item['pose_3d'][self.joint_mapper_gt, :-1].float()

        # Get 2D keypoints and apply augmentation transforms
        keypoints = self.j2d_processing(keypoints, center, sc * scale, rot, flip)

        # if self.options.USE_HEATMAPS:
        #     joints_vis = np.ones((keypoints.shape[0], 1))
        #     item['heatmaps_2d'], item['joint_vis'] = generate_heatmaps_2d(
        #         joints=keypoints[:24], joints_vis=joints_vis, num_joints=keypoints.shape[0],
        #         heatmap_size=self.options.IMG_RES//4, image_size=self.options.IMG_RES,
        #         sigma=(2./64.)*self.options.IMG_RES//4,
        #     )

        item['keypoints'] = torch.from_numpy(keypoints).float()
        item['has_smpl'] = self.has_smpl[index]
        item['has_pose_3d'] = self.has_pose_3d
        item['scale'] = float(sc * scale)
        item['center'] = center.astype(np.float32)
        item['orig_shape'] = orig_shape
        item['is_flipped'] = flip
        item['rot_angle'] = np.float32(rot)
        item['gender'] = self.gender[index]
        item['sample_index'] = index
        item['dataset_name'] = self.dataset

        try:
            item['maskname'] = self.maskname[index]
        except AttributeError:
            item['maskname'] = ''
        try:
            item['partname'] = self.partname[index]
        except AttributeError:
            item['partname'] = ''

        item['load_time'] = load_time
        item['proc_time'] = proc_time
        return item

    def __len__(self):
        return len(self.imgname)

datasets/base_dataset_m.py
sha2nkt/ipman-r
"""
This file contains the definition of different heterogeneous datasets used for training
"""
import torch
import numpy as np

from .base_dataset import BaseDataset

class MixedDataset(torch.utils.data.Dataset):

    def __init__(self, options, **kwargs):
        # self.dataset_list = ['h36m', 'lsp-orig', 'mpii', 'lspet', 'coco', 'mpi-inf-3dhp', 'rich']
        # self.dataset_dict = {'h36m': 0, 'lsp-orig': 1, 'mpii': 2, 'lspet': 3, 'coco': 4, 'mpi-inf-3dhp': 5, 'rich': 6}
        # self.dataset_list = ['lsp-orig', 'mpii', 'lspet', 'coco', 'mpi-inf-3dhp', 'rich']
        # self.dataset_dict = {'lsp-orig': 0, 'mpii': 1, 'lspet': 2, 'coco': 3, 'mpi-inf-3dhp': 4, 'rich': 5}
        # self.dataset_list = ['h36m']
        # self.dataset_dict = {'h36m': 0}
        # self.dataset_list = ['rich']
        # self.dataset_dict = {'rich': 0}
        # self.dataset_list = ['h36m', 'rich']
        # self.dataset_dict = {'h36m': 0, 'rich': 1}
        self.dataset_list = ['mpi-inf-3dhp', 'rich']
        self.dataset_dict = {'mpi-inf-3dhp': 0, 'rich': 1}
        # self.dataset_list = ['agora1', 'agora2', 'agora3', 'agora4']
        # self.dataset_dict = {'agora1': 0, 'agora2': 1, 'agora3': 2, 'agora4': 3, 'agora5':4}
        self.datasets = [BaseDataset(options, ds, **kwargs) for ds in self.dataset_list]
        total_length = sum([len(ds) for ds in self.datasets])
        length_itw = sum([len(ds) for ds in self.datasets])
        # length_itw = sum([len(ds) for ds in self.datasets[1:-2]]) # only when training with 2d datasets
        self.length = max([len(ds) for ds in self.datasets])
        """
        Data distribution inside each batch:
        # 30% H36M - 60% ITW - 10% MPI-INF
        30% H36M - 40% ITW - 10% MPI-INF - 20% RICH
        """
        # self.partition = [0.3, .4*len(self.datasets[1])/length_itw,
        #                   .4*len(self.datasets[2])/length_itw,
        #                   .4*len(self.datasets[3])/length_itw,
        #                   .4*len(self.datasets[4])/length_itw,
        #                   0.1, 0.2]
        # self.partition = [.4*len(self.datasets[0])/length_itw,
        #                   .4*len(self.datasets[1])/length_itw,
        #                   .4*len(self.datasets[2])/length_itw,
        #                   .4*len(self.datasets[3])/length_itw,
        #                   0.2, 0.4]
        # self.partition = [1.0]
        self.partition = [0.7, 0.3]
        # self.partition = [0.3, .4*len(self.datasets[1])/length_itw,
        #                   .4*len(self.datasets[2])/length_itw,
        #                   .4*len(self.datasets[3])/length_itw,
        #                   .4*len(self.datasets[4])/length_itw,
        #                   0.1, 0.2]
        self.partition = np.array(self.partition).cumsum()

    def __getitem__(self, index):
        p = np.random.rand()
        # for i in range(6):
        for i in range(2):
            if p <= self.partition[i]:
                return self.datasets[i][index % len(self.datasets[i])]

    def __len__(self):
        return self.length

datasets/mixed_dataset.py
sha2nkt/ipman-r
"""
This file contains the definition of different heterogeneous datasets used for training
"""
import torch
import numpy as np

from .base_dataset_eval import BaseDatasetEval

class MixedDatasetEval(torch.utils.data.Dataset):

    def __init__(self, options, **kwargs):
        self.dataset_list = ['h36m', 'lsp-orig', 'mpii', 'lspet', 'coco', 'mpi-inf-3dhp']
        self.dataset_dict = {'h36m': 0, 'lsp-orig': 1, 'mpii': 2, 'lspet': 3, 'coco': 4, 'mpi-inf-3dhp': 5}
        self.datasets = [BaseDatasetEval(options, ds, **kwargs) for ds in self.dataset_list]
        total_length = sum([len(ds) for ds in self.datasets])
        length_itw = sum([len(ds) for ds in self.datasets[1:-1]])
        self.length = max([len(ds) for ds in self.datasets])
        """
        Data distribution inside each batch:
        30% H36M - 60% ITW - 10% MPI-INF
        """
        self.partition = [.3, .6*len(self.datasets[1])/length_itw,
                          .6*len(self.datasets[2])/length_itw,
                          .6*len(self.datasets[3])/length_itw,
                          .6*len(self.datasets[4])/length_itw,
                          0.1]
        self.partition = np.array(self.partition).cumsum()

    def __getitem__(self, index):
        p = np.random.rand()
        for i in range(6):
            if p <= self.partition[i]:
                return self.datasets[i][index % len(self.datasets[i])]

    def __len__(self):
        return self.length
datasets/mixed_dataset_eval.py
sha2nkt/ipman-r
import os
import sys
import stat
import shutil
import subprocess

from loguru import logger

GPUS = {
    'v100-v16': ('\"Tesla V100-PCIE-16GB\"', 'tesla', 16000),
    'v100-p32': ('\"Tesla V100-PCIE-32GB\"', 'tesla', 32000),
    'v100-s32': ('\"Tesla V100-SXM2-32GB\"', 'tesla', 32000),
    'v100-p16': ('\"Tesla P100-PCIE-16GB\"', 'tesla', 16000),
    'rtx2080ti': ('\"NVIDIA GeForce RTX 2080 Ti\"', 'rtx', 11000),
    'quadro6000': ('\"Quadro RTX 6000\"', 'quadro', 24000),
}

def get_gpus(min_mem=10000, arch=('tesla', 'quadro', 'rtx')):
    gpu_names = []
    for k, (gpu_name, gpu_arch, gpu_mem) in GPUS.items():
        if gpu_mem >= min_mem and gpu_arch in arch:
            gpu_names.append(gpu_name)

    assert len(gpu_names) > 0, 'Suitable GPU model could not be found'

    return gpu_names


def execute_task_on_cluster(
        script,
        exp_name,
        input_folder,
        cluster_bs,
        num_queue,
        condor_dir,
        cfg_file,
        exp_opts=None,
        bid_amount=5,
        num_workers=8,
        memory=64000,
        gpu_min_mem=10000,
        gpu_arch=('tesla', 'quadro', 'rtx'),
        num_gpus=1
):
    # copy config to a new experiment directory and source from there
    temp_config_dir = os.path.join(os.path.dirname(condor_dir), 'temp_configs')
    logdir = os.path.join(temp_config_dir, f'{exp_name}')
    os.makedirs(logdir, exist_ok=True)
    new_cfg_file = os.path.join(logdir, 'config.yaml')
    shutil.copy(src=cfg_file, dst=new_cfg_file)

    gpus = get_gpus(min_mem=gpu_min_mem, arch=gpu_arch)

    gpus = ' || '.join([f'CUDADeviceName=={x}' for x in gpus])

    os.makedirs(os.path.join(condor_dir, 'condorlog', exp_name), exist_ok=True)
    submission = f'executable = {condor_dir}/condorlog/{exp_name}/{exp_name}_run.sh\n' \
                 'arguments = $(Process) $(Cluster)\n' \
                 f'error = {condor_dir}/condorlog/{exp_name}/$(Cluster).$(Process).err\n' \
                 f'output = {condor_dir}/condorlog/{exp_name}/$(Cluster).$(Process).out\n' \
                 f'log = {condor_dir}/condorlog/{exp_name}/$(Cluster).$(Process).log\n' \
                 f'request_memory = {memory}\n' \
                 f'request_cpus={int(num_workers/2)}\n' \
                 f'request_gpus={num_gpus}\n' \
                 f'requirements={gpus}\n' \
                 f'+MaxRunningPrice = 500\n' \
                 f'queue {num_queue}'
                 # f'request_cpus={int(num_workers/2)}\n' \
                 # f'+RunningPriceExceededAction = \"kill\"\n' \
    print('<<< Condor Submission >>> ')
    print(submission)

    with open('submit.sub', 'w') as f:
        f.write(submission)

    # output_dir = os.path.join(output_dir, exp_name)
    logger.info(f'The logs for this experiments can be found under: {condor_dir}/condorlog/{exp_name}')
    bash = 'export PYTHONBUFFERED=1\n export PATH=$PATH\n ' \
           f'{sys.executable} ../ipman_regr/datasets/preprocess/RICH/{script} ' \
           f'--config {new_cfg_file} ' \
           f'--input_folder {input_folder} ' \
           f' --condor_dir {condor_dir} ' \
           f' --cluster_bs {cluster_bs} ' \
           f' --ds_start_idx $1'  ## This is the trick. Notice there is no --cluster here
    if exp_opts is not None:
        bash += ' --opts '
        for opt in exp_opts:
            bash += f'{opt} '
        bash += 'SYSTEM.CLUSTER_NODE $2.$1'
    else:
        bash += ' --opts SYSTEM.CLUSTER_NODE $2.$1'

    executable_path =f'{condor_dir}/condorlog/{exp_name}/{exp_name}_run.sh'

    with open(executable_path, 'w') as f:
        f.write(bash)

    os.chmod(executable_path, stat.S_IRWXU)

    cmd = ['condor_submit_bid', f'{bid_amount}', 'submit.sub']
    logger.info('Executing ' + ' '.join(cmd))
    subprocess.call(cmd)

datasets/preprocess/RICH/rich_cluster.py
sha2nkt/ipman-r
import os
import time
import shutil
import yaml
import itertools
# from flatten_dict import flatten, unflatten
from typing import Dict, List, Union, Any
from loguru import logger
import operator
from functools import reduce
from rich_cluster import execute_task_on_cluster
from yacs.config import CfgNode as CN
SMPL_MODEL_DIR = '/ps/project/common/smplifyx/models/'

def get_hparams_defaults():
    """Get a yacs hparamsNode object with default values for my_project."""
    # Return a clone so that the defaults will not be altered
    # This is for the "local variable" use pattern
    return hparams.clone()

def update_hparams_from_dict(cfg_dict):
    hparams = get_hparams_defaults()
    cfg = hparams.load_cfg(str(cfg_dict))
    hparams.merge_from_other_cfg(cfg)
    return hparams.clone()

def get_grid_search_configs(config, excluded_keys=[]):
    """
    :param config: dictionary with the configurations
    :return: The different configurations
    """

    def bool_to_string(x: Union[List[bool], bool]) -> Union[List[str], str]:
        """
        boolean to string conversion
        :param x: list or bool to be converted
        :return: string converted thinghat
        """
        if isinstance(x, bool):
            return [str(x)]
        for i, j in enumerate(x):
            x[i] = str(j)
        return x

    # exclude from grid search

    flattened_config_dict = flatten(config, reducer='path')
    hyper_params = []
    for k,v in flattened_config_dict.items():
        if isinstance(v,list):
            if k in excluded_keys:
                flattened_config_dict[k] = ['+'.join(v)]
            elif len(v) > 1:
                hyper_params += [k]

        if isinstance(v, list) and isinstance(v[0], bool) :
            flattened_config_dict[k] = bool_to_string(v)

        if not isinstance(v,list):
            if isinstance(v, bool):
                flattened_config_dict[k] = bool_to_string(v)
            else:
                flattened_config_dict[k] = [v]

    keys, values = zip(*flattened_config_dict.items())
    experiments = [dict(zip(keys, v)) for v in itertools.product(*values)]

    for exp_id, exp in enumerate(experiments):
        for param in excluded_keys:
            exp[param] = exp[param].strip().split('+')
        for param_name, param_value in exp.items():
            # print(param_name,type(param_value))
            if isinstance(param_value, list) and (param_value[0] in ['True', 'False']):
                exp[param_name] = [True if x == 'True' else False for x in param_value]
            if param_value in ['True', 'False']:
                if param_value == 'True':
                    exp[param_name] = True
                else:
                    exp[param_name] = False


        experiments[exp_id] = unflatten(exp, splitter='path')

    return experiments, hyper_params

def run_grid_search_experiments(
        args,
    script='main.py'):
    # cfg = yaml.safe_load(open(cfg_file))
    # parse config file to get a list of configs and the name of hyperparameters that change
    # different_configs, hyperparams = get_grid_search_configs(
    #     cfg,
    #     excluded_keys=['TRAIN/DATASETS_2D', 'TRAIN/DATASETS_3D', 'TRAIN/DATASET_EVAL'],
    # )
    # logger.info(f'Grid search hparams: \n {hyperparams}')

    # different_configs = [update_hparams_from_dict(c) for c in different_configs]
    # logger.info(f'======> Number of experiment configurations is {len(different_configs)}')


    config_to_run = CN(args)

    if config_to_run.cluster: # if running on cluster
        execute_task_on_cluster(
            script=script,
            exp_name=config_to_run.exp_name,
            input_folder=config_to_run.input_folder,
            cluster_bs=config_to_run.cluster_bs,
            num_queue=config_to_run.num_queue,
            condor_dir=config_to_run.condor_dir,
            cfg_file=config_to_run.config,
            bid_amount=config_to_run.bid,
            num_workers=4,
            memory=config_to_run.memory,
            exp_opts=config_to_run.opts,
            gpu_min_mem=config_to_run.gpu_min_mem,
            gpu_arch=config_to_run.gpu_arch,
        )
        exit()

    return config_to_run
datasets/preprocess/RICH/rich_config.py
sha2nkt/ipman-r
### this toy demo script shows how to use RICH data in the subfolders

import os
import cv2
import json
import numpy as np
import trimesh
import os.path as osp

def extract_cam_param_xml(xml_path:str='', dtype=np.float):
    
    import xml.etree.ElementTree as ET
    tree = ET.parse(xml_path)

    extrinsics_mat = [float(s) for s in tree.find('./CameraMatrix/data').text.split()]
    intrinsics_mat = [float(s) for s in tree.find('./Intrinsics/data').text.split()]
    distortion_vec = [float(s) for s in tree.find('./Distortion/data').text.split()]

    focal_length_x = intrinsics_mat[0]
    focal_length_y = intrinsics_mat[4]
    center = np.array([[intrinsics_mat[2], intrinsics_mat[5]]], dtype=dtype)
    
    rotation = np.array([[extrinsics_mat[0], extrinsics_mat[1], extrinsics_mat[2]], 
                            [extrinsics_mat[4], extrinsics_mat[5], extrinsics_mat[6]], 
                            [extrinsics_mat[8], extrinsics_mat[9], extrinsics_mat[10]]], dtype=dtype)

    translation = np.array([[extrinsics_mat[3], extrinsics_mat[7], extrinsics_mat[11]]], dtype=dtype)

    # t = -Rc --> c = -R^Tt
    cam_center = [  -extrinsics_mat[0]*extrinsics_mat[3] - extrinsics_mat[4]*extrinsics_mat[7] - extrinsics_mat[8]*extrinsics_mat[11],
                    -extrinsics_mat[1]*extrinsics_mat[3] - extrinsics_mat[5]*extrinsics_mat[7] - extrinsics_mat[9]*extrinsics_mat[11], 
                    -extrinsics_mat[2]*extrinsics_mat[3] - extrinsics_mat[6]*extrinsics_mat[7] - extrinsics_mat[10]*extrinsics_mat[11]]

    cam_center =np.array([cam_center], dtype=dtype)

    k1 = np.array([distortion_vec[0]], dtype=dtype)
    k2 = np.array([distortion_vec[1]], dtype=dtype)

    return focal_length_x, focal_length_y, center, rotation, translation, cam_center, k1, k2


################## meta data #################
SPLIT = 'val'
SEQ_NAME ='2021-06-15_Multi_IOI_ID_03588_Yoga1'
ROOT_DIR = '/ps/scratch/ps_shared/stripathi/4yogi/RICH'
cam_id = 2
seq_dir = osp.join(ROOT_DIR,SPLIT,SEQ_NAME)
frame_id = 100 ## randomly chosen frame


################## visualize body in scan coordidate #################

## body resides in multi-ioi coordidate, where camera 0 is world zero.
body_mesh = trimesh.load(os.path.join(seq_dir,'params', f'{frame_id:05d}', '00/meshes/000.ply'), process=False)

## ground resides in Leica scan coordidate, which is (roughly) axis aligned. 
ground_mesh = trimesh.load(os.path.join(seq_dir,'ground_mesh.ply'), process=False)

## rigid transformation between multi-ioi and Leica scan
ioi2scan_fn = os.path.join(seq_dir,'cam2scan.json')
with open(ioi2scan_fn, 'r') as f:
    ioi2scan_dict = json.load(f)
    R = np.array(ioi2scan_dict['R'])
    t = np.array(ioi2scan_dict['t']).reshape(1, 3)

vertices_scan = body_mesh.vertices @ R + t # row vector --> right multiplication
body_mesh.vertices = vertices_scan

ground_eq = np.mean(ground_mesh.vertices, axis=0)
print(f'Ground plane equation: z = {ground_eq[2]}')
(body_mesh + ground_mesh).show()



################## visualize body kpts in images #################

## load the body (in ioi coordinate) again
body_mesh = trimesh.load(os.path.join(seq_dir,'params', f'{frame_id:05d}', '00/meshes/000.ply'), process=False)

## load and build camera params
camera_fn = os.path.join(seq_dir,f'calibration/{cam_id:03d}.xml')
focal_length_x, focal_length_y, center, rotation, translation, _, _, _ = extract_cam_param_xml(camera_fn)

K = np.eye(3, dtype=np.float)
K[0, 0] = focal_length_x
K[1, 1] = focal_length_y
K[:2, 2:] = center.T

## transform to each camera coordinate and project to the image domain
vertices_cam = body_mesh.vertices @ rotation.T + translation
projected_points = (K @ vertices_cam.T).T
img_points = projected_points[:, :2] / np.hstack((projected_points[:, 2:], projected_points[:, 2:]))

## visualize projected points
img = cv2.imread(osp.join(seq_dir,f'data/{frame_id:05d}','00/images_orig', f'{frame_id:05d}_{cam_id:02d}.png'))
img_points = img_points.astype(np.int)
img[img_points[:,1], img_points[:,0], :] = [0, 255, 0]
cv2.imshow('Vis', img)
cv2.waitKey()
##






datasets/preprocess/RICH/rich_demo.py
sha2nkt/ipman-r
### this toy demo script shows how to use RICH data in the subfolders

import cv2
import json
import numpy as np



with open('/ps/scratch/ps_shared/stripathi/4yogi/RICH/test/MultiIOI_201019_ID03581_parkingLot_Calibration06_Settings06_PushUp__2/data/00068/00/bbox_refine/00068_00.json', 'r') as f:
    bbox = json.load(f)

bbox_top_left = np.array([bbox['x1'], bbox['y1']])

with open('/ps/scratch/ps_shared/stripathi/4yogi/RICH/test/MultiIOI_201019_ID03581_parkingLot_Calibration06_Settings06_PushUp__2/data/00068/00/keypoints_refine/00068_00_keypoints.json', 'r') as f:
    keypoints = json.load(f)

keypoints = np.array(keypoints['people'][0]['pose_keypoints_2d']).reshape([-1, 3])

keypoints[:,:2] = keypoints[:,:2] + bbox_top_left

img = cv2.imread('/ps/scratch/ps_shared/stripathi/4yogi/RICH/test/MultiIOI_201019_ID03581_parkingLot_Calibration06_Settings06_PushUp__2/data/00068/00/images_orig/00068_00.bmp')

img[keypoints[:,1].astype(int), keypoints[:,0].astype(int), :] = [0, 255, 0] # green
cv2.circle(img, (int(bbox['x1']), int(bbox['y1'])), 4, (255, 0, 0), -1) #blue
cv2.circle(img, (int(bbox['x1_rect']), int(bbox['y1_rect'])), 4, (0, 0, 255), -1) #red

cv2.imshow('Vis', img)
cv2.waitKey()
##






datasets/preprocess/RICH/rich_demo_kpy.py
sha2nkt/ipman-r
"""
python ../ipman_regr/datasets/preprocess//RICH/rich_smplx2smpl_conversion.py \
 --condor_dir ../ipman_regr/datasets/preprocess//RICH/condor \
 --input_folder /ps/scratch/ps_shared/stripathi/4yogi/RICH/val/2021-06-15_Multi_IOI_ID_03588_Yoga1/params \
 --config ../ipman_regr/datasets/preprocess//RICH/rich_config.yaml \
  --exp_name 2021-06-15_Multi_IOI_ID_03588_Yoga1 \
  --ds_start_idx 0 \
  --num_queue 19 \
  --cluster_bs 50 \
  --cluster
"""

import os
import sys
import subprocess
import argparse
import glob
from rich_config import run_grid_search_experiments
import shutil
import time
import numpy as np
from PIL import Image
import multiprocessing as mp




def main(args, downscale_factor=4):
    in_dir = args.input_folder
    # recursively get all pkl files in the directory
    types = ('*.png', '*.bmp')  # the tuple of file types
    in_images = []
    for files in types:
        in_images.extend(glob.glob(os.path.join(in_dir, '*/*/images_orig/'+files)))
    # select index for cluster usage
    print('Input Folder has {} objects'.format(len(in_images)))

    idxs = np.arange(len(in_images))
    cbs = len(in_images) if args.get('cluster_bs') is None else args.get('cluster_bs')
    sidx = args.ds_start_idx
    for idx in idxs[sidx * cbs: sidx * cbs + cbs]:
        in_img = in_images[idx]
        out_img_dir = os.path.dirname(in_img).replace('images_orig', 'images_orig_720p')
        os.makedirs(out_img_dir, exist_ok=True)
        out_img = os.path.join(out_img_dir, os.path.basename(in_img).replace('.bmp', '.png'))
        # bash = 'export PYTHONBUFFERED=1\n export PATH=$PATH\n' \
        #        f'{sys.executable} {script} --cfg {new_cfg_file} --cfg_id $1'
        img = Image.open(in_img)
        new_w = img.width // downscale_factor
        new_h = img.height // downscale_factor
        img = img.resize((new_w, new_h), Image.ANTIALIAS)
        print(f'Saving {out_img}')
        img.save(out_img)

def add_common_cmdline_args(parser):
    parser.add_argument('--opts', default=[], nargs='*', help='additional options to update config')
    parser.add_argument('--cfg_id', type=int, default=0, help='cfg id to run when multiple experiments are spawned')
    parser.add_argument('--cluster', default=False, action='store_true', help='creates submission files for cluster')
    parser.add_argument('--bid', type=int, default=30, help='amount of bid for cluster')
    parser.add_argument('--memory', type=int, default=10000, help='memory amount for cluster')
    parser.add_argument('--gpu_min_mem', type=int, default=10000, help='minimum amount of GPU memory')
    parser.add_argument('--gpu_arch', default=['tesla', 'quadro', 'rtx'],
                        nargs='*', help='additional options to update config')
    parser.add_argument('--num_cpus', type=int, default=1, help='num cpus for cluster')
    return parser


if __name__ == '__main__':
    # load argparse and take data_path as input
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_folder', type=str, default='/ps/scratch/ps_shared/stripathi/frmKocabas/human36m',
                        help='path to the h36m dataset')
    # config file added in common agrugments below
    parser.add_argument('-c', '--config',
                        required=True,
                        help='config file path')
    parser.add_argument('--exp_name',
                        default=None,
                        help='Custom name of the experiment')
    parser.add_argument('--ds_start_idx', type=int, default=0,
                        help='set index at which to start processing dataset')
    parser.add_argument('--cluster_bs', type=int, default=1,
                        help='number of dataset objects to process')
    parser.add_argument('--condor_dir',
                        default='./condorlog',
                        type=str,
                        help='The folder where the condor logs and configs are stored.')
    parser.add_argument('--num_queue',
                        default=1,
                        help='The number of jobs in db_file to run in cluster')
    parser = add_common_cmdline_args(parser)
    args = parser.parse_args()
    args = vars(args)
    args = run_grid_search_experiments(
        args,
        script='rich_downsample_images.py',
    )
    main(args)


datasets/preprocess/RICH/rich_downsample_images.py
sha2nkt/ipman-r
"""
python ../ipman_regr/datasets/preprocess//RICH/rich_smplx2smpl_conversion.py \
 --condor_dir ../ipman_regr/datasets/preprocess//RICH/condor \
 --input_folder /ps/scratch/ps_shared/stripathi/4yogi/RICH/val/2021-06-15_Multi_IOI_ID_03588_Yoga1/params \
 --config ../ipman_regr/datasets/preprocess//RICH/rich_config.yaml \
  --exp_name 2021-06-15_Multi_IOI_ID_03588_Yoga1 \
  --ds_start_idx 0 \
  --num_queue 19 \
  --cluster_bs 50 \
  --cluster
"""

import os
import sys
import subprocess
import argparse
import glob
from rich_config import run_grid_search_experiments
import shutil
import time
import numpy as np
import multiprocessing as mp


def main(args):
    in_dir = args.input_folder
    # recursively get all pkl files in the directory
    in_mesh_dirs = sorted(glob.glob(os.path.join(in_dir, '*/*/meshes')))
    # select index for cluster usage
    print('Input Folder has {} objects'.format(len(in_mesh_dirs)))
    idxs = np.arange(len(in_mesh_dirs))
    cbs = len(in_mesh_dirs) if args.get('cluster_bs') is None else args.get('cluster_bs')
    sidx = args.ds_start_idx
    for idx in idxs[sidx * cbs: sidx * cbs + cbs]:
        in_mesh_dir = in_mesh_dirs[idx]
        out_mesh_dir = in_mesh_dir.replace('meshes', 'results_smpl')
        os.makedirs(out_mesh_dir, exist_ok=True)
        # bash = 'export PYTHONBUFFERED=1\n export PATH=$PATH\n' \
        #        f'{sys.executable} {script} --cfg {new_cfg_file} --cfg_id $1'
        print(f'Saving {out_mesh_dir}')
        cmd = 'export PYTHONBUFFERED=1\n export PATH=$PATH\n ' \
              f'{sys.executable} -m transfer_model ' \
              '--exp-cfg /is/cluster/scratch/stripathi/pycharm_remote/smplx/config_files/smplx2smpl.yaml ' \
              '--exp-opts ' \
              f'output_folder={out_mesh_dir} ' \
              f'datasets.mesh_folder.data_folder={in_mesh_dir}'
        os.system(cmd)

def add_common_cmdline_args(parser):
    parser.add_argument('--opts', default=[], nargs='*', help='additional options to update config')
    parser.add_argument('--cfg_id', type=int, default=0, help='cfg id to run when multiple experiments are spawned')
    parser.add_argument('--cluster', default=False, action='store_true', help='creates submission files for cluster')
    parser.add_argument('--bid', type=int, default=30, help='amount of bid for cluster')
    parser.add_argument('--memory', type=int, default=10000, help='memory amount for cluster')
    parser.add_argument('--gpu_min_mem', type=int, default=10000, help='minimum amount of GPU memory')
    parser.add_argument('--gpu_arch', default=['tesla', 'quadro', 'rtx'],
                        nargs='*', help='additional options to update config')
    parser.add_argument('--num_cpus', type=int, default=1, help='num cpus for cluster')
    return parser


if __name__ == '__main__':
    # load argparse and take data_path as input
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_folder', type=str, default='/ps/scratch/ps_shared/stripathi/frmKocabas/human36m',
                        help='path to the h36m dataset')
    # config file added in common agrugments below
    parser.add_argument('-c', '--config',
                        required=True,
                        help='config file path')
    parser.add_argument('--exp_name',
                        default=None,
                        help='Custom name of the experiment')
    parser.add_argument('--ds_start_idx', type=int, default=0,
                        help='set index at which to start processing dataset')
    parser.add_argument('--cluster_bs', type=int, default=1,
                        help='number of dataset objects to process')
    parser.add_argument('--condor_dir',
                        default='./condorlog',
                        type=str,
                        help='The folder where the condor logs and configs are stored.')
    parser.add_argument('--num_queue',
                        default=1,
                        help='The number of jobs in db_file to run in cluster')
    parser = add_common_cmdline_args(parser)
    args = parser.parse_args()
    args = vars(args)
    args = run_grid_search_experiments(
        args,
        script='rich_smplx2smpl_conversion.py',
    )
    main(args)


datasets/preprocess/RICH/rich_smplx2smpl_conversion.py
sha2nkt/ipman-r
from .h36m import h36m_extract
from .rich import rich_extract
from .pw3d import pw3d_extract
from .mpi_inf_3dhp import mpi_inf_3dhp_extract
from .lsp_dataset import lsp_dataset_extract
from .lsp_dataset_original import lsp_dataset_original_extract
from .hr_lspet import hr_lspet_extract
from .mpii import mpii_extract
from .coco import coco_extract

datasets/preprocess/__init__.py
sha2nkt/ipman-r
import json
import os
os.environ["CDF_LIB"] = "data/cdf37_1-dist/src/lib"
import sys
import cv2
import glob
import h5py
import torch
import trimesh
import constants
import numpy as np
import pickle as pkl
from tqdm import tqdm
from spacepy import pycdf
# from .read_openpose import read_openpose
import sys
sys.path.append('../../../')
from utils.geometry import batch_rodrigues, batch_rot2aa, ea2rm
from vis_utils.world_vis import overlay_mesh, vis_smpl_with_ground

joints_to_use = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 37])
joints_to_use = np.arange(0,156).reshape((-1,3))[joints_to_use].reshape(-1)

def h36m_extract(dataset_path, out_path, protocol=1, extract_img=False, vis_path=None, visualize=False):

    # convert joints to global order
    h36m_idx = [11, 6, 7, 8, 1, 2, 3, 12, 24, 14, 15, 17, 18, 19, 25, 26, 27]
    global_idx = [14, 3, 4, 5, 2, 1, 0, 16, 12, 17, 18, 9, 10, 11, 8, 7, 6]

    # structs we use
    imgnames_, scales_, centers_, parts_, Ss_, Ss_world_,  = [], [], [], [], [], []
    poses_, shapes_, transls_ = [], [], []
    poses_world_, transls_world_, cams_r_, cams_t_, cams_k_ = [], [], [], [], []
    in_bos_label_, contact_label_ = [], []
    # users in validation set
    user_list = [9, 11]

    # go over each user
    for user_i in user_list:
        user_name = 'S%d' % user_i
        print(user_name)
        # path with GT bounding boxes
        bbox_path = os.path.join(dataset_path, 'annotations_temp', user_name, 'MySegmentsMat', 'ground_truth_bb')
        # path with GT 3D pose in world coordinates
        pose_world_path = os.path.join(dataset_path, 'annotations_temp', user_name, 'MyPoseFeatures', 'D3_Positions')
        # path with GT 3D pose
        pose_path = os.path.join(dataset_path, 'annotations_temp', user_name, 'MyPoseFeatures', 'D3_Positions_mono')
        # path with GT 2D pose
        pose2d_path = os.path.join(dataset_path, 'annotations_temp', user_name, 'MyPoseFeatures', 'D2_Positions')
        # path with videos
        # vid_path = os.path.join(dataset_path, user_name, 'Videos')

        # path with videos
        imgs_path = os.path.join(dataset_path, 'Images')
        ## path with mosh data - there is no mosh for test sequences
        # mosh_path = os.path.join(dataset_path, 'mosh', user_name)

        # go over all the sequences of each user
        seq_list = glob.glob(os.path.join(pose_path, '*.cdf'))
        seq_list.sort()
        for i, seq_i in enumerate(tqdm(seq_list)):
            # sequence info
            seq_name = seq_i.split('/')[-1]
            action, camera, _ = seq_name.split('.')
            seq_world_i = os.path.join(pose_world_path, action+'.cdf')
            seq_world_name = seq_world_i.split('/')[-1]
            action = action.replace(' ', '_')

            # irrelevant sequences
            if action == '_ALL':
                continue

            # 3D pose file
            poses_3d = pycdf.CDF(seq_i)['Pose'][0]
            poses_world_3d = pycdf.CDF(seq_world_i)['Pose'][0]

            # 2D pose file
            pose2d_file = os.path.join(pose2d_path, seq_name)
            poses_2d = pycdf.CDF(pose2d_file)['Pose'][0]

            # bbox file
            bbox_file = os.path.join(bbox_path, seq_name.replace('cdf', 'mat'))
            bbox_h5py = h5py.File(bbox_file, 'r')

            # mosh_file = os.path.join(mosh_path, f'{action.replace("_", " ")}_poses.pkl')

            # try:
            #     thetas, transls, thetas_world, transls_world, betas = read_mosh_data(mosh_file, cam=camera, user=user_name)
            # except FileNotFoundError:
            #     print(f'{mosh_file} not found!!!')
            #     continue

            # camera info
            camR, camT, camK, cam_id = read_cam(cam=camera, user=user_name)

            # get stability labels: 1: stable, 0: unstable but in contact, -1: unstable and not in contact
            # in_bos_labels, contact_labels = vis_smpl_with_ground(thetas_world, transls_world, betas, seq_world_name, vis_path,
            #                      sub_sample=1,
            #                      ground_offset=constants.GROUND_OFFSETS['h36m'],
            #                      smpl_batch_size=3000,
            #                      visualize=False)
            # in_bos_labels = in_bos_labels.detach().cpu().numpy()
            # contact_labels = contact_labels.detach().cpu().numpy()

            # visualize world smpl on ground plane
            if visualize:
                if user_name == 'S1' and cam_id == 1:
                    vis_smpl_with_ground(thetas_world, transls_world, betas, seq_world_name, vis_path)

            # video file
            # if extract_img:
            #     vid_file = os.path.join(vid_path, seq_name.replace('cdf', 'mp4'))
            #     imgs_path = os.path.join(dataset_path, 'images')
            #     vidcap = cv2.VideoCapture(vid_file)
            #     success, image = vidcap.read()

            # go over each frame of the sequence
            for frame_i in range(poses_3d.shape[0]):
                # read video frame
                # if extract_img:
                #     success, image = vidcap.read()
                #     if not success:
                #         break

                # check if you can keep this frame
                if frame_i % 5 == 0 and (protocol == 1 or camera == '60457274'):
                    # image name
                    imgname = os.path.join(dataset_path, 'Images', user_name, f'{action}.{camera}',
                                           f'{frame_i + 1:06d}.jpg')
                    if user_name == 'S1' and 'TakingPhoto' in action:
                        temp_action = action.replace('TakingPhoto', 'Photo')
                        imgname = os.path.join(dataset_path, 'Images', user_name, f'{temp_action}.{camera}',
                                               f'{frame_i + 1:06d}.jpg')

                    if user_name == 'S1' and 'WalkingDog' in action:
                        temp_action = action.replace('WalkingDog', 'WalkDog')
                        imgname = os.path.join(dataset_path, 'Images', user_name, f'{temp_action}.{camera}',
                                               f'{frame_i + 1:06d}.jpg')

                    if not os.path.isfile(imgname):
                        print(imgname)
                        raise FileNotFoundError
                    # save image
                    if extract_img:
                        img_out = os.path.join(imgs_path, imgname)
                        cv2.imwrite(img_out, image)

                    # read GT bounding box
                    mask = bbox_h5py[bbox_h5py['Masks'][frame_i,0]].value.T
                    ys, xs = np.where(mask==1)
                    bbox = np.array([np.min(xs), np.min(ys), np.max(xs)+1, np.max(ys)+1])
                    center = [(bbox[2]+bbox[0])/2, (bbox[3]+bbox[1])/2]
                    scale = 0.9*max(bbox[2]-bbox[0], bbox[3]-bbox[1])/200.

                    # read GT 3D pose
                    partall = np.reshape(poses_2d[frame_i, :], [-1, 2])
                    part17 = partall[h36m_idx]
                    part = np.zeros([24, 3])
                    part[global_idx, :2] = part17
                    part[global_idx, 2] = 1

                    # read GT 3D pose
                    Sall = np.reshape(poses_3d[frame_i,:], [-1,3])/1000.
                    S17 = Sall[h36m_idx]
                    S17 -= S17[0] # root-centered
                    S24 = np.zeros([24,4])
                    S24[global_idx, :3] = S17
                    S24[global_idx, 3] = 1

                    # read GT 3D pose in world coordinates
                    Sall_world = np.reshape(poses_world_3d[frame_i, :], [-1, 3]) / 1000.
                    S17_world = Sall_world[h36m_idx]
                    # S17 -= S17[0]  # root-centered
                    S24_world = np.zeros([24, 4])
                    S24_world[global_idx, :3] = S17_world
                    S24_world[global_idx, 3] = 1

                    # read mosh data
                    # beta = betas[frame_i]
                    # theta = thetas[frame_i]
                    # transl = transls[frame_i]
                    # theta_world = thetas_world[frame_i]
                    # transl_world = transls_world[frame_i]
                    cam_r = camR
                    cam_t = camT
                    cam_k = camK
                    # in_bos_label = in_bos_labels[frame_i]
                    # contact_label = contact_labels[frame_i]

                    # store data
                    imgnames_.append(imgname)
                    centers_.append(center)
                    scales_.append(scale)
                    parts_.append(part)
                    Ss_.append(S24)
                    Ss_world_.append(S24_world)
                    # poses_.append(theta)
                    # transls_.append(transl)
                    # poses_world_.append(theta_world)
                    # transls_world_.append(transl_world)
                    # shapes_.append(beta)
                    cams_r_.append(cam_r)
                    cams_t_.append(cam_t)
                    cams_k_.append(cam_k)
                    # in_bos_label_.append(in_bos_label)
                    # contact_label_.append(contact_label)

    # store the data struct
    if not os.path.isdir(out_path):
        os.makedirs(out_path)
    out_file = os.path.join(out_path,
        'h36m_valid_protocol%d.npz' % protocol)
    np.savez(out_file, imgname=imgnames_,
                       center=centers_,
                       scale=scales_,
                       S=Ss_,
                       S_world=Ss_world_,
                       part=parts_,
                       # pose=poses_,
                       # transl=transls_,
                       # shape=shapes_,
                       # pose_world=poses_world_,
                       # transl_world=transls_world_,
                       cam_r=cams_r_,
                       cam_t=cams_t_,
                       cam_k=cams_k_,
                       # in_bos_label=in_bos_label_,
                       # contact_label=contact_label_,
             )
    print('Saved to ', out_file)

def load_camera_params( hf, path ):
    """Load h36m camera parameters
    Args
    hf: hdf5 open file with h36m cameras data
    path: path or key inside hf to the camera we are interested in
    Returns
    R: 3x3 Camera rotation matrix
    T: 3x1 Camera translation parameters
    f: (scalar) Camera focal length
    c: 2x1 Camera center
    k: 3x1 Camera radial distortion coefficients
    p: 2x1 Camera tangential distortion coefficients
    name: String with camera id
    """

    R = hf[ path.format('R') ][:]
    R = R.T

    T = hf[ path.format('T') ][:]
    f = hf[ path.format('f') ][:]
    c = hf[ path.format('c') ][:]
    k = hf[ path.format('k') ][:]
    p = hf[ path.format('p') ][:]

    name = hf[ path.format('Name') ][:]
    name = "".join( [chr(item) for item in name] )

    return R, T, f, c, k, p, name


def load_cameras(bpath='/ps/scratch/ps_shared/stripathi/frmKocabas/human36m/cameras.h5', subjects=[1,5,6,7,8,9,11]):
    """Loads the cameras of h36m
    Args
    bpath: path to hdf5 file with h36m camera data
    subjects: List of ints representing the subject IDs for which cameras are requested
    Returns
    rcams: dictionary of 4 tuples per subject ID containing its camera parameters for the 4 h36m cams
    """
    rcams = {}

    with h5py.File(bpath,'r') as hf:
        for s in subjects:
            for c in range(4): # There are 4 cameras in human3.6m
                rcams[c+1] = load_camera_params(hf, 'subject%d/camera%d/{0}' % (s,c+1) )
    return rcams

def load_cameras_json(bpath='/ps/scratch/ps_shared/stripathi/frmKocabas/human36m/h36m-camera-parameters.json'):
    """
    Loads the cameras from convenient json file at https://github.com/karfly/human36m-camera-parameters/blob/master/camera-parameters.json
    """
    with open(bpath) as f:
        data = json.load(f)
    return data

def read_pkl(f):
    return pkl.load(open(f, 'rb'), encoding='latin1')

def read_cam(cam=None, user=None):
    camera_ids = {
        '54138969': 1,
        '55011271': 2,
        '58860488': 3,
        '60457274': 4,
    }
    rcams_new = load_cameras_json()
    R = rcams_new['extrinsics'][user][cam]['R']
    T = rcams_new['extrinsics'][user][cam]['t']
    K = rcams_new['intrinsics'][cam]['calibration_matrix']
    cam_id = camera_ids[cam]
    # rcams = load_cameras()
    # cam_full = rcams[camera_ids[cam]] # order: R.T, T, f, c, k, p, name
    return R, T, K, cam_id

def read_mosh_data(f, frame_skip=4, cam=None, user=None):
    m2mm = 1000
    data = read_pkl(f)
    poses = data['pose_est_fullposes'][0::frame_skip,joints_to_use]
    poses_world = poses.copy()
    shape = data['shape_est_betas'][:10]
    shape = np.tile(shape, [poses.shape[0], 1])
    transls_world = data['pose_est_trans'][0::frame_skip]

    camera_ids = {
        '54138969': 1,
        '55011271': 2,
        '58860488': 3,
        '60457274': 4,
    }
    rcams_new = load_cameras_json()
    camR = rcams_new['extrinsics'][user][cam]['R']
    camT = rcams_new['extrinsics'][user][cam]['t']
    camPos = -np.matmul(np.array(camR).T, np.array(camT)) / m2mm  # C = -R.T x t

    # rcams = load_cameras()
    # camR = rcams[camera_ids[cam]][0]

    poses = torch.FloatTensor(poses)
    camR = torch.FloatTensor(camR).unsqueeze(0)

    poses[:,:3] = rectify_pose(camR, poses[:,:3])
    transl = transls_world - camPos.T
    return poses.numpy(), transl, poses_world, transls_world, shape


def rectify_pose(camera_r, body_aa):
    body_r = batch_rodrigues(body_aa).reshape(-1,3,3)
    final_r = camera_r @ body_r
    body_aa = batch_rot2aa(final_r)
    return body_aa

datasets/preprocess/h36m.py
sha2nkt/ipman-r
import os
os.environ["CDF_LIB"] = "data/cdf37_1-dist/src/lib"

import cv2
import json
import glob
import h5py
import torch
import trimesh
import constants
import numpy as np
import pickle as pkl
from tqdm import tqdm
from spacepy import pycdf
# from .read_openpose import read_openpose
import sys
sys.path.append('../../../')
from utils.geometry import batch_rodrigues, batch_rot2aa, ea2rm
from vis_utils.world_vis import overlay_mesh, vis_smpl_with_ground

joints_to_use = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 37])
joints_to_use = np.arange(0,156).reshape((-1,3))[joints_to_use].reshape(-1)

# Illustrative script for training data extraction
# No SMPL parameters will be included in the .npz file.
def h36m_train_extract(dataset_path, openpose_path, out_path, extract_img=False, vis_path=None, visualize=False):
    # convert joints to global order
    h36m_idx = [11, 6, 7, 8, 1, 2, 3, 12, 24, 14, 15, 17, 18, 19, 25, 26, 27]
    global_idx = [14, 3, 4, 5, 2, 1, 0, 16, 12, 17, 18, 9, 10, 11, 8, 7, 6]

    # structs we use
    imgnames_, scales_, centers_, parts_, Ss_, Ss_world_, openposes_  = [], [], [], [], [], [], []
    poses_, shapes_, transls_ = [], [], []
    poses_world_, transls_world_, cams_r_, cams_t_, cams_k_ = [], [], [], [], []
    in_bos_label_, contact_label_ = [], []
    # users in validation set
    user_list = [1, 5, 6, 7, 8]

    # go over each user
    for user_i in tqdm(user_list):
        user_name = 'S%d' % user_i
        print(user_name)
        # path with GT bounding boxes
        bbox_path = os.path.join(dataset_path, 'annotations_temp', user_name, 'MySegmentsMat', 'ground_truth_bb')
        # path with GT 3D pose in world coordinates
        pose_world_path = os.path.join(dataset_path, 'annotations_temp', user_name, 'MyPoseFeatures', 'D3_Positions')
        # path with GT 3D pose
        pose_path = os.path.join(dataset_path, 'annotations_temp', user_name, 'MyPoseFeatures', 'D3_Positions_mono')
        # path with GT 2D pose
        pose2d_path = os.path.join(dataset_path, 'annotations_temp', user_name, 'MyPoseFeatures', 'D2_Positions')
        # path with videos
        # vid_path = os.path.join(dataset_path, user_name, 'Videos')

        # path with videos
        imgs_path = os.path.join(dataset_path, 'Images')
        # path with mosh data
        mosh_path = os.path.join(dataset_path, 'mosh', user_name)

        # go over all the sequences of each user
        seq_list = glob.glob(os.path.join(pose_path, '*.cdf'))
        seq_list.sort()
        for i, seq_i in enumerate(tqdm(seq_list)):
            # sequence info
            seq_name = seq_i.split('/')[-1]
            action, camera, _ = seq_name.split('.')
            seq_world_i = os.path.join(pose_world_path, action+'.cdf')
            seq_world_name = seq_world_i.split('/')[-1]
            action = action.replace(' ', '_')

            # irrelevant sequences
            if action == '_ALL':
                continue

            # 3D pose file
            poses_3d = pycdf.CDF(seq_i)['Pose'][0]
            poses_world_3d = pycdf.CDF(seq_world_i)['Pose'][0]

            # 2D pose file
            pose2d_file = os.path.join(pose2d_path, seq_name)
            poses_2d = pycdf.CDF(pose2d_file)['Pose'][0]

            # bbox file
            bbox_file = os.path.join(bbox_path, seq_name.replace('cdf', 'mat'))
            bbox_h5py = h5py.File(bbox_file, 'r')

            mosh_file = os.path.join(mosh_path, f'{action.replace("_", " ")}_poses.pkl')

            try:
                thetas, transls, thetas_world, transls_world, betas = read_mosh_data(mosh_file, cam=camera, user=user_name)
            except FileNotFoundError:
                print(f'{mosh_file} not found!!!')
                continue

            # camera info
            camR, camT, camK, cam_id = read_cam(cam=camera, user=user_name)

            # get stability labels: 1: stable, 0: unstable but in contact, -1: unstable and not in contact
            in_bos_labels, contact_labels = vis_smpl_with_ground(thetas_world, transls_world, betas, seq_world_name, vis_path,
                                 sub_sample=1,
                                 ground_offset=constants.GROUND_OFFSETS['h36m'],
                                 smpl_batch_size=9000,
                                 visualize=False)
            in_bos_labels = in_bos_labels.detach().cpu().numpy()
            contact_labels = contact_labels.detach().cpu().numpy()

            # visualize world smpl on ground plane
            if visualize:
                if user_name == 'S1' and cam_id == 1:
                    vis_smpl_with_ground(thetas_world, transls_world, betas, seq_world_name, vis_path)

            # video file
            # if extract_img:
            #     vid_file = os.path.join(vid_path, seq_name.replace('cdf', 'mp4'))
            #     imgs_path = os.path.join(dataset_path, 'images')
            #     vidcap = cv2.VideoCapture(vid_file)

            # go over each frame of the sequence
            for frame_i in tqdm(range(poses_3d.shape[0])):
                # read video frame
                # if extract_img:
                #     success, image = vidcap.read()
                #     if not success:
                #         break

                # check if you can keep this frame
                # if frame_i % 5 == 0 and (protocol == 1 or camera == '60457274'):
                # image name
                # imgname = '%s_%s.%s_%06d.jpg' % (user_name, action, camera, frame_i+1)
                imgname = os.path.join(dataset_path, 'Images', user_name, f'{action}.{camera}',
                                       f'{frame_i + 1:06d}.jpg')
                if user_name == 'S1' and 'TakingPhoto' in action:
                    temp_action = action.replace('TakingPhoto', 'Photo')
                    imgname = os.path.join(dataset_path, 'Images', user_name, f'{temp_action}.{camera}',
                                           f'{frame_i + 1:06d}.jpg')

                if user_name == 'S1' and 'WalkingDog' in action:
                    temp_action = action.replace('WalkingDog', 'WalkDog')
                    imgname = os.path.join(dataset_path, 'Images', user_name, f'{temp_action}.{camera}',
                                           f'{frame_i + 1:06d}.jpg')


                if not os.path.isfile(imgname):
                    print(imgname)
                    raise FileNotFoundError
                # save image
                if extract_img:
                    img_out = os.path.join(imgs_path, imgname)
                    cv2.imwrite(img_out, image)

                # read GT bounding box
                mask = bbox_h5py[bbox_h5py['Masks'][frame_i,0]].value.T
                ys, xs = np.where(mask==1)
                bbox = np.array([np.min(xs), np.min(ys), np.max(xs)+1, np.max(ys)+1])
                center = [(bbox[2]+bbox[0])/2, (bbox[3]+bbox[1])/2]
                scale = 0.9*max(bbox[2]-bbox[0], bbox[3]-bbox[1])/200.

                # read GT 3D pose
                partall = np.reshape(poses_2d[frame_i,:], [-1,2])
                part17 = partall[h36m_idx]
                part = np.zeros([24,3])
                part[global_idx, :2] = part17
                part[global_idx, 2] = 1

                # read GT 3D pose
                Sall = np.reshape(poses_3d[frame_i,:], [-1,3])/1000.
                S17 = Sall[h36m_idx]
                S17 -= S17[0] # root-centered
                S24 = np.zeros([24,4])
                S24[global_idx, :3] = S17
                S24[global_idx, 3] = 1

                # read GT 3D pose in world coordinates
                Sall_world = np.reshape(poses_world_3d[frame_i, :], [-1, 3]) / 1000.
                S17_world = Sall_world[h36m_idx]
                # S17 -= S17[0]  # root-centered
                S24_world = np.zeros([24, 4])
                S24_world[global_idx, :3] = S17_world
                S24_world[global_idx, 3] = 1
                # if frame_i % 100 == 0:
                #     plot_3d_andCam(S24_world, np.array([[0,0,0]]))
                # TODO: run openpose for h36m
                # read openpose detections
                # json_file = os.path.join(openpose_path, 'coco',
                #     imgname.replace('.jpg', '_keypoints.json'))
                # openpose = read_openpose(json_file, part, 'h36m')
                openpose = np.zeros([25, 3])

                # read mosh data
                beta = betas[frame_i]
                theta = thetas[frame_i]
                transl = transls[frame_i]
                theta_world = thetas_world[frame_i]
                transl_world = transls_world[frame_i]
                cam_r = camR
                cam_t = camT
                cam_k = camK
                in_bos_label = in_bos_labels[frame_i]
                contact_label = contact_labels[frame_i]

                # store data
                imgnames_.append(imgname)
                centers_.append(center)
                scales_.append(scale)
                parts_.append(part)
                Ss_.append(S24)
                Ss_world_.append(S24_world)
                openposes_.append(openpose)
                poses_.append(theta)
                transls_.append(transl)
                poses_world_.append(theta_world)
                transls_world_.append(transl_world)
                shapes_.append(beta)
                cams_r_.append(cam_r)
                cams_t_.append(cam_t)
                cams_k_.append(cam_k)
                in_bos_label_.append(in_bos_label)
                contact_label_.append(contact_label)

    # store the data struct
    if not os.path.isdir(out_path):
        os.makedirs(out_path)
    out_file = os.path.join(out_path, 'h36m_train_world.npz')
    np.savez(out_file, imgname=imgnames_,
                       center=centers_,
                       scale=scales_,
                       part=parts_,
                       S=Ss_,
                       S_world=Ss_world_,
                       pose=poses_,
                       transl=transls_,
                       shape=shapes_,
                       openpose=openposes_,
                       pose_world=poses_world_,
                       transl_world=transls_world_,
                       cam_r=cams_r_,
                       cam_t=cams_t_,
                       cam_k=cams_k_,
                       in_bos_label=in_bos_label_,
                       contact_label=contact_label_,
             )
    print('Saved to ', out_file)

def load_camera_params( hf, path ):
    """Load h36m camera parameters
    Args
    hf: hdf5 open file with h36m cameras data
    path: path or key inside hf to the camera we are interested in
    Returns
    R: 3x3 Camera rotation matrix
    T: 3x1 Camera translation parameters
    f: (scalar) Camera focal length
    c: 2x1 Camera center
    k: 3x1 Camera radial distortion coefficients
    p: 2x1 Camera tangential distortion coefficients
    name: String with camera id
    """

    R = hf[ path.format('R') ][:]
    R = R.T

    T = hf[ path.format('T') ][:]
    f = hf[ path.format('f') ][:]
    c = hf[ path.format('c') ][:]
    k = hf[ path.format('k') ][:]
    p = hf[ path.format('p') ][:]

    name = hf[ path.format('Name') ][:]
    name = "".join( [chr(item) for item in name] )

    return R, T, f, c, k, p, name


def load_cameras(bpath='/ps/scratch/ps_shared/stripathi/frmKocabas/human36m/cameras.h5', subjects=[1,5,6,7,8,9,11]):
    """Loads the cameras of h36m
    Args
    bpath: path to hdf5 file with h36m camera data
    subjects: List of ints representing the subject IDs for which cameras are requested
    Returns
    rcams: dictionary of 4 tuples per subject ID containing its camera parameters for the 4 h36m cams
    """
    rcams = {}

    with h5py.File(bpath,'r') as hf:
        for s in subjects:
            for c in range(4): # There are 4 cameras in human3.6m
                rcams[c+1] = load_camera_params(hf, 'subject%d/camera%d/{0}' % (s,c+1) )
    return rcams

def load_cameras_json(bpath='/ps/scratch/ps_shared/stripathi/frmKocabas/human36m/h36m-camera-parameters.json'):
    """
    Loads the cameras from convenient json file at https://github.com/karfly/human36m-camera-parameters/blob/master/camera-parameters.json
    """
    with open(bpath) as f:
        data = json.load(f)
    return data

def read_pkl(f):
    return pkl.load(open(f, 'rb'), encoding='latin1')

def read_cam(cam=None, user=None):
    camera_ids = {
        '54138969': 1,
        '55011271': 2,
        '58860488': 3,
        '60457274': 4,
    }
    rcams_new = load_cameras_json()
    R = rcams_new['extrinsics'][user][cam]['R']
    T = rcams_new['extrinsics'][user][cam]['t']
    K = rcams_new['intrinsics'][cam]['calibration_matrix']
    cam_id = camera_ids[cam]
    # rcams = load_cameras()
    # cam_full = rcams[camera_ids[cam]] # order: R.T, T, f, c, k, p, name
    return R, T, K, cam_id

def read_mosh_data(f, frame_skip=4, cam=None, user=None):
    m2mm = 1000
    data = read_pkl(f)
    poses = data['pose_est_fullposes'][0::frame_skip,joints_to_use]
    poses_world = poses.copy()
    shape = data['shape_est_betas'][:10]
    shape = np.tile(shape, [poses.shape[0], 1])
    transls_world = data['pose_est_trans'][0::frame_skip]

    camera_ids = {
        '54138969': 1,
        '55011271': 2,
        '58860488': 3,
        '60457274': 4,
    }
    rcams_new = load_cameras_json()
    camR = rcams_new['extrinsics'][user][cam]['R']
    camT = rcams_new['extrinsics'][user][cam]['t']
    camPos = -np.matmul(np.array(camR).T, np.array(camT)) / m2mm  # C = -R.T x t

    # rcams = load_cameras()
    # camR = rcams[camera_ids[cam]][0]

    poses = torch.FloatTensor(poses)
    camR = torch.FloatTensor(camR).unsqueeze(0)

    poses[:,:3] = rectify_pose(camR, poses[:,:3])
    transl = transls_world - camPos.T
    return poses.numpy(), transl, poses_world, transls_world, shape


def rectify_pose(camera_r, body_aa):
    body_r = batch_rodrigues(body_aa).reshape(-1,3,3)
    final_r = camera_r @ body_r
    body_aa = batch_rot2aa(final_r)
    return body_aa
datasets/preprocess/h36m_train.py
sha2nkt/ipman-r
import os
import sys
import cv2
import glob
import h5py
import json
import numpy as np
import scipy.io as sio
import pandas as pd
import scipy.misc
from .read_openpose import read_openpose

def read_calibration(calib_file, vid_list):
    Ks, Rs, Ts = [], [], []
    file = open(calib_file, 'r')
    content = file.readlines()
    for vid_i in vid_list:
        K = np.array([float(s) for s in content[vid_i*7+5][11:-2].split()])
        K = np.reshape(K, (4, 4))
        RT = np.array([float(s) for s in content[vid_i*7+6][11:-2].split()])
        RT = np.reshape(RT, (4, 4))
        R = RT[:3,:3]
        T = RT[:3,3]/1000
        Ks.append(K)
        Rs.append(R)
        Ts.append(T)
    return Ks, Rs, Ts
    
def train_data(dataset_path, openpose_path, out_path, joints_idx, scaleFactor, extract_img=False, fits_3d=None, vis_path=None, visualize=None):

    joints17_idx = [4, 18, 19, 20, 23, 24, 25, 3, 5, 6, 7, 9, 10, 11, 14, 15, 16]

    h, w = 2048, 2048
    imgnames_, scales_, centers_ = [], [], []
    parts_, Ss_, Ss_world_, openposes_ = [], [], [], []
    poses_, shapes_, transls_ = [], [], []
    poses_world_, transls_world_, cams_r_, cams_t_, cams_k_ = [], [], [], [], []
    in_bos_label_, contact_label_ = [], []


    # training data
    user_list = range(1,9)
    seq_list = range(1,3)
    vid_list = list(range(3)) + list(range(4,9))

    counter = 0

    # load already prepared npz for 3d fits
    npz_path = '/is/cluster/scratch/stripathi/pycharm_remote/ipman_regr/data/dataset_extras/mpi_inf_3dhp_train.npz'
    mpi_npz = np.load(npz_path)
    mpi_df = pd.DataFrame.from_dict({item: mpi_npz[item] for item in mpi_npz.files}, orient='index').T

    for user_i in user_list:
        for seq_i in seq_list:
            seq_path = os.path.join(dataset_path,
                                    'S' + str(user_i),
                                    'Seq' + str(seq_i))
            # mat file with annotations
            annot_file = os.path.join(seq_path, 'annot.mat')
            annot2 = sio.loadmat(annot_file)['annot2']
            annot3 = sio.loadmat(annot_file)['annot3']
            # calibration file and camera parameters
            calib_file = os.path.join(seq_path, 'camera.calibration')
            Ks, Rs, Ts = read_calibration(calib_file, vid_list)

            for j, vid_i in enumerate(vid_list):

                # image folder
                imgs_path = os.path.join(seq_path,
                                         'imageFrames',
                                         'video_' + str(vid_i))

                # extract frames from video file
                if extract_img:

                    # if doesn't exist
                    if not os.path.isdir(imgs_path):
                        os.makedirs(imgs_path)

                    # video file
                    vid_file = os.path.join(seq_path,
                                            'imageSequence',
                                            'video_' + str(vid_i) + '.avi')
                    vidcap = cv2.VideoCapture(vid_file)

                    # process video
                    frame = 0
                    while 1:
                        # extract all frames
                        success, image = vidcap.read()
                        if not success:
                            break
                        frame += 1
                        # image name
                        imgname = os.path.join(imgs_path,
                            'frame_%06d.jpg' % frame)
                        # save image
                        cv2.imwrite(imgname, image)

                # per frame
                cam_aa = cv2.Rodrigues(Rs[j])[0].T[0]
                pattern = os.path.join(imgs_path, '*.jpg')
                img_list = sorted(glob.glob(pattern))

                K, R, T = Ks[j][:3, :3], Rs[j], Ts[j]
                for i, img_i in enumerate(img_list[::10]):

                    # for each image we store the relevant annotations
                    img_name = img_i.split('/')[-1]
                    img_view = os.path.join('S' + str(user_i),
                                            'Seq' + str(seq_i),
                                            'imageFrames',
                                            'video_' + str(vid_i),
                                            img_name)
                    df_img_view = os.path.join('S' + str(user_i),
                                            'Seq' + str(seq_i),
                                            'imageFrames',
                                            'video_' + str(vid_i),
                                            'frame_' + img_name)
                    joints = np.reshape(annot2[vid_i][0][i], (28, 2))[joints17_idx]
                    S17 = np.reshape(annot3[vid_i][0][i], (28, 3))/1000
                    S17_world = np.matmul(R.T, S17.T).T - np.matmul(R.T, T)
                    S17 = S17[joints17_idx] - S17[4] # 4 is the root

                    # # reproject joints into the image
                    # projected_points = (K @ S17.T).T
                    # joints2d = projected_points[:, :2] / np.hstack((projected_points[:, 2:], projected_points[:, 2:]))
                    # img = cv2.imread(os.path.join(imgs_path, img_name))
                    # img_points = joints2d.astype(np.int)
                    # for pt in img_points:
                    #     cv2.circle(img, (int(pt[0]), int(pt[1])), radius=6, color=(0, 0, 255), thickness=-1)
                    # cv2.imshow('Vis', img)
                    # cv2.waitKey()

                    bbox = [min(joints[:,0]), min(joints[:,1]),
                            max(joints[:,0]), max(joints[:,1])]
                    center = [(bbox[2]+bbox[0])/2, (bbox[3]+bbox[1])/2]
                    scale = scaleFactor*max(bbox[2]-bbox[0], bbox[3]-bbox[1])/200

                    # get SMPL Parameters
                    import ipdb; ipdb.set_trace()
                    theta = mpi_df[df_img_view == mpi_df['imgname']]['pose'][0]


                    # check that all joints are visible
                    x_in = np.logical_and(joints[:, 0] < w, joints[:, 0] >= 0)
                    y_in = np.logical_and(joints[:, 1] < h, joints[:, 1] >= 0)
                    ok_pts = np.logical_and(x_in, y_in)
                    if np.sum(ok_pts) < len(joints_idx):
                        continue
                        
                    part = np.zeros([24,3])
                    part[joints_idx] = np.hstack([joints, np.ones([17,1])])
                    json_file = os.path.join(openpose_path, 'mpi_inf_3dhp',
                                             img_view.replace('.jpg', '_keypoints.json'))
                    openpose = read_openpose(json_file, part, 'mpi_inf_3dhp')

                    S = np.zeros([24,4])
                    S[joints_idx] = np.hstack([S17, np.ones([17,1])])
                    S_world = np.zeros([24, 4])
                    S_world[joints_idx] = np.hstack([S17_world, np.ones([17,1])])

                    # because of the dataset size, we only keep every 10th frame
                    counter += 1
                    if counter % 10 != 1:
                        continue

                    # store the data
                    imgnames_.append(img_view)
                    centers_.append(center)
                    scales_.append(scale)
                    parts_.append(part)
                    Ss_.append(S)
                    Ss_world_.append(S_world)
                    pose=poses_,
                    transl=transls_,
                    shape=shapes_,
                    openposes_.append(openpose)
                    pose_world=poses_world_,
                    transl_world=transls_world_,
                    cam_r=camr_r_,
                    cam_t=cam_t_,
                    cam_k=cam_k_,
                    in_bos_label=in_bos_label_,
                    contact_label=contact_label_,
                       
    # store the data struct
    if not os.path.isdir(out_path):
        os.makedirs(out_path)
    out_file = os.path.join(out_path, 'mpi_inf_3dhp_train_world.npz')
    if fits_3d is not None:
        fits_3d = np.load(fits_3d)
        np.savez(out_file, imgname=imgnames_,
                           center=centers_,
                           scale=scales_,
                           part=parts_,
                           pose=fits_3d['pose'],
                           shape=fits_3d['shape'],
                           has_smpl=fits_3d['has_smpl'],
                           S=Ss_,
                           openpose=openposes_)
    else:
        np.savez(out_file, imgname=imgnames_,
                           center=centers_,
                           scale=scales_,
                           part=parts_,
                           S=Ss_,
                           openpose=openposes_)        
        
        
def test_data(dataset_path, out_path, joints_idx, scaleFactor, vis_path=None, visualize=None):

    joints17_idx = [14, 11, 12, 13, 8, 9, 10, 15, 1, 16, 0, 5, 6, 7, 2, 3, 4]

    imgnames_, scales_, centers_, parts_,  Ss_ = [], [], [], [], []

    # training data
    user_list = range(1,7)

    for user_i in tqdm(user_list):
        seq_path = os.path.join(dataset_path,
                                'mpi_inf_3dhp_test_set',
                                'TS' + str(user_i))
        # mat file with annotations
        annot_file = os.path.join(seq_path, 'annot_data.mat')
        mat_as_h5 = h5py.File(annot_file, 'r')
        annot2 = np.array(mat_as_h5['annot2'])
        annot3 = np.array(mat_as_h5['univ_annot3'])
        valid = np.array(mat_as_h5['valid_frame'])
        for frame_i, valid_i in enumerate(valid):
            if valid_i == 0:
                continue
            img_name = os.path.join('mpi_inf_3dhp_test_set',
                                   'TS' + str(user_i),
                                   'imageSequence',
                                   'img_' + str(frame_i+1).zfill(6) + '.jpg')

            joints = annot2[frame_i,0,joints17_idx,:]
            S17 = annot3[frame_i,0,joints17_idx,:]/1000
            S17 = S17 - S17[0]

            bbox = [min(joints[:,0]), min(joints[:,1]),
                    max(joints[:,0]), max(joints[:,1])]
            center = [(bbox[2]+bbox[0])/2, (bbox[3]+bbox[1])/2]
            scale = scaleFactor*max(bbox[2]-bbox[0], bbox[3]-bbox[1])/200

            # check that all joints are visible
            img_file = os.path.join(dataset_path, img_name)
            I = scipy.misc.imread(img_file)
            h, w, _ = I.shape
            x_in = np.logical_and(joints[:, 0] < w, joints[:, 0] >= 0)
            y_in = np.logical_and(joints[:, 1] < h, joints[:, 1] >= 0)
            ok_pts = np.logical_and(x_in, y_in)
            if np.sum(ok_pts) < len(joints_idx):
                continue

            part = np.zeros([24,3])
            part[joints_idx] = np.hstack([joints, np.ones([17,1])])

            S = np.zeros([24,4])
            S[joints_idx] = np.hstack([S17, np.ones([17,1])])

            # store the data
            imgnames_.append(img_name)
            centers_.append(center)
            scales_.append(scale)
            parts_.append(part)
            Ss_.append(S)

    # store the data struct
    if not os.path.isdir(out_path):
        os.makedirs(out_path)
    out_file = os.path.join(out_path, 'mpi_inf_3dhp_test_world.npz')
    np.savez(out_file, imgname=imgnames_,
                       center=centers_,
                       scale=scales_,
                       part=parts_,
                       S=Ss_)    

def mpi_inf_3dhp_extract(dataset_path, openpose_path, out_path, mode, extract_img=False, static_fits=None, vis_path=None, visualize=False):

    scaleFactor = 1.2
    joints_idx = [14, 3, 4, 5, 2, 1, 0, 16, 12, 17, 18, 9, 10, 11, 8, 7, 6]
    
    if static_fits is not None:
        fits_3d = os.path.join(static_fits, 
                               'mpi-inf-3dhp_mview_fits.npz')
    else:
        fits_3d = None
    
    if mode == 'train':
        train_data(dataset_path, openpose_path, out_path, 
                   joints_idx, scaleFactor, extract_img=extract_img, fits_3d=fits_3d, vis_path=vis_path, visualize=visualize)
    elif mode == 'test':
        test_data(dataset_path, out_path, joints_idx, scaleFactor, vis_path=vis_path, visualize=visualize)


# dataset_path = '/ps/project/datasets/mpi_inf_3dhp/'
# out_path = '/is/ps2/ppatel/SPIN_Finetuning2/spin_ft_on_agora/data/dataset_extras'
# mpi_inf_3dhp_extract(dataset_path,openpose_path=None,mode='train',out_path=out_path)

datasets/preprocess/mpi_inf_3dhp.py
sha2nkt/ipman-r
import os
import cv2
import numpy as np
import pickle

def pw3d_extract(dataset_path, out_path):

    # scale factor
    scaleFactor = 1.2

    # structs we use
    imgnames_, scales_, centers_, parts_ = [], [], [], []
    poses_, shapes_, genders_ = [], [], [] 

    # get a list of .pkl files in the directory
    dataset_path = os.path.join(dataset_path, 'sequenceFiles', 'test')
    files = [os.path.join(dataset_path, f) 
        for f in os.listdir(dataset_path) if f.endswith('.pkl')]
    # go through all the .pkl files
    for filename in files:
        with open(filename, 'rb') as f:
            data = pickle.load(f, encoding='latin1')
            smpl_pose = data['poses']
            smpl_betas = data['betas']
            poses2d = data['poses2d']
            global_poses = data['cam_poses']
            genders = data['genders']
            valid = np.array(data['campose_valid']).astype(np.bool)
            num_people = len(smpl_pose)
            num_frames = len(smpl_pose[0])
            seq_name = str(data['sequence'])
            img_names = np.array(['imageFiles/' + seq_name + '/image_%s.jpg' % str(i).zfill(5) for i in range(num_frames)])
            print(data['cam_intrinsics'])
            # get through all the people in the sequence
            for i in range(num_people):
                valid_pose = smpl_pose[i][valid[i]]
                valid_betas = np.tile(smpl_betas[i][:10].reshape(1,-1), (num_frames, 1))
                valid_betas = valid_betas[valid[i]]
                valid_keypoints_2d = poses2d[i][valid[i]]
                valid_img_names = img_names[valid[i]]
                valid_global_poses = global_poses[valid[i]]
                gender = genders[i]
                # consider only valid frames
                for valid_i in range(valid_pose.shape[0]):

                    part = valid_keypoints_2d[valid_i,:,:].T
                    part = part[part[:,2]>0,:]
                    bbox = [min(part[:,0]), min(part[:,1]),
                        max(part[:,0]), max(part[:,1])]
                    center = [(bbox[2]+bbox[0])/2, (bbox[3]+bbox[1])/2]
                    scale = scaleFactor*max(bbox[2]-bbox[0], bbox[3]-bbox[1])/200
                    
                    # transform global pose
                    pose = valid_pose[valid_i]
                    extrinsics = valid_global_poses[valid_i][:3,:3]
                    pose[:3] = cv2.Rodrigues(np.dot(extrinsics, cv2.Rodrigues(pose[:3])[0]))[0].T[0]                      

                    imgnames_.append(valid_img_names[valid_i])
                    parts_.append(part)
                    centers_.append(center)
                    scales_.append(scale)
                    poses_.append(pose)
                    shapes_.append(valid_betas[valid_i])
                    genders_.append(gender)

    # store data
    if not os.path.isdir(out_path):
        os.makedirs(out_path)
    out_file = os.path.join(out_path,
        '3dpw_test_sha.npz')
    np.savez(out_file, imgname=imgnames_,
                       part=parts_,
                       center=centers_,
                       scale=scales_,
                       pose=poses_,
                       shape=shapes_,
                       gender=genders_)


# dataset_path = '/ps/project/datasets/3DPW___v1/'
# out_path = '/is/ps2/ppatel/SPIN_Finetuning2/spin_ft_on_agora/data/dataset_extras'
# pw3d_extract(dataset_path, out_path)
datasets/preprocess/pw3d.py
sha2nkt/ipman-r
import os
os.environ["CDF_LIB"] = "data/cdf37_1-dist/src/lib"

import cv2
import json
import glob
import h5py
import torch
import trimesh
import numpy as np
import pickle as pkl
from xml.dom import minidom
import xml.etree.ElementTree as ET
from tqdm import tqdm
from spacepy import pycdf
# from .read_openpose import read_openpose
import sys
sys.path.append('../../../')
from models import hmr, SMPL
import config
import constants

from smplx import SMPL as SMPL_orig



from utils.geometry import batch_rodrigues, batch_rot2aa, ea2rm
from vis_utils.world_vis import overlay_mesh, vis_smpl_with_ground

smpl_obj = SMPL(config.SMPL_MODEL_DIR,
                         batch_size=1,
                         create_transl=True).to('cuda')

smpl_obj_orig = SMPL_orig(config.SMPL_MODEL_DIR,
                         batch_size=1,
                         create_transl=True).to('cuda')

def rich_extract(dataset_path, out_path, split=None, vis_path=None, visualize=False, downsample_factor=4):

    # structs we use
    imgnames_, scales_, centers_, parts_, Ss_, Ss_world_, openposes_ = [], [], [], [], [], [], []
    poses_, shapes_, transls_ = [], [], []
    poses_world_, transls_world_, cams_r_, cams_t_, cams_k_ = [], [], [], [], []
    ground_offset_ = []
    in_bos_label_, contact_label_ = [], []

    # seqs in validation set
    if split == 'val':
        seq_list = ['2021-06-15_Multi_IOI_ID_00176_Yoga2',
                    '2021-06-15_Multi_IOI_ID_00228_Yoga1',
                    '2021-06-15_Multi_IOI_ID_03588_Yoga1',
                    '2021-06-15_Multi_IOI_ID_00176_Yoga1']

    # seqs in testing set
    if split == 'test':
        seq_list = ['2021-06-15_Multi_IOI_ID_00186_Yoga1',
                    '2021-06-15_Multi_IOI_ID_03588_Yoga2',
                    'MultiIOI_201019_ID03581_parkingLot_Calibration06_Settings06_PushUp__2',
                    'Multi-IOI_ID00227_Scene_ParkingLot_Calibration_03_CameraSettings_4_pushup_1']


    for seq_i in tqdm(seq_list):
        print(f'Processing sequence: {seq_i}')

        # path with GT bounding boxes
        params_path = os.path.join(dataset_path, seq_i, 'params')

        # path to metadata for files
        md_path = os.path.join(dataset_path, seq_i, 'data')

        # glob all folders in params path
        frame_param_paths = sorted(glob.glob(os.path.join(params_path, '*')))
        frame_param_paths = [p for p in frame_param_paths if '.yaml' not in p]

        # get ioi2scan transformation per sequence
        ioi2scan_fn = os.path.join(dataset_path, seq_i, 'cam2scan.json')

        ## ground resides in Leica scan coordidate, which is (roughly) axis aligned.
        ground_mesh = trimesh.load(os.path.join(dataset_path, seq_i, 'ground_mesh.ply'), process=False)
        ground_eq = np.mean(ground_mesh.vertices, axis=0)

        # list all files in the folder
        cam_files = os.listdir(os.path.join(dataset_path, seq_i, f'calibration'))
        cam_list = sorted([int(os.path.splitext(f)[0]) for f in cam_files if '.xml' in f])

        if split == 'val':
            cam_list = cam_list[1:] # remove first camera in val
        for cam_num in cam_list:
            camera_fn = os.path.join(dataset_path, seq_i, f'calibration/{cam_num:03d}.xml')
            focal_length_x, focal_length_y, camC, camR, camT, _, _, _ = extract_cam_param_xml(camera_fn)

            for frame_param_path in tqdm(frame_param_paths):
                frame_id = os.path.basename(frame_param_path)
                frame_num = int(frame_id)

                # path to smpl params
                try:
                    smpl_param = os.path.join(frame_param_path, '00', 'results_smpl/000.pkl')
                except:
                    import ipdb; ipdb.set_trace()

                # path to GT bounding boxes
                bbox_path = os.path.join(md_path, frame_id, '00', 'bbox_refine', f'{frame_id}_{cam_num:02d}.json')
                # path with 2D openpose keypoints
                openpose_path = os.path.join(md_path, frame_id, '00', 'keypoints_refine', f'{frame_id}_{str(cam_num).zfill(2)}_keypoints.json')
                # path to image crops
                if downsample_factor == 1:
                    img_path = os.path.join(md_path, frame_id, '00', 'images_orig', f'{frame_id}_{cam_num:02d}.png')
                else:
                    img_path = os.path.join(md_path, frame_id, '00', 'images_orig_720p', f'{frame_id}_{cam_num:02d}.png')

                if not os.path.isfile(img_path):
                    print(f'image not found: {img_path}')
                    continue
                    # raise FileNotFoundError

                # bbox file
                try:
                    with open(bbox_path, 'r') as f:
                        bbox_dict = json.load(f)
                except:
                    print(f'bbox file not found: {bbox_path}')
                    continue
                # read GT bounding box
                x1_ul = bbox_dict['x1'] // downsample_factor
                y1_ul = bbox_dict['y1'] // downsample_factor
                x2_br = bbox_dict['x2'] // downsample_factor
                y2_br = bbox_dict['y2'] // downsample_factor
                bbox = np.array([x1_ul, y1_ul, x2_br, y2_br])
                center = [(bbox[2]+bbox[0])/2, (bbox[3]+bbox[1])/2]
                scale = 0.9 * max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 200.

                # get smpl parameters
                ## body resides in multi-ioi coordidate, where camera 0 is world zero.
                with open(smpl_param, 'rb') as f:
                    body_params = pkl.load(f)
                    # in ioi coordinates: cam 0
                    beta = body_params['betas'].detach().cpu().numpy()
                    pose_rotmat = body_params['body_pose'].detach().cpu().numpy()
                    transl = body_params['transl'].detach().cpu().numpy()
                    global_orient = body_params['global_orient'].detach().cpu().numpy()

                smpl_body_cam0 = smpl_obj_orig(betas=torch.FloatTensor(beta).to('cuda')) # canonical body with shape
                vertices_cam0 = smpl_body_cam0.vertices.detach().cpu().numpy().squeeze()
                joints_cam0 = smpl_body_cam0.joints.detach().cpu().numpy()
                pelvis_cam0 = joints_cam0[:, 0, :]

                ## rigid transformation between multi-ioi and Leica scan (world)
                with open(ioi2scan_fn, 'r') as f:
                    ioi2scan_dict = json.load(f)
                    R_ioi2world = np.array(ioi2scan_dict['R']) # Note: R is transposed
                    t_ioi2world= np.array(ioi2scan_dict['t']).reshape(1, 3)

                # get SMPL params in world coordinates
                global_orient_world = np.matmul(R_ioi2world.T, global_orient)
                transl_world = np.matmul((pelvis_cam0+transl), R_ioi2world) + t_ioi2world - pelvis_cam0 # right multiplication to avoid transpose
                full_pose_rotmat_world = np.concatenate((global_orient_world, pose_rotmat), axis=1).squeeze()
                theta_world = batch_rot2aa(torch.FloatTensor(full_pose_rotmat_world)).reshape(-1, 72).cpu().numpy()

                smpl_body_world = smpl_obj(betas=torch.FloatTensor(beta).to('cuda'),
                                    body_pose=torch.FloatTensor(pose_rotmat).to('cuda'),
                                    transl=torch.FloatTensor(transl_world).to('cuda'),
                                    global_orient=torch.FloatTensor(global_orient_world).to('cuda'),
                                      pose2rot=False)
                vertices_world = smpl_body_world.vertices.detach().cpu().numpy().squeeze()
                joints3d_world = smpl_body_world.joints[:, 25:, :].detach().cpu().numpy().squeeze()

                # get SMPL params in camera coordinates
                global_orient_cam = np.matmul(camR, global_orient)
                transl_cam = np.matmul(camR, (pelvis_cam0 + transl).T).T + camT - pelvis_cam0
                full_pose_rotmat_cam = np.concatenate((global_orient_cam, pose_rotmat), axis=1).squeeze()
                theta_cam = batch_rot2aa(torch.FloatTensor(full_pose_rotmat_cam)).reshape(-1, 72).cpu().numpy()
                smpl_body_cam = smpl_obj(betas=torch.FloatTensor(beta).to('cuda'),
                                         body_pose=torch.FloatTensor(pose_rotmat).to('cuda'),
                                         transl=torch.FloatTensor(transl_cam).to('cuda'),
                                           global_orient=torch.FloatTensor(global_orient_cam).to('cuda'),
                                           pose2rot=False)
                vertices_cam = smpl_body_cam.vertices.detach().cpu().numpy().squeeze()
                joints3d_cam = smpl_body_cam.joints[:, 25:, :].detach().cpu().numpy().squeeze()
                # mesh = trimesh.Trimesh(vertices_cam, smpl_obj.faces,
                #                        process=False,
                #                        maintain_order=True)
                # mesh.export('mesh_in_cam0.obj')

                # read GT 2D keypoints
                K = np.eye(3, dtype=np.float)
                K[0, 0] = focal_length_x / downsample_factor
                K[1, 1] = focal_length_y / downsample_factor
                K[:2, 2:] = camC.T / downsample_factor
                projected_points = (K @ joints3d_cam.T).T
                joints2d = projected_points[:, :2] / np.hstack((projected_points[:, 2:], projected_points[:, 2:]))
                part = np.hstack((joints2d, np.ones((joints2d.shape[0], 1))))

                # get openpose 2D keypoints
                try:
                    with open(openpose_path, 'r') as f:
                        openpose = json.load(f)
                    openpose = np.array(openpose['people'][0]['pose_keypoints_2d']).reshape([-1, 3])
                except:
                    print(f'No openpose !! Missing {openpose_path}')
                    continue

                # get camera parameters wrt to scan
                R_worldtocam = np.matmul(camR, R_ioi2world) # Note: R_ioi2world is transposed
                T_worldtocam = -t_ioi2world + camT

                # ground offset
                ground_offset = ground_eq[2]

                # get stability labels: 1: stable, 0: unstable but in contact, -1: unstable and not `in contact
                in_bos_label, contact_label, contact_mask = vis_smpl_with_ground(theta_world, transl_world, beta, seq_i,
                                                                     vis_path,
                                                                     start_idx=frame_num,
                                                                     sub_sample=1,
                                                                     ground_offset=ground_offset,
                                                                     smpl_batch_size=1,
                                                                     visualize=False)
                in_bos_label = in_bos_label.detach().cpu().numpy()
                contact_label = contact_label.detach().cpu().numpy()
                contact_mask = contact_mask.detach().cpu().numpy()

                # visualize world smpl on ground plane
                if visualize:
                    if cam_num == 0:
                        vis_smpl_with_ground(theta_world, transl_world, beta, split+'_'+seq_i, vis_path,
                                             start_idx=frame_num,
                                             sub_sample=1,
                                             ground_offset=ground_offset,
                                             smpl_batch_size=1,
                                             visualize=True)


                # ## visualize projected points
                # img = cv2.imread(img_path)
                # joints2d = joints2d.astype(np.int)
                # img[joints2d[:, 1], joints2d[:, 0], :] = [0, 255, 0]

                # read GT 3D pose in cam coordinates
                S24 = joints3d_cam
                pelvis_cam = (S24[[2], :] + S24[[3], :]) / 2
                S24 -= pelvis_cam
                S24 = np.hstack([S24, np.ones((S24.shape[0], 1))])

                # read GT 3D pose in world coordinates
                S24_world = joints3d_world
                S24_world = np.hstack([S24_world, np.ones((S24_world.shape[0], 1))])

                # store data
                imgnames_.append(img_path)
                centers_.append(center)
                scales_.append(scale)
                parts_.append(part)
                Ss_.append(S24)
                Ss_world_.append(S24_world)
                openposes_.append(openpose)
                poses_.append(theta_cam.squeeze())
                transls_.append(transl.squeeze())
                poses_world_.append(theta_world.squeeze())
                transls_world_.append(transl_world.squeeze())
                shapes_.append(beta.squeeze())
                cams_r_.append(R_worldtocam.tolist())
                # Todo: note that T_worldtocam here is (1,3) whereas in h36m T_worldtocam is (1,3)
                cams_t_.append(T_worldtocam.tolist())
                cams_k_.append(K.tolist())
                in_bos_label_.append(in_bos_label)
                contact_label_.append(contact_label)
                ground_offset_.append(ground_offset)

    # store the data struct
    if not os.path.isdir(out_path):
        os.makedirs(out_path)
    out_file = os.path.join(out_path, f'rich_world_{split}.npz')
    np.savez(out_file, imgname=imgnames_,
                       center=centers_,
                       scale=scales_,
                       part=parts_,
                       S=Ss_,
                       S_world=Ss_world_,
                       pose=poses_,
                       transl=transls_,
                       shape=shapes_,
                       openpose=openposes_,
                       pose_world=poses_world_,
                       transl_world=transls_world_,
                       cam_r=cams_r_,
                       cam_t=cams_t_,
                       cam_k=cams_k_,
                       in_bos_label=in_bos_label_,
                       contact_label=contact_label_,
                       ground_offset=ground_offset_
             )
    print('Saved to ', out_file)

def rectify_pose(camera_r, body_aa):
    body_r = batch_rodrigues(body_aa).reshape(-1,3,3)
    final_r = camera_r @ body_r
    body_aa = batch_rot2aa(final_r)
    return body_aa


def extract_cam_param_xml(xml_path: str = '', dtype=np.float):
    import xml.etree.ElementTree as ET
    tree = ET.parse(xml_path)

    extrinsics_mat = [float(s) for s in tree.find('./CameraMatrix/data').text.split()]
    intrinsics_mat = [float(s) for s in tree.find('./Intrinsics/data').text.split()]
    distortion_vec = [float(s) for s in tree.find('./Distortion/data').text.split()]

    focal_length_x = intrinsics_mat[0]
    focal_length_y = intrinsics_mat[4]
    center = np.array([[intrinsics_mat[2], intrinsics_mat[5]]], dtype=dtype)

    rotation = np.array([[extrinsics_mat[0], extrinsics_mat[1], extrinsics_mat[2]],
                         [extrinsics_mat[4], extrinsics_mat[5], extrinsics_mat[6]],
                         [extrinsics_mat[8], extrinsics_mat[9], extrinsics_mat[10]]], dtype=dtype)

    translation = np.array([[extrinsics_mat[3], extrinsics_mat[7], extrinsics_mat[11]]], dtype=dtype)

    # t = -Rc --> c = -R^Tt
    cam_center = [-extrinsics_mat[0] * extrinsics_mat[3] - extrinsics_mat[4] * extrinsics_mat[7] - extrinsics_mat[8] *
                  extrinsics_mat[11],
                  -extrinsics_mat[1] * extrinsics_mat[3] - extrinsics_mat[5] * extrinsics_mat[7] - extrinsics_mat[9] *
                  extrinsics_mat[11],
                  -extrinsics_mat[2] * extrinsics_mat[3] - extrinsics_mat[6] * extrinsics_mat[7] - extrinsics_mat[10] *
                  extrinsics_mat[11]]

    cam_center = np.array([cam_center], dtype=dtype)

    k1 = np.array([distortion_vec[0]], dtype=dtype)
    k2 = np.array([distortion_vec[1]], dtype=dtype)

    return focal_length_x, focal_length_y, center, rotation, translation, cam_center, k1, k2







datasets/preprocess/rich.py
sha2nkt/ipman-r
import os
os.environ["CDF_LIB"] = "data/cdf37_1-dist/src/lib"

import cv2
import json
import glob
import h5py
import torch
import trimesh
import numpy as np
import pickle as pkl
from xml.dom import minidom
import xml.etree.ElementTree as ET
from tqdm import tqdm
from spacepy import pycdf
# from .read_openpose import read_openpose
import sys
sys.path.append('../../../')
from models import hmr, SMPL
import config
import constants

import smplx
import pytorch3d.transforms as p3dt



from utils.geometry import batch_rodrigues, batch_rot2aa, ea2rm
from vis_utils.world_vis import overlay_mesh, vis_smpl_with_ground


model_type = 'smplx'
model_folder = '/ps/project/common/smplifyx/models/'
body_model_params = dict(model_path=model_folder,
                         model_type=model_type,
                         create_global_orient=True,
                         create_body_pose=True,
                         create_betas=True,
                         num_betas=10,
                         create_left_hand_pose=True,
                         create_right_hand_pose=True,
                         create_expression=True,
                         create_jaw_pose=True,
                         create_leye_pose=True,
                         create_reye_pose=True,
                         create_transl=True,
                         use_pca=False)
body_model = smplx.create(gender='neutral', **body_model_params).to('cuda')

def rich_extract(dataset_path, out_path, split=None, vis_path=None, visualize=False, downsample_factor=4):

    # structs we use
    imgnames_, scales_, centers_, parts_, Ss_, Ss_world_, openposes_ = [], [], [], [], [], [], []
    poses_, shapes_, transls_ = [], [], []
    poses_world_, transls_world_, cams_r_, cams_t_, cams_k_ = [], [], [], [], []
    ground_offset_ = []
    in_bos_label_, contact_label_ = [], []

    # seqs in validation set
    if split == 'val':
        seq_list = ['2021-06-15_Multi_IOI_ID_00176_Yoga2',
                    '2021-06-15_Multi_IOI_ID_00228_Yoga1',
                    '2021-06-15_Multi_IOI_ID_03588_Yoga1',
                    '2021-06-15_Multi_IOI_ID_00176_Yoga1']

    # seqs in testing set
    if split == 'test':
        seq_list = ['2021-06-15_Multi_IOI_ID_00186_Yoga1',
                    '2021-06-15_Multi_IOI_ID_03588_Yoga2',
                    'MultiIOI_201019_ID03581_parkingLot_Calibration06_Settings06_PushUp__2',
                    'Multi-IOI_ID00227_Scene_ParkingLot_Calibration_03_CameraSettings_4_pushup_1']


    for seq_i in tqdm(seq_list):
        print(f'Processing sequence: {seq_i}')

        # path with GT bounding boxes
        params_path = os.path.join(dataset_path, seq_i, 'params')

        # path to metadata for files
        md_path = os.path.join(dataset_path, seq_i, 'data')

        # glob all folders in params path
        frame_param_paths = sorted(glob.glob(os.path.join(params_path, '*')))
        frame_param_paths = [p for p in frame_param_paths if '.yaml' not in p]

        # get ioi2scan transformation per sequence
        ioi2scan_fn = os.path.join(dataset_path, seq_i, 'cam2scan.json')

        ## ground resides in Leica scan coordidate, which is (roughly) axis aligned.
        ground_mesh = trimesh.load(os.path.join(dataset_path, seq_i, 'ground_mesh.ply'), process=False)
        ground_eq = np.mean(ground_mesh.vertices, axis=0)

        # list all files in the folder
        cam_files = os.listdir(os.path.join(dataset_path, seq_i, f'calibration'))
        cam_list = sorted([int(os.path.splitext(f)[0]) for f in cam_files if '.xml' in f])

        # if split == 'val':
        #     cam_list = cam_list[1:] # remove first camera in val
        for cam_num in cam_list[:1]:
            camera_fn = os.path.join(dataset_path, seq_i, f'calibration/{cam_num:03d}.xml')
            focal_length_x, focal_length_y, camC, camR, camT, _, _, _ = extract_cam_param_xml(camera_fn)

            for frame_param_path in tqdm(frame_param_paths):
                frame_id = os.path.basename(frame_param_path)
                frame_num = int(frame_id)

                # path to smpl params
                try:
                    smplx_param = os.path.join(frame_param_path, '00', 'results/000.pkl')
                except:
                    import ipdb; ipdb.set_trace()

                # path to GT bounding boxes
                bbox_path = os.path.join(md_path, frame_id, '00', 'bbox_refine', f'{frame_id}_{cam_num:02d}.json')
                # path with 2D openpose keypoints
                openpose_path = os.path.join(md_path, frame_id, '00', 'keypoints_refine', f'{frame_id}_{str(cam_num).zfill(2)}_keypoints.json')
                # path to image crops
                if downsample_factor == 1:
                    img_path = os.path.join(md_path, frame_id, '00', 'images_orig', f'{frame_id}_{cam_num:02d}.png')
                else:
                    img_path = os.path.join(md_path, frame_id, '00', 'images_orig_720p', f'{frame_id}_{cam_num:02d}.png')

                if not os.path.isfile(img_path):
                    print(f'image not found: {img_path}')
                    continue
                    # raise FileNotFoundError

                # bbox file
                try:
                    with open(bbox_path, 'r') as f:
                        bbox_dict = json.load(f)
                except:
                    print(f'bbox file not found: {bbox_path}')
                    continue
                # read GT bounding box
                x1_ul = bbox_dict['x1'] // downsample_factor
                y1_ul = bbox_dict['y1'] // downsample_factor
                x2_br = bbox_dict['x2'] // downsample_factor
                y2_br = bbox_dict['y2'] // downsample_factor
                bbox = np.array([x1_ul, y1_ul, x2_br, y2_br])
                center = [(bbox[2]+bbox[0])/2, (bbox[3]+bbox[1])/2]
                scale = 0.9 * max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 200.

                # get smpl parameters
                ## body resides in multi-ioi coordidate, where camera 0 is world zero.
                with open(smplx_param, 'rb') as f:
                    body_params = pkl.load(f)
                    # in ioi coordinates: cam 0
                    beta = body_params['betas']
                    pose_aa = body_params['body_pose']
                    pose_rotmat = p3dt.axis_angle_to_matrix(torch.FloatTensor(pose_aa.reshape(-1,3))).numpy()

                    transl = body_params['transl']
                    global_orient = body_params['global_orient']
                    global_orient = p3dt.axis_angle_to_matrix(torch.FloatTensor(global_orient.reshape(-1,3))).numpy()

                smpl_body_cam0 = body_model(betas=torch.FloatTensor(beta).to('cuda')) # canonical body with shape
                vertices_cam0 = smpl_body_cam0.vertices.detach().cpu().numpy().squeeze()
                joints_cam0 = smpl_body_cam0.joints.detach().cpu().numpy()
                pelvis_cam0 = joints_cam0[:, 0, :]


                ## rigid transformation between multi-ioi and Leica scan (world)
                with open(ioi2scan_fn, 'r') as f:
                    ioi2scan_dict = json.load(f)
                    R_ioi2world = np.array(ioi2scan_dict['R']) # Note: R is transposed
                    t_ioi2world= np.array(ioi2scan_dict['t']).reshape(1, 3)

                # get SMPL params in world coordinates
                # import ipdb; ipdb.set_trace()
                global_orient_world = np.matmul(R_ioi2world.T, global_orient)
                transl_world = np.matmul((pelvis_cam0+transl), R_ioi2world) + t_ioi2world - pelvis_cam0 # right multiplication to avoid transpose
                full_pose_rotmat_world = np.concatenate((global_orient_world, pose_rotmat), axis=0).squeeze()
                theta_world = batch_rot2aa(torch.FloatTensor(full_pose_rotmat_world)).reshape(-1, 66).cpu().numpy()

                smpl_body_world = body_model(betas=torch.FloatTensor(beta).to('cuda'),
                                    body_pose=torch.FloatTensor(theta_world[:, 3:]).to('cuda'),
                                    transl=torch.FloatTensor(transl_world).to('cuda'),
                                    global_orient=torch.FloatTensor(theta_world[:, :3]).to('cuda'))
                vertices_world = smpl_body_world.vertices.detach().cpu().numpy().squeeze()
                joints3d_world = smpl_body_world.joints[:, 25:, :].detach().cpu().numpy().squeeze()

                mesh = trimesh.Trimesh(vertices_world, body_model.faces,
                                       process=False,
                                       maintain_order=True)
                mesh.export('gt_mesh_world_smplx.obj')



                # smpl_body_world = body_model(betas=torch.FloatTensor(beta).to('cuda'),
                #                     body_pose=torch.FloatTensor(pose_rotmat[None, ...]).to('cuda'),
                #                     transl=torch.FloatTensor(transl_world[None, ...]).to('cuda'),
                #                     global_orient=torch.FloatTensor(global_orient_world[None, ...]).to('cuda'),
                #                     left_hand_pose=torch.eye(3).reshape(1, 1, 3, 3).expand(1, 15, -1, -1).to('cuda'),
                #                     right_hand_pose=torch.eye(3).reshape(1, 1, 3, 3).expand(1, 15, -1, -1).to('cuda'),
                #                     leye_pose= torch.eye(3).reshape(1, 1, 3, 3).expand(batch_size, -1, -1, -1),
                #
                #                       pose2rot=False)
                # vertices_world = smpl_body_world.vertices.detach().cpu().numpy().squeeze()
                # joints3d_world = smpl_body_world.joints[:, 25:, :].detach().cpu().numpy().squeeze()

                # mesh = trimesh.Trimesh(vertices_world, body_model.faces,
                #                        process=False,
                #                        maintain_order=True)
                # mesh.export('gt_mesh_world_smplx.obj')

                # get SMPL params in camera coordinates
                global_orient_cam = np.matmul(camR, global_orient)
                transl_cam = np.matmul(camR, (pelvis_cam0 + transl).T).T + camT - pelvis_cam0
                full_pose_rotmat_cam = np.concatenate((global_orient_cam, pose_rotmat), axis=0).squeeze()
                theta_cam = batch_rot2aa(torch.FloatTensor(full_pose_rotmat_cam)).reshape(-1, 66).cpu().numpy()
                # smpl_body_cam = body_model(betas=torch.FloatTensor(beta).to('cuda'),
                #                          body_pose=torch.FloatTensor(pose_rotmat).to('cuda'),
                #                          transl=torch.FloatTensor(transl_cam).to('cuda'),
                #                            global_orient=torch.FloatTensor(global_orient_cam).to('cuda'),
                #                            pose2rot=False)
                # vertices_cam = smpl_body_cam.vertices.detach().cpu().numpy().squeeze()
                # joints3d_cam = smpl_body_cam.joints[:, 25:, :].detach().cpu().numpy().squeeze()
                #
                # mesh = trimesh.Trimesh(vertices_cam, body_model.faces,
                #                        process=False,
                #                        maintain_order=True)
                # mesh.export('mesh_in_cam0.obj')

                # read GT 2D keypoints
                K = np.eye(3, dtype=np.float)
                K[0, 0] = focal_length_x / downsample_factor
                K[1, 1] = focal_length_y / downsample_factor
                K[:2, 2:] = camC.T / downsample_factor
                # projected_points = (K @ joints3d_cam.T).T
                # joints2d = projected_points[:, :2] / np.hstack((projected_points[:, 2:], projected_points[:, 2:]))
                # part = np.hstack((joints2d, np.ones((joints2d.shape[0], 1))))

                # get openpose 2D keypoints
                try:
                    with open(openpose_path, 'r') as f:
                        openpose = json.load(f)
                    openpose = np.array(openpose['people'][0]['pose_keypoints_2d']).reshape([-1, 3])
                except:
                    print(f'No openpose !! Missing {openpose_path}')
                    continue

                # get camera parameters wrt to scan
                R_worldtocam = np.matmul(camR, R_ioi2world) # Note: R_ioi2world is transposed
                T_worldtocam = -t_ioi2world + camT

                # ground offset
                ground_offset = ground_eq[2]

                # # get stability labels: 1: stable, 0: unstable but in contact, -1: unstable and not `in contact
                # in_bos_label, contact_label, contact_mask = vis_smpl_with_ground(theta_world, transl_world, beta, seq_i,
                #                                                      vis_path,
                #                                                      start_idx=frame_num,
                #                                                      sub_sample=1,
                #                                                      ground_offset=ground_offset,
                #                                                      smpl_batch_size=1,
                #                                                      visualize=False)
                # in_bos_label = in_bos_label.detach().cpu().numpy()
                # contact_label = contact_label.detach().cpu().numpy()
                # contact_mask = contact_mask.detach().cpu().numpy()

                # visualize world smpl on ground plane
                if visualize:
                    if cam_num == 0:
                        vis_smpl_with_ground(theta_world, transl_world, beta, split+'_'+seq_i, vis_path,
                                             start_idx=frame_num,
                                             sub_sample=1,
                                             ground_offset=ground_offset,
                                             smpl_batch_size=1,
                                             visualize=True)


                # ## visualize projected points
                # img = cv2.imread(img_path)
                # joints2d = joints2d.astype(np.int)
                # img[joints2d[:, 1], joints2d[:, 0], :] = [0, 255, 0]

                # read GT 3D pose in cam coordinates
                # S24 = joints3d_cam
                # pelvis_cam = (S24[[2], :] + S24[[3], :]) / 2
                # S24 -= pelvis_cam
                # S24 = np.hstack([S24, np.ones((S24.shape[0], 1))])

                # read GT 3D pose in world coordinates
                # S24_world = joints3d_world
                # S24_world = np.hstack([S24_world, np.ones((S24_world.shape[0], 1))])

                # store data
                imgnames_.append(img_path)
                centers_.append(center)
                scales_.append(scale)
                # parts_.append(part)
                # Ss_.append(S24)
                # Ss_world_.append(S24_world)
                openposes_.append(openpose)
                poses_.append(theta_cam.squeeze())
                transls_.append(transl.squeeze())
                poses_world_.append(theta_world.squeeze())
                transls_world_.append(transl_world.squeeze())
                shapes_.append(beta.squeeze())
                cams_r_.append(R_worldtocam.tolist())
                # Todo: note that T_worldtocam here is (1,3) whereas in h36m T_worldtocam is (1,3)
                cams_t_.append(T_worldtocam.tolist())
                cams_k_.append(K.tolist())
                # in_bos_label_.append(in_bos_label)
                # contact_label_.append(contact_label)
                ground_offset_.append(ground_offset)

    # store the data struct
    if not os.path.isdir(out_path):
        os.makedirs(out_path)
    out_file = os.path.join(out_path, f'rich_world_{split}_only_cam0_smplx.npz')
    np.savez(out_file, imgname=imgnames_,
                       center=centers_,
                       scale=scales_,
                       part=parts_,
                       S=Ss_,
                       S_world=Ss_world_,
                       pose=poses_,
                       transl=transls_,
                       shape=shapes_,
                       openpose=openposes_,
                       pose_world=poses_world_,
                       transl_world=transls_world_,
                       cam_r=cams_r_,
                       cam_t=cams_t_,
                       cam_k=cams_k_,
                       in_bos_label=in_bos_label_,
                       contact_label=contact_label_,
                       ground_offset=ground_offset_
             )
    print('Saved to ', out_file)

def rectify_pose(camera_r, body_aa):
    body_r = batch_rodrigues(body_aa).reshape(-1,3,3)
    final_r = camera_r @ body_r
    body_aa = batch_rot2aa(final_r)
    return body_aa


def extract_cam_param_xml(xml_path: str = '', dtype=np.float):
    import xml.etree.ElementTree as ET
    tree = ET.parse(xml_path)

    extrinsics_mat = [float(s) for s in tree.find('./CameraMatrix/data').text.split()]
    intrinsics_mat = [float(s) for s in tree.find('./Intrinsics/data').text.split()]
    distortion_vec = [float(s) for s in tree.find('./Distortion/data').text.split()]

    focal_length_x = intrinsics_mat[0]
    focal_length_y = intrinsics_mat[4]
    center = np.array([[intrinsics_mat[2], intrinsics_mat[5]]], dtype=dtype)

    rotation = np.array([[extrinsics_mat[0], extrinsics_mat[1], extrinsics_mat[2]],
                         [extrinsics_mat[4], extrinsics_mat[5], extrinsics_mat[6]],
                         [extrinsics_mat[8], extrinsics_mat[9], extrinsics_mat[10]]], dtype=dtype)

    translation = np.array([[extrinsics_mat[3], extrinsics_mat[7], extrinsics_mat[11]]], dtype=dtype)

    # t = -Rc --> c = -R^Tt
    cam_center = [-extrinsics_mat[0] * extrinsics_mat[3] - extrinsics_mat[4] * extrinsics_mat[7] - extrinsics_mat[8] *
                  extrinsics_mat[11],
                  -extrinsics_mat[1] * extrinsics_mat[3] - extrinsics_mat[5] * extrinsics_mat[7] - extrinsics_mat[9] *
                  extrinsics_mat[11],
                  -extrinsics_mat[2] * extrinsics_mat[3] - extrinsics_mat[6] * extrinsics_mat[7] - extrinsics_mat[10] *
                  extrinsics_mat[11]]

    cam_center = np.array([cam_center], dtype=dtype)

    k1 = np.array([distortion_vec[0]], dtype=dtype)
    k2 = np.array([distortion_vec[1]], dtype=dtype)

    return focal_length_x, focal_length_y, center, rotation, translation, cam_center, k1, k2







datasets/preprocess/rich_smplx.py
sha2nkt/ipman-r
"""
This script can be used to evaluate a trained model on 3D pose/shape and masks/part segmentation. You first need to download the datasets and preprocess them.
Example usage:
```
python3 eval.py --checkpoint=data/model_checkpoint.pt --dataset=h36m-p1 --log_freq=20
```
Running the above command will compute the MPJPE and Reconstruction Error on the Human3.6M dataset (Protocol I). The ```--dataset``` option can take different values based on the type of evaluation you want to perform:
1. Human3.6M Protocol 1 ```--dataset=h36m-p1```
2. Human3.6M Protocol 2 ```--dataset=h36m-p2```
3. 3DPW ```--dataset=3dpw```
4. LSP ```--dataset=lsp```
5. MPI-INF-3DHP ```--dataset=mpi-inf-3dhp```
"""

import torch
from torch.utils.data import DataLoader
import numpy as np
import cv2
import os
# for cluster rendering
import argparse
from tqdm import tqdm
import torchgeometry as tgm


import config
import constants
from models import hmr, SMPL
from datasets import BaseDatasetEval
from utils.imutils import uncrop
from utils.pose_utils import reconstruction_error, compute_accel, compute_error_accel
from utils.geometry import batch_rot2aa, batch_rectify_pose

# from utils.part_utils import PartRenderer
from vis_utils.world_vis import vis_vert_with_ground

# Define command-line arguments
parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint', default=None, help='Path to network checkpoint')
parser.add_argument('--dataset', default='h36m-p1', choices=['h36m-p1', 'h36m-p2', 'h36m-test-s1', 'h36m-train-small', 'lsp', '3dpw', 'mpi-inf-3dhp',
                                                             'rich-val-onlycam0', 'rich-val-onlycam1', 'rich-val-onlycam5', 'rich-test-last2seq-onlycam0',
                                                             'rich-test-onlycam0', 'rich-val', 'rich-test'], help='Choose evaluation dataset')
parser.add_argument('--log_freq', default=50, type=int, help='Frequency of printing intermediate results')
parser.add_argument('--batch_size', default=32, help='Batch size for testing')
parser.add_argument('--shuffle', default=False, action='store_true', help='Shuffle data')
parser.add_argument('--num_workers', default=8, type=int, help='Number of processes for data loading')
parser.add_argument('--result_file', default=None, help='If set, save detections to a .npz file')
parser.add_argument('--vis_path', default=None, help='save the world frame visualizations here')
parser.add_argument('--visualize', default=False, action='store_true', help='generate visualization?')
parser.add_argument('--eval_stability', default=False, action='store_true', help='compute stability metrics?')

def cam_to_world(vertices, translation, cam_r, cam_t):
    """
    Convert vertices from camera coordinates to world coordinates.
    """
    # apply predicted translation to the mesh (as tb = -tc)
    vertices = vertices + translation[:, None, :]

    cam_r = cam_r.to(torch.float32)
    cam_t = cam_t.to(torch.float32)

    # cam extrinsics
    mm2m = 1000
    R = cam_r
    t = -torch.bmm(R, cam_t) / mm2m  # t= -RC
    # t = cam_t.to(torch.float32) / mm2m

    # reverse camera to go from camera to world
    R_T = R.permute(0, 2, 1)
    t_w = -torch.bmm(R_T, t)

    # apply extrinsics
    vertices_world = torch.einsum('bij,bkj->bki', R_T, vertices)
    vertices_world = vertices_world + t_w.squeeze()[:, None, :]
    return vertices_world

def run_evaluation(model, dataset_name, dataset, result_file, vis_path,
                   batch_size=32, img_res=224, 
                   num_workers=32, shuffle=False, log_freq=50, visualize=False, eval_stability=False):
    """Run evaluation on the datasets and metrics we report in the paper. """

    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    # Transfer model to the GPU
    model.to(device)

    # Load SMPL model
    smpl_neutral = SMPL(config.SMPL_MODEL_DIR,
                        create_transl=False).to(device)
    smpl_male = SMPL(config.SMPL_MODEL_DIR,
                     gender='male',
                     create_transl=False).to(device)
    smpl_female = SMPL(config.SMPL_MODEL_DIR,
                       gender='female',
                       create_transl=False).to(device)

    # renderer = PartRenderer()

    # Regressor for H36m joints
    J_regressor = torch.from_numpy(np.load(config.JOINT_REGRESSOR_H36M)).float()

    save_results = result_file is not None
    # Disable shuffling if you want to save the results
    if save_results:
        shuffle=False
    # Create dataloader for the dataset
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)

    # Pose metrics
    # MPJPE and Reconstruction error for the non-parametric and parametric shapes
    mpjpe = np.zeros(len(dataset))
    recon_err = np.zeros(len(dataset))
    mpjpe_smpl = np.zeros(len(dataset))
    recon_err_smpl = np.zeros(len(dataset))
    v2v_err = np.zeros(len(dataset))

    # Acceleration metrics
    accel_ = np.zeros(len(dataset))
    accel_err_ = np.zeros(len(dataset))

    gt_bos_accumulator = []
    gt_contact_accumulator = []
    gt_contact_mask_accumulator = []
    pred_bos_accumulator = []
    pred_contact_accumulator = []
    pred_contact_mask_accumulator = []

    # Shape metrics
    # Mean per-vertex error
    shape_err = np.zeros(len(dataset))
    shape_err_smpl = np.zeros(len(dataset))

    # Mask and part metrics
    # Accuracy
    accuracy = 0.
    parts_accuracy = 0.
    # True positive, false positive and false negative
    tp = np.zeros((2,1))
    fp = np.zeros((2,1))
    fn = np.zeros((2,1))
    parts_tp = np.zeros((7,1))
    parts_fp = np.zeros((7,1))
    parts_fn = np.zeros((7,1))
    # Pixel count accumulators
    pixel_count = 0
    parts_pixel_count = 0

    # Store SMPL parameters
    smpl_pose = np.zeros((len(dataset), 72))
    smpl_betas = np.zeros((len(dataset), 10))
    smpl_camera = np.zeros((len(dataset), 3))
    pred_joints = np.zeros((len(dataset), 17, 3))

    eval_pose = False
    eval_masks = False
    eval_parts = False
    # Choose appropriate evaluation for each dataset
    if dataset_name == 'h36m-p1' or dataset_name == 'h36m-p2'  \
            or dataset_name == 'h36m-test-s1' or dataset_name == '3dpw' \
            or dataset_name == 'mpi-inf-3dhp' \
            or dataset_name == 'rich-val-onlycam0' or dataset_name == 'rich-test-onlycam0'\
            or dataset_name == 'rich-val' or dataset_name == 'rich-test':
        eval_pose = True
    elif dataset_name == 'lsp':
        eval_masks = True
        eval_parts = True
        annot_path = config.DATASET_FOLDERS['upi-s1h']

    joint_mapper_h36m = constants.H36M_TO_J17 if dataset_name == 'mpi-inf-3dhp' else constants.H36M_TO_J14
    joint_mapper_gt = constants.J24_TO_J17 if dataset_name == 'mpi-inf-3dhp' else constants.J24_TO_J14
    # Iterate over the entire dataset
    for step, batch in enumerate(tqdm(data_loader, desc='Eval', total=len(data_loader))):
        # Get ground truth annotations from the batch
        if len(torch.unique(batch['has_smpl'])) == 1 and torch.unique(batch['has_smpl'])[0] == 1:
            gt_pose = batch['pose'].to(device)
            gt_betas = batch['betas'].to(device)
            gt_vertices = smpl_neutral(betas=gt_betas, body_pose=gt_pose[:, 3:], global_orient=gt_pose[:, :3]).vertices
            has_smpl = True

        images = batch['img'].to(device)
        imgnames = batch['imgname']
        gender = batch['gender'].to(device)
        curr_batch_size = images.shape[0]
        if 'rich' in dataset_name:
            ground_offset = batch['ground_offset'][:, None].to(device)
        elif 'h36m' in dataset_name:
            ground_offset = constants.GROUND_OFFSETS['h36m']
        else:
            ground_offset = 0
        # visualize mesh in world coordinates
        if len(torch.unique(batch['has_smpl_world'])) == 1 and torch.unique(batch['has_smpl_world'])[0] == 1:
            gt_pose_world = batch['pose_world'].to(device)
            gt_transl_world = batch['transl_world'].to(device)
            gt_body_world = smpl_neutral(betas=gt_betas,
                                         body_pose=gt_pose_world[:, 3:],
                                         global_orient=gt_pose_world[:, :3],
                                         transl=gt_transl_world)

            gt_joints_world = gt_body_world.joints[:, 25:, :]
            if eval_stability:
                # visualize the gt vertices in world coords
                gt_in_bos_label, gt_contact_metric, gt_contact_mask = vis_vert_with_ground(gt_body_world.vertices,
                                                                              gt_transl_world[:, None, :]*0.0, # since already included in smpl
                                                                              seq_name='test-gt_vert',
                                                                              vis_path=vis_path,
                                                                              start_idx=step * curr_batch_size,
                                                                              sub_sample=1,
                                                                              smpl_batch_size=curr_batch_size,
                                                                              imgnames = imgnames,
                                                                              ground_offset=ground_offset,
                                                                              visualize=False)
                # gt_in_bos_label, gt_contact_metric = vis_smpl_with_ground(gt_pose_world, gt_transl_world, gt_betas,
                #                                                           seq_name='test-gt_pose',
                #                                                           vis_path=vis_path,
                #                                                           start_idx= step * curr_batch_size,
                #                                                           sub_sample=1,
                #                                                           smpl_batch_size=curr_batch_size,
                #                                                           ground_offset = ground_offset,
                #                                                           visualize=visualize)
                gt_bos_accumulator.append(gt_in_bos_label)
                gt_contact_accumulator.append(gt_contact_metric)
                gt_contact_mask_accumulator.append(gt_contact_mask)
            has_smpl_world = True
        else:
            gt_joints_world = batch['pose_3d_world'][:, :, :-1].to(device)
            has_smpl_world = False



        # Get ground truth camera
        gt_cam_r = batch['cam_r'].to(device).to(torch.float32)
        gt_cam_t = batch['cam_t'].to(device).to(torch.float32)
        gt_cam_k = batch['cam_k'].to(device).to(torch.float32)

        # # visualize mesh by going from camera to world coordinates
        # # apply reverse rotation to pose in camera coordinates
        # gt_pose_new = gt_pose.clone()
        # gt_pose_new[:,:3] = batch_rectify_pose(gt_cam_r.permute(0, 2, 1), gt_pose_new[:,:3])
        # vis_smpl_with_ground(gt_pose_new, gt_transl_world, gt_betas,
        #                      seq_name='test-gt-recovered',
        #                      vis_path=vis_path,
        #                      sub_sample=1,
        #                      ground_offset = ground_offset,
        #                      smpl_batch_size=curr_batch_size)

        with torch.no_grad():
            pred_rotmat, pred_betas, pred_camera = model(images)
            pred_pose = batch_rot2aa(pred_rotmat.view(-1, 3, 3)).view(curr_batch_size, -1)
            pred_output = smpl_neutral(betas=pred_betas, body_pose=pred_rotmat[:,1:], global_orient=pred_rotmat[:,0].unsqueeze(1), pose2rot=False)
            pred_vertices = pred_output.vertices
            pred_joints = pred_output.joints

            ## move to world coordinates - first correct the predicted pose using gt camera rotation, then correct the translation by pelvis alignment
            pred_pose_world = pred_pose.clone()
            pred_pose_world[:, :3] = batch_rectify_pose(gt_cam_r.permute(0, 2, 1), pred_pose_world[:, :3])
            pred_output_world = smpl_neutral(betas=pred_betas,
                                             body_pose=pred_pose_world[:, 3:],
                                             global_orient=pred_pose_world[:, :3])
            pred_vertices_world = pred_output_world.vertices
            pred_joints_world = pred_output_world.joints[:, 25:, :]
            # compute pred transl world by equating it to the gt world pelvis
            gt_world_pelvis = (gt_joints_world[:, 2, :] + gt_joints_world[:, 3, :]) / 2
            pred_pelvis_world = (pred_joints_world[:, 2, :] + pred_joints_world[:, 3, :]) / 2
            pred_transl_world = gt_world_pelvis - pred_pelvis_world
            if eval_stability:
                pred_in_bos_label, pred_contact_metric, pred_contact_mask = vis_vert_with_ground(pred_vertices_world, pred_transl_world[:, None, :],
                                                                              seq_name='test-pred-offset',
                                                                              vis_path=vis_path,
                                                                              start_idx=step * curr_batch_size,
                                                                              sub_sample=1,
                                                                              imgnames = imgnames,
                                                                              smpl_batch_size=curr_batch_size,
                                                                              ground_offset=ground_offset,
                                                                              visualize=visualize)
                pred_bos_accumulator.append(pred_in_bos_label)
                pred_contact_accumulator.append(pred_contact_metric)
                pred_contact_mask_accumulator.append(pred_contact_mask)

        if save_results:
            rot_pad = torch.tensor([0,0,1], dtype=torch.float32, device=device).view(1,3,1)
            rotmat = torch.cat((pred_rotmat.view(-1, 3, 3), rot_pad.expand(curr_batch_size * 24, -1, -1)), dim=-1)
            pred_pose = tgm.rotation_matrix_to_angle_axis(rotmat).contiguous().view(-1, 72)
            smpl_pose[step * batch_size:step * batch_size + curr_batch_size, :] = pred_pose.cpu().numpy()
            smpl_betas[step * batch_size:step * batch_size + curr_batch_size, :]  = pred_betas.cpu().numpy()
            smpl_camera[step * batch_size:step * batch_size + curr_batch_size, :]  = pred_camera.cpu().numpy()

        # 3D pose evaluation
        if eval_pose:
            # Regressor broadcasting
            J_regressor_batch = J_regressor[None, :].expand(pred_vertices.shape[0], -1, -1).to(device)
            # Get 14 ground truth joints
            if 'h36m' in dataset_name or 'mpi-inf' in dataset_name:
                gt_keypoints_3d = batch['pose_3d'].cuda()
                gt_keypoints_3d = gt_keypoints_3d[:, joint_mapper_gt, :-1]
            # For 3DPW get the 14 common joints from the rendered shape
            if '3dpw' in dataset_name:
                gt_vertices = smpl_male(global_orient=gt_pose[:,:3], body_pose=gt_pose[:,3:], betas=gt_betas).vertices
                gt_vertices_female = smpl_female(global_orient=gt_pose[:,:3], body_pose=gt_pose[:,3:], betas=gt_betas).vertices
                gt_vertices[gender==1, :, :] = gt_vertices_female[gender==1, :, :]
                gt_keypoints_3d = torch.matmul(J_regressor_batch, gt_vertices)
                gt_pelvis = gt_keypoints_3d[:, [0],:].clone()
                gt_keypoints_3d = gt_keypoints_3d[:, joint_mapper_h36m, :]
                gt_keypoints_3d = gt_keypoints_3d - gt_pelvis
            if 'rich' in dataset_name:
                gt_keypoints_3d = torch.matmul(J_regressor_batch, gt_vertices)
                gt_pelvis = gt_keypoints_3d[:, [0], :].clone()
                gt_keypoints_3d = gt_keypoints_3d[:, joint_mapper_h36m, :]
                gt_keypoints_3d = gt_keypoints_3d - gt_pelvis
                gt_vertices = gt_vertices - gt_pelvis

                # Get 14 predicted joints from the mesh
            pred_keypoints_3d = torch.matmul(J_regressor_batch, pred_vertices)
            if save_results:
                pred_joints[step * batch_size:step * batch_size + curr_batch_size, :, :]  = pred_keypoints_3d.cpu().numpy()
            pred_pelvis = pred_keypoints_3d[:, [0],:].clone()
            pred_keypoints_3d = pred_keypoints_3d[:, joint_mapper_h36m, :]
            pred_keypoints_3d = pred_keypoints_3d - pred_pelvis
            pred_vertices = pred_vertices - pred_pelvis

            # Absolute error (MPJPE)
            error = torch.sqrt(((pred_keypoints_3d - gt_keypoints_3d) ** 2).sum(dim=-1)).mean(dim=-1).cpu().numpy()
            mpjpe[step * batch_size:step * batch_size + curr_batch_size] = error

            # Reconstuction_error
            r_error = reconstruction_error(pred_keypoints_3d.cpu().numpy(), gt_keypoints_3d.cpu().numpy(), reduction=None)
            recon_err[step * batch_size:step * batch_size + curr_batch_size] = r_error

            # V2V error in RICH
            if 'rich' in dataset_name:
                v2v = torch.sqrt(((gt_vertices - pred_vertices) ** 2).sum(dim=-1)).mean(dim=-1).cpu().numpy()
                v2v_err[step * batch_size:step * batch_size + curr_batch_size] = v2v

            # Acceleration Error
            if pred_keypoints_3d.shape[0] >= 3: # can only calculate accel if there are at least 3 frames
                accel = np.mean(compute_accel(pred_keypoints_3d.cpu().numpy()))
                if not np.isnan(accel):
                    accel_[step * batch_size:step * batch_size + curr_batch_size] = accel
                accel_err = np.mean(compute_error_accel(joints_pred=pred_keypoints_3d.cpu().numpy(), joints_gt=gt_keypoints_3d.cpu().numpy()))
                if not np.isnan(accel_err):
                    accel_err_[step * batch_size:step * batch_size + curr_batch_size] = accel_err

        # If mask or part evaluation, render the mask and part images
        # if eval_masks or eval_parts:
        #     mask, parts = renderer(pred_vertices, pred_camera)

        # Mask evaluation (for LSP)
        if eval_masks:
            center = batch['center'].cpu().numpy()
            scale = batch['scale'].cpu().numpy()
            # Dimensions of original image
            orig_shape = batch['orig_shape'].cpu().numpy()
            for i in range(curr_batch_size):
                # After rendering, convert imate back to original resolution
                pred_mask = uncrop(mask[i].cpu().numpy(), center[i], scale[i], orig_shape[i]) > 0
                # Load gt mask
                gt_mask = cv2.imread(os.path.join(annot_path, batch['maskname'][i]), 0) > 0
                # Evaluation consistent with the original UP-3D code
                accuracy += (gt_mask == pred_mask).sum()
                pixel_count += np.prod(np.array(gt_mask.shape))
                for c in range(2):
                    cgt = gt_mask == c
                    cpred = pred_mask == c
                    tp[c] += (cgt & cpred).sum()
                    fp[c] +=  (~cgt & cpred).sum()
                    fn[c] +=  (cgt & ~cpred).sum()
                f1 = 2 * tp / (2 * tp + fp + fn)

        # Part evaluation (for LSP)
        if eval_parts:
            center = batch['center'].cpu().numpy()
            scale = batch['scale'].cpu().numpy()
            orig_shape = batch['orig_shape'].cpu().numpy()
            for i in range(curr_batch_size):
                pred_parts = uncrop(parts[i].cpu().numpy().astype(np.uint8), center[i], scale[i], orig_shape[i])
                # Load gt part segmentation
                gt_parts = cv2.imread(os.path.join(annot_path, batch['partname'][i]), 0)
                # Evaluation consistent with the original UP-3D code
                # 6 parts + background
                for c in range(7):
                    cgt = gt_parts == c
                    cpred = pred_parts == c
                    cpred[gt_parts == 255] = 0
                    parts_tp[c] += (cgt & cpred).sum()
                    parts_fp[c] +=  (~cgt & cpred).sum()
                    parts_fn[c] +=  (cgt & ~cpred).sum()
                gt_parts[gt_parts == 255] = 0
                pred_parts[pred_parts == 255] = 0
                parts_f1 = 2 * parts_tp / (2 * parts_tp + parts_fp + parts_fn)
                parts_accuracy += (gt_parts == pred_parts).sum()
                parts_pixel_count += np.prod(np.array(gt_parts.shape))

        # # Print intermediate results during evaluation
        # if step % log_freq == log_freq - 1:
        #     if eval_pose:
        #         print('MPJPE: ' + str(1000 * mpjpe[:step * batch_size].mean()))
        #         print('Reconstruction Error: ' + str(1000 * recon_err[:step * batch_size].mean()))
        #         print('V2V Error: ' + str(1000 * v2v_err[:step * batch_size].mean()))
        #         print()
        #     if eval_masks:
        #         print('Accuracy: ', accuracy / pixel_count)
        #         print('F1: ', f1.mean())
        #         print()
        #     if eval_parts:
        #         print('Parts Accuracy: ', parts_accuracy / parts_pixel_count)
        #         print('Parts F1 (BG): ', parts_f1[[0,1,2,3,4,5,6]].mean())
        #         print()


    print('*** Stability Metrics ***')
    print()
    if eval_stability:
        # find accuracy of bos metric
        pred_bos_accumulator = torch.cat(pred_bos_accumulator, dim=0).to(torch.float32)
        pred_contact_accumulator = torch.cat(pred_contact_accumulator, dim=0).to(torch.float32)
        pred_contact_mask_accumulator = torch.cat(pred_contact_mask_accumulator, dim=0).to(torch.float32)

        if has_smpl_world:
            gt_bos_accumulator = torch.cat(gt_bos_accumulator, dim=0)
            pred_bos_accuracy = (gt_bos_accumulator == pred_bos_accumulator).sum().item() / gt_bos_accumulator.shape[0]
            gt_contact_accumulator = torch.cat(gt_contact_accumulator, dim=0)
            gt_contact_mask_accumulator = torch.cat(gt_contact_mask_accumulator, dim=0)
            pred_contact_accuracy = (gt_contact_accumulator == pred_contact_accumulator).sum().item() / gt_contact_accumulator.shape[0]
            dense_contact_accuracy = (gt_contact_mask_accumulator == pred_contact_mask_accumulator).sum().item() / \
                                     gt_contact_mask_accumulator.numel()
            print('BOS Accuracy: ', pred_bos_accuracy)
            print('Contact Accuracy: ', pred_contact_accuracy)
            print(f'Contact Mask Accuracy: ', dense_contact_accuracy)
        else:
            bos_mean = torch.mean(pred_bos_accumulator)
            contact_mean = torch.mean(pred_contact_accumulator)
            print('BOS Mean: ', bos_mean.cpu().numpy())
            print('Contact Mean: ', contact_mean.cpu().numpy())

    # Save reconstructions to a file for further processing
    if save_results:
        np.savez(result_file, pred_joints=pred_joints, pose=smpl_pose, betas=smpl_betas, camera=smpl_camera)
    # Print final results during evaluation
    print('*** Final Results ***')
    print()
    if eval_pose:
        print('MPJPE: ' + str(1000 * mpjpe.mean()))
        print('Reconstruction Error: ' + str(1000 * recon_err.mean()))
        print('V2V Error: ' + str(1000 * v2v_err.mean()))
        print()
    if eval_masks:
        print('Accuracy: ', accuracy / pixel_count)
        print('F1: ', f1.mean())
        print()
    if eval_parts:
        print('Parts Accuracy: ', parts_accuracy / parts_pixel_count)
        print('Parts F1 (BG): ', parts_f1[[0,1,2,3,4,5,6]].mean())
        print()

if __name__ == '__main__':
    args = parser.parse_args()
    model = hmr(config.SMPL_MEAN_PARAMS)
    checkpoint = torch.load(args.checkpoint)
    model.load_state_dict(checkpoint['model'], strict=False)
    model.eval()

    # Setup evaluation dataset
    dataset = BaseDatasetEval(None, args.dataset, is_train=False)
    # Run evaluation
    run_evaluation(model, args.dataset, dataset, args.result_file, args.vis_path,
                   batch_size=args.batch_size,
                   shuffle=args.shuffle,
                   log_freq=args.log_freq,
                   num_workers=args.num_workers,
                   visualize=args.visualize,
                   eval_stability=args.eval_stability)

eval.py
sha2nkt/ipman-r
"""
This script can be used to evaluate a trained model on 3D pose/shape and masks/part segmentation. You first need to download the datasets and preprocess them.
Example usage:
```
python3 eval.py --checkpoint=data/model_checkpoint.pt --dataset=h36m-p1 --log_freq=20
```
Running the above command will compute the MPJPE and Reconstruction Error on the Human3.6M dataset (Protocol I). The ```--dataset``` option can take different values based on the type of evaluation you want to perform:
1. Human3.6M Protocol 1 ```--dataset=h36m-p1```
2. Human3.6M Protocol 2 ```--dataset=h36m-p2```
3. 3DPW ```--dataset=3dpw```
4. LSP ```--dataset=lsp```
5. MPI-INF-3DHP ```--dataset=mpi-inf-3dhp```
"""

import torch
from torch.utils.data import DataLoader
import numpy as np
import cv2
import os
# for cluster rendering
import argparse
import trimesh
import json
from collections import namedtuple
from tqdm import tqdm
import torchgeometry as tgm
import pprint


import config
import constants
from models import hmr, SMPL
from datasets import BaseDatasetEval
from utils.imutils import uncrop
from utils.pose_utils import reconstruction_error, compute_accel, compute_error_accel
from utils.geometry import batch_rodrigues, batch_rot2aa, ea2rm, rot6d_to_rotmat, batch_rectify_pose

# from utils.part_utils import PartRenderer
from vis_utils.world_vis import vis_smpl_with_ground, vis_vert_with_ground

# Define command-line arguments
parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint', default=None, help='Path to network checkpoint')
parser.add_argument('--dataset', default='h36m-p1', choices=['h36m-p1', 'h36m-p2', 'h36m-test-s1', 'h36m-train-small', 'lsp', '3dpw', 'mpi-inf-3dhp',
                                                             'rich-val-onlycam0', 'rich-val-onlycam1', 'rich-val-onlycam5', 'rich-test-last2seq-onlycam0',
                                                             'rich-test-onlycam0', 'rich-val', 'rich-test'], help='Choose evaluation dataset')
parser.add_argument('--log_freq', default=50, type=int, help='Frequency of printing intermediate results')
parser.add_argument('--batch_size', default=32, help='Batch size for testing')
parser.add_argument('--shuffle', default=False, action='store_true', help='Shuffle data')
parser.add_argument('--num_workers', default=8, type=int, help='Number of processes for data loading')
parser.add_argument('--result_file', default=None, help='If set, save detections to a .npz file')
parser.add_argument('--vis_path', default=None, help='save the world frame visualizations here')
parser.add_argument('--visualize', default=False, action='store_true', help='generate visualization?')

def cam_to_world(vertices, translation, cam_r, cam_t):
    """
    Convert vertices from camera coordinates to world coordinates.
    """
    # apply predicted translation to the mesh (as tb = -tc)
    vertices = vertices + translation[:, None, :]

    cam_r = cam_r.to(torch.float32)
    cam_t = cam_t.to(torch.float32)

    # cam extrinsics
    mm2m = 1000
    R = cam_r
    t = -torch.bmm(R, cam_t) / mm2m  # t= -RC
    # t = cam_t.to(torch.float32) / mm2m

    # reverse camera to go from camera to world
    R_T = R.permute(0, 2, 1)
    t_w = -torch.bmm(R_T, t)

    # apply extrinsics
    vertices_world = torch.einsum('bij,bkj->bki', R_T, vertices)
    vertices_world = vertices_world + t_w.squeeze()[:, None, :]
    return vertices_world

def run_evaluation(model, dataset_name, dataset, result_file, vis_path,
                   batch_size=32, img_res=224, 
                   num_workers=32, shuffle=False, log_freq=50, visualize=False):
    """Run evaluation on the datasets and metrics we report in the paper. """

    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    # Transfer model to the GPU
    model.to(device)

    # Load SMPL model
    smpl_neutral = SMPL(config.SMPL_MODEL_DIR,
                        create_transl=False).to(device)
    smpl_male = SMPL(config.SMPL_MODEL_DIR,
                     gender='male',
                     create_transl=False).to(device)
    smpl_female = SMPL(config.SMPL_MODEL_DIR,
                       gender='female',
                       create_transl=False).to(device)

    # renderer = PartRenderer()

    # Regressor for H36m joints
    J_regressor = torch.from_numpy(np.load(config.JOINT_REGRESSOR_H36M)).float()

    save_results = result_file is not None
    # Disable shuffling if you want to save the results
    if save_results:
        shuffle=False
    # Create dataloader for the dataset
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)

    # Pose metrics
    # MPJPE and Reconstruction error for the non-parametric and parametric shapes
    mpjpe = np.zeros(len(dataset))
    recon_err = np.zeros(len(dataset))
    mpjpe_smpl = np.zeros(len(dataset))
    recon_err_smpl = np.zeros(len(dataset))
    v2v_err = np.zeros(len(dataset))

    # Acceleration metrics
    accel_ = np.zeros(len(dataset))
    accel_err_ = np.zeros(len(dataset))

    gt_bos_accumulator = []
    gt_contact_accumulator = []
    gt_contact_mask_accumulator = []
    pred_bos_accumulator = []
    pred_contact_accumulator = []
    pred_contact_mask_accumulator = []

    # Shape metrics
    # Mean per-vertex error
    shape_err = np.zeros(len(dataset))
    shape_err_smpl = np.zeros(len(dataset))

    # Mask and part metrics
    # Accuracy
    accuracy = 0.
    parts_accuracy = 0.
    # True positive, false positive and false negative
    tp = np.zeros((2,1))
    fp = np.zeros((2,1))
    fn = np.zeros((2,1))
    parts_tp = np.zeros((7,1))
    parts_fp = np.zeros((7,1))
    parts_fn = np.zeros((7,1))
    # Pixel count accumulators
    pixel_count = 0
    parts_pixel_count = 0

    # classwise error Metrics
    classwise_mpjpe = {}

    # Store SMPL parameters
    smpl_pose = np.zeros((len(dataset), 72))
    smpl_betas = np.zeros((len(dataset), 10))
    smpl_camera = np.zeros((len(dataset), 3))
    pred_joints = np.zeros((len(dataset), 17, 3))

    eval_pose = False
    eval_masks = False
    eval_parts = False
    # Choose appropriate evaluation for each dataset
    if dataset_name == 'h36m-p1' or dataset_name == 'h36m-p2'  \
            or dataset_name == 'h36m-test-s1' or dataset_name == '3dpw' \
            or dataset_name == 'mpi-inf-3dhp' \
            or dataset_name == 'rich-val-onlycam0' or dataset_name == 'rich-test-onlycam0'\
            or dataset_name == 'rich-val' or dataset_name == 'rich-test':
        eval_pose = True
    elif dataset_name == 'lsp':
        eval_masks = True
        eval_parts = True
        annot_path = config.DATASET_FOLDERS['upi-s1h']

    joint_mapper_h36m = constants.H36M_TO_J17 if dataset_name == 'mpi-inf-3dhp' else constants.H36M_TO_J14
    joint_mapper_gt = constants.J24_TO_J17 if dataset_name == 'mpi-inf-3dhp' else constants.J24_TO_J14
    # Iterate over the entire dataset
    for step, batch in enumerate(tqdm(data_loader, desc='Eval', total=len(data_loader))):
        # Get ground truth annotations from the batch
        if len(torch.unique(batch['has_smpl'])) == 1 and torch.unique(batch['has_smpl'])[0] == 1:
            gt_pose = batch['pose'].to(device)
            gt_betas = batch['betas'].to(device)
            gt_vertices = smpl_neutral(betas=gt_betas, body_pose=gt_pose[:, 3:], global_orient=gt_pose[:, :3]).vertices
            has_smpl = True

        images = batch['img'].to(device)
        imgnames = batch['imgname']
        gender = batch['gender'].to(device)
        curr_batch_size = images.shape[0]
        if 'rich' in dataset_name:
            ground_offset = batch['ground_offset'][:, None].to(device)
        elif 'h36m' in dataset_name:
            ground_offset = constants.GROUND_OFFSETS['h36m']
        else:
            ground_offset = 0
        # visualize mesh in world coordinates
        if len(torch.unique(batch['has_smpl_world'])) == 1 and torch.unique(batch['has_smpl_world'])[0] == 1:
            gt_pose_world = batch['pose_world'].to(device)
            gt_transl_world = batch['transl_world'].to(device)
            gt_body_world = smpl_neutral(betas=gt_betas,
                                         body_pose=gt_pose_world[:, 3:],
                                         global_orient=gt_pose_world[:, :3],
                                         transl=gt_transl_world)

            gt_joints_world = gt_body_world.joints[:, 25:, :]
            # # visualize the gt vertices in world coords
            # gt_in_bos_label, gt_contact_metric, gt_contact_mask = vis_vert_with_ground(gt_body_world.vertices,
            #                                                               gt_transl_world[:, None, :]*0.0, # since already included in smpl
            #                                                               seq_name='test-gt_vert',
            #                                                               vis_path=vis_path,
            #                                                               start_idx=step * curr_batch_size,
            #                                                               sub_sample=1,
            #                                                               smpl_batch_size=curr_batch_size,
            #                                                               imgnames = imgnames,
            #                                                               ground_offset=ground_offset,
            #                                                               visualize=False)
            # # gt_in_bos_label, gt_contact_metric = vis_smpl_with_ground(gt_pose_world, gt_transl_world, gt_betas,
            # #                                                           seq_name='test-gt_pose',
            # #                                                           vis_path=vis_path,
            # #                                                           start_idx= step * curr_batch_size,
            # #                                                           sub_sample=1,
            # #                                                           smpl_batch_size=curr_batch_size,
            # #                                                           ground_offset = ground_offset,
            # #                                                           visualize=visualize)
            # gt_bos_accumulator.append(gt_in_bos_label)
            # gt_contact_accumulator.append(gt_contact_metric)
            # gt_contact_mask_accumulator.append(gt_contact_mask)
            has_smpl_world = True
        else:
            gt_joints_world = batch['pose_3d_world'][:, :, :-1].to(device)
            has_smpl_world = False



        # Get ground truth camera
        gt_cam_r = batch['cam_r'].to(device).to(torch.float32)
        gt_cam_t = batch['cam_t'].to(device).to(torch.float32)
        gt_cam_k = batch['cam_k'].to(device).to(torch.float32)

        # # visualize mesh by going from camera to world coordinates
        # # apply reverse rotation to pose in camera coordinates
        # gt_pose_new = gt_pose.clone()
        # gt_pose_new[:,:3] = batch_rectify_pose(gt_cam_r.permute(0, 2, 1), gt_pose_new[:,:3])
        # vis_smpl_with_ground(gt_pose_new, gt_transl_world, gt_betas,
        #                      seq_name='test-gt-recovered',
        #                      vis_path=vis_path,
        #                      sub_sample=1,
        #                      ground_offset = ground_offset,
        #                      smpl_batch_size=curr_batch_size)

        with torch.no_grad():
            pred_rotmat, pred_betas, pred_camera = model(images)
            pred_pose = batch_rot2aa(pred_rotmat.view(-1, 3, 3)).view(curr_batch_size, -1)
            pred_output = smpl_neutral(betas=pred_betas, body_pose=pred_rotmat[:,1:], global_orient=pred_rotmat[:,0].unsqueeze(1), pose2rot=False)
            pred_vertices = pred_output.vertices
            pred_joints = pred_output.joints

            ## move to world coordinates - first correct the predicted pose using gt camera rotation, then correct the translation by pelvis alignment
            pred_pose_world = pred_pose.clone()
            pred_pose_world[:, :3] = batch_rectify_pose(gt_cam_r.permute(0, 2, 1), pred_pose_world[:, :3])
            pred_output_world = smpl_neutral(betas=pred_betas,
                                             body_pose=pred_pose_world[:, 3:],
                                             global_orient=pred_pose_world[:, :3])
            pred_vertices_world = pred_output_world.vertices
            pred_joints_world = pred_output_world.joints[:, 25:, :]
            # compute pred transl world by equating it to the gt world pelvis
            gt_world_pelvis = (gt_joints_world[:, 2, :] + gt_joints_world[:, 3, :]) / 2
            pred_pelvis_world = (pred_joints_world[:, 2, :] + pred_joints_world[:, 3, :]) / 2
            pred_transl_world = gt_world_pelvis - pred_pelvis_world
            pred_in_bos_label, pred_contact_metric, pred_contact_mask = vis_vert_with_ground(pred_vertices_world, pred_transl_world[:, None, :],
                                                                          seq_name='test-pred-offset',
                                                                          vis_path=vis_path,
                                                                          start_idx=step * curr_batch_size,
                                                                          sub_sample=1,
                                                                          imgnames = imgnames,
                                                                          smpl_batch_size=curr_batch_size,
                                                                          ground_offset=ground_offset,
                                                                          visualize=visualize)
            pred_bos_accumulator.append(pred_in_bos_label)
            pred_contact_accumulator.append(pred_contact_metric)
            pred_contact_mask_accumulator.append(pred_contact_mask)

        if save_results:
            rot_pad = torch.tensor([0,0,1], dtype=torch.float32, device=device).view(1,3,1)
            rotmat = torch.cat((pred_rotmat.view(-1, 3, 3), rot_pad.expand(curr_batch_size * 24, -1, -1)), dim=-1)
            pred_pose = tgm.rotation_matrix_to_angle_axis(rotmat).contiguous().view(-1, 72)
            smpl_pose[step * batch_size:step * batch_size + curr_batch_size, :] = pred_pose.cpu().numpy()
            smpl_betas[step * batch_size:step * batch_size + curr_batch_size, :]  = pred_betas.cpu().numpy()
            smpl_camera[step * batch_size:step * batch_size + curr_batch_size, :]  = pred_camera.cpu().numpy()

        # 3D pose evaluation
        if eval_pose:
            # Regressor broadcasting
            J_regressor_batch = J_regressor[None, :].expand(pred_vertices.shape[0], -1, -1).to(device)
            # Get 14 ground truth joints
            if 'h36m' in dataset_name or 'mpi-inf' in dataset_name:
                gt_keypoints_3d = batch['pose_3d'].cuda()
                gt_keypoints_3d = gt_keypoints_3d[:, joint_mapper_gt, :-1]
            # For 3DPW get the 14 common joints from the rendered shape
            if '3dpw' in dataset_name:
                gt_vertices = smpl_male(global_orient=gt_pose[:,:3], body_pose=gt_pose[:,3:], betas=gt_betas).vertices
                gt_vertices_female = smpl_female(global_orient=gt_pose[:,:3], body_pose=gt_pose[:,3:], betas=gt_betas).vertices
                gt_vertices[gender==1, :, :] = gt_vertices_female[gender==1, :, :]
                gt_keypoints_3d = torch.matmul(J_regressor_batch, gt_vertices)
                gt_pelvis = gt_keypoints_3d[:, [0],:].clone()
                gt_keypoints_3d = gt_keypoints_3d[:, joint_mapper_h36m, :]
                gt_keypoints_3d = gt_keypoints_3d - gt_pelvis
            if 'rich' in dataset_name:
                gt_keypoints_3d = torch.matmul(J_regressor_batch, gt_vertices)
                gt_pelvis = gt_keypoints_3d[:, [0], :].clone()
                gt_keypoints_3d = gt_keypoints_3d[:, joint_mapper_h36m, :]
                gt_keypoints_3d = gt_keypoints_3d - gt_pelvis
                gt_vertices = gt_vertices - gt_pelvis

                # Get 14 predicted joints from the mesh
            pred_keypoints_3d = torch.matmul(J_regressor_batch, pred_vertices)
            if save_results:
                pred_joints[step * batch_size:step * batch_size + curr_batch_size, :, :]  = pred_keypoints_3d.cpu().numpy()
            pred_pelvis = pred_keypoints_3d[:, [0],:].clone()
            pred_keypoints_3d = pred_keypoints_3d[:, joint_mapper_h36m, :]
            pred_keypoints_3d = pred_keypoints_3d - pred_pelvis
            pred_vertices = pred_vertices - pred_pelvis

            # Absolute error (MPJPE)
            error = torch.sqrt(((pred_keypoints_3d - gt_keypoints_3d) ** 2).sum(dim=-1)).mean(dim=-1).cpu().numpy()
            mpjpe[step * batch_size:step * batch_size + curr_batch_size] = error
            for err, imgname in zip(error, imgnames):
                class_name = imgname.split('/')[-2].split('.')[0].split('_')[0]
                classwise_mpjpe.setdefault(class_name, []).append(err)

            # Reconstuction_error
            r_error = reconstruction_error(pred_keypoints_3d.cpu().numpy(), gt_keypoints_3d.cpu().numpy(), reduction=None)
            recon_err[step * batch_size:step * batch_size + curr_batch_size] = r_error

            # V2V error in RICH
            if 'rich' in dataset_name:
                v2v = torch.sqrt(((gt_vertices - pred_vertices) ** 2).sum(dim=-1)).mean(dim=-1).cpu().numpy()
                v2v_err[step * batch_size:step * batch_size + curr_batch_size] = v2v

            # Acceleration Error
            if pred_keypoints_3d.shape[0] >= 3: # can only calculate accel if there are at least 3 frames
                accel = np.mean(compute_accel(pred_keypoints_3d.cpu().numpy()))
                if not np.isnan(accel):
                    accel_[step * batch_size:step * batch_size + curr_batch_size] = accel
                accel_err = np.mean(compute_error_accel(joints_pred=pred_keypoints_3d.cpu().numpy(), joints_gt=gt_keypoints_3d.cpu().numpy()))
                if not np.isnan(accel_err):
                    accel_err_[step * batch_size:step * batch_size + curr_batch_size] = accel_err

        # If mask or part evaluation, render the mask and part images
        # if eval_masks or eval_parts:
        #     mask, parts = renderer(pred_vertices, pred_camera)

        # Mask evaluation (for LSP)
        if eval_masks:
            center = batch['center'].cpu().numpy()
            scale = batch['scale'].cpu().numpy()
            # Dimensions of original image
            orig_shape = batch['orig_shape'].cpu().numpy()
            for i in range(curr_batch_size):
                # After rendering, convert imate back to original resolution
                pred_mask = uncrop(mask[i].cpu().numpy(), center[i], scale[i], orig_shape[i]) > 0
                # Load gt mask
                gt_mask = cv2.imread(os.path.join(annot_path, batch['maskname'][i]), 0) > 0
                # Evaluation consistent with the original UP-3D code
                accuracy += (gt_mask == pred_mask).sum()
                pixel_count += np.prod(np.array(gt_mask.shape))
                for c in range(2):
                    cgt = gt_mask == c
                    cpred = pred_mask == c
                    tp[c] += (cgt & cpred).sum()
                    fp[c] +=  (~cgt & cpred).sum()
                    fn[c] +=  (cgt & ~cpred).sum()
                f1 = 2 * tp / (2 * tp + fp + fn)

        # Part evaluation (for LSP)
        if eval_parts:
            center = batch['center'].cpu().numpy()
            scale = batch['scale'].cpu().numpy()
            orig_shape = batch['orig_shape'].cpu().numpy()
            for i in range(curr_batch_size):
                pred_parts = uncrop(parts[i].cpu().numpy().astype(np.uint8), center[i], scale[i], orig_shape[i])
                # Load gt part segmentation
                gt_parts = cv2.imread(os.path.join(annot_path, batch['partname'][i]), 0)
                # Evaluation consistent with the original UP-3D code
                # 6 parts + background
                for c in range(7):
                    cgt = gt_parts == c
                    cpred = pred_parts == c
                    cpred[gt_parts == 255] = 0
                    parts_tp[c] += (cgt & cpred).sum()
                    parts_fp[c] +=  (~cgt & cpred).sum()
                    parts_fn[c] +=  (cgt & ~cpred).sum()
                gt_parts[gt_parts == 255] = 0
                pred_parts[pred_parts == 255] = 0
                parts_f1 = 2 * parts_tp / (2 * parts_tp + parts_fp + parts_fn)
                parts_accuracy += (gt_parts == pred_parts).sum()
                parts_pixel_count += np.prod(np.array(gt_parts.shape))

        # # Print intermediate results during evaluation
        # if step % log_freq == log_freq - 1:
        #     if eval_pose:
        #         print('MPJPE: ' + str(1000 * mpjpe[:step * batch_size].mean()))
        #         print('Reconstruction Error: ' + str(1000 * recon_err[:step * batch_size].mean()))
        #         print('V2V Error: ' + str(1000 * v2v_err[:step * batch_size].mean()))
        #         print()
        #     if eval_masks:
        #         print('Accuracy: ', accuracy / pixel_count)
        #         print('F1: ', f1.mean())
        #         print()
        #     if eval_parts:
        #         print('Parts Accuracy: ', parts_accuracy / parts_pixel_count)
        #         print('Parts F1 (BG): ', parts_f1[[0,1,2,3,4,5,6]].mean())
        #         print()

    # find accuracy of bos metric
    pred_bos_accumulator = torch.cat(pred_bos_accumulator, dim=0).to(torch.float32)
    pred_contact_accumulator = torch.cat(pred_contact_accumulator, dim=0).to(torch.float32)
    pred_contact_mask_accumulator = torch.cat(pred_contact_mask_accumulator, dim=0).to(torch.float32)
    print('*** Stability Metrics ***')
    print()
    if has_smpl_world:
        gt_bos_accumulator = torch.cat(gt_bos_accumulator, dim=0)
        pred_bos_accuracy = (gt_bos_accumulator == pred_bos_accumulator).sum().item() / gt_bos_accumulator.shape[0]
        gt_contact_accumulator = torch.cat(gt_contact_accumulator, dim=0)
        gt_contact_mask_accumulator = torch.cat(gt_contact_mask_accumulator, dim=0)
        pred_contact_accuracy = (gt_contact_accumulator == pred_contact_accumulator).sum().item() / gt_contact_accumulator.shape[0]
        dense_contact_accuracy = (gt_contact_mask_accumulator == pred_contact_mask_accumulator).sum().item() / \
                                 gt_contact_mask_accumulator.numel()
        print('BOS Accuracy: ', pred_bos_accuracy)
        print('Contact Accuracy: ', pred_contact_accuracy)
        print(f'Contact Mask Accuracy: ', dense_contact_accuracy)
    else:
        bos_mean = torch.mean(pred_bos_accumulator)
        contact_mean = torch.mean(pred_contact_accumulator)
        print('BOS Mean: ', bos_mean.cpu().numpy())
        print('Contact Mean: ', contact_mean.cpu().numpy())

    # Save reconstructions to a file for further processing
    if save_results:
        np.savez(result_file, pred_joints=pred_joints, pose=smpl_pose, betas=smpl_betas, camera=smpl_camera)
    # Print final results during evaluation
    print('*** Final Results ***')
    print()
    if eval_pose:
        print('MPJPE: ' + str(1000 * mpjpe.mean()))
        print('Reconstruction Error: ' + str(1000 * recon_err.mean()))
        print('V2V Error: ' + str(1000 * v2v_err.mean()))
        print('---Classwise MPJPE---')
        for k, v in classwise_mpjpe.items():
            classwise_mpjpe[k] = str(1000 * np.array(v).mean())
        pprint.pprint(classwise_mpjpe)
        print()
    if eval_masks:
        print('Accuracy: ', accuracy / pixel_count)
        print('F1: ', f1.mean())
        print()
    if eval_parts:
        print('Parts Accuracy: ', parts_accuracy / parts_pixel_count)
        print('Parts F1 (BG): ', parts_f1[[0,1,2,3,4,5,6]].mean())
        print()

if __name__ == '__main__':
    args = parser.parse_args()
    model = hmr(config.SMPL_MEAN_PARAMS)
    checkpoint = torch.load(args.checkpoint)
    model.load_state_dict(checkpoint['model'], strict=False)
    model.eval()

    # Setup evaluation dataset
    dataset = BaseDatasetEval(None, args.dataset, is_train=False)
    # Run evaluation
    run_evaluation(model, args.dataset, dataset, args.result_file, args.vis_path,
                   batch_size=args.batch_size,
                   shuffle=args.shuffle,
                   log_freq=args.log_freq,
                   num_workers=args.num_workers,
                   visualize=args.visualize)

eval_h36m_per_class.py
sha2nkt/ipman-r
import torch
import numpy as np
import smplx
from smplx import SMPL as _SMPL
from smplx.body_models import ModelOutput
from smplx.lbs import vertices2joints

import config
import constants

class SMPL(_SMPL):
    """ Extension of the official SMPL implementation to support more joints """

    def __init__(self, *args, **kwargs):
        super(SMPL, self).__init__(*args, **kwargs)
        #joints = [constants.JOINT_MAP[i] for i in constants.JOINT_NAMES]
        #J_regressor_extra = np.load(config.JOINT_REGRESSOR_TRAIN_EXTRA)
        #self.register_buffer('J_regressor_extra', torch.tensor(J_regressor_extra, dtype=torch.float32))
        #self.joint_map = torch.tensor(joints, dtype=torch.long)

    def forward(self, *args, **kwargs):
        kwargs['get_skin'] = True
        smpl_output = super(SMPL, self).forward(*args, **kwargs)
        #extra_joints = vertices2joints(self.J_regressor_extra, smpl_output.vertices)
        #joints = torch.cat([smpl_output.joints, extra_joints], dim=1)
        #joints = joints[:, self.joint_map, :]
        output = ModelOutput(vertices=smpl_output.vertices,
                             global_orient=smpl_output.global_orient,
                             body_pose=smpl_output.body_pose,
                             joints=smpl_output.joints,
                             betas=smpl_output.betas,
                             full_pose=smpl_output.full_pose)
        
        return output

models/smpl_orig.py
sha2nkt/ipman-r
import numpy as np
import constants
from utils.geometry import perspective_projection
import torch

def get_2d_joints(joints_3d, pred_cam_t, focal_length,dev, batch_size ):
    rotation = torch.eye(3, device=dev).unsqueeze(0).expand(batch_size, -1, -1)
    camera_center = torch.zeros(batch_size, 2, device=dev)
    print('cam)t',pred_cam_t)
    print('rotaion',rotation)

    pred_keypoints_2d = perspective_projection(joints_3d,
                                               rotation=rotation,
                                               translation=pred_cam_t,
                                               focal_length=focal_length,
                                               camera_center=camera_center)

    return pred_keypoints_2d

def get_original(proc_parm, verts, cam, joints):
    img_size = proc_parm['img_size'][0]
    # undo_scale = np.array(proc_parm['scale'])#1. / np.array(proc_parm['scale'])
    undo_scale = img_size / constants.IMG_RES

    cam_pos = cam[1:]
    principal_pt = np.array([img_size, img_size]) / 2.
    flength = constants.FOCAL_LENGTH
    # tz = flength / (0.5 * constants.IMG_RES * cam_s)
    tz = (2 * flength) / (constants.IMG_RES * cam[0] + 1e-9)
    trans = np.hstack([cam_pos, tz])
    vert_shifted = (verts + trans)

    start_pt = proc_parm['start_pt']
    final_principal_pt = principal_pt+ start_pt
    cam_for_render = np.hstack(
        [np.mean(flength * undo_scale), final_principal_pt[1], final_principal_pt[0]])

    # This is in padded image.
    # kp_original = (joints + proc_param['start_pt']) * undo_scale
    # Subtract padding from joints.
    # import ipdb; ipdb.set_trace()
    # kp_original = (joints + proc_parm['start_pt']) * undo_scale
    kp_original = (joints) * undo_scale + proc_parm['start_pt'][::-1]- (proc_parm['pad_x'][0],proc_parm['pad_y'][0])

    return cam_for_render, vert_shifted, kp_original

# def get_original(proc_parm, verts, cam, joints):
#
#     img_size = np.array(224)
#     undo_scale = 1.
#     print('und scale',undo_scale)
#    # undo_scale = img_size / constants.IMG_RES
#
#
#     cam_pos = cam[1:]
#     principal_pt = np.array([img_size, img_size]) / 2.
#     flength = 500
#     # tz = flength / (0.5 * constants.IMG_RES * cam_s)
#     tz = (2 * flength) / (constants.IMG_RES * cam[ 0] + 1e-9)
#     trans = np.hstack([cam_pos, tz])
#     vert_shifted = (verts + trans)
#     print('img size', img_size)
#     start_pt = proc_parm['start_pt']
#     print('start pt', proc_parm['start_pt'] )
#     print('principal pt', principal_pt)
#     final_principal_pt = (principal_pt  + start_pt)*undo_scale
#     cam_for_render = np.hstack(
#         [np.mean(flength * undo_scale), 1280/2,720/2])
#
#     # This is in padded image.
#     # kp_original = (joints + proc_param['start_pt']) * undo_scale
#     # Subtract padding from joints.
#     # import ipdb; ipdb.set_trace()
#     # kp_original = (joints + proc_parm['start_pt']) * undo_scale
#
#     kp_original = (joints ) * undo_scale + proc_parm['start_pt'][::-1] - (proc_parm['pad_x'][0],proc_parm['pad_y'][0])
#
#     return cam_for_render, vert_shifted, kp_original

render_original.py
sha2nkt/ipman-r
# take a human3.6 file and visualize skeletons in world coordinates

import argparse
import numpy as np
import matplotlib.pyplot as plt
from datasets.preprocess.rich import rich_extract
import os
# for cluster rendering
os.environ['PYOPENGL_PLATFORM'] = 'egl'

if __name__ == '__main__':
    # load argparse and take data_path as input
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_path', type=str, default='/ps/scratch/ps_shared/stripathi/frmKocabas/human36m',
                        help='path to the h36m dataset')
    parser.add_argument('--split', type=str, default='val',
                        help='test or val split')
    parser.add_argument('--out_path', type=str, default='data/dataset_extras/',
                        help='path to the h36m dataset')
    parser.add_argument('--vis_path', type=str, default='data/debug_vis/',
                        help='path to the data where visualizations are stored')
    parser.add_argument('--visualize', default=False, action='store_true',
                        help='generate visualization?')
    args = parser.parse_args()

    rich_extract(dataset_path=args.data_path,
                       split=args.split,
                       out_path=args.out_path,
                       vis_path=args.vis_path,
                       visualize=args.visualize)


rich_world_vis.py
sha2nkt/ipman-r
import torch
import os

from models.smpl import SMPL
from .losses import camera_fitting_loss, body_fitting_loss
import config
import constants

# For the GMM prior, we use the GMM implementation of SMPLify-X
# https://github.com/vchoutas/smplify-x/blob/master/smplifyx/prior.py
from .prior import MaxMixturePrior

class SMPLify():
    """Implementation of single-stage SMPLify.""" 
    def __init__(self, 
                 step_size=1e-2,
                 batch_size=66,
                 num_iters=100,
                 focal_length=5000,
                 device=torch.device('cuda')):

        # Store options
        self.device = device
        self.focal_length = focal_length
        self.step_size = step_size

        # Ignore the the following joints for the fitting process
        ign_joints = ['OP Neck', 'OP RHip', 'OP LHip', 'Right Hip', 'Left Hip']
        self.ign_joints = [constants.JOINT_IDS[i] for i in ign_joints]
        self.num_iters = num_iters
        # GMM pose prior
        self.pose_prior = MaxMixturePrior(prior_folder='/is/cluster/scratch/stripathi/pycharm_remote/ipman_regr/data',
                                          num_gaussians=8,
                                          dtype=torch.float32).to(device)
        # Load SMPL model
        self.smpl = SMPL(config.SMPL_MODEL_DIR,
                         batch_size=batch_size,
                         create_transl=False).to(self.device)

    def __call__(self, init_pose, init_betas, init_cam_t, camera_center, keypoints_2d):
        """Perform body fitting.
        Input:
            init_pose: SMPL pose estimate
            init_betas: SMPL betas estimate
            init_cam_t: Camera translation estimate
            camera_center: Camera center location
            keypoints_2d: Keypoints used for the optimization
        Returns:
            vertices: Vertices of optimized shape
            joints: 3D joints of optimized shape
            pose: SMPL pose parameters of optimized shape
            betas: SMPL beta parameters of optimized shape
            camera_translation: Camera translation
            reprojection_loss: Final joint reprojection loss
        """

        batch_size = init_pose.shape[0]

        # Make camera translation a learnable parameter
        camera_translation = init_cam_t.clone()

        # Get joint confidence
        joints_2d = keypoints_2d[:, :, :2]
        joints_conf = keypoints_2d[:, :, -1]

        # Split SMPL pose to body pose and global orientation
        body_pose = init_pose[:, 3:].detach().clone()
        global_orient = init_pose[:, :3].detach().clone()
        betas = init_betas.detach().clone()

        # Step 1: Optimize camera translation and body orientation
        # Optimize only camera translation and body orientation
        body_pose.requires_grad=False
        betas.requires_grad=False
        global_orient.requires_grad=True
        camera_translation.requires_grad = True

        camera_opt_params = [global_orient, camera_translation]
        camera_optimizer = torch.optim.Adam(camera_opt_params, lr=self.step_size, betas=(0.9, 0.999))

        for i in range(self.num_iters):
            smpl_output = self.smpl(global_orient=global_orient,
                                    body_pose=body_pose,
                                    betas=betas)
            model_joints = smpl_output.joints
            loss = camera_fitting_loss(model_joints, camera_translation,
                                       init_cam_t, camera_center,
                                       joints_2d, joints_conf, focal_length=self.focal_length)
            camera_optimizer.zero_grad()
            loss.backward()
            camera_optimizer.step()

        # Fix camera translation after optimizing camera
        camera_translation.requires_grad = False

        # Step 2: Optimize body joints
        # Optimize only the body pose and global orientation of the body
        body_pose.requires_grad=True
        betas.requires_grad=True
        global_orient.requires_grad=True
        camera_translation.requires_grad = False
        body_opt_params = [body_pose, betas, global_orient]

        # For joints ignored during fitting, set the confidence to 0
        joints_conf[:, self.ign_joints] = 0.
        
        body_optimizer = torch.optim.Adam(body_opt_params, lr=self.step_size, betas=(0.9, 0.999))
        for i in range(self.num_iters):
            smpl_output = self.smpl(global_orient=global_orient,
                                    body_pose=body_pose,
                                    betas=betas)
            model_joints = smpl_output.joints
            loss = body_fitting_loss(body_pose, betas, model_joints, camera_translation, camera_center,
                                     joints_2d, joints_conf, self.pose_prior,
                                     focal_length=self.focal_length)
            body_optimizer.zero_grad()
            loss.backward()
            body_optimizer.step()

        # Get final loss value
        with torch.no_grad():
            smpl_output = self.smpl(global_orient=global_orient,
                                    body_pose=body_pose,
                                    betas=betas, return_full_pose=True)
            model_joints = smpl_output.joints
            reprojection_loss = body_fitting_loss(body_pose, betas, model_joints, camera_translation, camera_center,
                                                  joints_2d, joints_conf, self.pose_prior,
                                                  focal_length=self.focal_length,
                                                  output='reprojection')

        vertices = smpl_output.vertices.detach()
        joints = smpl_output.joints.detach()
        pose = torch.cat([global_orient, body_pose], dim=-1).detach()
        betas = betas.detach()

        return vertices, joints, pose, betas, camera_translation, reprojection_loss

    def get_fitting_loss(self, pose, betas, cam_t, camera_center, keypoints_2d):
        """Given body and camera parameters, compute reprojection loss value.
        Input:
            pose: SMPL pose parameters
            betas: SMPL beta parameters
            cam_t: Camera translation
            camera_center: Camera center location
            keypoints_2d: Keypoints used for the optimization
        Returns:
            reprojection_loss: Final joint reprojection loss
        """

        batch_size = pose.shape[0]

        # Get joint confidence
        joints_2d = keypoints_2d[:, :, :2]
        joints_conf = keypoints_2d[:, :, -1]
        # For joints ignored during fitting, set the confidence to 0
        joints_conf[:, self.ign_joints] = 0.

        # Split SMPL pose to body pose and global orientation
        body_pose = pose[:, 3:]
        global_orient = pose[:, :3]

        with torch.no_grad():
            smpl_output = self.smpl(global_orient=global_orient,
                                    body_pose=body_pose,
                                    betas=betas, return_full_pose=True)
            model_joints = smpl_output.joints
            reprojection_loss = body_fitting_loss(body_pose, betas, model_joints, cam_t, camera_center,
                                                  joints_2d, joints_conf, self.pose_prior,
                                                  focal_length=self.focal_length,
                                                  output='reprojection')

        return reprojection_loss

smplify/smplify.py
sha2nkt/ipman-r
# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de

import os.path as osp
import argparse

import numpy as np
import torch

import smplx


def main(model_folder, model_type='smplx', ext='npz',
         gender='neutral', plot_joints=False,
         plotting_module='pyrender',
         use_face_contour=False):

    model = smplx.create(model_folder, model_type=model_type,
                         gender=gender, use_face_contour=use_face_contour,
                         ext=ext)
    print(model)

    betas = torch.randn([1, 10], dtype=torch.float32)
    expression = torch.randn([1, 10], dtype=torch.float32)

    output = model(betas=betas, expression=expression,
                   return_verts=True)
    vertices = output.vertices.detach().cpu().numpy().squeeze()
    joints = output.joints.detach().cpu().numpy().squeeze()

    print('Vertices shape =', vertices.shape)
    print('Joints shape =', joints.shape)

    if plotting_module == 'pyrender':
        import pyrender
        import trimesh
        vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8]
        tri_mesh = trimesh.Trimesh(vertices, model.faces,
                                   vertex_colors=vertex_colors)

        mesh = pyrender.Mesh.from_trimesh(tri_mesh)

        scene = pyrender.Scene()
        scene.add(mesh)

        if plot_joints:
            sm = trimesh.creation.uv_sphere(radius=0.005)
            sm.visual.vertex_colors = [0.9, 0.1, 0.1, 1.0]
            tfs = np.tile(np.eye(4), (len(joints), 1, 1))
            tfs[:, :3, 3] = joints
            joints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
            scene.add(joints_pcl)

        pyrender.Viewer(scene, use_raymond_lighting=True)
    elif plotting_module == 'matplotlib':
        from matplotlib import pyplot as plt
        from mpl_toolkits.mplot3d import Axes3D
        from mpl_toolkits.mplot3d.art3d import Poly3DCollection

        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')

        mesh = Poly3DCollection(vertices[model.faces], alpha=0.1)
        face_color = (1.0, 1.0, 0.9)
        edge_color = (0, 0, 0)
        mesh.set_edgecolor(edge_color)
        mesh.set_facecolor(face_color)
        ax.add_collection3d(mesh)
        ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], color='r')

        if plot_joints:
            ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], alpha=0.1)
        plt.show()
    elif plotting_module == 'open3d':
        import open3d as o3d

        mesh = o3d.geometry.TriangleMesh()
        mesh.vertices = o3d.utility.Vector3dVector(
            vertices)
        mesh.triangles = o3d.utility.Vector3iVector(model.faces)
        mesh.compute_vertex_normals()
        mesh.paint_uniform_color([0.3, 0.3, 0.3])

        o3d.visualization.draw_geometries([mesh])
    else:
        raise ValueError('Unknown plotting_module: {}'.format(plotting_module))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='SMPL-X Demo')

    parser.add_argument('--model-folder', required=True, type=str,
                        help='The path to the model folder')
    parser.add_argument('--model-type', default='smplx', type=str,
                        choices=['smpl', 'smplh', 'smplx'],
                        help='The type of model to load')
    parser.add_argument('--gender', type=str, default='neutral',
                        help='The gender of the model')
    parser.add_argument('--plotting-module', type=str, default='pyrender',
                        dest='plotting_module',
                        choices=['pyrender', 'matplotlib', 'open3d'],
                        help='The module to use for plotting the result')
    parser.add_argument('--ext', type=str, default='npz',
                        help='Which extension to use for loading')
    parser.add_argument('--plot-joints', default=False,
                        type=lambda arg: arg.lower() in ['true', '1'],
                        help='The path to the model folder')
    parser.add_argument('--use-face-contour', default=False,
                        type=lambda arg: arg.lower() in ['true', '1'],
                        help='Compute the contour of the face')

    args = parser.parse_args()

    model_folder = osp.expanduser(osp.expandvars(args.model_folder))
    model_type = args.model_type
    plot_joints = args.plot_joints
    use_face_contour = args.use_face_contour
    gender = args.gender
    ext = args.ext
    plotting_module = args.plotting_module

    main(model_folder, model_type, ext=ext,
         gender=gender, plot_joints=plot_joints,
         plotting_module=plotting_module,
         use_face_contour=use_face_contour)

smplx/examples/demo.py
sha2nkt/ipman-r
# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import os
import os.path as osp

try:
    import cPickle as pickle
except ImportError:
    import pickle

import numpy as np

from collections import namedtuple

import torch
import torch.nn as nn

from .lbs import (
    lbs, vertices2landmarks, find_dynamic_lmk_idx_and_bcoords)

from .vertex_ids import vertex_ids as VERTEX_IDS
from .utils import Struct, to_np, to_tensor
from .vertex_joint_selector import VertexJointSelector


ModelOutput = namedtuple('ModelOutput',
                         ['vertices', 'joints', 'full_pose', 'betas',
                          'global_orient',
                          'body_pose', 'expression',
                          'left_hand_pose', 'right_hand_pose',
                          'jaw_pose'])
ModelOutput.__new__.__defaults__ = (None,) * len(ModelOutput._fields)


def create(model_path, model_type='smpl',
           **kwargs):
    ''' Method for creating a model from a path and a model type

        Parameters
        ----------
        model_path: str
            Either the path to the model you wish to load or a folder,
            where each subfolder contains the differents types, i.e.:
            model_path:
            |
            |-- smpl
                |-- SMPL_FEMALE
                |-- SMPL_NEUTRAL
                |-- SMPL_MALE
            |-- smplh
                |-- SMPLH_FEMALE
                |-- SMPLH_MALE
            |-- smplx
                |-- SMPLX_FEMALE
                |-- SMPLX_NEUTRAL
                |-- SMPLX_MALE
        model_type: str, optional
            When model_path is a folder, then this parameter specifies  the
            type of model to be loaded
        **kwargs: dict
            Keyword arguments

        Returns
        -------
            body_model: nn.Module
                The PyTorch module that implements the corresponding body model
        Raises
        ------
            ValueError: In case the model type is not one of SMPL, SMPLH or
            SMPLX
    '''

    # If it's a folder, assume
    if osp.isdir(model_path):
        model_path = os.path.join(model_path, model_type)

    if model_type.lower() == 'smpl':
        return SMPL(model_path, **kwargs)
    elif model_type.lower() == 'smplh':
        return SMPLH(model_path, **kwargs)
    elif model_type.lower() == 'smplx':
        return SMPLX(model_path, **kwargs)
    else:
        raise ValueError('Unknown model type {}, exiting!'.format(model_type))


class SMPL(nn.Module):

    NUM_JOINTS = 23
    NUM_BODY_JOINTS = 23
    NUM_BETAS = 10

    def __init__(self, model_path, data_struct=None,
                 create_betas=True,
                 betas=None,
                 create_global_orient=True,
                 global_orient=None,
                 create_body_pose=True,
                 body_pose=None,
                 create_transl=True,
                 transl=None,
                 dtype=torch.float32,
                 batch_size=1,
                 joint_mapper=None, gender='neutral',
                 vertex_ids=None,
                 **kwargs):
        ''' SMPL model constructor

            Parameters
            ----------
            model_path: str
                The path to the folder or to the file where the model
                parameters are stored
            data_struct: Strct
                A struct object. If given, then the parameters of the model are
                read from the object. Otherwise, the model tries to read the
                parameters from the given `model_path`. (default = None)
            create_global_orient: bool, optional
                Flag for creating a member variable for the global orientation
                of the body. (default = True)
            global_orient: torch.tensor, optional, Bx3
                The default value for the global orientation variable.
                (default = None)
            create_body_pose: bool, optional
                Flag for creating a member variable for the pose of the body.
                (default = True)
            body_pose: torch.tensor, optional, Bx(Body Joints * 3)
                The default value for the body pose variable.
                (default = None)
            create_betas: bool, optional
                Flag for creating a member variable for the shape space
                (default = True).
            betas: torch.tensor, optional, Bx10
                The default value for the shape member variable.
                (default = None)
            create_transl: bool, optional
                Flag for creating a member variable for the translation
                of the body. (default = True)
            transl: torch.tensor, optional, Bx3
                The default value for the transl variable.
                (default = None)
            dtype: torch.dtype, optional
                The data type for the created variables
            batch_size: int, optional
                The batch size used for creating the member variables
            joint_mapper: object, optional
                An object that re-maps the joints. Useful if one wants to
                re-order the SMPL joints to some other convention (e.g. MSCOCO)
                (default = None)
            gender: str, optional
                Which gender to load
            vertex_ids: dict, optional
                A dictionary containing the indices of the extra vertices that
                will be selected
        '''

        self.gender = gender

        if data_struct is None:
            if osp.isdir(model_path):
                model_fn = 'SMPL_{}.{ext}'.format(gender.upper(), ext='pkl')
                smpl_path = os.path.join(model_path, model_fn)
            else:
                smpl_path = model_path
            assert osp.exists(smpl_path), 'Path {} does not exist!'.format(
                smpl_path)

            with open(smpl_path, 'rb') as smpl_file:
                data_struct = Struct(**pickle.load(smpl_file,
                                                   encoding='latin1'))

        super(SMPL, self).__init__()
        self.batch_size = batch_size

        if vertex_ids is None:
            # SMPL and SMPL-H share the same topology, so any extra joints can
            # be drawn from the same place
            vertex_ids = VERTEX_IDS['smplh']

        self.dtype = dtype

        self.joint_mapper = joint_mapper

        self.vertex_joint_selector = VertexJointSelector(
            vertex_ids=vertex_ids, **kwargs)

        self.faces = data_struct.f
        self.register_buffer('faces_tensor',
                             to_tensor(to_np(self.faces, dtype=np.int64),
                                       dtype=torch.long))

        if create_betas:
            if betas is None:
                default_betas = torch.zeros([batch_size, self.NUM_BETAS],
                                            dtype=dtype)
            else:
                if 'torch.Tensor' in str(type(betas)):
                    default_betas = betas.clone().detach()
                else:
                    default_betas = torch.tensor(betas,
                                                 dtype=dtype)

            self.register_parameter('betas', nn.Parameter(default_betas,
                                                          requires_grad=True))

        # The tensor that contains the global rotation of the model
        # It is separated from the pose of the joints in case we wish to
        # optimize only over one of them
        if create_global_orient:
            if global_orient is None:
                default_global_orient = torch.zeros([batch_size, 3],
                                                    dtype=dtype)
            else:
                if 'torch.Tensor' in str(type(global_orient)):
                    default_global_orient = global_orient.clone().detach()
                else:
                    default_global_orient = torch.tensor(global_orient,
                                                         dtype=dtype)

            global_orient = nn.Parameter(default_global_orient,
                                         requires_grad=True)
            self.register_parameter('global_orient', global_orient)

        if create_body_pose:
            if body_pose is None:
                default_body_pose = torch.zeros(
                    [batch_size, self.NUM_BODY_JOINTS * 3], dtype=dtype)
            else:
                if 'torch.Tensor' in str(type(body_pose)):
                    default_body_pose = body_pose.clone().detach()
                else:
                    default_body_pose = torch.tensor(body_pose,
                                                     dtype=dtype)
            self.register_parameter(
                'body_pose',
                nn.Parameter(default_body_pose, requires_grad=True))

        if create_transl:
            if transl is None:
                default_transl = torch.zeros([batch_size, 3],
                                             dtype=dtype,
                                             requires_grad=True)
            else:
                default_transl = torch.tensor(transl, dtype=dtype)
            self.register_parameter(
                'transl',
                nn.Parameter(default_transl, requires_grad=True))

        # The vertices of the template model
        self.register_buffer('v_template',
                             to_tensor(to_np(data_struct.v_template),
                                       dtype=dtype))

        # The shape components
        shapedirs = data_struct.shapedirs
        # The shape components
        self.register_buffer(
            'shapedirs',
            to_tensor(to_np(shapedirs), dtype=dtype))

        j_regressor = to_tensor(to_np(
            data_struct.J_regressor), dtype=dtype)
        self.register_buffer('J_regressor', j_regressor)

        # Pose blend shape basis: 6890 x 3 x 207, reshaped to 6890*3 x 207
        num_pose_basis = data_struct.posedirs.shape[-1]
        # 207 x 20670
        posedirs = np.reshape(data_struct.posedirs, [-1, num_pose_basis]).T
        self.register_buffer('posedirs',
                             to_tensor(to_np(posedirs), dtype=dtype))

        # indices of parents for each joints
        parents = to_tensor(to_np(data_struct.kintree_table[0])).long()
        parents[0] = -1
        self.register_buffer('parents', parents)

        self.register_buffer('lbs_weights',
                             to_tensor(to_np(data_struct.weights), dtype=dtype))

    def create_mean_pose(self, data_struct):
        pass

    @torch.no_grad()
    def reset_params(self, **params_dict):
        for param_name, param in self.named_parameters():
            if param_name in params_dict:
                param[:] = torch.tensor(params_dict[param_name])
            else:
                param.fill_(0)

    def get_num_verts(self):
        return self.v_template.shape[0]

    def get_num_faces(self):
        return self.faces.shape[0]

    def extra_repr(self):
        return 'Number of betas: {}'.format(self.NUM_BETAS)

    def forward(self, betas=None, body_pose=None, global_orient=None,
                transl=None, return_verts=True, return_full_pose=False, pose2rot=True,
                **kwargs):
        ''' Forward pass for the SMPL model

            Parameters
            ----------
            global_orient: torch.tensor, optional, shape Bx3
                If given, ignore the member variable and use it as the global
                rotation of the body. Useful if someone wishes to predicts this
                with an external model. (default=None)
            betas: torch.tensor, optional, shape Bx10
                If given, ignore the member variable `betas` and use it
                instead. For example, it can used if shape parameters
                `betas` are predicted from some external model.
                (default=None)
            body_pose: torch.tensor, optional, shape Bx(J*3)
                If given, ignore the member variable `body_pose` and use it
                instead. For example, it can used if someone predicts the
                pose of the body joints are predicted from some external model.
                It should be a tensor that contains joint rotations in
                axis-angle format. (default=None)
            transl: torch.tensor, optional, shape Bx3
                If given, ignore the member variable `transl` and use it
                instead. For example, it can used if the translation
                `transl` is predicted from some external model.
                (default=None)
            return_verts: bool, optional
                Return the vertices. (default=True)
            return_full_pose: bool, optional
                Returns the full axis-angle pose vector (default=False)

            Returns
            -------
        '''
        # If no shape and pose parameters are passed along, then use the
        # ones from the module
        global_orient = (global_orient if global_orient is not None else
                         self.global_orient)
        body_pose = body_pose if body_pose is not None else self.body_pose
        betas = betas if betas is not None else self.betas

        apply_trans = transl is not None or hasattr(self, 'transl')
        if transl is None and hasattr(self, 'transl'):
            transl = self.transl

        full_pose = torch.cat([global_orient, body_pose], dim=1)

        batch_size = max(betas.shape[0], global_orient.shape[0],
                         body_pose.shape[0])

        if betas.shape[0] != batch_size:
            num_repeats = int(batch_size / betas.shape[0])
            betas = betas.expand(num_repeats, -1)

        vertices, joints = lbs(betas, full_pose, self.v_template,
                               self.shapedirs, self.posedirs,
                               self.J_regressor, self.parents,
                               self.lbs_weights, pose2rot=pose2rot, dtype=self.dtype)

        joints = self.vertex_joint_selector(vertices, joints)
        # Map the joints to the current dataset
        if self.joint_mapper is not None:
            joints = self.joint_mapper(joints)

        if apply_trans:
            joints += transl.unsqueeze(dim=1)
            vertices += transl.unsqueeze(dim=1)

        output = ModelOutput(vertices=vertices if return_verts else None,
                             global_orient=global_orient,
                             body_pose=body_pose,
                             joints=joints,
                             betas=betas,
                             full_pose=full_pose if return_full_pose else None)

        return output


class SMPLH(SMPL):

    # The hand joints are replaced by MANO
    NUM_BODY_JOINTS = SMPL.NUM_JOINTS - 2
    NUM_HAND_JOINTS = 15
    NUM_JOINTS = NUM_BODY_JOINTS + 2 * NUM_HAND_JOINTS

    def __init__(self, model_path,
                 data_struct=None,
                 create_left_hand_pose=True,
                 left_hand_pose=None,
                 create_right_hand_pose=True,
                 right_hand_pose=None,
                 use_pca=True,
                 num_pca_comps=6,
                 flat_hand_mean=False,
                 batch_size=1,
                 gender='neutral',
                 dtype=torch.float32,
                 vertex_ids=None,
                 use_compressed=True,
                 ext='pkl',
                 **kwargs):
        ''' SMPLH model constructor

            Parameters
            ----------
            model_path: str
                The path to the folder or to the file where the model
                parameters are stored
            data_struct: Strct
                A struct object. If given, then the parameters of the model are
                read from the object. Otherwise, the model tries to read the
                parameters from the given `model_path`. (default = None)
            create_left_hand_pose: bool, optional
                Flag for creating a member variable for the pose of the left
                hand. (default = True)
            left_hand_pose: torch.tensor, optional, BxP
                The default value for the left hand pose member variable.
                (default = None)
            create_right_hand_pose: bool, optional
                Flag for creating a member variable for the pose of the right
                hand. (default = True)
            right_hand_pose: torch.tensor, optional, BxP
                The default value for the right hand pose member variable.
                (default = None)
            num_pca_comps: int, optional
                The number of PCA components to use for each hand.
                (default = 6)
            flat_hand_mean: bool, optional
                If False, then the pose of the hand is initialized to False.
            batch_size: int, optional
                The batch size used for creating the member variables
            gender: str, optional
                Which gender to load
            dtype: torch.dtype, optional
                The data type for the created variables
            vertex_ids: dict, optional
                A dictionary containing the indices of the extra vertices that
                will be selected
        '''

        self.num_pca_comps = num_pca_comps
        # If no data structure is passed, then load the data from the given
        # model folder
        if data_struct is None:
            # Load the model
            if osp.isdir(model_path):
                model_fn = 'SMPLH_{}.{ext}'.format(gender.upper(), ext=ext)
                smplh_path = os.path.join(model_path, model_fn)
            else:
                smplh_path = model_path
            assert osp.exists(smplh_path), 'Path {} does not exist!'.format(
                smplh_path)

            if ext == 'pkl':
                with open(smplh_path, 'rb') as smplh_file:
                    model_data = pickle.load(smplh_file, encoding='latin1')
            elif ext == 'npz':
                model_data = np.load(smplh_path, allow_pickle=True)
            else:
                raise ValueError('Unknown extension: {}'.format(ext))
            data_struct = Struct(**model_data)

        if vertex_ids is None:
            vertex_ids = VERTEX_IDS['smplh']

        super(SMPLH, self).__init__(
            model_path=model_path, data_struct=data_struct,
            batch_size=batch_size, vertex_ids=vertex_ids, gender=gender,
            use_compressed=use_compressed, dtype=dtype, ext=ext, **kwargs)

        self.use_pca = use_pca
        self.num_pca_comps = num_pca_comps
        self.flat_hand_mean = flat_hand_mean

        left_hand_components = data_struct.hands_componentsl[:num_pca_comps]
        right_hand_components = data_struct.hands_componentsr[:num_pca_comps]

        self.np_left_hand_components = left_hand_components
        self.np_right_hand_components = right_hand_components
        if self.use_pca:
            self.register_buffer(
                'left_hand_components',
                torch.tensor(left_hand_components, dtype=dtype))
            self.register_buffer(
                'right_hand_components',
                torch.tensor(right_hand_components, dtype=dtype))

        if self.flat_hand_mean:
            left_hand_mean = np.zeros_like(data_struct.hands_meanl)
        else:
            left_hand_mean = data_struct.hands_meanl

        if self.flat_hand_mean:
            right_hand_mean = np.zeros_like(data_struct.hands_meanr)
        else:
            right_hand_mean = data_struct.hands_meanr

        self.register_buffer('left_hand_mean',
                             to_tensor(left_hand_mean, dtype=self.dtype))
        self.register_buffer('right_hand_mean',
                             to_tensor(right_hand_mean, dtype=self.dtype))

        # Create the buffers for the pose of the left hand
        hand_pose_dim = num_pca_comps if use_pca else 3 * self.NUM_HAND_JOINTS
        if create_left_hand_pose:
            if left_hand_pose is None:
                default_lhand_pose = torch.zeros([batch_size, hand_pose_dim],
                                                 dtype=dtype)
            else:
                default_lhand_pose = torch.tensor(left_hand_pose, dtype=dtype)

            left_hand_pose_param = nn.Parameter(default_lhand_pose,
                                                requires_grad=True)
            self.register_parameter('left_hand_pose',
                                    left_hand_pose_param)

        if create_right_hand_pose:
            if right_hand_pose is None:
                default_rhand_pose = torch.zeros([batch_size, hand_pose_dim],
                                                 dtype=dtype)
            else:
                default_rhand_pose = torch.tensor(right_hand_pose, dtype=dtype)

            right_hand_pose_param = nn.Parameter(default_rhand_pose,
                                                 requires_grad=True)
            self.register_parameter('right_hand_pose',
                                    right_hand_pose_param)

        # Create the buffer for the mean pose.
        pose_mean = self.create_mean_pose(data_struct,
                                          flat_hand_mean=flat_hand_mean)
        pose_mean_tensor = torch.tensor(pose_mean, dtype=dtype)
        self.register_buffer('pose_mean', pose_mean_tensor)

    def create_mean_pose(self, data_struct, flat_hand_mean=False):
        # Create the array for the mean pose. If flat_hand is false, then use
        # the mean that is given by the data, rather than the flat open hand
        global_orient_mean = torch.zeros([3], dtype=self.dtype)
        body_pose_mean = torch.zeros([self.NUM_BODY_JOINTS * 3],
                                     dtype=self.dtype)

        pose_mean = torch.cat([global_orient_mean, body_pose_mean,
                               self.left_hand_mean,
                               self.right_hand_mean], dim=0)
        return pose_mean

    def extra_repr(self):
        msg = super(SMPLH, self).extra_repr()
        if self.use_pca:
            msg += '\nNumber of PCA components: {}'.format(self.num_pca_comps)
        msg += '\nFlat hand mean: {}'.format(self.flat_hand_mean)
        return msg

    def forward(self, betas=None, global_orient=None, body_pose=None,
                left_hand_pose=None, right_hand_pose=None, transl=None,
                return_verts=True, return_full_pose=False, pose2rot=True,
                **kwargs):
        '''
        '''
        # If no shape and pose parameters are passed along, then use the
        # ones from the module
        global_orient = (global_orient if global_orient is not None else
                         self.global_orient)
        body_pose = body_pose if body_pose is not None else self.body_pose
        betas = betas if betas is not None else self.betas
        left_hand_pose = (left_hand_pose if left_hand_pose is not None else
                          self.left_hand_pose)
        right_hand_pose = (right_hand_pose if right_hand_pose is not None else
                           self.right_hand_pose)

        apply_trans = transl is not None or hasattr(self, 'transl')
        if transl is None:
            if hasattr(self, 'transl'):
                transl = self.transl

        if self.use_pca:
            left_hand_pose = torch.einsum(
                'bi,ij->bj', [left_hand_pose, self.left_hand_components])
            right_hand_pose = torch.einsum(
                'bi,ij->bj', [right_hand_pose, self.right_hand_components])

        full_pose = torch.cat([global_orient, body_pose,
                               left_hand_pose,
                               right_hand_pose], dim=1)
        full_pose += self.pose_mean

        vertices, joints = lbs(self.betas, full_pose, self.v_template,
                               self.shapedirs, self.posedirs,
                               self.J_regressor, self.parents,
                               self.lbs_weights, pose2rot=pose2rot,
                               dtype=self.dtype)

        # Add any extra joints that might be needed
        joints = self.vertex_joint_selector(vertices, joints)
        if self.joint_mapper is not None:
            joints = self.joint_mapper(joints)

        if apply_trans:
            joints += transl.unsqueeze(dim=1)
            vertices += transl.unsqueeze(dim=1)

        output = ModelOutput(vertices=vertices if return_verts else None,
                             joints=joints,
                             betas=betas,
                             global_orient=global_orient,
                             body_pose=body_pose,
                             left_hand_pose=left_hand_pose,
                             right_hand_pose=right_hand_pose,
                             full_pose=full_pose if return_full_pose else None)

        return output


class SMPLX(SMPLH):
    '''
    SMPL-X (SMPL eXpressive) is a unified body model, with shape parameters
    trained jointly for the face, hands and body.
    SMPL-X uses standard vertex based linear blend skinning with learned
    corrective blend shapes, has N=10475 vertices and K=54 joints,
    which includes joints for the neck, jaw, eyeballs and fingers.
    '''

    NUM_BODY_JOINTS = SMPLH.NUM_BODY_JOINTS
    NUM_HAND_JOINTS = 15
    NUM_FACE_JOINTS = 3
    NUM_JOINTS = NUM_BODY_JOINTS + 2 * NUM_HAND_JOINTS + NUM_FACE_JOINTS
    NUM_EXPR_COEFFS = 10
    NECK_IDX = 12

    def __init__(self, model_path,
                 create_expression=True, expression=None,
                 create_jaw_pose=True, jaw_pose=None,
                 create_leye_pose=True, leye_pose=None,
                 create_reye_pose=True, reye_pose=None,
                 use_face_contour=False,
                 batch_size=1, gender='neutral',
                 dtype=torch.float32,
                 ext='npz',
                 **kwargs):
        ''' SMPLX model constructor

            Parameters
            ----------
            model_path: str
                The path to the folder or to the file where the model
                parameters are stored
            create_expression: bool, optional
                Flag for creating a member variable for the expression space
                (default = True).
            expression: torch.tensor, optional, Bx10
                The default value for the expression member variable.
                (default = None)
            create_jaw_pose: bool, optional
                Flag for creating a member variable for the jaw pose.
                (default = False)
            jaw_pose: torch.tensor, optional, Bx3
                The default value for the jaw pose variable.
                (default = None)
            create_leye_pose: bool, optional
                Flag for creating a member variable for the left eye pose.
                (default = False)
            leye_pose: torch.tensor, optional, Bx10
                The default value for the left eye pose variable.
                (default = None)
            create_reye_pose: bool, optional
                Flag for creating a member variable for the right eye pose.
                (default = False)
            reye_pose: torch.tensor, optional, Bx10
                The default value for the right eye pose variable.
                (default = None)
            use_face_contour: bool, optional
                Whether to compute the keypoints that form the facial contour
            batch_size: int, optional
                The batch size used for creating the member variables
            gender: str, optional
                Which gender to load
            dtype: torch.dtype
                The data type for the created variables
        '''

        # Load the model
        if osp.isdir(model_path):
            model_fn = 'SMPLX_{}.{ext}'.format(gender.upper(), ext=ext)
            smplx_path = os.path.join(model_path, model_fn)
        else:
            smplx_path = model_path
        assert osp.exists(smplx_path), 'Path {} does not exist!'.format(
            smplx_path)

        if ext == 'pkl':
            with open(smplx_path, 'rb') as smplx_file:
                model_data = pickle.load(smplx_file, encoding='latin1')
        elif ext == 'npz':
            model_data = np.load(smplx_path, allow_pickle=True)
        else:
            raise ValueError('Unknown extension: {}'.format(ext))

        data_struct = Struct(**model_data)

        super(SMPLX, self).__init__(
            model_path=model_path,
            data_struct=data_struct,
            dtype=dtype,
            batch_size=batch_size,
            vertex_ids=VERTEX_IDS['smplx'],
            gender=gender, ext=ext,
            **kwargs)

        lmk_faces_idx = data_struct.lmk_faces_idx
        self.register_buffer('lmk_faces_idx',
                             torch.tensor(lmk_faces_idx, dtype=torch.long))
        lmk_bary_coords = data_struct.lmk_bary_coords
        self.register_buffer('lmk_bary_coords',
                             torch.tensor(lmk_bary_coords, dtype=dtype))

        self.use_face_contour = use_face_contour
        if self.use_face_contour:
            dynamic_lmk_faces_idx = data_struct.dynamic_lmk_faces_idx
            dynamic_lmk_faces_idx = torch.tensor(
                dynamic_lmk_faces_idx,
                dtype=torch.long)
            self.register_buffer('dynamic_lmk_faces_idx',
                                 dynamic_lmk_faces_idx)

            dynamic_lmk_bary_coords = data_struct.dynamic_lmk_bary_coords
            dynamic_lmk_bary_coords = torch.tensor(
                dynamic_lmk_bary_coords, dtype=dtype)
            self.register_buffer('dynamic_lmk_bary_coords',
                                 dynamic_lmk_bary_coords)

            neck_kin_chain = []
            curr_idx = torch.tensor(self.NECK_IDX, dtype=torch.long)
            while curr_idx != -1:
                neck_kin_chain.append(curr_idx)
                curr_idx = self.parents[curr_idx]
            self.register_buffer('neck_kin_chain',
                                 torch.stack(neck_kin_chain))

        if create_jaw_pose:
            if jaw_pose is None:
                default_jaw_pose = torch.zeros([batch_size, 3], dtype=dtype)
            else:
                default_jaw_pose = torch.tensor(jaw_pose, dtype=dtype)
            jaw_pose_param = nn.Parameter(default_jaw_pose,
                                          requires_grad=True)
            self.register_parameter('jaw_pose', jaw_pose_param)

        if create_leye_pose:
            if leye_pose is None:
                default_leye_pose = torch.zeros([batch_size, 3], dtype=dtype)
            else:
                default_leye_pose = torch.tensor(leye_pose, dtype=dtype)
            leye_pose_param = nn.Parameter(default_leye_pose,
                                           requires_grad=True)
            self.register_parameter('leye_pose', leye_pose_param)

        if create_reye_pose:
            if reye_pose is None:
                default_reye_pose = torch.zeros([batch_size, 3], dtype=dtype)
            else:
                default_reye_pose = torch.tensor(reye_pose, dtype=dtype)
            reye_pose_param = nn.Parameter(default_reye_pose,
                                           requires_grad=True)
            self.register_parameter('reye_pose', reye_pose_param)

        if create_expression:
            if expression is None:
                default_expression = torch.zeros(
                    [batch_size, self.NUM_EXPR_COEFFS], dtype=dtype)
            else:
                default_expression = torch.tensor(expression, dtype=dtype)
            expression_param = nn.Parameter(default_expression,
                                            requires_grad=True)
            self.register_parameter('expression', expression_param)

    def create_mean_pose(self, data_struct, flat_hand_mean=False):
        # Create the array for the mean pose. If flat_hand is false, then use
        # the mean that is given by the data, rather than the flat open hand
        global_orient_mean = torch.zeros([3], dtype=self.dtype)
        body_pose_mean = torch.zeros([self.NUM_BODY_JOINTS * 3],
                                     dtype=self.dtype)
        jaw_pose_mean = torch.zeros([3], dtype=self.dtype)
        leye_pose_mean = torch.zeros([3], dtype=self.dtype)
        reye_pose_mean = torch.zeros([3], dtype=self.dtype)

        pose_mean = np.concatenate([global_orient_mean, body_pose_mean,
                                    jaw_pose_mean,
                                    leye_pose_mean, reye_pose_mean,
                                    self.left_hand_mean, self.right_hand_mean],
                                   axis=0)

        return pose_mean

    def extra_repr(self):
        msg = super(SMPLX, self).extra_repr()
        msg += '\nGender: {}'.format(self.gender.title())
        msg += '\nExpression Coefficients: {}'.format(
            self.NUM_EXPR_COEFFS)
        msg += '\nUse face contour: {}'.format(self.use_face_contour)
        return msg

    def forward(self, betas=None, global_orient=None, body_pose=None,
                left_hand_pose=None, right_hand_pose=None, transl=None,
                expression=None, jaw_pose=None, leye_pose=None, reye_pose=None,
                return_verts=True, return_full_pose=False, pose2rot=True, **kwargs):
        '''
        Forward pass for the SMPLX model

            Parameters
            ----------
            global_orient: torch.tensor, optional, shape Bx3
                If given, ignore the member variable and use it as the global
                rotation of the body. Useful if someone wishes to predicts this
                with an external model. (default=None)
            betas: torch.tensor, optional, shape Bx10
                If given, ignore the member variable `betas` and use it
                instead. For example, it can used if shape parameters
                `betas` are predicted from some external model.
                (default=None)
            expression: torch.tensor, optional, shape Bx10
                If given, ignore the member variable `expression` and use it
                instead. For example, it can used if expression parameters
                `expression` are predicted from some external model.
            body_pose: torch.tensor, optional, shape Bx(J*3)
                If given, ignore the member variable `body_pose` and use it
                instead. For example, it can used if someone predicts the
                pose of the body joints are predicted from some external model.
                It should be a tensor that contains joint rotations in
                axis-angle format. (default=None)
            left_hand_pose: torch.tensor, optional, shape BxP
                If given, ignore the member variable `left_hand_pose` and
                use this instead. It should either contain PCA coefficients or
                joint rotations in axis-angle format.
            right_hand_pose: torch.tensor, optional, shape BxP
                If given, ignore the member variable `right_hand_pose` and
                use this instead. It should either contain PCA coefficients or
                joint rotations in axis-angle format.
            jaw_pose: torch.tensor, optional, shape Bx3
                If given, ignore the member variable `jaw_pose` and
                use this instead. It should either joint rotations in
                axis-angle format.
            transl: torch.tensor, optional, shape Bx3
                If given, ignore the member variable `transl` and use it
                instead. For example, it can used if the translation
                `transl` is predicted from some external model.
                (default=None)
            return_verts: bool, optional
                Return the vertices. (default=True)
            return_full_pose: bool, optional
                Returns the full axis-angle pose vector (default=False)

            Returns
            -------
                output: ModelOutput
                A named tuple of type `ModelOutput`
        '''

        # If no shape and pose parameters are passed along, then use the
        # ones from the module
        global_orient = (global_orient if global_orient is not None else
                         self.global_orient)
        body_pose = body_pose if body_pose is not None else self.body_pose
        betas = betas if betas is not None else self.betas

        left_hand_pose = (left_hand_pose if left_hand_pose is not None else
                          self.left_hand_pose)
        right_hand_pose = (right_hand_pose if right_hand_pose is not None else
                           self.right_hand_pose)
        jaw_pose = jaw_pose if jaw_pose is not None else self.jaw_pose
        leye_pose = leye_pose if leye_pose is not None else self.leye_pose
        reye_pose = reye_pose if reye_pose is not None else self.reye_pose
        expression = expression if expression is not None else self.expression

        apply_trans = transl is not None or hasattr(self, 'transl')
        if transl is None:
            if hasattr(self, 'transl'):
                transl = self.transl

        if self.use_pca:
            left_hand_pose = torch.einsum(
                'bi,ij->bj', [left_hand_pose, self.left_hand_components])
            right_hand_pose = torch.einsum(
                'bi,ij->bj', [right_hand_pose, self.right_hand_components])

        full_pose = torch.cat([global_orient, body_pose,
                               jaw_pose, leye_pose, reye_pose,
                               left_hand_pose,
                               right_hand_pose], dim=1)

        # Add the mean pose of the model. Does not affect the body, only the
        # hands when flat_hand_mean == False
        full_pose += self.pose_mean

        batch_size = max(betas.shape[0], global_orient.shape[0],
                         body_pose.shape[0])
        # Concatenate the shape and expression coefficients
        scale = int(batch_size / betas.shape[0])
        if scale > 1:
            betas = betas.expand(scale, -1)
        shape_components = torch.cat([betas, expression], dim=-1)

        vertices, joints = lbs(shape_components, full_pose, self.v_template,
                               self.shapedirs, self.posedirs,
                               self.J_regressor, self.parents,
                               self.lbs_weights, pose2rot=pose2rot,
                               dtype=self.dtype)

        lmk_faces_idx = self.lmk_faces_idx.unsqueeze(
            dim=0).expand(batch_size, -1).contiguous()
        lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat(
            self.batch_size, 1, 1)
        if self.use_face_contour:
            dyn_lmk_faces_idx, dyn_lmk_bary_coords = find_dynamic_lmk_idx_and_bcoords(
                vertices, full_pose, self.dynamic_lmk_faces_idx,
                self.dynamic_lmk_bary_coords,
                self.neck_kin_chain, dtype=self.dtype)

            lmk_faces_idx = torch.cat([lmk_faces_idx,
                                       dyn_lmk_faces_idx], 1)
            lmk_bary_coords = torch.cat(
                [lmk_bary_coords.expand(batch_size, -1, -1),
                 dyn_lmk_bary_coords], 1)

        landmarks = vertices2landmarks(vertices, self.faces_tensor,
                                       lmk_faces_idx,
                                       lmk_bary_coords)

        # Add any extra joints that might be needed
        joints = self.vertex_joint_selector(vertices, joints)
        # Add the landmarks to the joints
        joints = torch.cat([joints, landmarks], dim=1)
        # Map the joints to the current dataset

        if self.joint_mapper is not None:
            joints = self.joint_mapper(joints=joints, vertices=vertices)

        if apply_trans:
            joints += transl.unsqueeze(dim=1)
            vertices += transl.unsqueeze(dim=1)

        output = ModelOutput(vertices=vertices if return_verts else None,
                             joints=joints,
                             betas=betas,
                             expression=expression,
                             global_orient=global_orient,
                             body_pose=body_pose,
                             left_hand_pose=left_hand_pose,
                             right_hand_pose=right_hand_pose,
                             jaw_pose=jaw_pose,
                             full_pose=full_pose if return_full_pose else None)
        return output

smplx/smplx/body_models.py
sha2nkt/ipman-r
# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import numpy as np

import torch
import torch.nn.functional as F

from .utils import rot_mat_to_euler


def find_dynamic_lmk_idx_and_bcoords(vertices, pose, dynamic_lmk_faces_idx,
                                     dynamic_lmk_b_coords,
                                     neck_kin_chain, dtype=torch.float32):
    ''' Compute the faces, barycentric coordinates for the dynamic landmarks


        To do so, we first compute the rotation of the neck around the y-axis
        and then use a pre-computed look-up table to find the faces and the
        barycentric coordinates that will be used.

        Special thanks to Soubhik Sanyal (soubhik.sanyal@tuebingen.mpg.de)
        for providing the original TensorFlow implementation and for the LUT.

        Parameters
        ----------
        vertices: torch.tensor BxVx3, dtype = torch.float32
            The tensor of input vertices
        pose: torch.tensor Bx(Jx3), dtype = torch.float32
            The current pose of the body model
        dynamic_lmk_faces_idx: torch.tensor L, dtype = torch.long
            The look-up table from neck rotation to faces
        dynamic_lmk_b_coords: torch.tensor Lx3, dtype = torch.float32
            The look-up table from neck rotation to barycentric coordinates
        neck_kin_chain: list
            A python list that contains the indices of the joints that form the
            kinematic chain of the neck.
        dtype: torch.dtype, optional

        Returns
        -------
        dyn_lmk_faces_idx: torch.tensor, dtype = torch.long
            A tensor of size BxL that contains the indices of the faces that
            will be used to compute the current dynamic landmarks.
        dyn_lmk_b_coords: torch.tensor, dtype = torch.float32
            A tensor of size BxL that contains the indices of the faces that
            will be used to compute the current dynamic landmarks.
    '''

    batch_size = vertices.shape[0]

    aa_pose = torch.index_select(pose.view(batch_size, -1, 3), 1,
                                 neck_kin_chain)
    rot_mats = batch_rodrigues(
        aa_pose.view(-1, 3), dtype=dtype).view(batch_size, -1, 3, 3)

    rel_rot_mat = torch.eye(
        3, device=vertices.device, dtype=dtype).unsqueeze_(dim=0).repeat(
            batch_size, 1, 1)
    for idx in range(len(neck_kin_chain)):
        rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat)

    y_rot_angle = torch.round(
        torch.clamp(-rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi,
                    max=39)).to(dtype=torch.long)
    neg_mask = y_rot_angle.lt(0).to(dtype=torch.long)
    mask = y_rot_angle.lt(-39).to(dtype=torch.long)
    neg_vals = mask * 78 + (1 - mask) * (39 - y_rot_angle)
    y_rot_angle = (neg_mask * neg_vals +
                   (1 - neg_mask) * y_rot_angle)

    dyn_lmk_faces_idx = torch.index_select(dynamic_lmk_faces_idx,
                                           0, y_rot_angle)
    dyn_lmk_b_coords = torch.index_select(dynamic_lmk_b_coords,
                                          0, y_rot_angle)

    return dyn_lmk_faces_idx, dyn_lmk_b_coords


def vertices2landmarks(vertices, faces, lmk_faces_idx, lmk_bary_coords):
    ''' Calculates landmarks by barycentric interpolation

        Parameters
        ----------
        vertices: torch.tensor BxVx3, dtype = torch.float32
            The tensor of input vertices
        faces: torch.tensor Fx3, dtype = torch.long
            The faces of the mesh
        lmk_faces_idx: torch.tensor L, dtype = torch.long
            The tensor with the indices of the faces used to calculate the
            landmarks.
        lmk_bary_coords: torch.tensor Lx3, dtype = torch.float32
            The tensor of barycentric coordinates that are used to interpolate
            the landmarks

        Returns
        -------
        landmarks: torch.tensor BxLx3, dtype = torch.float32
            The coordinates of the landmarks for each mesh in the batch
    '''
    # Extract the indices of the vertices for each face
    # BxLx3
    batch_size, num_verts = vertices.shape[:2]
    device = vertices.device

    lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).view(
        batch_size, -1, 3)

    lmk_faces += torch.arange(
        batch_size, dtype=torch.long, device=device).view(-1, 1, 1) * num_verts

    lmk_vertices = vertices.view(-1, 3)[lmk_faces].view(
        batch_size, -1, 3, 3)

    landmarks = torch.einsum('blfi,blf->bli', [lmk_vertices, lmk_bary_coords])
    return landmarks


def lbs(betas, pose, v_template, shapedirs, posedirs, J_regressor, parents,
        lbs_weights, pose2rot=True, dtype=torch.float32):
    ''' Performs Linear Blend Skinning with the given shape and pose parameters

        Parameters
        ----------
        betas : torch.tensor BxNB
            The tensor of shape parameters
        pose : torch.tensor Bx(J + 1) * 3
            The pose parameters in axis-angle format
        v_template torch.tensor BxVx3
            The template mesh that will be deformed
        shapedirs : torch.tensor 1xNB
            The tensor of PCA shape displacements
        posedirs : torch.tensor Px(V * 3)
            The pose PCA coefficients
        J_regressor : torch.tensor JxV
            The regressor array that is used to calculate the joints from
            the position of the vertices
        parents: torch.tensor J
            The array that describes the kinematic tree for the model
        lbs_weights: torch.tensor N x V x (J + 1)
            The linear blend skinning weights that represent how much the
            rotation matrix of each part affects each vertex
        pose2rot: bool, optional
            Flag on whether to convert the input pose tensor to rotation
            matrices. The default value is True. If False, then the pose tensor
            should already contain rotation matrices and have a size of
            Bx(J + 1)x9
        dtype: torch.dtype, optional

        Returns
        -------
        verts: torch.tensor BxVx3
            The vertices of the mesh after applying the shape and pose
            displacements.
        joints: torch.tensor BxJx3
            The joints of the model
    '''

    batch_size = max(betas.shape[0], pose.shape[0])
    device = betas.device

    # Add shape contribution
    v_shaped = v_template + blend_shapes(betas, shapedirs)

    # Get the joints
    # NxJx3 array
    J = vertices2joints(J_regressor, v_shaped)

    # 3. Add pose blend shapes
    # N x J x 3 x 3
    ident = torch.eye(3, dtype=dtype, device=device)
    if pose2rot:
        rot_mats = batch_rodrigues(
            pose.view(-1, 3), dtype=dtype).view([batch_size, -1, 3, 3])

        pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
        # (N x P) x (P, V * 3) -> N x V x 3
        pose_offsets = torch.matmul(pose_feature, posedirs) \
            .view(batch_size, -1, 3)
    else:
        pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident
        rot_mats = pose.view(batch_size, -1, 3, 3)

        pose_offsets = torch.matmul(pose_feature.view(batch_size, -1),
                                    posedirs).view(batch_size, -1, 3)

    v_posed = pose_offsets + v_shaped
    # 4. Get the global joint location
    J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype)

    # 5. Do skinning:
    # W is N x V x (J + 1)
    W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1])
    # (N x V x (J + 1)) x (N x (J + 1) x 16)
    num_joints = J_regressor.shape[0]
    T = torch.matmul(W, A.view(batch_size, num_joints, 16)) \
        .view(batch_size, -1, 4, 4)

    homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1],
                               dtype=dtype, device=device)
    v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2)
    v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))

    verts = v_homo[:, :, :3, 0]

    return verts, J_transformed


def vertices2joints(J_regressor, vertices):
    ''' Calculates the 3D joint locations from the vertices

    Parameters
    ----------
    J_regressor : torch.tensor JxV
        The regressor array that is used to calculate the joints from the
        position of the vertices
    vertices : torch.tensor BxVx3
        The tensor of mesh vertices

    Returns
    -------
    torch.tensor BxJx3
        The location of the joints
    '''

    return torch.einsum('bik,ji->bjk', [vertices, J_regressor])


def blend_shapes(betas, shape_disps):
    ''' Calculates the per vertex displacement due to the blend shapes


    Parameters
    ----------
    betas : torch.tensor Bx(num_betas)
        Blend shape coefficients
    shape_disps: torch.tensor Vx3x(num_betas)
        Blend shapes

    Returns
    -------
    torch.tensor BxVx3
        The per-vertex displacement due to shape deformation
    '''

    # Displacement[b, m, k] = sum_{l} betas[b, l] * shape_disps[m, k, l]
    # i.e. Multiply each shape displacement by its corresponding beta and
    # then sum them.
    blend_shape = torch.einsum('bl,mkl->bmk', [betas, shape_disps])
    return blend_shape


def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
    ''' Calculates the rotation matrices for a batch of rotation vectors
        Parameters
        ----------
        rot_vecs: torch.tensor Nx3
            array of N axis-angle vectors
        Returns
        -------
        R: torch.tensor Nx3x3
            The rotation matrices for the given axis-angle parameters
    '''

    batch_size = rot_vecs.shape[0]
    device = rot_vecs.device

    angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
    rot_dir = rot_vecs / angle

    cos = torch.unsqueeze(torch.cos(angle), dim=1)
    sin = torch.unsqueeze(torch.sin(angle), dim=1)

    # Bx1 arrays
    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
    K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)

    zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))

    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
    return rot_mat


def transform_mat(R, t):
    ''' Creates a batch of transformation matrices
        Args:
            - R: Bx3x3 array of a batch of rotation matrices
            - t: Bx3x1 array of a batch of translation vectors
        Returns:
            - T: Bx4x4 Transformation matrix
    '''
    # No padding left or right, only add an extra row
    return torch.cat([F.pad(R, [0, 0, 0, 1]),
                      F.pad(t, [0, 0, 0, 1], value=1)], dim=2)


def batch_rigid_transform(rot_mats, joints, parents, dtype=torch.float32):
    """
    Applies a batch of rigid transformations to the joints

    Parameters
    ----------
    rot_mats : torch.tensor BxNx3x3
        Tensor of rotation matrices
    joints : torch.tensor BxNx3
        Locations of joints
    parents : torch.tensor BxN
        The kinematic tree of each object
    dtype : torch.dtype, optional:
        The data type of the created tensors, the default is torch.float32

    Returns
    -------
    posed_joints : torch.tensor BxNx3
        The locations of the joints after applying the pose rotations
    rel_transforms : torch.tensor BxNx4x4
        The relative (with respect to the root joint) rigid transformations
        for all the joints
    """

    joints = torch.unsqueeze(joints, dim=-1)

    rel_joints = joints.clone()
    rel_joints[:, 1:] -= joints[:, parents[1:]]

    transforms_mat = transform_mat(
        rot_mats.reshape(-1, 3, 3),
        rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4)

    transform_chain = [transforms_mat[:, 0]]
    for i in range(1, parents.shape[0]):
        # Subtract the joint location at the rest pose
        # No need for rotation, since it's identity when at rest
        curr_res = torch.matmul(transform_chain[parents[i]],
                                transforms_mat[:, i])
        transform_chain.append(curr_res)

    transforms = torch.stack(transform_chain, dim=1)

    # The last column of the transformations contains the posed joints
    posed_joints = transforms[:, :, :3, 3]

    # The last column of the transformations contains the posed joints
    posed_joints = transforms[:, :, :3, 3]

    joints_homogen = F.pad(joints, [0, 0, 0, 1])

    rel_transforms = transforms - F.pad(
        torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0])

    return posed_joints, rel_transforms

smplx/smplx/lbs.py
sha2nkt/ipman-r
# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de

from __future__ import print_function
from __future__ import absolute_import
from __future__ import division

# Joint name to vertex mapping. SMPL/SMPL-H/SMPL-X vertices that correspond to
# MSCOCO and OpenPose joints
vertex_ids = {
    'smplh': {
        'nose':		    332,
        'reye':		    6260,
        'leye':		    2800,
        'rear':		    4071,
        'lear':		    583,
        'rthumb':		6191,
        'rindex':		5782,
        'rmiddle':		5905,
        'rring':		6016,
        'rpinky':		6133,
        'lthumb':		2746,
        'lindex':		2319,
        'lmiddle':		2445,
        'lring':		2556,
        'lpinky':		2673,
        'LBigToe':		3216,
        'LSmallToe':	3226,
        'LHeel':		3387,
        'RBigToe':		6617,
        'RSmallToe':    6624,
        'RHeel':		6787
    },
    'smplx': {
        'nose':		    9120,
        'reye':		    9929,
        'leye':		    9448,
        'rear':		    616,
        'lear':		    6,
        'rthumb':		8079,
        'rindex':		7669,
        'rmiddle':		7794,
        'rring':		7905,
        'rpinky':		8022,
        'lthumb':		5361,
        'lindex':		4933,
        'lmiddle':		5058,
        'lring':		5169,
        'lpinky':		5286,
        'LBigToe':		5770,
        'LSmallToe':    5780,
        'LHeel':		8846,
        'RBigToe':		8463,
        'RSmallToe': 	8474,
        'RHeel':  		8635
    }
}

smplx/smplx/vertex_ids.py
sha2nkt/ipman-r
# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2021 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
import sys

import torch
import torch.nn as nn
import pickle as pkl
import numpy as np
from vis_utils.mesh_utils import GMoF_unscaled, batch_face_normals, HDfier
from constants import CONTACT_THRESH
from stability.part_volumes import PartVolume

SMPL_PART_BOUNDS = 'data/essentials/yogi_segments/smpl/part_meshes_ply/smpl_segments_bounds.pkl'
FID_TO_PART = 'data/essentials/yogi_segments/smpl/part_meshes_ply/fid_to_part.pkl'
PART_VID_FID = 'data/essentials/yogi_segments/smpl/part_meshes_ply/smpl_part_vid_fid.pkl'
HD_SMPL_MAP  = 'data/essentials/hd_model/smpl/smpl_neutral_hd_sample_from_mesh_out.pkl'

class StabilityLossCoP(nn.Module):
    def __init__(self,
                 faces,
                 cop_w = 10,
                 cop_k = 100,
                 contact_thresh=CONTACT_THRESH,
                 model_type='smpl',
                 device='cuda',
    ):
        super().__init__()
        """
        Loss that ensures that the COM of the SMPL mesh is close to the center of support 
        """
        if model_type == 'smpl':
            num_faces = 13776
            num_verts_hd = 20000

        assert faces is not None, 'Faces tensor is none'
        if type(faces) is not torch.Tensor:
            faces = torch.tensor(faces.astype(np.int64), dtype=torch.long).to(device)
        self.register_buffer('faces', faces)

        self.cop_w = cop_w
        self.cop_k = cop_k
        self.contact_thresh = contact_thresh

        self.hdfy_op = HDfier(model_type=model_type)

        with open(SMPL_PART_BOUNDS, 'rb') as f:
            d = pkl.load(f)
            self.part_bounds = {k: d[k] for k in sorted(d)}
        self.part_order = sorted(self.part_bounds)

        with open(PART_VID_FID, 'rb') as f:
            self.part_vid_fid = pkl.load(f)

        # mapping between vid_hd and fid
        with open(HD_SMPL_MAP, 'rb') as f:
            faces_vert_is_sampled_from = pkl.load(f)['faces_vert_is_sampled_from']
        index_row_col = torch.stack(
            [torch.LongTensor(np.arange(0, num_verts_hd)), torch.LongTensor(faces_vert_is_sampled_from)], dim=0)
        values = torch.ones(num_verts_hd, dtype=torch.float)
        size = torch.Size([num_verts_hd, num_faces])
        hd_vert_on_fid = torch.sparse.FloatTensor(index_row_col, values, size)

        # mapping between fid and part label
        with open(FID_TO_PART, 'rb') as f:
            fid_to_part_dict = pkl.load(f)
        fid_to_part = torch.zeros([len(fid_to_part_dict.keys()), len(self.part_order)], dtype=torch.float32)
        for fid, partname in fid_to_part_dict.items():
            part_idx = self.part_order.index(partname)
            fid_to_part[fid, part_idx] = 1.

        # mapping between vid_hd and part label
        self.hd_vid_in_part = self.vertex_id_to_part_mapping(hd_vert_on_fid, fid_to_part)

    def compute_triangle_area(self, triangles):
        ### Compute the area of each triangle in the mesh
        # Compute the cross product of the two vectors of each triangle
        # Then compute the length of the cross product
        # Finally, divide by 2 to get the area of each triangle

        vectors = torch.diff(triangles, dim=2)
        crosses = torch.cross(vectors[:, :, 0], vectors[:, :, 1])
        area = torch.norm(crosses, dim=2) / 2
        return area

    def compute_per_part_volume(self, vertices):
        """
        Compute the volume of each part in the reposed mesh
        """
        part_volume = []
        for part_name, part_bounds in self.part_bounds.items():
            # get part vid and fid
            part_vid = torch.LongTensor(self.part_vid_fid[part_name]['vert_id']).to(vertices.device)
            part_fid = torch.LongTensor(self.part_vid_fid[part_name]['face_id']).to(vertices.device)
            pv = PartVolume(part_name, vertices, self.faces)
            for bound_name, bound_vids in part_bounds.items():
                pv.close_mesh(bound_vids)
            # add extra vids and fids to original part ids
            new_vert_ids = torch.LongTensor(pv.new_vert_ids).to(vertices.device)
            new_face_ids = torch.LongTensor(pv.new_face_ids).to(vertices.device)
            part_vid = torch.cat((part_vid, new_vert_ids), dim=0)
            part_fid = torch.cat((part_fid, new_face_ids), dim=0)
            pv.extract_part_triangles(part_vid, part_fid)
            part_volume.append(pv.part_volume())
        return torch.vstack(part_volume).permute(1,0).to(vertices.device)

    def vertex_id_to_part_volume_mapping(self, per_part_volume, device):
        batch_size = per_part_volume.shape[0]
        self.hd_vid_in_part = self.hd_vid_in_part.to(device)
        hd_vid_in_part = self.hd_vid_in_part[None, :, :].repeat(batch_size, 1, 1)
        vid_to_vol = torch.bmm(hd_vid_in_part, per_part_volume[:, :, None])
        return vid_to_vol

    def vertex_id_to_part_mapping(self, hd_vert_on_fid, fid_to_part):
        vid_to_part = torch.mm(hd_vert_on_fid, fid_to_part)
        return vid_to_part

    def forward(self, vertices):
        # Note: the vertices should be aligned along y-axis and in world coordinates
        batch_size = vertices.shape[0]
        # calculate per part volume
        per_part_volume = self.compute_per_part_volume(vertices)
        # sample 20k vertices uniformly on the smpl mesh
        vertices_hd = self.hdfy_op.hdfy_mesh(vertices)
        # get volume per vertex id in the hd mesh
        volume_per_vert_hd = self.vertex_id_to_part_volume_mapping(per_part_volume, vertices.device)
        # calculate com using volume weighted mean
        com = torch.sum(vertices_hd * volume_per_vert_hd, dim=1) / torch.sum(volume_per_vert_hd, dim=1)

        # # get COM of the SMPLX mesh
        # triangles = torch.index_select(vertices, 1, self.faces.view(-1)).reshape(batch_size, -1, 3, 3)
        # triangle_centroids = torch.mean(triangles, dim=2)
        # triangle_area = self.compute_triangle_area(triangles)
        # com_naive = torch.einsum('bij,bi->bj', triangle_centroids, triangle_area) / torch.sum(triangle_area, dim=1)

        # pressure based center of support
        ground_plane_height = 0.0
        eps = 1e-6
        vertex_height = (vertices_hd[:, :, 1] - ground_plane_height)
        inside_mask = (vertex_height < 0.0).float()
        outside_mask = (vertex_height >= 0.0).float()
        pressure_weights = inside_mask * (1-self.cop_k*vertex_height) + outside_mask *  torch.exp(-self.cop_w * vertex_height)
        cop = torch.sum(vertices_hd * pressure_weights.unsqueeze(-1), dim=1) / (torch.sum(pressure_weights, dim=1, keepdim=True) +eps)

        # naive center of support
        # vertex_height_robustified = GMoF_unscaled(rho=self.gmof_rho)(vertex_height)
        contact_confidence = torch.sum(pressure_weights, dim=1)
        # contact_mask = (vertex_height < self.contact_thresh).float()
        # num_contact_verts = torch.sum(contact_mask, dim=1)
        # contact_centroid_naive = torch.sum(vertices_hd * contact_mask[:, :, None], dim=1) / (torch.sum(contact_mask, dim=1) + eps)

        # project com, cop to ground plane (x-z plane)
        # weight loss by number of contact vertices to zero out if zero vertices in contact
        com_xz = torch.stack([com[:, 0], torch.zeros_like(com)[:, 0], com[:, 2]], dim=1)
        contact_centroid_xz = torch.stack([cop[:, 0], torch.zeros_like(cop)[:, 0], cop[:, 2]], dim=1)
        # stability_loss = (contact_confidence * torch.norm(com_xz - contact_centroid_xz, dim=1)).sum(dim=-1)
        stability_loss = (torch.norm(com_xz - contact_centroid_xz, dim=1))
        return stability_loss


class GroundLoss(nn.Module):
    def __init__(self,
                 faces,
                 device='cuda',
                 model_type='smpl',
                 out_alpha1=1,
                 out_alpha2=0.5,
                 in_alpha1=1,
                 in_alpha2=0.15,
    ):
        super().__init__()
        assert faces is not None, 'Faces tensor is none'
        if type(faces) is not torch.Tensor:
            faces = torch.tensor(faces.astype(np.int64), dtype=torch.long).to(device)
        self.register_buffer('faces', faces)

        self.device = device

        # loss specifications

        # push, pull, and contact weights
        # self.inside_w = inside_loss_weight
        # self.outside_w = outside_loss_weight

        # hyper params
        self.a1 = out_alpha1
        self.a2 = out_alpha2
        self.b1 = in_alpha1
        self.b2 = in_alpha2

        self.hdfy_op = HDfier(model_type=model_type)



    def forward(self, vertices):
        """
                Loss that ensures that the ground plane pulls body mesh vertices towards it till contact and resolves
                ground-plane intersections
                Using assymetric loss such that the
                - loss is 0 when the body is in contact with the ground plane
                - loss is >> 0 when the plane interpenetrates the body (completely implausible)
                - loss >0 if the plane is close to nearby vertices that are not touching
                vertices:the vertices should be aligned along y-axis and in world coordinates
        """
        bs = vertices.shape[0]
        vertices_hd = self.hdfy_op.hdfy_mesh(vertices)
        # get vertices under the ground plane
        ground_plane_height = 0.0  # obtained by visualization on the presented pose
        vertex_height = (vertices_hd[:, :, 1] - ground_plane_height)
        inside_mask = vertex_height < 0.00
        outside_mask = vertex_height >= 0.00
        # pull closeby outside vertices
        v2v_pull = (self.a1 * torch.tanh((vertex_height * outside_mask) / self.a2) ** 2)
        # # apply loss to inside vertices to remove intersection
        v2v_push = (self.b1 * torch.tanh((vertex_height * inside_mask) / self.b2)**2)
        return v2v_pull, v2v_push

# class MimickedSupportLoss(nn.Module):
#     def __init__(self,
#                  faces,
#                  inside_loss_weight=1.0,
#                  outside_loss_weight=1.0,
#                  contact_loss_weight=2.5,
#                  align_faces=False,
#                  gmof_rho=0.02,
#                  use_hd=True,
#                  device='cuda',
#                  model_type='smplx',
#                  alpha1=0.005,
#                  alpha2=0.005,
#                  beta1=600,
#     ):
#         super().__init__()
#         assert faces is not None, 'Faces tensor is none'
#         if type(faces) is not torch.Tensor:
#             faces = torch.tensor(faces.astype(np.int64), dtype=torch.long)
#         self.register_buffer('faces', faces)
#
#         self.device = device
#         self.model_type = model_type
#
#         # loss specifications
#         self.use_hd = use_hd
#         self.align_faces = align_faces
#         # self.gmof_rho = gmof_rho
#
#         # push, pull, and contact weights
#         self.inside_w = inside_loss_weight
#         self.contact_w = contact_loss_weight
#
#         # hyper params
#         self.a1 = alpha1
#         self.a2 = alpha2
#         self.b1 = beta1
#
#
#
#     def forward(self, vertices):
#         """
#                 Loss that ensures that the ground plane pulls body mesh vertices towards it till contact and resolves
#                 ground-plane intersections
#                 Using assymetric loss such that the
#                 - loss is 0 when the body is in contact with the ground plane
#                 - loss is >> 0 when the plane interpenetrates the body (completely implausible)
#                 - loss >0 if the plane is close to nearby vertices that are not touching
#         """
#         bs = vertices.shape[0]
#
#         import ipdb; ipdb.set_trace()
#         # get vertices under the ground plane
#         ground_plane_height = 0.0  # obtained by visualization on the presented pose
#         vertex_height = (vertices[:, :, 1] - ground_plane_height)
#         inside_mask = (vertex_height < 0.00).float()
#
#         # pull closeby outside vertices
#         v2v_pull = self.contact_w * (self.a1 * torch.tanh(vertex_height * [~inside_mask] / self.a2)**2).sum()
#         # # apply loss to inside vertices to remove intersection
#         v2v_push = self.inside_w * (self.b1 * (vertex_height * inside_mask)**2).sum()
#         contactloss = v2v_push + v2v_pull
#
#         # now align faces that are close
#         # dot product should be -1 for faces in contact
#         if self.align_faces:
#             triangles = torch.index_select(vertices, 1, self.faces.view(-1)).reshape(bs, -1, 3, 3)
#             face_normals = batch_face_normals(triangles)
#
#         face_angle_loss = torch.zeros(bs, device=vertices.device)
#         if self.align_faces:
#             if self.use_hd:
#                 if hd_faces_in_contact[idx] is not None:
#                     hd_fn_in_contact = face_normals[idx][hd_faces_in_contact[idx]]
#                     dotprod_normals = 1 + (hd_fn_in_contact[0] * hd_fn_in_contact[1]).sum(1)
#                     face_angle_loss[idx] = dotprod_normals.sum()
#             else:
#                 sys.exit('You can only align vertices when use_hd=True')
#
#         return contactloss, face_angle_loss
stability/ground_losses.py
sha2nkt/ipman-r
# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2021 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
"""
This script is used to close part meshes.
"""
from tempfile import gettempdir

import torch
import trimesh
import torch.nn as nn
import numpy as np
import pickle
# from .utils.mesh import winding_numbers

class SMPLMesh(nn.Module):
    def __init__(self, vertices, faces):
        super(SMPLMesh, self).__init__()

        self.vertices = vertices
        self.faces = faces

class PartVolume(nn.Module):
    def __init__(self,
                 part_name,
                 vertices,
                 faces):
        super(PartVolume, self).__init__()

        self.part_name = part_name
        self.smpl_mesh = SMPLMesh(vertices, faces)

        self.part_triangles = None
        self.device = vertices.device

        self.new_vert_ids = []
        self.new_face_ids = []

    def close_mesh(self, boundary_vids):
        # find the center of the boundary
        mean_vert = self.smpl_mesh.vertices[:, boundary_vids, :].mean(dim=1, keepdim=True)
        self.smpl_mesh.vertices = torch.cat([self.smpl_mesh.vertices, mean_vert], dim=1)
        new_vert_idx = self.smpl_mesh.vertices.shape[1]-1
        self.new_vert_ids.append(new_vert_idx)
        # add faces
        new_faces = [[boundary_vids[i + 1], boundary_vids[i], new_vert_idx] for i in range(len(boundary_vids) - 1)]
        self.smpl_mesh.faces = torch.cat([self.smpl_mesh.faces, torch.tensor(new_faces, dtype=torch.long, device=self.device)], dim=0)
        self.new_face_ids += list(range(self.smpl_mesh.faces.shape[0]-len(new_faces), self.smpl_mesh.faces.shape[0]))

    def extract_part_triangles(self, part_vids, part_fid):
        # make sure that inputs are from a watertight part mesh
        batch_size = self.smpl_mesh.vertices.shape[0]

        part_vertices = self.smpl_mesh.vertices[:, part_vids, :]
        part_faces = self.smpl_mesh.faces[part_fid, :]

        part_mean = part_vertices.mean(dim=1, keepdim=True)

        # subtract vert mean because volume computation only applies if origin is inside the triangles
        self.smpl_mesh.vertices = self.smpl_mesh.vertices - part_mean

        # compute triangle
        if self.part_triangles is None:
            # self.part_triangles = torch.index_select(self.smpl_mesh.vertices, 1, self.smpl_mesh.faces.view(-1)).reshape(batch_size, -1, 3, 3)
            self.part_triangles = torch.index_select(self.smpl_mesh.vertices, 1, part_faces.view(-1)).reshape(batch_size, -1, 3, 3)
        else:
            self.part_triangles = torch.cat([self.part_triangles,
                                             torch.index_select(self.smpl_mesh.vertices, 1,
                                                     part_faces.view(-1)).reshape(batch_size, -1, 3, 3)], dim=1)
        # add back vert mean
        self.smpl_mesh.vertices = self.smpl_mesh.vertices + part_mean

    def part_volume(self):
        # Note: the mesh should be enclosing the origin (mean-subtracted)
        # compute volume of triangles by drawing tetrahedrons
        # https://stackoverflow.com/questions/1406029/how-to-calculate-the-volume-of-a-3d-mesh-object-the-surface-of-which-is-made-up
        x = self.part_triangles[:, :, :, 0]
        y = self.part_triangles[:, :, :, 1]
        z = self.part_triangles[:, :, :, 2]
        volume = (
                         -x[:, :, 2] * y[:, :, 1] * z[:, :, 0] +
                         x[:, :, 1] * y[:, :, 2] * z[:, :, 0] +
                         x[:, :, 2] * y[:, :, 0] * z[:, :, 1] -
                         x[:, :, 0] * y[:, :, 2] * z[:, :, 1] -
                         x[:, :, 1] * y[:, :, 0] * z[:, :, 2] +
                         x[:, :, 0] * y[:, :, 1] * z[:, :, 2]
                 ).sum(dim=1).abs() / 6.0
        return volume

class BodySegment(nn.Module):
    def __init__(self,
                 name,
                 faces,
                 segments_folder,
                 model_type='smpl',
                 append_idx=None):
        super(BodySegment, self).__init__()

        self.name = name
        self.append_idx = faces.max().item() if append_idx is None \
            else append_idx

        self.model_type = model_type
        sb_path = f'{segments_folder}/{model_type}_segments_bounds.pkl'
        sxseg = pickle.load(open(sb_path, 'rb'))

        # read mesh and find faces of segment
        segment_path = f'{segments_folder}/{model_type}_segment_{name}.ply'
        bandmesh = trimesh.load(segment_path, process=False)
        segment_vidx = torch.from_numpy(np.where(
            np.array(bandmesh.visual.vertex_colors[:,0]) == 255)[0])
        self.register_buffer('segment_vidx', segment_vidx)

        # read boundary information
        self.bands = [x for x in sxseg[name].keys()]
        self.bands_verts = [x for x in sxseg[name].values()]
        self.num_bounds = len(self.bands_verts)
        for idx, bv in enumerate(self.bands_verts):
            self.register_buffer(f'bands_verts_{idx}', torch.tensor(bv))
        self.bands_faces = self.create_band_faces()

        # read mesh and find
        segment_faces_ids = np.where(np.isin(faces.cpu().numpy(),
            segment_vidx).sum(1) == 3)[0]
        segment_faces = faces[segment_faces_ids,:]
        segment_faces = torch.cat((faces[segment_faces_ids,:],
            self.bands_faces), 0)
        self.register_buffer('segment_faces', segment_faces)

        # create vector to select vertices form faces
        tri_vidx = []
        for ii in range(faces.max().item()+1):
            tri_vidx += [torch.nonzero(faces==ii)[0].tolist()]
        self.register_buffer('tri_vidx', torch.tensor(tri_vidx))

    def create_band_faces(self):
        """
            create the faces that close the segment.
        """
        bands_faces = []
        for idx, k in enumerate(self.bands):
            new_vert_idx = self.append_idx + 1 + idx
            new_faces = [[self.bands_verts[idx][i+1], \
                self.bands_verts[idx][i], new_vert_idx] \
                for i in range(len(self.bands_verts[idx])-1)]
            bands_faces += new_faces

        bands_faces_tensor = torch.tensor(
            np.array(bands_faces).astype(np.int64), dtype=torch.long)

        return bands_faces_tensor

    def get_closed_segment(self, vertices):
        """
            create the closed segment mesh from SMPL-X vertices.
        """
        vertices = vertices.detach().clone()
        # append vertices to smpl, that close the segment and compute faces
        for idx in range(self.num_bounds):
            bv = eval(f'self.bands_verts_{idx}')
            close_segment_vertices = torch.mean(vertices[:, bv,:], 1,
                                    keepdim=True)
            vertices = torch.cat((vertices, close_segment_vertices), 1)
        segm_triangles = vertices[:, self.segment_faces, :]

        return segm_triangles

    # def has_self_isect_verts(self, vertices, thres=0.99):
    #     """
    #         check if segment (its vertices) are self intersecting.
    #     """
    #     segm_triangles = self.get_closed_segment(vertices)
    #     segm_verts = vertices[:,self.segment_vidx,:]
    #
    #     # do inside outside segmentation
    #     exterior = winding_numbers(segm_verts, segm_triangles) \
    #                 .le(thres)
    #
    #     return exterior
    #
    # def has_self_isect_points(self, points, triangles, thres=0.99):
    #     """
    #         check if points on segment are self intersecting.
    #     """
    #     smpl_verts = triangles[:,self.tri_vidx[:,0], self.tri_vidx[:,1],:]
    #     segm_triangles = self.get_closed_segment(smpl_verts)
    #
    #     # do inside outside segmentation
    #     exterior = winding_numbers(points, segm_triangles) \
    #                 .le(thres)
    #
    #     return exterior

class BatchBodySegment(nn.Module):
    def __init__(self,
                 names,
                 faces,
                 segments_folder,
                 model_type='smpl',
                 device='cuda'):
        super(BatchBodySegment, self).__init__()
        self.names = names
        self.num_segments = len(names)
        self.nv = faces.max().item()

        self.model_type = model_type
        sb_path = f'{segments_folder}/{model_type}_segments_bounds.pkl'
        sxseg = pickle.load(open(sb_path, 'rb'))

        self.append_idx = [len(b) for a,b in sxseg.items() \
            for c,d in b.items() if a in self.names]
        self.append_idx = np.cumsum(np.array([self.nv] + self.append_idx))

        self.segmentation = {}
        for idx, name in enumerate(names):
            self.segmentation[name] = BodySegment(name, faces, segments_folder,
                model_type).to('cuda')

    def batch_has_self_isec_verts(self, vertices):
        """
            check is mesh is intersecting with itself
        """
        exteriors = []
        for k, segm in self.segmentation.items():
            exteriors += [segm.has_self_isect_verts(vertices)]
        return exteriors

stability/part_volumes.py
sha2nkt/ipman-r
import torch
import sys
import numpy as np
import trimesh
from scipy.optimize import linprog
from stability.part_volumes import PartVolume
from vis_utils.mesh_utils import GMoF_unscaled, HDfier
import pickle as pkl
import time

SMPL_PART_BOUNDS = 'data/essentials/yogi_segments/smpl/part_meshes_ply/smpl_segments_bounds.pkl'
FID_TO_PART = 'data/essentials/yogi_segments/smpl/part_meshes_ply/fid_to_part.pkl'
PART_VID_FID = 'data/essentials/yogi_segments/smpl/part_meshes_ply/smpl_part_vid_fid.pkl'
HD_SMPL_MAP  = 'data/essentials/hd_model/smpl/smpl_neutral_hd_sample_from_mesh_out.pkl'

def stability_error(vertices, faces, contact_thresh, cop_w=10, cop_k=100):
    with torch.no_grad():
        batch_size = vertices.shape[0]
        # calculate per part volume
        per_part_volume = compute_per_part_volume(vertices, faces)
        # sample 20k vertices uniformly on the smpl mesh
        vertices_hd = HDfier().hdfy_mesh(vertices, vertices.device)

        # create face_id to part label map
        with open(PART_VID_FID, 'rb') as f:
            part_vid_fid = pkl.load(f)
        face_id_to_part_label = face_id_to_part_mapping(part_vid_fid, faces)
        with open(HD_SMPL_MAP, 'rb') as f:
            faces_vert_is_sampled_from = pkl.load(f)['faces_vert_is_sampled_from']
        all_fids = torch.arange(len(faces), dtype=torch.long)
        hd_vert_on_fid = all_fids[faces_vert_is_sampled_from]
        hd_vert_label = [face_id_to_part_label[fid.item()] for fid in hd_vert_on_fid]

        # calculate volume per vert
        volume_per_vert_hd = torch.FloatTensor([per_part_volume[part] for part in hd_vert_label]).to(vertices.device)
        if volume_per_vert_hd.dim() == 1:
            volume_per_vert_hd = volume_per_vert_hd.unsqueeze(0)
        # calculate com using volume weighted mean
        com = torch.sum(vertices_hd * volume_per_vert_hd.unsqueeze(-1), dim=1) / torch.sum(volume_per_vert_hd, dim=1,
                                                                                           keepdim=True)

        # get vertices in contact
        ground_plane_height = 0.0
        vertex_height = (vertices_hd[:, :, 1] - ground_plane_height)
        # vertex_height_robustified = GMoF_unscaled(rho=gmof_rho)(vertex_height)
        contact_mask = (vertex_height < contact_thresh).float()
        num_contact_verts = torch.sum(contact_mask, dim=1)
        contact_metric = torch.ones_like(num_contact_verts)
        contact_metric[torch.nonzero(num_contact_verts == 0)] = 0
        # if torch.nonzero(num_contact_verts).shape[0] != len(num_contact_verts):
        #     print('An element with zero contact vertices found in the batch. Removing that element from stability metric and adding to no_contact metric')
        #     contact_metric[torch.nonzero(num_contact_verts==0)] = 0
            # remove no_contact pose for further computation
            # vertices_hd = torch.index_select(vertices_hd, dim=0, index=torch.nonzero(contact_metric).squeeze())
            # contact_mask = torch.index_select(contact_mask, dim=0, index=torch.nonzero(contact_metric).squeeze())
            # com = torch.index_select(com, dim=0, index=torch.nonzero(contact_metric).squeeze())

        # calculate pressure based cos
        inside_mask = (vertex_height < 0.0).float()
        outside_mask = (vertex_height >= 0.0).float()
        pressure_weights = inside_mask * (1-cop_k*vertex_height) + outside_mask *  torch.exp(-cop_w * vertex_height)
        cos = torch.sum(vertices_hd * pressure_weights.unsqueeze(-1), dim=1) / (
            torch.sum(pressure_weights, dim=1, keepdim=True))

        # get distance of COM to center of support
        # project com, cos to ground plane (x-z plane)
        com_xz = torch.stack([com[:, 0], torch.zeros_like(com)[:, 0], com[:, 2]], dim=1)
        cos_xz = torch.stack(
            [cos[:, 0], torch.zeros_like(cos)[:, 0], cos[:, 2]], dim=1)
        stability_loss = torch.norm(com_xz - cos_xz, dim=1)

    # if non contact, set nan
    stability_loss[torch.nonzero(contact_metric == 0)] = float('nan')
    return stability_loss, contact_metric

class StabilityMetrics():
    def __init__(self, model_type='smpl'):
        if model_type == 'smpl':
            num_faces = 13776
            num_verts_hd = 20000

        with open(SMPL_PART_BOUNDS, 'rb') as f:
            d = pkl.load(f)
            self.part_bounds = {k: d[k] for k in sorted(d)}
        self.part_order = sorted(self.part_bounds)

        with open(PART_VID_FID, 'rb') as f:
            self.part_vid_fid = pkl.load(f)

        # mapping between vid_hd and fid
        with open(HD_SMPL_MAP, 'rb') as f:
            faces_vert_is_sampled_from = pkl.load(f)['faces_vert_is_sampled_from']
        index_row_col = torch.stack([torch.LongTensor(np.arange(0, num_verts_hd)), torch.LongTensor(faces_vert_is_sampled_from)], dim=0)
        values = torch.ones(num_verts_hd, dtype=torch.float)
        size = torch.Size([num_verts_hd, num_faces])
        hd_vert_on_fid = torch.sparse.FloatTensor(index_row_col, values, size)

        # mapping between fid and part label
        with open(FID_TO_PART, 'rb') as f:
            fid_to_part_dict = pkl.load(f)
        fid_to_part = torch.zeros([len(fid_to_part_dict.keys()), len(self.part_order)], dtype=torch.float32)
        for fid, partname in fid_to_part_dict.items():
            part_idx = self.part_order.index(partname)
            fid_to_part[fid, part_idx] = 1.

         # mapping between vid_hd and part label
        self.hd_vid_in_part = self.vertex_id_to_part_mapping(hd_vert_on_fid, fid_to_part)

    def check_bos_interior(self, vertices, faces, contact_thresh, no_contact_vert_tol=10):
        """
        Check if the projection of COM lies inside the 2D convex hull/base-of-support of the contact vertices.
        Test this by finding if the COM projectection is a convex combination of the contact vertices.
        Ref: https://stackoverflow.com/questions/16750618/whats-an-efficient-way-to-find-if-a-point-lies-in-the-convex-hull-of-a-point-cl#answer-43564754
        Args:
            vertices:
            faces:
            gmof_rho:
            contact_thresh:

        Returns:

        """
        with torch.no_grad():
            batch_size = vertices.shape[0]
            # calculate per part volume
            per_part_volume = self.compute_per_part_volume(vertices, faces)
            # sample 20k vertices uniformly on the smpl mesh
            vertices_hd = HDfier(model_type='smpl').hdfy_mesh(vertices)
            # get volume per vertex id in the hd mesh
            volume_per_vert_hd = self.vertex_id_to_part_volume_mapping(per_part_volume, vertices.device)
            # calculate com using volume weighted mean
            com = torch.sum(vertices_hd * volume_per_vert_hd, dim=1) / torch.sum(volume_per_vert_hd, dim=1)

            # get vertices in contact
            ground_plane_height = 0.0
            vertex_height = (vertices_hd[:, :, 1] - ground_plane_height)
            # vertex_height_robustified = GMoF_unscaled(rho=gmof_rho)(vertex_height)
            contact_mask = (vertex_height < contact_thresh)


            num_contact_verts = torch.sum(contact_mask, dim=1, keepdim=True)
            contact_metric = torch.ones_like(num_contact_verts)
            contact_metric[num_contact_verts<no_contact_vert_tol] = 0
            # if torch.nonzero(num_contact_verts).shape[0] != len(num_contact_verts):
            #     print('An element with zero contact vertices found in the batch. Removing that element from stability metric and adding to no_contact metric')
            #     # no_contact_metric is a boolean mask of poses with zero contact vertices
            #     contact_metric = torch.ones_like(num_contact_verts)
            #     contact_metric[torch.nonzero(num_contact_verts == 0)] = 0
                # # remove no_contact pose for further computation
                # vertices_hd = torch.index_select(vertices_hd, dim=0, index=torch.nonzero(contact_metric).squeeze())
                # contact_mask = torch.index_select(contact_mask, dim=0, index=torch.nonzero(contact_metric).squeeze())
                # com = torch.index_select(com, dim=0, index=torch.nonzero(contact_metric).squeeze())

            # filter out the vertices that are not in contact
            # contact_vertices_hd = vertices_hd * contact_mask[:, :, None]
            in_hull_label = torch.zeros(batch_size, dtype=torch.float)
            for i, com_el in enumerate(com):
                # filter out the vertices that are not in contact
                c_el= torch.masked_select(vertices_hd[i], contact_mask[i, :, None]).view(-1, 3)
                if num_contact_verts[i] <= no_contact_vert_tol:
                    in_hull_label[i] = -1 # no contact
                    continue
                # project to ground plane y=0
                c_el[:, 1] = torch.zeros_like(c_el)[:, 1]
                com_el[1] = 0.0
                label = self.in_hull(c_el.cpu().numpy(), com_el.cpu().numpy())
                in_hull_label[i] = float(label)
        return torch.tensor(in_hull_label), com, contact_metric, contact_mask

    def in_hull(self, points, x):
        n_points = len(points)
        n_dim = len(x)
        c = np.zeros(n_points)
        A = np.r_[points.T, np.ones((1, n_points))]
        b = np.r_[x, np.ones(1)]
        try:
            lp = linprog(c, A_eq=A, b_eq=b, method='interior-point')
            return lp.success
        except:
            print('Linprog failed. Problem is infeasible')
            return False


    def compute_triangle_area(self, triangles):
        ### Compute the area of each triangle in the mesh
        # Compute the cross product of the two vectors of each triangle
        # Then compute the length of the cross product
        # Finally, divide by 2 to get the area of each triangle

        vectors = torch.diff(triangles, dim=2)
        crosses = torch.cross(vectors[:, :, 0], vectors[:, :, 1])
        area = torch.norm(crosses, dim=2) / 2
        return area

    def compute_per_part_volume(self, vertices, faces):
        """
        Compute the volume of each part in the reposed mesh
        """
        part_volume = []

        for part_name, part_bounds in self.part_bounds.items():
            # get part vid and fid
            part_vid = torch.LongTensor(self.part_vid_fid[part_name]['vert_id']).to(vertices.device)
            part_fid = torch.LongTensor(self.part_vid_fid[part_name]['face_id']).to(vertices.device)
            pv = PartVolume(part_name, vertices, faces)
            for bound_name, bound_vids in part_bounds.items():
                pv.close_mesh(bound_vids)
            # add extra vids and fids to original part ids
            new_vert_ids = torch.LongTensor(pv.new_vert_ids).to(vertices.device)
            new_face_ids = torch.LongTensor(pv.new_face_ids).to(vertices.device)
            part_vid = torch.cat((part_vid, new_vert_ids), dim=0)
            part_fid = torch.cat((part_fid, new_face_ids), dim=0)
            pv.extract_part_triangles(part_vid, part_fid)
            part_volume.append(pv.part_volume())
        return torch.vstack(part_volume).permute(1,0).to(vertices.device)

    def face_id_to_part_mapping(self, part_vid_fid, faces):
        fipm = dict({k: [] for k in range(len(faces))})

        # get mapping from face_id to part label
        for part_label, part_vid_fid in part_vid_fid.items():
            for fid in part_vid_fid['face_id']:
                if part_label not in fipm[fid]:
                    fipm[fid].append(part_label)
        # check if at most two mappings exist. If two, it is a boundary face, take the first one
        for fid, part_labels in fipm.items():
            if len(part_labels) > 2:
                print('Warning: more than two part labels for face {}'.format(fid))
                import ipdb;
                ipdb.set_trace()
            if len(part_labels) <= 2:
                fipm[fid] = part_labels[0]
        return fipm

    def vertex_id_to_part_mapping(self, hd_vert_on_fid, fid_to_part):
        vid_to_part = torch.mm(hd_vert_on_fid, fid_to_part)
        return vid_to_part

    def vertex_id_to_part_volume_mapping(self, per_part_volume, device):
        batch_size = per_part_volume.shape[0]
        self.hd_vid_in_part = self.hd_vid_in_part.to(device)
        hd_vid_in_part = self.hd_vid_in_part[None, :, :].repeat(batch_size, 1, 1)
        vid_to_vol = torch.bmm(hd_vid_in_part, per_part_volume[:, :, None])
        return vid_to_vol
stability/stability_metrics.py
sha2nkt/ipman-r
from utils import TrainOptions
from train import Trainer
from loguru import logger
import argparse
from config import run_grid_search_experiments

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--cfg', type=str, help='cfg file path')
    parser.add_argument('--opts', default=[], nargs='*', help='additional options to update config')
    parser.add_argument('--cfg_id', type=int, default=0, help='cfg id to run when multiple experiments are spawned')
    parser.add_argument('--cluster', default=False, action='store_true', help='creates submission files for cluster')
    parser.add_argument('--bid', type=int, default=30, help='amount of bid for cluster')
    parser.add_argument('--memory', type=int, default=20000, help='memory amount for cluster')
    parser.add_argument('--gpu_min_mem', type=int, default=11000, help='minimum amount of GPU memory')
    parser.add_argument('--gpu_arch', default=['tesla', 'quadro', 'rtx'],
                        nargs='*', help='additional options to update config')
    parser.add_argument('--num_cpus', type=int, default=8, help='num cpus for cluster')

    args = parser.parse_args()

    logger.info(f'Input arguments: \n {args}')

    hparams = run_grid_search_experiments(
        cfg_id=args.cfg_id,
        cfg_file=args.cfg,
        bid=args.bid,
        use_cluster=args.cluster,
        memory=args.memory,
        script='train.py',
        cmd_opts=args.opts,
        gpu_min_mem=args.gpu_min_mem,
        gpu_arch=args.gpu_arch,
    )
    trainer = Trainer(hparams)
    trainer.train()

train.py
sha2nkt/ipman-r
import torch
import numpy as np
import os
import cv2
from torchgeometry import angle_axis_to_rotation_matrix
import ipdb

import config
import constants

class FitsDict():
    """ Dictionary keeping track of the best fit per image in the training set """
    def __init__(self, options, train_dataset):
        self.options = options
        self.train_dataset = train_dataset
        self.fits_dict = {}
        # array used to flip SMPL pose parameters
        self.flipped_parts = torch.tensor(constants.SMPL_POSE_FLIP_PERM, dtype=torch.int64)
        # Load dictionary state
        for ds_name, ds in train_dataset.dataset_dict.items():
            ipdb.set_trace()
            try:
                dict_file = os.path.join(options.checkpoint_dir, ds_name + '_fits.npy')
                self.fits_dict[ds_name] = torch.from_numpy(np.load(dict_file))
            except IOError:
                # Dictionary does not exist, so populate with static fits
                dict_file = os.path.join(config.STATIC_FITS_DIR, ds_name + '_fits.npy')
                self.fits_dict[ds_name] = torch.from_numpy(np.load(dict_file))

    def save(self):
        """ Save dictionary state to disk """
        for ds_name in self.train_dataset.dataset_dict.keys():
            dict_file = os.path.join(self.options.checkpoint_dir, ds_name + '_fits.npy')
            np.save(dict_file, self.fits_dict[ds_name].cpu().numpy())

    def __getitem__(self, x):
        """ Retrieve dictionary entries """
        dataset_name, ind, rot, is_flipped = x
        batch_size = len(dataset_name)
        pose = torch.zeros((batch_size, 72))
        betas = torch.zeros((batch_size, 10))
        for ds, i, n in zip(dataset_name, ind, range(batch_size)):
            params = self.fits_dict[ds][i]
            pose[n, :] = params[:72]
            betas[n, :] = params[72:]
        pose = pose.clone()
        # Apply flipping and rotation
        pose = self.flip_pose(self.rotate_pose(pose, rot), is_flipped)
        betas = betas.clone()
        return pose, betas

    def __setitem__(self, x, val):
        """ Update dictionary entries """
        dataset_name, ind, rot, is_flipped, update = x
        pose, betas = val
        batch_size = len(dataset_name)
        # Undo flipping and rotation
        pose = self.rotate_pose(self.flip_pose(pose, is_flipped), -rot)
        params = torch.cat((pose, betas), dim=-1).cpu()
        for ds, i, n in zip(dataset_name, ind, range(batch_size)):
            if update[n]:
                self.fits_dict[ds][i] = params[n]

    def flip_pose(self, pose, is_flipped):
        """flip SMPL pose parameters"""
        is_flipped = is_flipped.byte()
        pose_f = pose.clone()
        pose_f[is_flipped, :] = pose[is_flipped][:, self.flipped_parts]
        # we also negate the second and the third dimension of the axis-angle representation
        pose_f[is_flipped, 1::3] *= -1
        pose_f[is_flipped, 2::3] *= -1
        return pose_f

    def rotate_pose(self, pose, rot):
        """Rotate SMPL pose parameters by rot degrees"""
        pose = pose.clone()
        cos = torch.cos(-np.pi * rot / 180.)
        sin = torch.sin(-np.pi * rot/ 180.)
        zeros = torch.zeros_like(cos)
        r3 = torch.zeros(cos.shape[0], 1, 3, device=cos.device)
        r3[:,0,-1] = 1
        R = torch.cat([torch.stack([cos, -sin, zeros], dim=-1).unsqueeze(1),
                       torch.stack([sin, cos, zeros], dim=-1).unsqueeze(1),
                       r3], dim=1)
        global_pose = pose[:, :3]
        global_pose_rotmat = angle_axis_to_rotation_matrix(global_pose)
        global_pose_rotmat_3b3 = global_pose_rotmat[:, :3, :3]
        global_pose_rotmat_3b3 = torch.matmul(R, global_pose_rotmat_3b3)
        global_pose_rotmat[:, :3, :3] = global_pose_rotmat_3b3
        global_pose_rotmat = global_pose_rotmat[:, :-1, :-1].cpu().numpy()
        global_pose_np = np.zeros((global_pose.shape[0], 3))
        for i in range(global_pose.shape[0]):
            aa, _ = cv2.Rodrigues(global_pose_rotmat[i])
            global_pose_np[i,:] = aa.squeeze()
        pose[:, :3] = torch.from_numpy(global_pose_np).to(pose.device)
        return pose

train/fits_dict.py
sha2nkt/ipman-r
import torch
import torch.nn as nn
import numpy as np
from torchgeometry import angle_axis_to_rotation_matrix, rotation_matrix_to_angle_axis
import cv2
import ipdb

from datasets import MixedDataset
from models import hmr, SMPL
from smplify import SMPLify
from utils.geometry import batch_rot2aa, batch_rodrigues, perspective_projection, estimate_translation, \
    batch_rectify_pose, ea2rm
# from utils.renderer import Renderer
from utils.base_trainer import BaseTrainer

# for eval
from tqdm import tqdm
from torch.utils.data import DataLoader
from utils.imutils import uncrop
from utils.pose_utils import reconstruction_error, compute_accel, compute_error_accel
from utils.renderer import Renderer
# from utils.part_utils import PartRenderer
from stability.ground_losses import StabilityLossCoP, GroundLoss
from vis_utils.world_vis import vis_vert_with_ground

import config
import constants
from .fits_dict import FitsDict


class Trainer(BaseTrainer):

    def init_fn(self):
        self.train_ds = MixedDataset(self.options, ignore_3d=self.options.ignore_3d, is_train=True)

        self.model = hmr(config.SMPL_MEAN_PARAMS, pretrained=True).to(self.device)
        self.optimizer = torch.optim.Adam(params=self.model.parameters(),
                                          lr=self.options.lr,
                                          weight_decay=0)
        self.smpl = SMPL(config.SMPL_MODEL_DIR,
                         batch_size=self.options.batch_size,
                         create_transl=False).to(self.device)
        # Per-vertex loss on the shape
        self.criterion_shape = nn.L1Loss().to(self.device)
        # Keypoint (2D and 3D) loss
        # No reduction because confidence weighting needs to be applied
        self.criterion_keypoints = nn.MSELoss(reduction='none').to(self.device)
        # Loss for SMPL parameter regression
        self.criterion_regr = nn.MSELoss().to(self.device)
        self.models_dict = {'model': self.model}
        self.optimizers_dict = {'optimizer': self.optimizer}
        self.focal_length = constants.FOCAL_LENGTH

        # Intuitive physics losses
        self.criterion_stability = StabilityLossCoP(
            faces = self.smpl.faces,
            cop_w = self.options.cop_w,
            cop_k = self.options.cop_k,
            contact_thresh = self.options.contact_thresh,
            model_type='smpl'
        )

        self.criterion_ground = GroundLoss(
            faces=self.smpl.faces,
            out_alpha1=self.options.out_alpha1,
            out_alpha2=self.options.out_alpha2,
            in_alpha1=self.options.in_alpha1,
            in_alpha2=self.options.in_alpha2,
            device=self.device,
            model_type='smpl'
        )

        # Initialize SMPLify fitting module
        self.smplify = SMPLify(step_size=1e-4, batch_size=self.options.batch_size,
                               num_iters=self.options.num_smplify_iters, focal_length=self.focal_length)
        if self.options.pretrained_checkpoint is not None:
            self.load_pretrained(checkpoint_file=self.options.pretrained_checkpoint)

        # Load dictionary of fits
        # self.fits_dict = FitsDict(self.options, self.train_ds)

        # Create renderer
        self.renderer = Renderer(focal_length=self.focal_length, img_res=self.options.img_res, faces=self.smpl.faces)

    def finalize(self):
        self.fits_dict.save()

    def keypoint_loss(self, pred_keypoints_2d, gt_keypoints_2d, openpose_weight, gt_weight):
        """ Compute 2D reprojection loss on the keypoints.
        The loss is weighted by the confidence.
        The available keypoints are different for each dataset.
        """
        conf = gt_keypoints_2d[:, :, -1].unsqueeze(-1).clone()
        conf[:, :25] *= openpose_weight
        conf[:, 25:] *= gt_weight
        loss = (conf * self.criterion_keypoints(pred_keypoints_2d, gt_keypoints_2d[:, :, :-1])).mean()
        return loss

    def keypoint_3d_loss(self, pred_keypoints_3d, gt_keypoints_3d, has_pose_3d):
        """Compute 3D keypoint loss for the examples that 3D keypoint annotations are available.
        The loss is weighted by the confidence.
        """
        pred_keypoints_3d = pred_keypoints_3d[:, 25:, :]
        conf = gt_keypoints_3d[:, :, -1].unsqueeze(-1).clone()
        gt_keypoints_3d = gt_keypoints_3d[:, :, :-1].clone()
        gt_keypoints_3d = gt_keypoints_3d[has_pose_3d == 1]
        conf = conf[has_pose_3d == 1]
        pred_keypoints_3d = pred_keypoints_3d[has_pose_3d == 1]
        if len(gt_keypoints_3d) > 0:
            gt_pelvis = (gt_keypoints_3d[:, 2, :] + gt_keypoints_3d[:, 3, :]) / 2
            gt_keypoints_3d = gt_keypoints_3d - gt_pelvis[:, None, :]
            pred_pelvis = (pred_keypoints_3d[:, 2, :] + pred_keypoints_3d[:, 3, :]) / 2
            pred_keypoints_3d = pred_keypoints_3d - pred_pelvis[:, None, :]
            return (conf * self.criterion_keypoints(pred_keypoints_3d, gt_keypoints_3d)).mean()
        else:
            return torch.FloatTensor(1).fill_(0.).to(self.device)

    def shape_loss(self, pred_vertices, gt_vertices, has_smpl):
        """Compute per-vertex loss on the shape for the examples that SMPL annotations are available."""
        pred_vertices_with_shape = pred_vertices[has_smpl == 1]
        gt_vertices_with_shape = gt_vertices[has_smpl == 1]
        if len(gt_vertices_with_shape) > 0:
            return self.criterion_shape(pred_vertices_with_shape, gt_vertices_with_shape)
        else:
            return torch.FloatTensor(1).fill_(0.).to(self.device)

    def smpl_losses(self, pred_rotmat, pred_betas, gt_pose, gt_betas, has_smpl):
        pred_rotmat_valid = pred_rotmat[has_smpl == 1]
        gt_rotmat_valid = batch_rodrigues(gt_pose.view(-1, 3)).view(-1, 24, 3, 3)[has_smpl == 1]
        pred_betas_valid = pred_betas[has_smpl == 1]
        gt_betas_valid = gt_betas[has_smpl == 1]
        if len(pred_rotmat_valid) > 0:
            loss_regr_pose = self.criterion_regr(pred_rotmat_valid, gt_rotmat_valid)
            loss_regr_betas = self.criterion_regr(pred_betas_valid, gt_betas_valid)
        else:
            loss_regr_pose = torch.FloatTensor(1).fill_(0.).to(self.device)
            loss_regr_betas = torch.FloatTensor(1).fill_(0.).to(self.device)
        return loss_regr_pose, loss_regr_betas

    def ip_losses(self, pred_rotmat, pred_betas, gt_cam_r, gt_joints_world, ground_offset, valid):
        """Intuitive physics losses"""

        ## bring predictions to world coordinates:
        # 1. apply GT cam on pred_pose
        # 2. compute tranlsation to align pred_pelvis to gt_pelvis
        # 3. rotate the vertices to align with y-axis (instead of z-axis in smpl)
        batch_size = pred_rotmat.shape[0]
        pred_rotmat_global_orient_world = torch.bmm(gt_cam_r.permute(0, 2, 1), pred_rotmat[:, 0])

        pred_output_world = self.smpl(betas=pred_betas,
                                    body_pose=pred_rotmat[:, 1:],
                                    global_orient=pred_rotmat_global_orient_world.unsqueeze(1),
                                      pose2rot=False)
        pred_vertices_world = pred_output_world.vertices
        pred_joints_world = pred_output_world.joints[:, 25:, :]

        # compute pred transl world by equating it to the gt world pelvis
        gt_joints_world = gt_joints_world[:, :, :-1] # removing the confidence values as always 1
        gt_world_pelvis = (gt_joints_world[:, 2, :] + gt_joints_world[:, 3, :]) / 2
        pred_pelvis_world = (pred_joints_world[:, 2, :] + pred_joints_world[:, 3, :]) / 2
        pred_transl_world = gt_world_pelvis - pred_pelvis_world

        pred_vertices_world += pred_transl_world[:, None, :]
        # rotate the vertices to align with the pyrender coordinates (yaxis-up), not aligned with y axis
        R1 = ea2rm(torch.tensor([[np.radians(270)]]), torch.tensor([[np.radians(0)]]),
                   torch.tensor([[np.radians(0)]])).float().to(self.device).expand(pred_vertices_world.shape[0], 3, 3)
        pred_vertices_world = torch.einsum('bki,bji->bjk', [R1, pred_vertices_world])
        ground_offset = torch.hstack([torch.zeros_like(ground_offset), ground_offset, torch.zeros_like(ground_offset)])
        pred_vertices_world += -ground_offset[:, None, :]

        # # rotate the vertices to align with the pyrender coordinates (yaxis-up), not aligned with y axis
        # gt_vertices_world = torch.einsum('bki,bji->bjk', [R1, gt_vertices_world])
        # gt_vertices_world += -ground_offset[:, None, :]
        #
        # import trimesh
        # mesh = trimesh.Trimesh(pred_vertices_world[0].detach().squeeze().cpu().numpy(), self.smpl.faces, process=False,
        #                        maintain_order=True)
        # mesh.export('pred_mesh_in_training.obj')
        # mesh = trimesh.Trimesh(gt_vertices_world[0].detach().squeeze().cpu().numpy(), self.smpl.faces, process=False,
        #                        maintain_order=True)
        # mesh.export('gt_mesh_in_training.obj')
        # import ipdb; ipdb.set_trace()


        # stability loss
        valid = valid.to(torch.bool)
        stability_loss = self.criterion_stability(pred_vertices_world)
        stability_loss = stability_loss[valid].mean()

        # ground loss
        out_pull_loss, in_push_loss  = self.criterion_ground(pred_vertices_world)
        out_pull_loss = out_pull_loss[valid].mean()
        in_push_loss = in_push_loss[valid].mean()
        return stability_loss, out_pull_loss, in_push_loss, pred_vertices_world


    def train_step(self, input_batch):
        self.model.train()

        # Get data from the batch
        images = input_batch['img']  # input image
        gt_keypoints_2d = input_batch['keypoints']  # 2D keypoints
        gt_pose = input_batch['pose']  # SMPL pose parameters
        gt_betas = input_batch['betas']  # SMPL beta parameters
        gt_joints = input_batch['pose_3d']  # 3D pose
        gt_cam_r = input_batch['cam_r'].to(torch.float32)
        gt_joints_world = input_batch['pose_3d_world'].to(torch.float32)
        gt_in_bos_label = input_batch['in_bos_label']
        gt_contact_label = input_batch['contact_label']
        has_smpl = input_batch['has_smpl'].byte()  # flag that indicates whether SMPL parameters are valid
        has_smpl_world = input_batch['has_smpl_world'].byte()  # flag that indicates whether SMPL world parameters are valid
        has_pose_3d = input_batch['has_pose_3d'].byte()  # flag that indicates whether 3D pose is valid
        is_flipped = input_batch['is_flipped']  # flag that indicates whether image was flipped during data augmentation
        rot_angle = input_batch['rot_angle']  # rotation angle used for data augmentation
        dataset_name = input_batch['dataset_name']  # name of the dataset the image comes from
        indices = input_batch['sample_index']  # index of example inside its dataset
        batch_size = images.shape[0]

        h36m_mask = torch.BoolTensor([1 if k == 'h36m' else 0 for k in dataset_name]).to(self.device)
        h36m_ground_offset = torch.FloatTensor([constants.GROUND_OFFSETS['h36m']]).to(self.device).to(torch.float32).expand(batch_size)
        ground_offset = input_batch['ground_offset'].to(torch.float32)
        ground_offset = ground_offset.masked_scatter_(h36m_mask, h36m_ground_offset)
        ground_offset = ground_offset[:, None]
        # Get GT vertices and model joints
        # Note that gt_model_joints is different from gt_joints as it comes from SMPL
        # gt_out = self.smpl(betas=gt_betas, body_pose=gt_pose[:, 3:], global_orient=gt_pose[:, :3])
        # gt_model_joints = gt_out.joints
        # gt_vertices = gt_out.vertices

        # # Get GT vertices and model joints in world coordinates
        # # Note that gt_model_joints is different from gt_joints as it comes from SMPL
        # gt_out_world = self.smpl(betas=gt_betas,
        #                    body_pose=gt_pose_world[:, 3:],
        #                    global_orient=gt_pose_world[:, :3],
        #                    transl=gt_transl_world)
        # gt_vertices_world = gt_out_world.vertices

        # Get current best fits from the dictionary
        # opt_pose, opt_betas = self.fits_dict[(dataset_name, indices.cpu(), rot_angle.cpu(), is_flipped.cpu())]
        # opt_pose = opt_pose.to(self.device)
        # opt_betas = opt_betas.to(self.device)
        # opt_output = self.smpl(betas=opt_betas, body_pose=opt_pose[:, 3:], global_orient=opt_pose[:, :3])
        # opt_vertices = opt_output.vertices
        # opt_joints = opt_output.joints

        # Get gt fits
        gt_pose = gt_pose.to(self.device)
        gt_betas = gt_betas.to(self.device)
        gt_output = self.smpl(betas=gt_betas,
                              body_pose=gt_pose[:, 3:],
                              global_orient=gt_pose[:, :3])
        gt_vertices = gt_output.vertices
        gt_model_joints = gt_output.joints

        # De-normalize 2D keypoints from [-1,1] to pixel space
        gt_keypoints_2d_orig = gt_keypoints_2d.clone()
        gt_keypoints_2d_orig[:, :, :-1] = 0.5 * self.options.img_res * (gt_keypoints_2d_orig[:, :, :-1] + 1)

        # # Estimate camera translation given the model joints and 2D keypoints
        # # by minimizing a weighted least squares loss
        # gt_cam_t = estimate_translation(gt_model_joints, gt_keypoints_2d_orig, focal_length=self.focal_length,
        #                                 img_size=self.options.img_res)


        gt_cam_t = estimate_translation(gt_model_joints,
                                        gt_keypoints_2d_orig,
                                        focal_length=self.focal_length,
                                        img_size=self.options.img_res,
                                        use_all_joints=True if '3dpw' in dataset_name else False)

        # Feed images in the network to predict camera and SMPL parameters
        pred_rotmat, pred_betas, pred_camera = self.model(images)
        pred_output = self.smpl(betas=pred_betas, body_pose=pred_rotmat[:, 1:],
                                global_orient=pred_rotmat[:, 0].unsqueeze(1), pose2rot=False)
        pred_vertices = pred_output.vertices
        pred_joints = pred_output.joints

        # Convert Weak Perspective Camera [s, tx, ty] to camera translation [tx, ty, tz] in 3D given the bounding box size
        # This camera translation can be used in a full perspective projection
        pred_cam_t = torch.stack([pred_camera[:, 1],
                                  pred_camera[:, 2],
                                  2 * self.focal_length / (self.options.img_res * pred_camera[:, 0] + 1e-9)], dim=-1)

        camera_center = torch.zeros(batch_size, 2, device=self.device)
        pred_keypoints_2d = perspective_projection(pred_joints,
                                                   rotation=torch.eye(3, device=self.device).unsqueeze(0).expand(
                                                       batch_size, -1, -1),
                                                   translation=pred_cam_t,
                                                   focal_length=self.focal_length,
                                                   camera_center=camera_center)
        # Normalize keypoints to [-1,1]
        pred_keypoints_2d = pred_keypoints_2d / (self.options.img_res / 2.)

        valid_fit = has_smpl
        # Add the examples with GT parameters to the list of valid fits
        valid_fit = valid_fit | has_smpl

        # Compute loss on SMPL parameters
        loss_regr_pose, loss_regr_betas = self.smpl_losses(pred_rotmat, pred_betas, gt_pose, gt_betas, valid_fit)

        # Compute 2D reprojection loss for the keypoints
        loss_keypoints = self.keypoint_loss(pred_keypoints_2d, gt_keypoints_2d,
                                            self.options.openpose_train_weight,
                                            self.options.gt_train_weight)

        # Compute 3D keypoint loss
        loss_keypoints_3d = self.keypoint_3d_loss(pred_joints, gt_joints, has_pose_3d)

        # Per-vertex loss for the shape
        loss_shape = self.shape_loss(pred_vertices, gt_vertices, valid_fit)

        # Intuitive physics losses
        if self.options.stability_loss_weight != 0 or \
            self.options.inside_push_loss_weight !=0 or\
            self.options.outside_pull_loss_weight !=0:
            # convert in_bos_label and contact label to one-hot (all unstable {0, -1} go to 0}
            gt_in_bos_label = (gt_in_bos_label == 1).byte()
            gt_contact_label = (gt_contact_label == 1).byte()
            valid_fit_stability = has_smpl & gt_in_bos_label & gt_contact_label
            loss_stability, loss_out_pull, loss_in_push, pred_vertices_world = self.ip_losses(pred_rotmat, pred_betas, gt_cam_r, gt_joints_world, ground_offset, valid_fit_stability)
        else:
            loss_stability = torch.FloatTensor(1).fill_(0.).to(self.device)
            loss_out_pull = torch.FloatTensor(1).fill_(0.).to(self.device)
            loss_in_push = torch.FloatTensor(1).fill_(0.).to(self.device)
            pred_vertices_world = pred_vertices

        # Compute total loss
        # The last component is a loss that forces the network to predict positive depth values
        loss = self.options.shape_loss_weight * loss_shape + \
               self.options.keypoint_loss_weight * loss_keypoints + \
               self.options.keypoint_loss_weight * loss_keypoints_3d + \
               self.options.pose_loss_weight * loss_regr_pose + \
               self.options.beta_loss_weight * loss_regr_betas + \
               self.options.stability_loss_weight * loss_stability + \
               self.options.inside_push_loss_weight * loss_in_push + \
               self.options.outside_pull_loss_weight * loss_out_pull + \
               ((torch.exp(-pred_camera[:, 0] * 10)) ** 2).mean()
        loss *= 60

        # Do backprop
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # Pack output arguments for tensorboard logging
        output = {'pred_vertices': pred_vertices.detach(),
                  'pred_vertices_world': pred_vertices_world.detach(),
                  'opt_vertices': gt_vertices,
                  'pred_cam_t': pred_cam_t.detach(),
                  'opt_cam_t': gt_cam_t}
        losses = {'loss': loss.detach().item(),
                  'loss_keypoints': self.options.keypoint_loss_weight * loss_keypoints.detach().item(),
                  'loss_keypoints_3d': self.options.keypoint_loss_weight * loss_keypoints_3d.detach().item(),
                  'loss_regr_pose': self.options.pose_loss_weight * loss_regr_pose.detach().item(),
                  'loss_regr_betas': self.options.beta_loss_weight * loss_regr_betas.detach().item(),
                  'loss_shape': self.options.shape_loss_weight * loss_shape.detach().item(),
                  'loss_stability': self.options.stability_loss_weight * loss_stability.detach().item(),
                  'loss_in_push': self.options.inside_push_loss_weight * loss_in_push.detach().item(),
                  'loss_out_pull': self.options.outside_pull_loss_weight * loss_out_pull.detach().item()}
        return output, losses

    def train_summaries(self, input_batch, output, losses):
        images = input_batch['img']
        images = images * torch.tensor([0.229, 0.224, 0.225], device=images.device).reshape(1, 3, 1, 1)
        images = images + torch.tensor([0.485, 0.456, 0.406], device=images.device).reshape(1, 3, 1, 1)

        pred_vertices = output['pred_vertices']
        opt_vertices = output['opt_vertices']
        pred_cam_t = output['pred_cam_t']
        opt_cam_t = output['opt_cam_t']
        images_pred = self.renderer.visualize_tb(pred_vertices, pred_cam_t, images)
        images_opt = self.renderer.visualize_tb(opt_vertices, opt_cam_t, images)
        self.summary_writer.add_image('pred_shape', images_pred, self.step_count)


    def eval_step(self, epoch, dataset_name, dataset, result_file,
                       batch_size=32, img_res=224,
                       num_workers=32, shuffle=False, log_freq=None, step_count=-1):
        """Run evaluation on the datasets and metrics we report in the paper. """
        self.model.eval()

        device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')


        # Load SMPL model
        smpl_neutral = SMPL(config.SMPL_MODEL_DIR,
                            create_transl=False).to(device)
        smpl_male = SMPL(config.SMPL_MODEL_DIR,
                         gender='male',
                         create_transl=False).to(device)
        smpl_female = SMPL(config.SMPL_MODEL_DIR,
                           gender='female',
                           create_transl=False).to(device)

        # renderer = PartRenderer()

        # Regressor for H36m joints
        J_regressor = torch.from_numpy(np.load(config.JOINT_REGRESSOR_H36M)).float()

        save_results = result_file is not None
        # Disable shuffling if you want to save the results
        if save_results:
            shuffle = False
        # Create dataloader for the dataset
        data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)

        # Pose metrics
        # MPJPE and Reconstruction error for the non-parametric and parametric shapes
        mpjpe = np.zeros(len(dataset))
        recon_err = np.zeros(len(dataset))
        mpjpe_smpl = np.zeros(len(dataset))
        recon_err_smpl = np.zeros(len(dataset))
        v2v_err = np.zeros(len(dataset))

        # Acceleration metrics
        accel_ = np.zeros(len(dataset))
        accel_err_ = np.zeros(len(dataset))

        # Stability metrics
        gt_bos_accumulator = []
        gt_contact_accumulator = []
        gt_contact_mask_accumulator = []
        pred_bos_accumulator = []
        pred_contact_accumulator = []
        pred_contact_mask_accumulator = []

        # Shape metrics
        # Mean per-vertex error
        shape_err = np.zeros(len(dataset))
        shape_err_smpl = np.zeros(len(dataset))

        # Mask and part metrics
        # Accuracy
        accuracy = 0.
        parts_accuracy = 0.
        # True positive, false positive and false negative
        tp = np.zeros((2, 1))
        fp = np.zeros((2, 1))
        fn = np.zeros((2, 1))
        parts_tp = np.zeros((7, 1))
        parts_fp = np.zeros((7, 1))
        parts_fn = np.zeros((7, 1))
        # Pixel count accumulators
        pixel_count = 0
        parts_pixel_count = 0

        # Store SMPL parameters
        smpl_pose = np.zeros((len(dataset), 72))
        smpl_betas = np.zeros((len(dataset), 10))
        smpl_camera = np.zeros((len(dataset), 3))
        pred_joints = np.zeros((len(dataset), 17, 3))

        eval_pose = False
        eval_masks = False
        eval_parts = False
        # Choose appropriate evaluation for each dataset
        if dataset_name == 'h36m-p1' or dataset_name == 'h36m-p2' \
                or dataset_name == 'h36m-test-s1' or dataset_name == '3dpw' \
                or dataset_name == 'mpi-inf-3dhp' \
                or dataset_name == 'rich-val-onlycam0' or dataset_name == 'rich-test-onlycam0' \
                or dataset_name == 'rich-val' or dataset_name == 'rich-test':
            eval_pose = True
        elif dataset_name == 'lsp':
            eval_masks = True
            eval_parts = True
            annot_path = config.DATASET_FOLDERS['upi-s1h']

        joint_mapper_h36m = constants.H36M_TO_J17 if dataset_name == 'mpi-inf-3dhp' else constants.H36M_TO_J14
        joint_mapper_gt = constants.J24_TO_J17 if dataset_name == 'mpi-inf-3dhp' else constants.J24_TO_J14
        # Iterate over the entire dataset
        for step, batch in enumerate(tqdm(data_loader, desc='Eval', total=len(data_loader))):
            # Get ground truth annotations from the batch
            gt_pose = batch['pose'].to(device)
            gt_betas = batch['betas'].to(device)
            gt_vertices = smpl_neutral(betas=gt_betas, body_pose=gt_pose[:, 3:], global_orient=gt_pose[:, :3]).vertices
            images = batch['img'].to(device)
            gender = batch['gender'].to(device)
            if 'rich' in dataset_name:
                ground_offset = batch['ground_offset'][:, None].to(device)
            elif 'h36m' in dataset_name:
                ground_offset = constants.GROUND_OFFSETS['h36m']
            else: #for 3dpw
                ground_offset = 0.0

            curr_batch_size = images.shape[0]

            # visualize mesh in world coordinates
            if len(torch.unique(batch['has_smpl_world'])) == 1 and torch.unique(batch['has_smpl_world'])[0] == 1:
                gt_pose_world = batch['pose_world'].to(device)
                gt_transl_world = batch['transl_world'].to(device)
                gt_body_world = smpl_neutral(betas=gt_betas,
                                             body_pose=gt_pose_world[:, 3:],
                                             global_orient=gt_pose_world[:, :3],
                                             transl=gt_transl_world)

                gt_joints_world = gt_body_world.joints[:, 25:, :]
                # visualize the gt vertices in world coords
                gt_in_bos_label, gt_contact_metric, gt_contact_mask = vis_vert_with_ground(gt_body_world.vertices,
                                                                              gt_transl_world[:, None, :]*0.0,
                                                                              seq_name='test-gt_vert',
                                                                              start_idx=step * curr_batch_size,
                                                                              sub_sample=1,
                                                                              smpl_batch_size=curr_batch_size,
                                                                              ground_offset=ground_offset,
                                                                              visualize=False)
                if gt_in_bos_label.dim() == 0:
                    gt_in_bos_label = gt_in_bos_label.view(-1)
                    gt_contact_metric = gt_contact_metric.view(-1)
                    gt_contact_mask = gt_contact_mask.view(1, -1)

                gt_bos_accumulator.append(gt_in_bos_label)
                gt_contact_accumulator.append(gt_contact_metric)
                gt_contact_mask_accumulator.append(gt_contact_mask)

                has_smpl_world = True
            else:
                gt_joints_world = batch['pose_3d_world'][:, :, :-1].to(device)
                has_smpl_world = False

            # Get ground truth camera
            gt_cam_r = batch['cam_r'].to(device).to(torch.float32)

            with torch.no_grad():
                pred_rotmat, pred_betas, pred_camera = self.model(images)
                pred_pose = batch_rot2aa(pred_rotmat.view(-1, 3, 3)).view(curr_batch_size, -1)
                pred_output = smpl_neutral(betas=pred_betas, body_pose=pred_rotmat[:, 1:],
                                           global_orient=pred_rotmat[:, 0].unsqueeze(1), pose2rot=False)
                pred_vertices = pred_output.vertices
                pred_joints = pred_output.joints

                ## move to world coordinates - first correct the predicted pose using gt camera rotation, then correct the translation by pelvis alignment
                pred_pose_world = pred_pose.clone()
                pred_pose_world[:, :3] = batch_rectify_pose(gt_cam_r.permute(0, 2, 1), pred_pose_world[:, :3])
                pred_output_world = smpl_neutral(betas=pred_betas,
                                                 body_pose=pred_pose_world[:, 3:],
                                                 global_orient=pred_pose_world[:, :3])
                pred_vertices_world = pred_output_world.vertices
                pred_joints_world = pred_output_world.joints[:, 25:, :]
                # compute pred transl world by equating it to the gt world pelvis
                gt_world_pelvis = (gt_joints_world[:, 2, :] + gt_joints_world[:, 3, :]) / 2
                pred_pelvis_world = (pred_joints_world[:, 2, :] + pred_joints_world[:, 3, :]) / 2
                pred_transl_world = gt_world_pelvis - pred_pelvis_world
                pred_in_bos_label, pred_contact_metric, pred_contact_mask = vis_vert_with_ground(pred_vertices_world,
                                                                              pred_transl_world[:, None, :],
                                                                              seq_name='test-pred',
                                                                              start_idx=step * curr_batch_size,
                                                                              sub_sample=1,
                                                                              smpl_batch_size=curr_batch_size,
                                                                              ground_offset=ground_offset,
                                                                              visualize=False)

                if pred_in_bos_label.dim() == 0:
                    pred_in_bos_label = pred_in_bos_label.view(-1)
                    pred_contact_metric = pred_contact_metric.view(-1)
                    pred_contact_mask = pred_contact_mask.view(1, -1)

                pred_bos_accumulator.append(pred_in_bos_label)
                pred_contact_accumulator.append(pred_contact_metric)
                pred_contact_mask_accumulator.append(pred_contact_mask)

            if save_results:
                rot_pad = torch.tensor([0, 0, 1], dtype=torch.float32, device=device).view(1, 3, 1)
                rotmat = torch.cat((pred_rotmat.view(-1, 3, 3), rot_pad.expand(curr_batch_size * 24, -1, -1)), dim=-1)
                pred_pose = tgm.rotation_matrix_to_angle_axis(rotmat).contiguous().view(-1, 72)
                smpl_pose[step * batch_size:step * batch_size + curr_batch_size, :] = pred_pose.cpu().numpy()
                smpl_betas[step * batch_size:step * batch_size + curr_batch_size, :] = pred_betas.cpu().numpy()
                smpl_camera[step * batch_size:step * batch_size + curr_batch_size, :] = pred_camera.cpu().numpy()

            # 3D pose evaluation
            if eval_pose:
                # Regressor broadcasting
                J_regressor_batch = J_regressor[None, :].expand(pred_vertices.shape[0], -1, -1).to(device)
                # Get 14 ground truth joints
                if 'h36m' in dataset_name or 'mpi-inf' in dataset_name:
                    gt_keypoints_3d = batch['pose_3d'].cuda()
                    gt_keypoints_3d = gt_keypoints_3d[:, joint_mapper_gt, :-1]
                # For 3DPW get the 14 common joints from the rendered shape
                if '3dpw' in dataset_name:
                    gt_vertices = smpl_male(global_orient=gt_pose[:, :3], body_pose=gt_pose[:, 3:],
                                            betas=gt_betas).vertices
                    gt_vertices_female = smpl_female(global_orient=gt_pose[:, :3], body_pose=gt_pose[:, 3:],
                                                     betas=gt_betas).vertices
                    gt_vertices[gender == 1, :, :] = gt_vertices_female[gender == 1, :, :]
                    gt_keypoints_3d = torch.matmul(J_regressor_batch, gt_vertices)
                    gt_pelvis = gt_keypoints_3d[:, [0], :].clone()
                    gt_keypoints_3d = gt_keypoints_3d[:, joint_mapper_h36m, :]
                    gt_keypoints_3d = gt_keypoints_3d - gt_pelvis
                if 'rich' in dataset_name:
                    gt_keypoints_3d = torch.matmul(J_regressor_batch, gt_vertices)
                    gt_pelvis = gt_keypoints_3d[:, [0], :].clone()
                    gt_keypoints_3d = gt_keypoints_3d[:, joint_mapper_h36m, :]
                    gt_keypoints_3d = gt_keypoints_3d - gt_pelvis
                    gt_vertices = gt_vertices - gt_pelvis

                    # Get 14 predicted joints from the mesh
                pred_keypoints_3d = torch.matmul(J_regressor_batch, pred_vertices)
                if save_results:
                    pred_joints[step * batch_size:step * batch_size + curr_batch_size, :,
                    :] = pred_keypoints_3d.cpu().numpy()
                pred_pelvis = pred_keypoints_3d[:, [0], :].clone()
                pred_keypoints_3d = pred_keypoints_3d[:, joint_mapper_h36m, :]
                pred_keypoints_3d = pred_keypoints_3d - pred_pelvis
                pred_vertices = pred_vertices - pred_pelvis

                # Absolute error (MPJPE)
                error = torch.sqrt(((pred_keypoints_3d - gt_keypoints_3d) ** 2).sum(dim=-1)).mean(dim=-1).cpu().numpy()
                mpjpe[step * batch_size:step * batch_size + curr_batch_size] = error

                # Reconstuction_error
                r_error = reconstruction_error(pred_keypoints_3d.cpu().numpy(), gt_keypoints_3d.cpu().numpy(),
                                               reduction=None)
                recon_err[step * batch_size:step * batch_size + curr_batch_size] = r_error

                # V2V error in RICH
                if 'rich' in dataset_name:
                    v2v = torch.sqrt(((gt_vertices - pred_vertices) ** 2).sum(dim=-1)).mean(dim=-1).cpu().numpy()
                    v2v_err[step * batch_size:step * batch_size + curr_batch_size] = v2v

                # Acceleration Error
                if pred_keypoints_3d.shape[0] >= 3: # can only calculate accel if there are at least 3 frames
                    accel = np.mean(compute_accel(pred_keypoints_3d.cpu().numpy()))
                    if not np.isnan(accel):
                        accel_[step * batch_size:step * batch_size + curr_batch_size] = accel
                    accel_err = np.mean(compute_error_accel(joints_pred=pred_keypoints_3d.cpu().numpy(), joints_gt=gt_keypoints_3d.cpu().numpy()))
                    if not np.isnan(accel_err):
                        accel_err_[step * batch_size:step * batch_size + curr_batch_size] = accel_err

            # If mask or part evaluation, render the mask and part images
            # if eval_masks or eval_parts:
                # mask, parts = renderer(pred_vertices, pred_camera)

            # Mask evaluation (for LSP)
            if eval_masks:
                center = batch['center'].cpu().numpy()
                scale = batch['scale'].cpu().numpy()
                # Dimensions of original image
                orig_shape = batch['orig_shape'].cpu().numpy()
                for i in range(curr_batch_size):
                    # After rendering, convert imate back to original resolution
                    pred_mask = uncrop(mask[i].cpu().numpy(), center[i], scale[i], orig_shape[i]) > 0
                    # Load gt mask
                    gt_mask = cv2.imread(os.path.join(annot_path, batch['maskname'][i]), 0) > 0
                    # Evaluation consistent with the original UP-3D code
                    accuracy += (gt_mask == pred_mask).sum()
                    pixel_count += np.prod(np.array(gt_mask.shape))
                    for c in range(2):
                        cgt = gt_mask == c
                        cpred = pred_mask == c
                        tp[c] += (cgt & cpred).sum()
                        fp[c] += (~cgt & cpred).sum()
                        fn[c] += (cgt & ~cpred).sum()
                    f1 = 2 * tp / (2 * tp + fp + fn)

            # Part evaluation (for LSP)
            if eval_parts:
                center = batch['center'].cpu().numpy()
                scale = batch['scale'].cpu().numpy()
                orig_shape = batch['orig_shape'].cpu().numpy()
                for i in range(curr_batch_size):
                    pred_parts = uncrop(parts[i].cpu().numpy().astype(np.uint8), center[i], scale[i], orig_shape[i])
                    # Load gt part segmentation
                    gt_parts = cv2.imread(os.path.join(annot_path, batch['partname'][i]), 0)
                    # Evaluation consistent with the original UP-3D code
                    # 6 parts + background
                    for c in range(7):
                        cgt = gt_parts == c
                        cpred = pred_parts == c
                        cpred[gt_parts == 255] = 0
                        parts_tp[c] += (cgt & cpred).sum()
                        parts_fp[c] += (~cgt & cpred).sum()
                        parts_fn[c] += (cgt & ~cpred).sum()
                    gt_parts[gt_parts == 255] = 0
                    pred_parts[pred_parts == 255] = 0
                    parts_f1 = 2 * parts_tp / (2 * parts_tp + parts_fp + parts_fn)
                    parts_accuracy += (gt_parts == pred_parts).sum()
                    parts_pixel_count += np.prod(np.array(gt_parts.shape))

            # # Print intermediate results during evaluation
            # if step % log_freq == log_freq - 1:
            #     if eval_pose:
            #         print('MPJPE: ' + str(1000 * mpjpe[:step * batch_size].mean()))
            #         print('Reconstruction Error: ' + str(1000 * recon_err[:step * batch_size].mean()))
            #         print()
            #     if eval_masks:
            #         print('Accuracy: ', accuracy / pixel_count)
            #         print('F1: ', f1.mean())
            #         print()
            #     if eval_parts:
            #         print('Parts Accuracy: ', parts_accuracy / parts_pixel_count)
            #         print('Parts F1 (BG): ', parts_f1[[0, 1, 2, 3, 4, 5, 6]].mean())
            #         print()

        # find accuracy of bos metric
        pred_bos_accumulator = torch.cat(pred_bos_accumulator, dim=0).to(torch.float32)
        pred_contact_accumulator = torch.cat(pred_contact_accumulator, dim=0).to(torch.float32)
        pred_contact_mask_accumulator = torch.cat(pred_contact_mask_accumulator, dim=0).to(torch.float32)
        print('*** Stability Metrics ***')
        print()
        if has_smpl_world:
            gt_bos_accumulator = torch.cat(gt_bos_accumulator, dim=0)
            pred_bos_accuracy = (gt_bos_accumulator == pred_bos_accumulator).sum().item() / \
                                gt_bos_accumulator.shape[0]
            gt_contact_accumulator = torch.cat(gt_contact_accumulator, dim=0)
            gt_contact_mask_accumulator = torch.cat(gt_contact_mask_accumulator, dim=0)
            pred_contact_accuracy = (gt_contact_accumulator == pred_contact_accumulator).sum().item() / \
                                    gt_contact_accumulator.shape[0]
            dense_contact_accuracy =  (gt_contact_mask_accumulator == pred_contact_mask_accumulator).sum().item() / \
                                      gt_contact_mask_accumulator.numel()
            print(f'{dataset_name} BOS Accuracy: ', pred_bos_accuracy)
            print(f'{dataset_name} Contact Accuracy: ', pred_contact_accuracy)
            print(f'{dataset_name} Contact Mask Accuracy: ', dense_contact_accuracy)
            self.summary_writer.add_scalar(f'{dataset_name} BOS Accuracy', pred_bos_accuracy, epoch)
            self.summary_writer.add_scalar(f'{dataset_name} Contact Accuracy', pred_contact_accuracy, epoch)
            self.summary_writer.add_scalar(f'{dataset_name} Contact Mask Accuracy', dense_contact_accuracy, epoch)
        else:
            bos_mean = torch.mean(pred_bos_accumulator)
            contact_mean = torch.mean(pred_contact_accumulator)
            print(f'{dataset_name} BOS Mean: ', bos_mean.cpu().numpy())
            print(f'{dataset_name} Contact Mean: ', contact_mean.cpu().numpy())
            self.summary_writer.add_scalar(f'{dataset_name} BOS Mean', bos_mean.cpu().numpy(), epoch)
            self.summary_writer.add_scalar(f'{dataset_name} Contact Mean', contact_mean.cpu().numpy(), epoch)

        # Save reconstructions to a file for further processing
        if save_results:
            np.savez(result_file, pred_joints=pred_joints, pose=smpl_pose, betas=smpl_betas, camera=smpl_camera)
        # Print final results during evaluation
        print('*** Final Results ***')
        print()
        if eval_pose:
            print(f'{dataset_name} MPJPE: ' + str(1000 * mpjpe.mean()))
            print(f'{dataset_name} Reconstruction Error: ' + str(1000 * recon_err.mean()))
            print(f'{dataset_name} V2V Error: ' + str(1000 * v2v_err.mean()))
            print(f'{dataset_name} Accel: ' + str(1000 * accel_.mean()))
            print(f'{dataset_name} Accel Error: ' + str(1000 * accel_err_.mean()))
            print()
            self.summary_writer.add_scalar(f'{dataset_name} MPJPE', 1000 * mpjpe.mean(), epoch)
            self.summary_writer.add_scalar(f'{dataset_name} Reconstruction Error', 1000 * recon_err.mean(), epoch)
            self.summary_writer.add_scalar(f'{dataset_name} V2V Error', 1000 * v2v_err.mean(), epoch)
            self.summary_writer.add_scalar(f'{dataset_name} Accel', 1000 * accel_.mean(), epoch)
            self.summary_writer.add_scalar(f'{dataset_name} Acceleration Error', 1000 * accel_err_.mean(), epoch)
        if eval_masks:
            print('Accuracy: ', accuracy / pixel_count)
            print('F1: ', f1.mean())
            print()
        if eval_parts:
            print('Parts Accuracy: ', parts_accuracy / parts_pixel_count)
            print('Parts F1 (BG): ', parts_f1[[0, 1, 2, 3, 4, 5, 6]].mean())




train/trainer.py
sha2nkt/ipman-r
from __future__ import division
import sys
import time
import ipdb

import torch
import random
from tqdm import tqdm
tqdm.monitor_interval = 0
from torch.utils.tensorboard import SummaryWriter
from datasets import BaseDataset, BaseDatasetEval


from utils import CheckpointDataLoader, CheckpointSaver

class BaseTrainer(object):
    """Base class for Trainer objects.
    Takes care of checkpointing/logging/resuming training.
    """
    def __init__(self, options):
        self.options = options
        self.endtime = time.time() + self.options.time_to_run
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        # override this function to define your model, optimizers etc.
        self.init_fn()
        self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
        self.summary_writer = SummaryWriter(self.options.summary_dir)
        self.run_eval_h36m = options.run_eval_h36m
        self.run_eval_rich_val = options.run_eval_rich_val
        self.run_eval_rich_test = options.run_eval_rich_test
        self.run_eval_3dpw = options.run_eval_3dpw

        self.is_agora = options.is_agora
        self.is_others = options.is_others
        self.run_smplify = options.run_smplify

        self.checkpoint = None
        if self.options.resume and self.saver.exists_checkpoint():
            self.checkpoint = self.saver.load_checkpoint(self.models_dict, self.optimizers_dict, checkpoint_file=self.options.checkpoint)

        if self.checkpoint is None:
            self.epoch_count = 0
            self.step_count = 0
        else:
            self.epoch_count = self.checkpoint['epoch']
            self.step_count = self.checkpoint['total_step_count']

        if self.run_eval_h36m:
            self.eval_data_loader = BaseDatasetEval(None, 'h36m-p2', is_train=False)
        if self.run_eval_rich_val:
            self.eval_data_loader_rich_val = BaseDatasetEval(None, 'rich-val', is_train=False)
        if self.run_eval_rich_test:
            self.eval_data_loader_rich_test = BaseDatasetEval(None, 'rich-test', is_train=False)
        if self.run_eval_3dpw:
            self.eval_data_loader_3dpw = BaseDatasetEval(None, '3dpw', is_train=False)
            # self.eval_data_loader_agora = BaseDataset(None, 'agora_val', is_train=False)

    def load_pretrained(self, checkpoint_file=None):
        """Load a pretrained checkpoint.
        This is different from resuming training using --resume.
        """
        if checkpoint_file is not None:
            checkpoint = torch.load(checkpoint_file)
            for model in self.models_dict:
                if model in checkpoint:
                    self.models_dict[model].load_state_dict(checkpoint[model], strict=False)
                    print('Checkpoint loaded')

    def train(self):
        """Training process."""
        # Run training for num_epochs epochs
        for epoch in tqdm(range(self.epoch_count, self.options.num_epochs), total=self.options.num_epochs, initial=self.epoch_count):
            # Run eval for off-the-shelf spin
            if epoch == 0:
                if self.run_eval_h36m:
                    self.eval_step(epoch, dataset_name='h36m-p2', dataset=self.eval_data_loader, result_file=None, num_workers=self.options.num_workers)
                if self.run_eval_rich_val:
                    self.eval_step(epoch, dataset_name='rich-val', dataset=self.eval_data_loader_rich_val, result_file=None, num_workers=self.options.num_workers)
                if self.run_eval_rich_test:
                    self.eval_step(epoch, dataset_name='rich-test', dataset=self.eval_data_loader_rich_test, result_file=None, num_workers=self.options.num_workers)
                if self.run_eval_3dpw:
                    self.eval_step(epoch, dataset_name='3dpw', dataset=self.eval_data_loader_3dpw, result_file=None, num_workers=self.options.num_workers)
                    # self.eval_step(epoch, dataset_name='agora_val', dataset=self.eval_data_loader_agora, result_file=None, num_workers=self.options.num_workers)

            # Create new DataLoader every epoch and (possibly) resume from an arbitrary step inside an epoch
            train_data_loader = CheckpointDataLoader(self.train_ds,checkpoint=self.checkpoint,
                                                     batch_size=self.options.batch_size,
                                                     num_workers=self.options.num_workers,
                                                     pin_memory=self.options.pin_memory,
                                                     shuffle=self.options.shuffle_train)

            # Iterate over all batches in an epoch
            for step, batch in enumerate(tqdm(train_data_loader, desc='Epoch '+str(epoch),
                                              total=len(self.train_ds) // self.options.batch_size,
                                              initial=train_data_loader.checkpoint_batch_idx),
                                         train_data_loader.checkpoint_batch_idx):
                if time.time() < self.endtime:
                    batch = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k,v in batch.items()}
                    out = self.train_step(batch)
                    self.step_count += 1
                    # Tensorboard logging every summary_steps steps
                    if self.step_count % self.options.summary_steps == 0:
                        self.train_summaries(batch, *out)
                    # Save checkpoint every checkpoint_steps steps
                    # if self.step_count % self.options.checkpoint_steps == 0:
                    if self.step_count % self.options.checkpoint_steps == 0:
                        self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, step+1, self.options.batch_size, train_data_loader.sampler.dataset_perm, self.step_count)
                        tqdm.write('Checkpoint saved')
                        if self.run_eval_h36m:
                            self.eval_step(epoch, dataset_name='h36m-p2', dataset=self.eval_data_loader,
                                           result_file=None, num_workers=self.options.num_workers, step_count=self.step_count)
                        if self.run_eval_rich_val:
                            self.eval_step(epoch, dataset_name='rich-val',
                                           dataset=self.eval_data_loader_rich_val, result_file=None, num_workers=self.options.num_workers, step_count=self.step_count)
                        if self.run_eval_rich_test:
                            self.eval_step(epoch, dataset_name='rich-test',
                                           dataset=self.eval_data_loader_rich_test, result_file=None, num_workers=self.options.num_workers, step_count=self.step_count)
                        # if self.run_eval_3dpw:
                        #     self.eval_step(epoch, dataset_name='3dpw', dataset=self.eval_data_loader_3dpw,
                        #                    result_file=None, num_workers=self.options.num_workers, step_count=self.step_count)

                    # Run validation every test_steps steps
                    if self.step_count % self.options.test_steps == 0:
                        self.test()
                else:
                    tqdm.write('Timeout reached')
                    self.finalize()
                    self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, step, self.options.batch_size, train_data_loader.sampler.dataset_perm, self.step_count)
                    tqdm.write('Checkpoint saved')
                    sys.exit(0)

            # load a checkpoint only on startup, for the next epochs
            # just iterate over the dataset as usual
            self.checkpoint=None
            # save checkpoint after each epoch
            # if (epoch+1) % 10 == 0:
            if (epoch + 1) % 2 == 0:
                # self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch+1, 0, self.step_count)
                self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, 0, self.options.batch_size, None, self.step_count)
                if self.run_eval_h36m:
                    self.eval_step(epoch, dataset_name='h36m-p2', dataset=self.eval_data_loader,
                                   result_file=None, num_workers=self.options.num_workers)
                if self.run_eval_rich_val:
                    self.eval_step(epoch, dataset_name='rich-val',
                                   dataset=self.eval_data_loader_rich_val, result_file=None, num_workers=self.options.num_workers)
                if self.run_eval_rich_test:
                    self.eval_step(epoch, dataset_name='rich-test',
                                   dataset=self.eval_data_loader_rich_test, result_file=None, num_workers=self.options.num_workers)
                if self.run_eval_3dpw:
                    self.eval_step(epoch, dataset_name='3dpw', dataset=self.eval_data_loader_3dpw,
                                   result_file=None, num_workers=self.options.num_workers)

                    # self.eval_step(epoch, dataset_name='agora_val', dataset=self.eval_data_loader_agora, result_file=None, num_workers=self.options.num_workers)
        return

    # The following methods (with the possible exception of test) have to be implemented in the derived classes
    def init_fn(self):
        raise NotImplementedError('You need to provide an _init_fn method')

    def train_step(self, input_batch):
        raise NotImplementedError('You need to provide a _train_step method')

    def train_summaries(self, input_batch):
        raise NotImplementedError('You need to provide a _train_summaries method')

    def eval_step(self):
        raise NotImplementedError('You need to provide a _eval_step method')

    def test(self):
        pass

class BaseTrainerAgoraPlusOther(object):
    """Base class for Trainer objects.
    Takes care of checkpointing/logging/resuming training.
    """
    def __init__(self, options):
        self.options = options
        self.endtime = time.time() + self.options.time_to_run
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        # override this function to define your model, optimizers etc.
        self.init_fn()
        self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
        self.summary_writer = SummaryWriter(self.options.summary_dir)
        self.run_eval = options.run_eval

        self.is_agora = options.is_agora
        self.is_others = options.is_others
        self.run_smplify = options.run_smplify

        self.checkpoint = None
        if self.options.resume and self.saver.exists_checkpoint():
            self.checkpoint = self.saver.load_checkpoint(self.models_dict, self.optimizers_dict, checkpoint_file=self.options.checkpoint)

        if self.checkpoint is None:
            self.epoch_count = 0
            self.step_count = 0
        else:
            self.epoch_count = self.checkpoint['epoch']
            self.step_count = self.checkpoint['total_step_count']

        if self.run_eval:
            self.eval_data_loader = BaseDataset(None, '3dpw', is_train=False)
            self.eval_data_loader_agora = BaseDataset(None, 'agora_val', is_train=False)

    def load_pretrained(self, checkpoint_file=None):
        """Load a pretrained checkpoint.
        This is different from resuming training using --resume.
        """
        if checkpoint_file is not None:
            checkpoint = torch.load(checkpoint_file)
            for model in self.models_dict:
                if model in checkpoint:
                    self.models_dict[model].load_state_dict(checkpoint[model], strict=False)
                    print('Checkpoint loaded')

    def train(self):
        """Training process."""
        # Run training for num_epochs epochs
        for epoch in tqdm(range(self.epoch_count, self.options.num_epochs), total=self.options.num_epochs, initial=self.epoch_count):
            # Run eval for off-the-shelf spin
            if epoch == 0:
                if self.run_eval:
                    self.eval_step(epoch, dataset_name='3dpw', dataset=self.eval_data_loader, result_file=None, num_workers=self.options.num_workers)
                    self.eval_step(epoch, dataset_name='agora_val', dataset=self.eval_data_loader_agora, result_file=None, num_workers=self.options.num_workers)

            # Create new DataLoader every epoch and (possibly) resume from an arbitrary step inside an epoch
            train_data_loader_agora = CheckpointDataLoader(self.train_ds_agora,checkpoint=self.checkpoint,
                                                     batch_size=self.options.batch_size,
                                                     num_workers=self.options.num_workers,
                                                     pin_memory=self.options.pin_memory,
                                                     shuffle=self.options.shuffle_train)

            train_data_loader_other = CheckpointDataLoader(self.train_ds_other, checkpoint=self.checkpoint,
                                                           batch_size=self.options.batch_size,
                                                           num_workers=self.options.num_workers,
                                                           pin_memory=self.options.pin_memory,
                                                           shuffle=self.options.shuffle_train)

            # Iterate over all batches in an epoch
            for step, batch_other in enumerate(tqdm(train_data_loader_other, desc='Epoch '+str(epoch),
                                              total=len(self.train_ds_other) // self.options.batch_size,
                                              initial=train_data_loader_other.checkpoint_batch_idx),
                                         train_data_loader_other.checkpoint_batch_idx):
                if time.time() < self.endtime:
                    batch_other = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k,v in batch_other.items()}
                    try:
                        batch_agora = next(train_iter_agora)
                    except:
                        train_iter_agora = iter(train_data_loader_agora)
                        try:
                            batch_agora = next(train_iter_agora)
                        except:
                            print('Agora data is done for this epoch, but others still left --- Restart from start')
                            train_iter_agora = sorted(iter(train_data_loader_agora), key=lambda k: random.random())
                            batch_agora = next(train_iter_agora)

                    batch_agora = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k, v in
                                   batch_agora.items()}

                    out_other = self.train_step_other(batch_other)
                    out_agora= self.train_step_agora(batch_agora)

                    self.step_count += 1
                    # Tensorboard logging every summary_steps steps
                    if self.step_count % self.options.summary_steps == 0:
                        self.train_summaries(batch_other, *out_other, dset='other')
                        self.train_summaries(batch_agora, *out_agora,  dset='agora')
                    # Save checkpoint every checkpoint_steps steps
                    # if self.step_count % self.options.checkpoint_steps == 0:
                    if self.step_count % self.options.checkpoint_steps == 0:
                        self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, step+1, self.options.batch_size, train_data_loader_other.sampler.dataset_perm, self.step_count)
                        tqdm.write('Checkpoint saved')

                    # Run validation every test_steps steps
                    if self.step_count % self.options.test_steps == 0:
                        self.test()
                else:
                    tqdm.write('Timeout reached')
                    self.finalize()
                    self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, step, self.options.batch_size, train_data_loader_other.sampler.dataset_perm, self.step_count)
                    tqdm.write('Checkpoint saved')
                    sys.exit(0)

            # load a checkpoint only on startup, for the next epochs
            # just iterate over the dataset as usual
            self.checkpoint=None
            # save checkpoint after each epoch
            # if (epoch+1) % 10 == 0:
            if (epoch + 1) % 2 == 0:
                # self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch+1, 0, self.step_count)
                self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, 0, self.options.batch_size, None, self.step_count)
                if self.run_eval:
                    self.eval_step(epoch, dataset_name='3dpw', dataset=self.eval_data_loader, result_file=None, num_workers=self.options.num_workers)
                    self.eval_step(epoch, dataset_name='agora_val', dataset=self.eval_data_loader_agora, result_file=None, num_workers=self.options.num_workers)
        return

    # The following methods (with the possible exception of test) have to be implemented in the derived classes
    def init_fn(self):
        raise NotImplementedError('You need to provide an _init_fn method')

    def train_step_agora(self, input_batch):
        raise NotImplementedError('You need to provide a _train_step_agora method')

    def train_step_other(self, input_batch):
        raise NotImplementedError('You need to provide a _train_step_other method')

    def train_summaries(self, input_batch):
        raise NotImplementedError('You need to provide a _train_summaries method')

    def eval_step(self):
        raise NotImplementedError('You need to provide a _eval_step method')

    def test(self):
        pass



utils/base_trainer.py
sha2nkt/ipman-r
import os
import sys
import stat
import shutil
import subprocess

from loguru import logger

GPUS = {
    'v100-v16': ('\"Tesla V100-PCIE-16GB\"', 'tesla', 16000),
    'v100-p32': ('\"Tesla V100-PCIE-32GB\"', 'tesla', 32000),
    'v100-s32': ('\"Tesla V100-SXM2-32GB\"', 'tesla', 32000),
    'v100-p16': ('\"Tesla P100-PCIE-16GB\"', 'tesla', 16000),
    'rtx2080ti': ('\"NVIDIA GeForce RTX 2080 Ti\"', 'rtx', 11000),
    # 'quadro6000': ('\"Quadro RTX 6000\"', 'quadro', 24000),
}

def get_gpus(min_mem=10000, arch=('tesla', 'quadro', 'rtx')):
    gpu_names = []
    for k, (gpu_name, gpu_arch, gpu_mem) in GPUS.items():
        if gpu_mem >= min_mem and gpu_arch in arch:
            gpu_names.append(gpu_name)

    assert len(gpu_names) > 0, 'Suitable GPU model could not be found'

    return gpu_names


def execute_task_on_cluster(
        script,
        exp_name,
        output_dir,
        cfg_file,
        num_exp=1,
        exp_opts=None,
        bid_amount=5,
        num_workers=8,
        memory=64000,
        gpu_min_mem=10000,
        gpu_arch=('tesla', 'quadro', 'rtx'),
        num_gpus=1
):
    # copy config to a new experiment directory and source from there
    temp_config_dir = os.path.join(os.path.dirname(output_dir), 'temp_configs')
    logdir = os.path.join(temp_config_dir, f'{exp_name}')
    os.makedirs(logdir, exist_ok=True)
    new_cfg_file = os.path.join(logdir, 'config.yaml')
    shutil.copy(src=cfg_file, dst=new_cfg_file)

    gpus = get_gpus(min_mem=gpu_min_mem, arch=gpu_arch)

    gpus = ' || '.join([f'CUDADeviceName=={x}' for x in gpus])

    os.makedirs(os.path.join('cluster_logs', exp_name), exist_ok=True)
    submission = f'executable = cluster_logs/{exp_name}/{exp_name}_run.sh\n' \
                 'arguments = $(Process) $(Cluster)\n' \
                 f'error = cluster_logs/{exp_name}/$(Cluster).$(Process).err\n' \
                 f'output = cluster_logs/{exp_name}/$(Cluster).$(Process).out\n' \
                 f'log = cluster_logs/{exp_name}/$(Cluster).$(Process).log\n' \
                 f'request_memory = {memory}\n' \
                 f'request_cpus={int(num_workers/2)}\n' \
                 f'request_gpus={num_gpus}\n' \
                 f'requirements={gpus}\n' \
                 f'requirements = UtsnameNodename =!= "g043"\n'\
                 f'+MaxRunningPrice = 200\n' \
                 f'queue {num_exp}'
                 # f'request_cpus={int(num_workers/2)}\n' \
                 # f'+RunningPriceExceededAction = \"kill\"\n' \
    print('<<< Condor Submission >>> ')
    print(submission)

    with open('submit.sub', 'w') as f:
        f.write(submission)

    logger.info(f'The logs for this experiments can be found under: cluster_logs/{exp_name}')
    bash = 'export PYTHONBUFFERED=1\n export PATH=$PATH\n' \
           f'{sys.executable} {script} --cfg {new_cfg_file} --cfg_id $1'  ## This is the trick. Notice there is no --cluster here

    if exp_opts is not None:
        bash += ' --opts '
        for opt in exp_opts:
            bash += f'{opt} '
        bash += 'SYSTEM.CLUSTER_NODE $2.$1'
    else:
        bash += ' --opts SYSTEM.CLUSTER_NODE $2.$1'

    executable_path =f'cluster_logs/{exp_name}/{exp_name}_run.sh'

    with open(executable_path, 'w') as f:
        f.write(bash)

    os.chmod(executable_path, stat.S_IRWXU)

    cmd = ['condor_submit_bid', f'{bid_amount}', 'submit.sub']
    logger.info('Executing ' + ' '.join(cmd))
    subprocess.call(cmd)

utils/cluster.py
sha2nkt/ipman-r
import torch
import numpy as np
from torch.nn import functional as F

"""
Useful geometric operations, e.g. Perspective projection and a differentiable Rodrigues formula
Parts of the code are taken from https://github.com/MandyMo/pytorch_HMR
"""


def batch_rot2aa(Rs):
    """
    Rs is B x 3 x 3
    void cMathUtil::RotMatToAxisAngle(const tMatrix& mat, tVector& out_axis,
                                      double& out_theta)
    {
        double c = 0.5 * (mat(0, 0) + mat(1, 1) + mat(2, 2) - 1);
        c = cMathUtil::Clamp(c, -1.0, 1.0);

        out_theta = std::acos(c);

        if (std::abs(out_theta) < 0.00001)
        {
            out_axis = tVector(0, 0, 1, 0);
        }
        else
        {
            double m21 = mat(2, 1) - mat(1, 2);
            double m02 = mat(0, 2) - mat(2, 0);
            double m10 = mat(1, 0) - mat(0, 1);
            double denom = std::sqrt(m21 * m21 + m02 * m02 + m10 * m10);
            out_axis[0] = m21 / denom;
            out_axis[1] = m02 / denom;
            out_axis[2] = m10 / denom;
            out_axis[3] = 0;
        }
    }
    """
    cos = 0.5 * (torch.stack([torch.trace(x) for x in Rs]) - 1)
    cos = torch.clamp(cos, -1, 1)

    theta = torch.acos(cos)

    m21 = Rs[:, 2, 1] - Rs[:, 1, 2]
    m02 = Rs[:, 0, 2] - Rs[:, 2, 0]
    m10 = Rs[:, 1, 0] - Rs[:, 0, 1]
    denom = torch.sqrt(m21 * m21 + m02 * m02 + m10 * m10)

    axis0 = torch.where(torch.abs(theta) < 0.00001, m21, m21 / denom)
    axis1 = torch.where(torch.abs(theta) < 0.00001, m02, m02 / denom)
    axis2 = torch.where(torch.abs(theta) < 0.00001, m10, m10 / denom)

    return theta.unsqueeze(1) * torch.stack([axis0, axis1, axis2], 1)


def batch_rodrigues(theta):
    """Convert axis-angle representation to rotation matrix.
    Args:
        theta: size = [B, 3]
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    l1norm = torch.norm(theta + 1e-8, p = 2, dim = 1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim = 1)
    return quat_to_rotmat(quat)


def quat_to_rotmat(quat):
    """Convert quaternion coefficients to rotation matrix.
    Args:
        quat: size = [B, 4] 4 <===>(w, x, y, z)
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    norm_quat = quat
    norm_quat = norm_quat/norm_quat.norm(p=2, dim=1, keepdim=True)
    w, x, y, z = norm_quat[:,0], norm_quat[:,1], norm_quat[:,2], norm_quat[:,3]

    B = quat.size(0)

    w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
    wx, wy, wz = w*x, w*y, w*z
    xy, xz, yz = x*y, x*z, y*z

    rotMat = torch.stack([w2 + x2 - y2 - z2, 2*xy - 2*wz, 2*wy + 2*xz,
                          2*wz + 2*xy, w2 - x2 + y2 - z2, 2*yz - 2*wx,
                          2*xz - 2*wy, 2*wx + 2*yz, w2 - x2 - y2 + z2], dim=1).view(B, 3, 3)
    return rotMat


def rot6d_to_rotmat(x):
    """Convert 6D rotation representation to 3x3 rotation matrix.
    Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019
    Input:
        (B,6) Batch of 6-D rotation representations
    Output:
        (B,3,3) Batch of corresponding rotation matrices
    """
    x = x.reshape(-1,3,2)
    a1 = x[:, :, 0]
    a2 = x[:, :, 1]
    b1 = F.normalize(a1)
    b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1)
    b3 = torch.cross(b1, b2)
    return torch.stack((b1, b2, b3), dim=-1)


def rotmat_to_rot6d(x):
    rotmat = x.reshape(-1, 3, 3)
    rot6d = rotmat[:, :, :2].reshape(x.shape[0], -1)
    return rot6d


def rotation_matrix_to_angle_axis(rotation_matrix):
    """
    This function is borrowed from https://github.com/kornia/kornia

    Convert 3x4 rotation matrix to Rodrigues vector

    Args:
        rotation_matrix (Tensor): rotation matrix.

    Returns:
        Tensor: Rodrigues vector transformation.

    Shape:
        - Input: :math:`(N, 3, 4)`
        - Output: :math:`(N, 3)`

    Example:
        >>> input = torch.rand(2, 3, 4)  # Nx4x4
        >>> output = tgm.rotation_matrix_to_angle_axis(input)  # Nx3
    """
    if rotation_matrix.shape[1:] == (3,3):
        rot_mat = rotation_matrix.reshape(-1, 3, 3)
        hom = torch.tensor([0, 0, 1], dtype=torch.float32,
                           device=rotation_matrix.device).reshape(1, 3, 1).expand(rot_mat.shape[0], -1, -1)
        rotation_matrix = torch.cat([rot_mat, hom], dim=-1)

    quaternion = rotation_matrix_to_quaternion(rotation_matrix)
    aa = quaternion_to_angle_axis(quaternion)
    aa[torch.isnan(aa)] = 0.0
    return aa


def quaternion_to_angle_axis(quaternion: torch.Tensor) -> torch.Tensor:
    """
    This function is borrowed from https://github.com/kornia/kornia

    Convert quaternion vector to angle axis of rotation.

    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h

    Args:
        quaternion (torch.Tensor): tensor with quaternions.

    Return:
        torch.Tensor: tensor with angle axis of rotation.

    Shape:
        - Input: :math:`(*, 4)` where `*` means, any number of dimensions
        - Output: :math:`(*, 3)`

    Example:
        >>> quaternion = torch.rand(2, 4)  # Nx4
        >>> angle_axis = tgm.quaternion_to_angle_axis(quaternion)  # Nx3
    """
    if not torch.is_tensor(quaternion):
        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
            type(quaternion)))

    if not quaternion.shape[-1] == 4:
        raise ValueError("Input must be a tensor of shape Nx4 or 4. Got {}"
                         .format(quaternion.shape))
    # unpack input and compute conversion
    q1: torch.Tensor = quaternion[..., 1]
    q2: torch.Tensor = quaternion[..., 2]
    q3: torch.Tensor = quaternion[..., 3]
    sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3

    sin_theta: torch.Tensor = torch.sqrt(sin_squared_theta)
    cos_theta: torch.Tensor = quaternion[..., 0]
    two_theta: torch.Tensor = 2.0 * torch.where(
        cos_theta < 0.0,
        torch.atan2(-sin_theta, -cos_theta),
        torch.atan2(sin_theta, cos_theta))

    k_pos: torch.Tensor = two_theta / sin_theta
    k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta)
    k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg)

    angle_axis: torch.Tensor = torch.zeros_like(quaternion)[..., :3]
    angle_axis[..., 0] += q1 * k
    angle_axis[..., 1] += q2 * k
    angle_axis[..., 2] += q3 * k
    return angle_axis


def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):
    """
    This function is borrowed from https://github.com/kornia/kornia

    Convert 3x4 rotation matrix to 4d quaternion vector

    This algorithm is based on algorithm described in
    https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201

    Args:
        rotation_matrix (Tensor): the rotation matrix to convert.

    Return:
        Tensor: the rotation in quaternion

    Shape:
        - Input: :math:`(N, 3, 4)`
        - Output: :math:`(N, 4)`

    Example:
        >>> input = torch.rand(4, 3, 4)  # Nx3x4
        >>> output = tgm.rotation_matrix_to_quaternion(input)  # Nx4
    """
    if not torch.is_tensor(rotation_matrix):
        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
            type(rotation_matrix)))

    if len(rotation_matrix.shape) > 3:
        raise ValueError(
            "Input size must be a three dimensional tensor. Got {}".format(
                rotation_matrix.shape))
    if not rotation_matrix.shape[-2:] == (3, 4):
        raise ValueError(
            "Input size must be a N x 3 x 4  tensor. Got {}".format(
                rotation_matrix.shape))

    rmat_t = torch.transpose(rotation_matrix, 1, 2)

    mask_d2 = rmat_t[:, 2, 2] < eps

    mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1]
    mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1]

    t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
    q0 = torch.stack([rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
                      t0, rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
                      rmat_t[:, 2, 0] + rmat_t[:, 0, 2]], -1)
    t0_rep = t0.repeat(4, 1).t()

    t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
    q1 = torch.stack([rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
                      rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
                      t1, rmat_t[:, 1, 2] + rmat_t[:, 2, 1]], -1)
    t1_rep = t1.repeat(4, 1).t()

    t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
    q2 = torch.stack([rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
                      rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
                      rmat_t[:, 1, 2] + rmat_t[:, 2, 1], t2], -1)
    t2_rep = t2.repeat(4, 1).t()

    t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
    q3 = torch.stack([t3, rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
                      rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
                      rmat_t[:, 0, 1] - rmat_t[:, 1, 0]], -1)
    t3_rep = t3.repeat(4, 1).t()

    mask_c0 = mask_d2 * mask_d0_d1
    mask_c1 = mask_d2 * ~mask_d0_d1
    mask_c2 = ~mask_d2 * mask_d0_nd1
    mask_c3 = ~mask_d2 * ~mask_d0_nd1
    mask_c0 = mask_c0.view(-1, 1).type_as(q0)
    mask_c1 = mask_c1.view(-1, 1).type_as(q1)
    mask_c2 = mask_c2.view(-1, 1).type_as(q2)
    mask_c3 = mask_c3.view(-1, 1).type_as(q3)

    q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3
    q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 +  # noqa
                    t2_rep * mask_c2 + t3_rep * mask_c3)  # noqa
    q *= 0.5
    return q

def convert_perspective_to_weak_perspective(
        perspective_camera,
        focal_length=5000.,
        img_res=224,
):
    # Convert Weak Perspective Camera [s, tx, ty] to camera translation [tx, ty, tz]
    # in 3D given the bounding box size
    # This camera translation can be used in a full perspective projection
    # if isinstance(focal_length, torch.Tensor):
    #     focal_length = focal_length[:, 0]

    weak_perspective_camera = torch.stack(
        [

            2 * focal_length / (img_res * perspective_camera[:, 2] + 1e-9),
            perspective_camera[:, 0],
            perspective_camera[:, 1],
        ],
        dim=-1
    )
    return weak_perspective_camera


def convert_weak_perspective_to_perspective(
        weak_perspective_camera,
        focal_length=5000.,
        img_res=224,
):
    # Convert Weak Perspective Camera [s, tx, ty] to camera translation [tx, ty, tz]
    # in 3D given the bounding box size
    # This camera translation can be used in a full perspective projection
    # if isinstance(focal_length, torch.Tensor):
    #     focal_length = focal_length[:, 0]

    perspective_camera = torch.stack(
        [
            weak_perspective_camera[:, 1],
            weak_perspective_camera[:, 2],
            2 * focal_length / (img_res * weak_perspective_camera[:, 0] + 1e-9)
        ],
        dim=-1
    )
    return perspective_camera


def perspective_projection(points, rotation, translation,
                           focal_length, camera_center):
    """
    This function computes the perspective projection of a set of points.
    Input:
        points (bs, N, 3): 3D points
        rotation (bs, 3, 3): Camera rotation
        translation (bs, 3): Camera translation
        focal_length (bs,) or scalar: Focal length
        camera_center (bs, 2): Camera center
    """
    batch_size = points.shape[0]
    K = torch.zeros([batch_size, 3, 3], device=points.device)
    K[:,0,0] = focal_length
    K[:,1,1] = focal_length
    K[:,2,2] = 1.
    K[:,:-1, -1] = camera_center

    # Transform points
    points = torch.einsum('bij,bkj->bki', rotation, points)
    points = points + translation.unsqueeze(1)

    # Apply perspective distortion
    projected_points = points / points[:,:,-1].unsqueeze(-1)

    # Apply camera intrinsics
    projected_points = torch.einsum('bij,bkj->bki', K, projected_points)

    return projected_points[:, :, :-1]


def weak_perspective_projection(points, rotation, weak_cam_params, focal_length, camera_center, img_res):
    """
    This function computes the perspective projection of a set of points.
    Input:
        points (bs, N, 3): 3D points
        rotation (bs, 3, 3): Camera rotation
        translation (bs, 3): Camera translation
        focal_length (bs,) or scalar: Focal length
        camera_center (bs, 2): Camera center
    """
    batch_size = points.shape[0]
    K = torch.zeros([batch_size, 3, 3], device=points.device)
    K[:,0,0] = focal_length
    K[:,1,1] = focal_length
    K[:,2,2] = 1.
    K[:,:-1, -1] = camera_center

    translation = convert_weak_perspective_to_perspective(weak_cam_params, focal_length, img_res)

    # Transform points
    points = torch.einsum('bij,bkj->bki', rotation, points)
    points = points + translation.unsqueeze(1)

    # Apply perspective distortion
    projected_points = points / points[:,:,-1].unsqueeze(-1)

    # Apply camera intrinsics
    projected_points = torch.einsum('bij,bkj->bki', K, projected_points)

    return projected_points[:, :, :-1]


def estimate_translation_np(S, joints_2d, joints_conf, focal_length=5000., img_size=224.):
    """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.
    Input:
        S: (25, 3) 3D joint locations
        joints: (25, 3) 2D joint locations and confidence
    Returns:
        (3,) camera translation vector
    """

    num_joints = S.shape[0]
    # focal length
    f = np.array([focal_length,focal_length])
    # optical center
    center = np.array([img_size/2., img_size/2.])

    # transformations
    Z = np.reshape(np.tile(S[:,2],(2,1)).T,-1)
    XY = np.reshape(S[:,0:2],-1)
    O = np.tile(center,num_joints)
    F = np.tile(f,num_joints)
    weight2 = np.reshape(np.tile(np.sqrt(joints_conf),(2,1)).T,-1)

    # least squares
    Q = np.array([F*np.tile(np.array([1,0]),num_joints), F*np.tile(np.array([0,1]),num_joints), O-np.reshape(joints_2d,-1)]).T
    c = (np.reshape(joints_2d,-1)-O)*Z - F*XY

    # weighted least squares
    W = np.diagflat(weight2)
    Q = np.dot(W,Q)
    c = np.dot(W,c)

    # square matrix
    A = np.dot(Q.T,Q)
    b = np.dot(Q.T,c)

    # solution
    trans = np.linalg.solve(A, b)

    return trans


def estimate_translation(S, joints_2d, focal_length=5000., img_size=224., use_all_joints=False, rotation=None):
    """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.
    Input:
        S: (B, 49, 3) 3D joint locations
        joints: (B, 49, 3) 2D joint locations and confidence
    Returns:
        (B, 3) camera translation vectors
    """

    device = S.device

    if rotation is not None:
        S = torch.einsum('bij,bkj->bki', rotation, S)

    # Use only joints 25:49 (GT joints)
    if use_all_joints:
        S = S.cpu().numpy()
        joints_2d = joints_2d.cpu().numpy()
    else:
        S = S[:, 25:, :].cpu().numpy()
        joints_2d = joints_2d[:, 25:, :].cpu().numpy()

    joints_conf = joints_2d[:, :, -1]
    joints_2d = joints_2d[:, :, :-1]
    trans = np.zeros((S.shape[0], 3), dtype=np.float32)
    # Find the translation for each example in the batch
    for i in range(S.shape[0]):
        S_i = S[i]
        joints_i = joints_2d[i]
        conf_i = joints_conf[i]
        trans[i] = estimate_translation_np(S_i, joints_i, conf_i, focal_length=focal_length, img_size=img_size)
    return torch.from_numpy(trans).to(device)


def estimate_translation_cam(S, joints_2d, focal_length=(5000., 5000.), img_size=(224., 224.),
                             use_all_joints=False, rotation=None):
    """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.
    Input:
        S: (B, 49, 3) 3D joint locations
        joints: (B, 49, 3) 2D joint locations and confidence
    Returns:
        (B, 3) camera translation vectors
    """

    def estimate_translation_np(S, joints_2d, joints_conf, focal_length=(5000., 5000.), img_size=(224., 224.)):
        """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.
        Input:
            S: (25, 3) 3D joint locations
            joints: (25, 3) 2D joint locations and confidence
        Returns:
            (3,) camera translation vector
        """

        num_joints = S.shape[0]
        # focal length
        f = np.array([focal_length[0], focal_length[1]])
        # optical center
        center = np.array([img_size[0] / 2., img_size[1] / 2.])

        # transformations
        Z = np.reshape(np.tile(S[:, 2], (2, 1)).T, -1)
        XY = np.reshape(S[:, 0:2], -1)
        O = np.tile(center, num_joints)
        F = np.tile(f, num_joints)
        weight2 = np.reshape(np.tile(np.sqrt(joints_conf), (2, 1)).T, -1)

        # least squares
        Q = np.array([F * np.tile(np.array([1, 0]), num_joints), F * np.tile(np.array([0, 1]), num_joints),
                      O - np.reshape(joints_2d, -1)]).T
        c = (np.reshape(joints_2d, -1) - O) * Z - F * XY

        # weighted least squares
        W = np.diagflat(weight2)
        Q = np.dot(W, Q)
        c = np.dot(W, c)

        # square matrix
        A = np.dot(Q.T, Q)
        b = np.dot(Q.T, c)

        # solution
        trans = np.linalg.solve(A, b)

        return trans

    device = S.device

    if rotation is not None:
        S = torch.einsum('bij,bkj->bki', rotation, S)

    # Use only joints 25:49 (GT joints)
    if use_all_joints:
        S = S.cpu().numpy()
        joints_2d = joints_2d.cpu().numpy()
    else:
        S = S[:, 25:, :].cpu().numpy()
        joints_2d = joints_2d[:, 25:, :].cpu().numpy()

    joints_conf = joints_2d[:, :, -1]
    joints_2d = joints_2d[:, :, :-1]
    trans = np.zeros((S.shape[0], 3), dtype=np.float32)
    # Find the translation for each example in the batch
    for i in range(S.shape[0]):
        S_i = S[i]
        joints_i = joints_2d[i]
        conf_i = joints_conf[i]
        trans[i] = estimate_translation_np(S_i, joints_i, conf_i, focal_length=focal_length, img_size=img_size)
    return torch.from_numpy(trans).to(device)


def get_coord_maps(size=56):
    xx_ones = torch.ones([1, size], dtype=torch.int32)
    xx_ones = xx_ones.unsqueeze(-1)

    xx_range = torch.arange(size, dtype=torch.int32).unsqueeze(0)
    xx_range = xx_range.unsqueeze(1)

    xx_channel = torch.matmul(xx_ones, xx_range)
    xx_channel = xx_channel.unsqueeze(-1)

    yy_ones = torch.ones([1, size], dtype=torch.int32)
    yy_ones = yy_ones.unsqueeze(1)

    yy_range = torch.arange(size, dtype=torch.int32).unsqueeze(0)
    yy_range = yy_range.unsqueeze(-1)

    yy_channel = torch.matmul(yy_range, yy_ones)
    yy_channel = yy_channel.unsqueeze(-1)

    xx_channel = xx_channel.permute(0, 3, 1, 2)
    yy_channel = yy_channel.permute(0, 3, 1, 2)

    xx_channel = xx_channel.float() / (size - 1)
    yy_channel = yy_channel.float() / (size - 1)

    xx_channel = xx_channel * 2 - 1
    yy_channel = yy_channel * 2 - 1

    out = torch.cat([xx_channel, yy_channel], dim=1)
    return out


def look_at(eye, at=np.array([0, 0, 0]), up=np.array([0, 0, 1]), eps=1e-5):
    at = at.astype(float).reshape(1, 3)
    up = up.astype(float).reshape(1, 3)

    eye = eye.reshape(-1, 3)
    up = up.repeat(eye.shape[0] // up.shape[0], axis=0)
    eps = np.array([eps]).reshape(1, 1).repeat(up.shape[0], axis=0)

    z_axis = eye - at
    z_axis /= np.max(np.stack([np.linalg.norm(z_axis, axis=1, keepdims=True), eps]))

    x_axis = np.cross(up, z_axis)
    x_axis /= np.max(np.stack([np.linalg.norm(x_axis, axis=1, keepdims=True), eps]))

    y_axis = np.cross(z_axis, x_axis)
    y_axis /= np.max(np.stack([np.linalg.norm(y_axis, axis=1, keepdims=True), eps]))

    r_mat = np.concatenate((x_axis.reshape(-1, 3, 1), y_axis.reshape(-1, 3, 1), z_axis.reshape(-1, 3, 1)), axis=2)

    return r_mat


def to_sphere(u, v):
    theta = 2 * np.pi * u
    phi = np.arccos(1 - 2 * v)
    cx = np.sin(phi) * np.cos(theta)
    cy = np.sin(phi) * np.sin(theta)
    cz = np.cos(phi)
    s = np.stack([cx, cy, cz])
    return s


def sample_on_sphere(range_u=(0, 1), range_v=(0, 1)):
    u = np.random.uniform(*range_u)
    v = np.random.uniform(*range_v)
    return to_sphere(u, v)


def sample_pose_on_sphere(range_v=(0,1), range_u=(0,1), radius=1, up=[0,1,0]):
    # sample location on unit sphere
    loc = sample_on_sphere(range_u, range_v)

    # sample radius if necessary
    if isinstance(radius, tuple):
        radius = np.random.uniform(*radius)

    loc = loc * radius
    R = look_at(loc, up=np.array(up))[0]

    RT = np.concatenate([R, loc.reshape(3, 1)], axis=1)
    RT = torch.Tensor(RT.astype(np.float32))
    return RT


def rectify_pose(camera_r, body_aa, rotate_x=False):
    body_r = batch_rodrigues(body_aa).reshape(-1,3,3)

    if rotate_x:
        rotate_x = torch.tensor([[[1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, -1.0]]])
        body_r = body_r @ rotate_x

    final_r = camera_r @ body_r
    body_aa = batch_rot2aa(final_r)
    return body_aa

def batch_rectify_pose(camera_r, body_aa):
    body_r = batch_rodrigues(body_aa).reshape(-1, 3, 3)
    final_r = torch.bmm(camera_r, body_r)
    body_aa = batch_rot2aa(final_r)
    return body_aa


def batch_euler2matrix(r):
    return quaternion_to_rotation_matrix(euler_to_quaternion(r))


def euler_to_quaternion(r):
    x = r[..., 0]
    y = r[..., 1]
    z = r[..., 2]

    z = z/2.0
    y = y/2.0
    x = x/2.0
    cz = torch.cos(z)
    sz = torch.sin(z)
    cy = torch.cos(y)
    sy = torch.sin(y)
    cx = torch.cos(x)
    sx = torch.sin(x)
    quaternion = torch.zeros_like(r.repeat(1,2))[..., :4].to(r.device)
    quaternion[..., 0] += cx*cy*cz - sx*sy*sz
    quaternion[..., 1] += cx*sy*sz + cy*cz*sx
    quaternion[..., 2] += cx*cz*sy - sx*cy*sz
    quaternion[..., 3] += cx*cy*sz + sx*cz*sy
    return quaternion


def quaternion_to_rotation_matrix(quat):
    """Convert quaternion coefficients to rotation matrix.
    Args:
        quat: size = [B, 4] 4 <===>(w, x, y, z)
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    norm_quat = quat
    norm_quat = norm_quat / norm_quat.norm(p=2, dim=1, keepdim=True)
    w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:, 2], norm_quat[:, 3]

    B = quat.size(0)

    w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
    wx, wy, wz = w * x, w * y, w * z
    xy, xz, yz = x * y, x * z, y * z

    rotMat = torch.stack([w2 + x2 - y2 - z2, 2 * xy - 2 * wz, 2 * wy + 2 * xz,
                          2 * wz + 2 * xy, w2 - x2 + y2 - z2, 2 * yz - 2 * wx,
                          2 * xz - 2 * wy, 2 * wx + 2 * yz, w2 - x2 - y2 + z2], dim=1).view(B, 3, 3)
    return rotMat

def ea2rm(x, y, z):
    cos_x, sin_x = torch.cos(x), torch.sin(x)
    cos_y, sin_y = torch.cos(y), torch.sin(y)
    cos_z, sin_z = torch.cos(z), torch.sin(z)

    R = torch.stack(
            [torch.cat([cos_y*cos_z, sin_x*sin_y*cos_z - cos_x*sin_z, cos_x*sin_y*cos_z + sin_x*sin_z], dim=1),
            torch.cat([cos_y*sin_z, sin_x*sin_y*sin_z + cos_x*cos_z, cos_x*sin_y*sin_z - sin_x*cos_z], dim=1),
            torch.cat([-sin_y, sin_x*cos_y, cos_x*cos_y], dim=1)], dim=1)
    return R

def euler_angles_from_rotmat(R):
    """
    computer euler angles for rotation around x, y, z axis
    from rotation amtrix
    R: 4x4 rotation matrix
    https://www.gregslabaugh.net/publications/euler.pdf
    """
    r21 = np.round(R[:, 2, 0].item(), 4)
    if abs(r21) != 1:
        y_angle1 = -1 * torch.asin(R[:, 2, 0])
        y_angle2 = math.pi + torch.asin(R[:, 2, 0])
        cy1, cy2 = torch.cos(y_angle1), torch.cos(y_angle2)

        x_angle1 = torch.atan2(R[:, 2, 1] / cy1, R[:, 2, 2] / cy1)
        x_angle2 = torch.atan2(R[:, 2, 1] / cy2, R[:, 2, 2] / cy2)
        z_angle1 = torch.atan2(R[:, 1, 0] / cy1, R[:, 0, 0] / cy1)
        z_angle2 = torch.atan2(R[:, 1, 0] / cy2, R[:, 0, 0] / cy2)

        s1 = (x_angle1, y_angle1, z_angle1)
        s2 = (x_angle2, y_angle2, z_angle2)
        s = (s1, s2)

    else:
        z_angle = torch.tensor([0], device=R.device).float()
        if r21 == -1:
            y_angle = torch.tensor([math.pi / 2], device=R.device).float()
            x_angle = z_angle + torch.atan2(R[:, 0, 1], R[:, 0, 2])
        else:
            y_angle = -torch.tensor([math.pi / 2], device=R.device).float()
            x_angle = -z_angle + torch.atan2(-R[:, 0, 1], R[:, 0, 2])
        s = ((x_angle, y_angle, z_angle),)
    return s

def get_view_matrix(pitch=0, yaw=0, roll=0, tx=0, ty=0, tz=0):
    camera_rot = create_rotmat(pitch, yaw, roll)

    view_matrix = np.eye(4)
    view_matrix[:3, :3] = camera_rot
    view_matrix[0, 3] = tx
    view_matrix[1, 3] = ty
    view_matrix[2, 3] = tz

    return view_matrix


def create_rotmat(pitch, yaw, roll):
    pitch = (180 + pitch) * np.pi / 180
    yaw = yaw * np.pi / 180
    roll = roll * np.pi / 180

    cos_p, sin_p = np.cos(pitch), np.sin(pitch)
    cos_y, sin_y = np.cos(yaw), np.sin(yaw)
    cos_r, sin_r = np.cos(roll), np.sin(roll)

    R = np.array(
        [[cos_y * cos_r, sin_p * sin_y * cos_r - cos_p * sin_r, cos_p * sin_y * cos_r + sin_p * sin_r],
         [cos_y * sin_r, sin_p * sin_y * sin_r + cos_p * cos_r, cos_p * sin_y * sin_r - sin_p * cos_r],
         [-sin_y, sin_p * cos_y, cos_p * cos_y]]
    )

    return R
utils/geometry.py
sha2nkt/ipman-r
"""
This file contains functions that are used to perform data augmentation.
"""
import os
import torch
import numpy as np
import scipy.misc
import skimage.transform as skt
from PIL import Image
import cv2

import constants

def get_transform(center, scale, res, rot=0):
    """Generate transformation matrix."""
    h = 200 * scale
    t = np.zeros((3, 3))
    t[0, 0] = float(res[1]) / h
    t[1, 1] = float(res[0]) / h
    t[0, 2] = res[1] * (-float(center[0]) / h + .5)
    t[1, 2] = res[0] * (-float(center[1]) / h + .5)
    t[2, 2] = 1
    if not rot == 0:
        rot = -rot # To match direction of rotation from cropping
        rot_mat = np.zeros((3,3))
        rot_rad = rot * np.pi / 180
        sn,cs = np.sin(rot_rad), np.cos(rot_rad)
        rot_mat[0,:2] = [cs, -sn]
        rot_mat[1,:2] = [sn, cs]
        rot_mat[2,2] = 1
        # Need to rotate around center
        t_mat = np.eye(3)
        t_mat[0,2] = -res[1]/2
        t_mat[1,2] = -res[0]/2
        t_inv = t_mat.copy()
        t_inv[:2,2] *= -1
        t = np.dot(t_inv,np.dot(rot_mat,np.dot(t_mat,t)))
    return t

def transform(pt, center, scale, res, invert=0, rot=0):
    """Transform pixel location to different reference."""
    t = get_transform(center, scale, res, rot=rot)
    if invert:
        t = np.linalg.inv(t)
    new_pt = np.array([pt[0]-1, pt[1]-1, 1.]).T
    new_pt = np.dot(t, new_pt)
    return new_pt[:2].astype(int)+1


def crop_bbox(img, center, scale, res,imagepath, i,rot=0, return_proc_param=False):
    """Crop image according to the supplied bounding box."""
    # Upper left point
    ul = np.array(transform([1, 1], center, scale, res, invert=1)) - 1
    # Bottom right point
    br = np.array(transform([res[0] + 1,
                             res[1] + 1], center, scale, res, invert=1)) - 1

    bbox=np.concatenate((ul,br),axis=0)
    openpose_bbox_pixie_dir ='/ps/scratch/priyanka/bbox_openpose'
    imagename = imagepath.split('/')[-1].split('.')[0]
    fname = os.path.join(openpose_bbox_pixie_dir,f'{imagename}_{i:02}.txt') 
    np.savetxt(fname, bbox)
    # Padding so that when rotated proper amount of context is included
    pad = int(np.linalg.norm(br - ul) / 2 - float(br[1] - ul[1]) / 2)
    if not rot == 0:
        ul -= pad
        br += pad

    new_shape = [br[1] - ul[1], br[0] - ul[0]]
    if len(img.shape) > 2:
        new_shape += [img.shape[2]]
    new_img = np.zeros(new_shape)

    # Range to fill new array
    new_x = max(0, -ul[0]), min(br[0], len(img[0])) - ul[0]
    new_y = max(0, -ul[1]), min(br[1], len(img)) - ul[1]
    # Range to sample from original image
    old_x = max(0, ul[0]), min(len(img[0]), br[0])
    old_y = max(0, ul[1]), min(len(img), br[1])
    pad_x = max(0,0-ul[0]), max(0,br[0]-len(img[0]))

    pad_y =max(0,0-ul[1]), max(0,br[1]-len(img))

    # todo: check for case when very close to border of image ( padding required)
    if new_x[1]<0 or new_y[1]<0 or old_x[1]<0 or old_y[1]<0:
        return None, None

    new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[0]:old_y[1],
                                                    old_x[0]:old_x[1]]
    # new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[1]:old_y[0],
    #                                                     old_x[1]:old_x[0]]

    if not rot == 0:
        # Remove padding
        new_img = np.array(Image.fromarray(new_img.astype(np.uint8)).rotate(rot))
        # new_img_old = scipy.misc.imrotate(new_img, rot)
        new_img = new_img[pad:-pad, pad:-pad]

    if res[0] == 0 or res[1] == 0 or new_img.shape[0] == 0 or new_img.shape[1] == 0:
        return None, None

    new_img = np.array(Image.fromarray(new_img.astype(np.uint8)).resize([res[1], res[0]]))# crop shape is h x w
    # new_img = scipy.misc.imresize(new_img, res)

    if not return_proc_param:
        return new_img
    else:
        # similar to HMR
        proc_param = {
            'scale': scale,
            'start_pt': [old_y[0], old_x[0]],
            'end_pt': [old_y[1], old_x[1]],
            'img_size': new_shape,  # todo: check if correct
            'pad_margin': pad,
            'pad_x':pad_x,
            'pad_y':pad_y,
            'center_orig': center
        }
        return new_img, proc_param


def crop(img, center, scale, res, rot=0, return_proc_param=False):
    """Crop image according to the supplied bounding box."""
    # Upper left point
    ul = np.array(transform([1, 1], center, scale, res, invert=1)) - 1
    # Bottom right point
    br = np.array(transform([res[0] + 1,
                             res[1] + 1], center, scale, res, invert=1)) - 1

    # Padding so that when rotated proper amount of context is included
    pad = int(np.linalg.norm(br - ul) / 2 - float(br[1] - ul[1]) / 2)
    if not rot == 0:
        ul -= pad
        br += pad

    new_shape = [br[1] - ul[1], br[0] - ul[0]]
    if len(img.shape) > 2:
        new_shape += [img.shape[2]]
    new_img = np.zeros(new_shape)

    # Range to fill new array
    new_x = max(0, -ul[0]), min(br[0], len(img[0])) - ul[0]
    new_y = max(0, -ul[1]), min(br[1], len(img)) - ul[1]
    # Range to sample from original image
    old_x = max(0, ul[0]), min(len(img[0]), br[0])
    old_y = max(0, ul[1]), min(len(img), br[1])
    pad_x = max(0,0-ul[0]), max(0,br[0]-len(img[0]))

    pad_y =max(0,0-ul[1]), max(0,br[1]-len(img))

    # todo: check for case when very close to border of image ( padding required)
    if new_x[1]<0 or new_y[1]<0 or old_x[1]<0 or old_y[1]<0:
        return None, None

    new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[0]:old_y[1],
                                                    old_x[0]:old_x[1]]
    # new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[1]:old_y[0],
    #                                                     old_x[1]:old_x[0]]

    if not rot == 0:
        # Remove padding
        new_img = np.array(Image.fromarray(new_img.astype(np.uint8)).rotate(rot))
        # new_img_old = scipy.misc.imrotate(new_img, rot)
        new_img = new_img[pad:-pad, pad:-pad]

    if res[0] == 0 or res[1] == 0 or new_img.shape[0] == 0 or new_img.shape[1] == 0:
        return None, None
    new_img = np.array(Image.fromarray(new_img.astype(np.uint8)).resize([res[1], res[0]])) # crop shape is h x w
    # new_img = scipy.misc.imresize(new_img, res)

    if not return_proc_param:
        return new_img
    else:
        # similar to HMR
        proc_param = {
            'scale': scale,
            'start_pt': [old_y[0], old_x[0]],
            'end_pt': [old_y[1], old_x[1]],
            'img_size': new_shape,  # todo: check if correct
            'pad_margin': pad,
            'pad_x':pad_x,
            'pad_y':pad_y,
            'center_orig': center
        }
        return new_img, proc_param


# def crop(img, center, scale, res, rot=0):
#     """Crop image according to the supplied bounding box."""
#     # Upper left point
#     ul = np.array(transform([1, 1], center, scale, res, invert=1))-1
#     # Bottom right point
#     br = np.array(transform([res[0]+1,
#                              res[1]+1], center, scale, res, invert=1))-1
#
#     # Padding so that when rotated proper amount of context is included
#     pad = int(np.linalg.norm(br - ul) / 2 - float(br[1] - ul[1]) / 2)
#     if not rot == 0:
#         ul -= pad
#         br += pad
#
#     new_shape = [br[1] - ul[1], br[0] - ul[0]]
#     if len(img.shape) > 2:
#         new_shape += [img.shape[2]]
#     new_img = np.zeros(new_shape)
#
#     # Range to fill new array
#     new_x = max(0, -ul[0]), min(br[0], len(img[0])) - ul[0]
#     new_y = max(0, -ul[1]), min(br[1], len(img)) - ul[1]
#     # Range to sample from original image
#     old_x = max(0, ul[0]), min(len(img[0]), br[0])
#     old_y = max(0, ul[1]), min(len(img), br[1])
#     new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[0]:old_y[1],
#                                                         old_x[0]:old_x[1]]
#
#     if not rot == 0:
#         # Remove padding
#         new_img = scipy.misc.imrotate(new_img, rot)
#         new_img = new_img[pad:-pad, pad:-pad]
#
#     # resize image
#     new_img = scipy.misc.imresize(new_img, res)
#     return new_img

# def uncrop(img, center, scale, orig_shape, rot=0, is_rgb=True):
#     """'Undo' the image cropping/resizing.
#     This function is used when evaluating mask/part segmentation.
#     """
#     res = img.shape[:2]
#     # Upper left point
#     ul = np.array(transform([1, 1], center, scale, res, invert=1))-1
#     # Bottom right point
#     br = np.array(transform([res[0]+1,res[1]+1], center, scale, res, invert=1))-1
#     # size of cropped image
#     crop_shape = [br[1] - ul[1], br[0] - ul[0]]
#
#     new_shape = [br[1] - ul[1], br[0] - ul[0]]
#     if len(img.shape) > 2:
#         new_shape += [img.shape[2]]
#     new_img = np.zeros(orig_shape, dtype=np.uint8)
#     # Range to fill new array
#     new_x = max(0, -ul[0]), min(br[0], orig_shape[1]) - ul[0]
#     new_y = max(0, -ul[1]), min(br[1], orig_shape[0]) - ul[1]
#     # Range to sample from original image
#     old_x = max(0, ul[0]), min(orig_shape[1], br[0])
#     old_y = max(0, ul[1]), min(orig_shape[0], br[1])
#     img = scipy.misc.imresize(img, crop_shape, interp='nearest')
#     new_img[old_y[0]:old_y[1], old_x[0]:old_x[1]] = img[new_y[0]:new_y[1], new_x[0]:new_x[1]]
#     return new_img

def uncrop(img, center, scale, orig_shape, rot=0, is_rgb=True, getCorner=False, orig_img=None):
    """'Undo' the image cropping/resizing.
    This function is used when evaluating mask/part segmentation.
    """
    res = img.shape[:2]
    # Upper left point
    ul = np.array(transform([1, 1], center, scale, res, invert=1))-1
    # Bottom right point
    br = np.array(transform([res[0]+1,res[1]+1], center, scale, res, invert=1))-1
    # size of cropped image
    crop_shape = [br[1] - ul[1], br[0] - ul[0]]

    new_shape = [br[1] - ul[1], br[0] - ul[0]]
    if len(img.shape) > 2:
        new_shape += [img.shape[2]]

    if orig_img is None:
        new_img = np.zeros(orig_shape, dtype=np.uint8)
    else:
        new_img = orig_img
    # Range to fill new array
    new_x = max(0, -ul[0]), min(br[0], orig_shape[1]) - ul[0]
    new_y = max(0, -ul[1]), min(br[1], orig_shape[0]) - ul[1]
    # Range to sample from original image
    old_x = max(0, ul[0]), min(orig_shape[1], br[0])
    old_y = max(0, ul[1]), min(orig_shape[0], br[1])
    img = np.array(Image.fromarray(img.astype(np.uint8)).resize([crop_shape[1], crop_shape[0]])) # crop shape is h x w
    # img = scipy.misc.imresize(img, crop_shape, interp='nearest')
    new_img[old_y[0]:old_y[1], old_x[0]:old_x[1]] = img[new_y[0]:new_y[1], new_x[0]:new_x[1]]
    if getCorner:
        return new_img, old_y[0], old_x[0], crop_shape
    else:
        return new_img

def rot_aa(aa, rot):
    """Rotate axis angle parameters."""
    # pose parameters
    R = np.array([[np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0],
                  [np.sin(np.deg2rad(-rot)), np.cos(np.deg2rad(-rot)), 0],
                  [0, 0, 1]])
    # find the rotation of the body in camera frame
    per_rdg, _ = cv2.Rodrigues(aa)
    # apply the global rotation to the global orientation
    resrot, _ = cv2.Rodrigues(np.dot(R,per_rdg))
    aa = (resrot.T)[0]
    return aa

def flip_img(img):
    """Flip rgb images or masks.
    channels come last, e.g. (256,256,3).
    """
    img = np.fliplr(img)
    return img

def flip_kp(kp):
    """Flip keypoints."""
    if len(kp) == 24:
        flipped_parts = constants.J24_FLIP_PERM
    elif len(kp) == 49:
        flipped_parts = constants.J49_FLIP_PERM
    kp = kp[flipped_parts]
    kp[:,0] = - kp[:,0]
    return kp

def flip_pose(pose):
    """Flip pose.
    The flipping is based on SMPL parameters.
    """
    flipped_parts = constants.SMPL_POSE_FLIP_PERM
    pose = pose[flipped_parts]
    # we also negate the second and the third dimension of the axis-angle
    pose[1::3] = -pose[1::3]
    pose[2::3] = -pose[2::3]
    return pose

utils/imutils.py
sha2nkt/ipman-r
import numpy as np


def keypoint_hflip(kp, img_width):
    # Flip a keypoint horizontally around the y-axis
    # kp N,2
    if len(kp.shape) == 2:
        kp[:,0] = (img_width - 1.) - kp[:,0]
    elif len(kp.shape) == 3:
        kp[:, :, 0] = (img_width - 1.) - kp[:, :, 0]
    return kp


def convert_kps(joints2d, src, dst):
    src_names = eval(f'get_{src}_joint_names')()
    dst_names = eval(f'get_{dst}_joint_names')()

    out_joints2d = np.zeros((joints2d.shape[0], len(dst_names), joints2d.shape[-1]))

    for idx, jn in enumerate(dst_names):
        if jn in src_names:
            out_joints2d[:, idx] = joints2d[:, src_names.index(jn)]

    return out_joints2d


def get_perm_idxs(src, dst):
    src_names = eval(f'get_{src}_joint_names')()
    dst_names = eval(f'get_{dst}_joint_names')()
    idxs = [src_names.index(h) for h in dst_names if h in src_names]
    return idxs


def get_mpii3d_test_joint_names():
    return [
        'headtop', # 'head_top',
        'neck',
        'rshoulder',# 'right_shoulder',
        'relbow',# 'right_elbow',
        'rwrist',# 'right_wrist',
        'lshoulder',# 'left_shoulder',
        'lelbow', # 'left_elbow',
        'lwrist', # 'left_wrist',
        'rhip', # 'right_hip',
        'rknee', # 'right_knee',
        'rankle',# 'right_ankle',
        'lhip',# 'left_hip',
        'lknee',# 'left_knee',
        'lankle',# 'left_ankle'
        'hip',# 'pelvis',
        'Spine (H36M)',# 'spine',
        'Head (H36M)',# 'head'
    ]


def get_mpii3d_joint_names():
    return [
        'spine3', # 0,
        'spine4', # 1,
        'spine2', # 2,
        'Spine (H36M)', #'spine', # 3,
        'hip', # 'pelvis', # 4,
        'neck', # 5,
        'Head (H36M)', # 'head', # 6,
        "headtop", # 'head_top', # 7,
        'left_clavicle', # 8,
        "lshoulder", # 'left_shoulder', # 9,
        "lelbow", # 'left_elbow',# 10,
        "lwrist", # 'left_wrist',# 11,
        'left_hand',# 12,
        'right_clavicle',# 13,
        'rshoulder',# 'right_shoulder',# 14,
        'relbow',# 'right_elbow',# 15,
        'rwrist',# 'right_wrist',# 16,
        'right_hand',# 17,
        'lhip', # left_hip',# 18,
        'lknee', # 'left_knee',# 19,
        'lankle', #left ankle # 20
        'left_foot', # 21
        'left_toe', # 22
        "rhip", # 'right_hip',# 23
        "rknee", # 'right_knee',# 24
        "rankle", #'right_ankle', # 25
        'right_foot',# 26
        'right_toe' # 27
    ]


# def get_insta_joint_names():
#     return [
#         'rheel'            ,   # 0
#         'rknee'            ,   # 1
#         'rhip'             ,   # 2
#         'lhip'             ,   # 3
#         'lknee'            ,   # 4
#         'lheel'            ,   # 5
#         'rwrist'           ,   # 6
#         'relbow'           ,   # 7
#         'rshoulder'        ,   # 8
#         'lshoulder'        ,   # 9
#         'lelbow'           ,   # 10
#         'lwrist'           ,   # 11
#         'neck'             ,   # 12
#         'headtop'          ,   # 13
#         'nose'             ,   # 14
#         'leye'             ,   # 15
#         'reye'             ,   # 16
#         'lear'             ,   # 17
#         'rear'             ,   # 18
#         'lbigtoe'          ,   # 19
#         'rbigtoe'          ,   # 20
#         'lsmalltoe'        ,   # 21
#         'rsmalltoe'        ,   # 22
#         'lankle'           ,   # 23
#         'rankle'           ,   # 24
#     ]


def get_insta_joint_names():
    return [
        'OP RHeel',
        'OP RKnee',
        'OP RHip',
        'OP LHip',
        'OP LKnee',
        'OP LHeel',
        'OP RWrist',
        'OP RElbow',
        'OP RShoulder',
        'OP LShoulder',
        'OP LElbow',
        'OP LWrist',
        'OP Neck',
        'headtop',
        'OP Nose',
        'OP LEye',
        'OP REye',
        'OP LEar',
        'OP REar',
        'OP LBigToe',
        'OP RBigToe',
        'OP LSmallToe',
        'OP RSmallToe',
        'OP LAnkle',
        'OP RAnkle',
    ]


def get_mmpose_joint_names():
    # this naming is for the first 23 joints of MMPose
    # does not include hands and face
    return [
        'OP Nose', # 1
        'OP LEye', # 2
        'OP REye', # 3
        'OP LEar', # 4
        'OP REar', # 5
        'OP LShoulder', # 6
        'OP RShoulder', # 7
        'OP LElbow', # 8
        'OP RElbow', # 9
        'OP LWrist', # 10
        'OP RWrist', # 11
        'OP LHip', # 12
        'OP RHip', # 13
        'OP LKnee', # 14
        'OP RKnee', # 15
        'OP LAnkle', # 16
        'OP RAnkle', # 17
        'OP LBigToe', # 18
        'OP LSmallToe', # 19
        'OP LHeel', # 20
        'OP RBigToe', # 21
        'OP RSmallToe', # 22
        'OP RHeel', # 23
    ]


def get_insta_skeleton():
    return np.array(
        [
            [0 , 1],
            [1 , 2],
            [2 , 3],
            [3 , 4],
            [4 , 5],
            [6 , 7],
            [7 , 8],
            [8 , 9],
            [9 ,10],
            [2 , 8],
            [3 , 9],
            [10,11],
            [8 ,12],
            [9 ,12],
            [12,13],
            [12,14],
            [14,15],
            [14,16],
            [15,17],
            [16,18],
            [0 ,20],
            [20,22],
            [5 ,19],
            [19,21],
            [5 ,23],
            [0 ,24],
        ])


def get_staf_skeleton():
    return np.array(
        [
            [0, 1],
            [1, 2],
            [2, 3],
            [3, 4],
            [1, 5],
            [5, 6],
            [6, 7],
            [1, 8],
            [8, 9],
            [9, 10],
            [10, 11],
            [8, 12],
            [12, 13],
            [13, 14],
            [0, 15],
            [0, 16],
            [15, 17],
            [16, 18],
            [2, 9],
            [5, 12],
            [1, 19],
            [20, 19],
        ]
    )


def get_staf_joint_names():
    return [
        'OP Nose', # 0,
        'OP Neck', # 1,
        'OP RShoulder', # 2,
        'OP RElbow', # 3,
        'OP RWrist', # 4,
        'OP LShoulder', # 5,
        'OP LElbow', # 6,
        'OP LWrist', # 7,
        'OP MidHip', # 8,
        'OP RHip', # 9,
        'OP RKnee', # 10,
        'OP RAnkle', # 11,
        'OP LHip', # 12,
        'OP LKnee', # 13,
        'OP LAnkle', # 14,
        'OP REye', # 15,
        'OP LEye', # 16,
        'OP REar', # 17,
        'OP LEar', # 18,
        'Neck (LSP)', # 19,
        'Top of Head (LSP)', # 20,
    ]


def get_spin_op_joint_names():
    return [
        'OP Nose',        # 0
        'OP Neck',        # 1
        'OP RShoulder',   # 2
        'OP RElbow',      # 3
        'OP RWrist',      # 4
        'OP LShoulder',   # 5
        'OP LElbow',      # 6
        'OP LWrist',      # 7
        'OP MidHip',      # 8
        'OP RHip',        # 9
        'OP RKnee',       # 10
        'OP RAnkle',      # 11
        'OP LHip',        # 12
        'OP LKnee',       # 13
        'OP LAnkle',      # 14
        'OP REye',        # 15
        'OP LEye',        # 16
        'OP REar',        # 17
        'OP LEar',        # 18
        'OP LBigToe',     # 19
        'OP LSmallToe',   # 20
        'OP LHeel',       # 21
        'OP RBigToe',     # 22
        'OP RSmallToe',   # 23
        'OP RHeel',       # 24
    ]


def get_openpose_joint_names():
    return [
        'OP Nose',        # 0
        'OP Neck',        # 1
        'OP RShoulder',   # 2
        'OP RElbow',      # 3
        'OP RWrist',      # 4
        'OP LShoulder',   # 5
        'OP LElbow',      # 6
        'OP LWrist',      # 7
        'OP MidHip',      # 8
        'OP RHip',        # 9
        'OP RKnee',       # 10
        'OP RAnkle',      # 11
        'OP LHip',        # 12
        'OP LKnee',       # 13
        'OP LAnkle',      # 14
        'OP REye',        # 15
        'OP LEye',        # 16
        'OP REar',        # 17
        'OP LEar',        # 18
        'OP LBigToe',     # 19
        'OP LSmallToe',   # 20
        'OP LHeel',       # 21
        'OP RBigToe',     # 22
        'OP RSmallToe',   # 23
        'OP RHeel',       # 24
    ]


def get_spin_joint_names():
    return [
        'OP Nose',        # 0
        'OP Neck',        # 1
        'OP RShoulder',   # 2
        'OP RElbow',      # 3
        'OP RWrist',      # 4
        'OP LShoulder',   # 5
        'OP LElbow',      # 6
        'OP LWrist',      # 7
        'OP MidHip',      # 8
        'OP RHip',        # 9
        'OP RKnee',       # 10
        'OP RAnkle',      # 11
        'OP LHip',        # 12
        'OP LKnee',       # 13
        'OP LAnkle',      # 14
        'OP REye',        # 15
        'OP LEye',        # 16
        'OP REar',        # 17
        'OP LEar',        # 18
        'OP LBigToe',     # 19
        'OP LSmallToe',   # 20
        'OP LHeel',       # 21
        'OP RBigToe',     # 22
        'OP RSmallToe',   # 23
        'OP RHeel',       # 24
        'rankle',         # 25
        'rknee',          # 26
        'rhip',           # 27
        'lhip',           # 28
        'lknee',          # 29
        'lankle',         # 30
        'rwrist',         # 31
        'relbow',         # 32
        'rshoulder',      # 33
        'lshoulder',      # 34
        'lelbow',         # 35
        'lwrist',         # 36
        'neck',           # 37
        'headtop',        # 38
        'hip',            # 39 'Pelvis (MPII)', # 39
        'thorax',         # 40 'Thorax (MPII)', # 40
        'Spine (H36M)',   # 41
        'Jaw (H36M)',     # 42
        'Head (H36M)',    # 43
        'nose',           # 44
        'leye',           # 45 'Left Eye', # 45
        'reye',           # 46 'Right Eye', # 46
        'lear',           # 47 'Left Ear', # 47
        'rear',           # 48 'Right Ear', # 48
    ]

def get_muco3dhp_joint_names():
    return [
        'headtop',
        'thorax',
        'rshoulder',
        'relbow',
        'rwrist',
        'lshoulder',
        'lelbow',
        'lwrist',
        'rhip',
        'rknee',
        'rankle',
        'lhip',
        'lknee',
        'lankle',
        'hip',
        'Spine (H36M)',
        'Head (H36M)',
        'R_Hand',
        'L_Hand',
        'R_Toe',
        'L_Toe'
    ]

def get_h36m_joint_names():
    return [
        'hip',  # 0
        'lhip',  # 1
        'lknee',  # 2
        'lankle',  # 3
        'rhip',  # 4
        'rknee',  # 5
        'rankle',  # 6
        'Spine (H36M)',  # 7
        'neck',  # 8
        'Head (H36M)',  # 9
        'headtop',  # 10
        'lshoulder',  # 11
        'lelbow',  # 12
        'lwrist',  # 13
        'rshoulder',  # 14
        'relbow',  # 15
        'rwrist',  # 16
    ]


def get_spin_skeleton():
    return np.array(
        [
            [0 , 1],
            [1 , 2],
            [2 , 3],
            [3 , 4],
            [1 , 5],
            [5 , 6],
            [6 , 7],
            [1 , 8],
            [8 , 9],
            [9 ,10],
            [10,11],
            [8 ,12],
            [12,13],
            [13,14],
            [0 ,15],
            [0 ,16],
            [15,17],
            [16,18],
            [21,19],
            [19,20],
            [14,21],
            [11,24],
            [24,22],
            [22,23],
            [0 ,38],
        ]
    )


def get_openpose_skeleton():
    return np.array(
        [
            [0 , 1],
            [1 , 2],
            [2 , 3],
            [3 , 4],
            [1 , 5],
            [5 , 6],
            [6 , 7],
            [1 , 8],
            [8 , 9],
            [9 ,10],
            [10,11],
            [8 ,12],
            [12,13],
            [13,14],
            [0 ,15],
            [0 ,16],
            [15,17],
            [16,18],
            [21,19],
            [19,20],
            [14,21],
            [11,24],
            [24,22],
            [22,23],
        ]
    )


def get_posetrack_joint_names():
    return [
        "nose",
        "neck",
        "headtop",
        "lear",
        "rear",
        "lshoulder",
        "rshoulder",
        "lelbow",
        "relbow",
        "lwrist",
        "rwrist",
        "lhip",
        "rhip",
        "lknee",
        "rknee",
        "lankle",
        "rankle"
    ]


def get_posetrack_original_kp_names():
    return [
        'nose',
        'head_bottom',
        'head_top',
        'left_ear',
        'right_ear',
        'left_shoulder',
        'right_shoulder',
        'left_elbow',
        'right_elbow',
        'left_wrist',
        'right_wrist',
        'left_hip',
        'right_hip',
        'left_knee',
        'right_knee',
        'left_ankle',
        'right_ankle'
    ]


def get_pennaction_joint_names():
   return [
       "headtop",   # 0
       "lshoulder", # 1
       "rshoulder", # 2
       "lelbow",    # 3
       "relbow",    # 4
       "lwrist",    # 5
       "rwrist",    # 6
       "lhip" ,     # 7
       "rhip" ,     # 8
       "lknee",     # 9
       "rknee" ,    # 10
       "lankle",    # 11
       "rankle"     # 12
   ]


def get_common_joint_names():
    return [
        "rankle",    # 0  "lankle",    # 0
        "rknee",     # 1  "lknee",     # 1
        "rhip",      # 2  "lhip",      # 2
        "lhip",      # 3  "rhip",      # 3
        "lknee",     # 4  "rknee",     # 4
        "lankle",    # 5  "rankle",    # 5
        "rwrist",    # 6  "lwrist",    # 6
        "relbow",    # 7  "lelbow",    # 7
        "rshoulder", # 8  "lshoulder", # 8
        "lshoulder", # 9  "rshoulder", # 9
        "lelbow",    # 10  "relbow",    # 10
        "lwrist",    # 11  "rwrist",    # 11
        "neck",      # 12  "neck",      # 12
        "headtop",   # 13  "headtop",   # 13
    ]


def get_common_paper_joint_names():
    return [
        "Right Ankle",    # 0  "lankle",    # 0
        "Right Knee",     # 1  "lknee",     # 1
        "Right Hip",      # 2  "lhip",      # 2
        "Left Hip",      # 3  "rhip",      # 3
        "Left Knee",     # 4  "rknee",     # 4
        "Left Ankle",    # 5  "rankle",    # 5
        "Right Wrist",    # 6  "lwrist",    # 6
        "Right Elbow",    # 7  "lelbow",    # 7
        "Right Shoulder", # 8  "lshoulder", # 8
        "Left Shoulder", # 9  "rshoulder", # 9
        "Left Elbow",    # 10  "relbow",    # 10
        "Left Wrist",    # 11  "rwrist",    # 11
        "Neck",      # 12  "neck",      # 12
        "Head",   # 13  "headtop",   # 13
    ]


def get_common_skeleton():
    return np.array(
        [
            [ 0, 1 ],
            [ 1, 2 ],
            [ 3, 4 ],
            [ 4, 5 ],
            [ 6, 7 ],
            [ 7, 8 ],
            [ 8, 2 ],
            [ 8, 9 ],
            [ 9, 3 ],
            [ 2, 3 ],
            [ 8, 12],
            [ 9, 10],
            [12, 9 ],
            [10, 11],
            [12, 13],
        ]
    )


def get_coco_joint_names():
    return [
        "nose",      # 0
        "leye",      # 1
        "reye",      # 2
        "lear",      # 3
        "rear",      # 4
        "lshoulder", # 5
        "rshoulder", # 6
        "lelbow",    # 7
        "relbow",    # 8
        "lwrist",    # 9
        "rwrist",    # 10
        "lhip",      # 11
        "rhip",      # 12
        "lknee",     # 13
        "rknee",     # 14
        "lankle",    # 15
        "rankle",    # 16
    ]


def get_ochuman_joint_names():
    return [
        'rshoulder',
        'relbow',
        'rwrist',
        'lshoulder',
        'lelbow',
        'lwrist',
        'rhip',
        'rknee',
        'rankle',
        'lhip',
        'lknee',
        'lankle',
        'headtop',
        'neck',
        'rear',
        'lear',
        'nose',
        'reye',
        'leye'
    ]


def get_crowdpose_joint_names():
    return [
        'lshoulder',
        'rshoulder',
        'lelbow',
        'relbow',
        'lwrist',
        'rwrist',
        'lhip',
        'rhip',
        'lknee',
        'rknee',
        'lankle',
        'rankle',
        'headtop',
        'neck'
    ]

def get_coco_skeleton():
    # 0  - nose,
    # 1  - leye,
    # 2  - reye,
    # 3  - lear,
    # 4  - rear,
    # 5  - lshoulder,
    # 6  - rshoulder,
    # 7  - lelbow,
    # 8  - relbow,
    # 9  - lwrist,
    # 10 - rwrist,
    # 11 - lhip,
    # 12 - rhip,
    # 13 - lknee,
    # 14 - rknee,
    # 15 - lankle,
    # 16 - rankle,
    return np.array(
        [
            [15, 13],
            [13, 11],
            [16, 14],
            [14, 12],
            [11, 12],
            [ 5, 11],
            [ 6, 12],
            [ 5, 6 ],
            [ 5, 7 ],
            [ 6, 8 ],
            [ 7, 9 ],
            [ 8, 10],
            [ 1, 2 ],
            [ 0, 1 ],
            [ 0, 2 ],
            [ 1, 3 ],
            [ 2, 4 ],
            [ 3, 5 ],
            [ 4, 6 ]
        ]
    )


def get_mpii_joint_names():
    return [
        "rankle",    # 0
        "rknee",     # 1
        "rhip",      # 2
        "lhip",      # 3
        "lknee",     # 4
        "lankle",    # 5
        "hip",       # 6
        "thorax",    # 7
        "neck",      # 8
        "headtop",   # 9
        "rwrist",    # 10
        "relbow",    # 11
        "rshoulder", # 12
        "lshoulder", # 13
        "lelbow",    # 14
        "lwrist",    # 15
    ]


def get_mpii_skeleton():
    # 0  - rankle,
    # 1  - rknee,
    # 2  - rhip,
    # 3  - lhip,
    # 4  - lknee,
    # 5  - lankle,
    # 6  - hip,
    # 7  - thorax,
    # 8  - neck,
    # 9  - headtop,
    # 10 - rwrist,
    # 11 - relbow,
    # 12 - rshoulder,
    # 13 - lshoulder,
    # 14 - lelbow,
    # 15 - lwrist,
    return np.array(
        [
            [ 0, 1 ],
            [ 1, 2 ],
            [ 2, 6 ],
            [ 6, 3 ],
            [ 3, 4 ],
            [ 4, 5 ],
            [ 6, 7 ],
            [ 7, 8 ],
            [ 8, 9 ],
            [ 7, 12],
            [12, 11],
            [11, 10],
            [ 7, 13],
            [13, 14],
            [14, 15]
        ]
    )


def get_aich_joint_names():
    return [
        "rshoulder", # 0
        "relbow",    # 1
        "rwrist",    # 2
        "lshoulder", # 3
        "lelbow",    # 4
        "lwrist",    # 5
        "rhip",      # 6
        "rknee",     # 7
        "rankle",    # 8
        "lhip",      # 9
        "lknee",     # 10
        "lankle",    # 11
        "headtop",   # 12
        "neck",      # 13
    ]


def get_aich_skeleton():
    # 0  - rshoulder,
    # 1  - relbow,
    # 2  - rwrist,
    # 3  - lshoulder,
    # 4  - lelbow,
    # 5  - lwrist,
    # 6  - rhip,
    # 7  - rknee,
    # 8  - rankle,
    # 9  - lhip,
    # 10 - lknee,
    # 11 - lankle,
    # 12 - headtop,
    # 13 - neck,
    return np.array(
        [
            [ 0, 1 ],
            [ 1, 2 ],
            [ 3, 4 ],
            [ 4, 5 ],
            [ 6, 7 ],
            [ 7, 8 ],
            [ 9, 10],
            [10, 11],
            [12, 13],
            [13, 0 ],
            [13, 3 ],
            [ 0, 6 ],
            [ 3, 9 ]
        ]
    )


def get_3dpw_joint_names():
    return [
        "nose",      # 0
        "thorax",    # 1
        "rshoulder", # 2
        "relbow",    # 3
        "rwrist",    # 4
        "lshoulder", # 5
        "lelbow",    # 6
        "lwrist",    # 7
        "rhip",      # 8
        "rknee",     # 9
        "rankle",    # 10
        "lhip",      # 11
        "lknee",     # 12
        "lankle",    # 13
    ]


def get_3dpw_skeleton():
    return np.array(
        [
            [ 0, 1 ],
            [ 1, 2 ],
            [ 2, 3 ],
            [ 3, 4 ],
            [ 1, 5 ],
            [ 5, 6 ],
            [ 6, 7 ],
            [ 2, 8 ],
            [ 5, 11],
            [ 8, 11],
            [ 8, 9 ],
            [ 9, 10],
            [11, 12],
            [12, 13]
        ]
    )


def get_smplcoco_joint_names():
    return [
        "rankle",    # 0
        "rknee",     # 1
        "rhip",      # 2
        "lhip",      # 3
        "lknee",     # 4
        "lankle",    # 5
        "rwrist",    # 6
        "relbow",    # 7
        "rshoulder", # 8
        "lshoulder", # 9
        "lelbow",    # 10
        "lwrist",    # 11
        "neck",      # 12
        "headtop",   # 13
        "nose",      # 14
        "leye",      # 15
        "reye",      # 16
        "lear",      # 17
        "rear",      # 18
    ]


def get_smplcoco_skeleton():
    return np.array(
        [
            [ 0, 1 ],
            [ 1, 2 ],
            [ 3, 4 ],
            [ 4, 5 ],
            [ 6, 7 ],
            [ 7, 8 ],
            [ 8, 12],
            [12, 9 ],
            [ 9, 10],
            [10, 11],
            [12, 13],
            [14, 15],
            [15, 17],
            [16, 18],
            [14, 16],
            [ 8, 2 ],
            [ 9, 3 ],
            [ 2, 3 ],
        ]
    )


def get_smpl_joint_names():
    return [
        'hips',            # 0
        'leftUpLeg',       # 1
        'rightUpLeg',      # 2
        'spine',           # 3
        'leftLeg',         # 4
        'rightLeg',        # 5
        'spine1',          # 6
        'leftFoot',        # 7
        'rightFoot',       # 8
        'spine2',          # 9
        'leftToeBase',     # 10
        'rightToeBase',    # 11
        'neck',            # 12
        'leftShoulder',    # 13
        'rightShoulder',   # 14
        'head',            # 15
        'leftArm',         # 16
        'rightArm',        # 17
        'leftForeArm',     # 18
        'rightForeArm',    # 19
        'leftHand',        # 20
        'rightHand',       # 21
        'leftHandIndex1',  # 22
        'rightHandIndex1', # 23
    ]


def get_smpl_paper_joint_names():
    return [
        'Hips',            # 0
        'Left Hip',       # 1
        'Right Hip',      # 2
        'Spine',           # 3
        'Left Knee',         # 4
        'Right Knee',        # 5
        'Spine_1',          # 6
        'Left Ankle',        # 7
        'Right Ankle',       # 8
        'Spine_2',          # 9
        'Left Toe',     # 10
        'Right Toe',    # 11
        'Neck',            # 12
        'Left Shoulder',    # 13
        'Right Shoulder',   # 14
        'Head',            # 15
        'Left Arm',         # 16
        'Right Arm',        # 17
        'Left Elbow',     # 18
        'Right Elbow',    # 19
        'Left Hand',        # 20
        'Right Hand',       # 21
        'Left Thumb',  # 22
        'Right Thumb', # 23
    ]


def get_smpl_neighbor_triplets():
    return [
        [ 0,  1, 2 ],  # 0
        [ 1,  4, 0 ],  # 1
        [ 2,  0, 5 ],  # 2
        [ 3,  0, 6 ],  # 3
        [ 4,  7, 1 ],  # 4
        [ 5,  2, 8 ],  # 5
        [ 6,  3, 9 ],  # 6
        [ 7, 10, 4 ],  # 7
        [ 8,  5, 11],  # 8
        [ 9, 13, 14],  # 9
        [10,  7, 4 ],  # 10
        [11,  8, 5 ],  # 11
        [12,  9, 15],  # 12
        [13, 16, 9 ],  # 13
        [14,  9, 17],  # 14
        [15,  9, 12],  # 15
        [16, 18, 13],  # 16
        [17, 14, 19],  # 17
        [18, 20, 16],  # 18
        [19, 17, 21],  # 19
        [20, 22, 18],  # 20
        [21, 19, 23],  # 21
        [22, 20, 18],  # 22
        [23, 19, 21],  # 23
    ]


def get_smpl_skeleton():
    return np.array(
        [
            [ 0, 1 ],
            [ 0, 2 ],
            [ 0, 3 ],
            [ 1, 4 ],
            [ 2, 5 ],
            [ 3, 6 ],
            [ 4, 7 ],
            [ 5, 8 ],
            [ 6, 9 ],
            [ 7, 10],
            [ 8, 11],
            [ 9, 12],
            [ 9, 13],
            [ 9, 14],
            [12, 15],
            [13, 16],
            [14, 17],
            [16, 18],
            [17, 19],
            [18, 20],
            [19, 21],
            [20, 22],
            [21, 23],
        ]
    )


def map_spin_joints_to_smpl():
    # this function primarily will be used to copy 2D keypoint
    # confidences to pose parameters
    return [
        [(39, 27, 28), 0],  # hip,lhip,rhip->hips
        [(28,), 1],  # lhip->leftUpLeg
        [(27,), 2],  # rhip->rightUpLeg
        [(41, 27, 28, 39), 3],  # Spine->spine
        [(29,), 4],  # lknee->leftLeg
        [(26,), 5],  # rknee->rightLeg
        [(41, 40, 33, 34,), 6],  # spine, thorax ->spine1
        [(30,), 7],  # lankle->leftFoot
        [(25,), 8],  # rankle->rightFoot
        [(40, 33, 34), 9],  # thorax,shoulders->spine2
        [(30,), 10],  # lankle -> leftToe
        [(25,), 11],  # rankle -> rightToe
        [(37, 42, 33, 34), 12],  # neck, shoulders -> neck
        [(34,), 13],  # lshoulder->leftShoulder
        [(33,), 14],  # rshoulder->rightShoulder
        [(33, 34, 38, 43, 44, 45, 46, 47, 48,), 15],  # nose, eyes, ears, headtop, shoulders->head
        [(34,), 16],  # lshoulder->leftArm
        [(33,), 17],  # rshoulder->rightArm
        [(35,), 18],  # lelbow->leftForeArm
        [(32,), 19],  # relbow->rightForeArm
        [(36,), 20],  # lwrist->leftHand
        [(31,), 21],  # rwrist->rightHand
        [(36,), 22],  # lhand -> leftHandIndex
        [(31,), 23],  # rhand -> rightHandIndex
    ]


def map_smpl_to_common():
    return [
        [(11, 8), 0], # rightToe, rightFoot -> rankle
        [(5,), 1], # rightleg -> rknee,
        [(2,), 2], # rhip
        [(1,), 3], # lhip
        [(4,), 4], # leftLeg -> lknee
        [(10, 7), 5], # lefttoe, leftfoot -> lankle
        [(21, 23), 6], # rwrist
        [(18,), 7], # relbow
        [(17, 14), 8],  # rshoulder
        [(16, 13), 9],  # lshoulder
        [(19,), 10],  # lelbow
        [(20, 22), 11],  # lwrist
        [(0, 3, 6, 9, 12), 12],  # neck
        [(15,), 13],  # headtop
    ]


def relation_among_spin_joints():
    # this function primarily will be used to copy 2D keypoint
    # confidences to 3D joints
    return [
        [(), 25],
        [(), 26],
        [(39,), 27],
        [(39,), 28],
        [(), 29],
        [(), 30],
        [(), 31],
        [(), 32],
        [(), 33],
        [(), 34],
        [(), 35],
        [(), 36],
        [(40,42,44,43,38,33,34,), 37],
        [(43,44,45,46,47,48,33,34,), 38],
        [(27,28,), 39],
        [(27,28,37,41,42,), 40],
        [(27,28,39,40,), 41],
        [(37,38,44,45,46,47,48,), 42],
        [(44,45,46,47,48,38,42,37,33,34,), 43],
        [(44,45,46,47,48,38,42,37,33,34), 44],
        [(44,45,46,47,48,38,42,37,33,34), 45],
        [(44,45,46,47,48,38,42,37,33,34), 46],
        [(44,45,46,47,48,38,42,37,33,34), 47],
        [(44,45,46,47,48,38,42,37,33,34), 48],
    ]
utils/kp_utils.py
sha2nkt/ipman-r
"""
Parts of the code are adapted from https://github.com/akanazawa/hmr
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import torch

def compute_similarity_transform(S1, S2):
    """
    Computes a similarity transform (sR, t) that takes
    a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
    where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
    i.e. solves the orthogonal Procrutes problem.
    """
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.T
        S2 = S2.T
        transposed = True
    assert(S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=1, keepdims=True)
    mu2 = S2.mean(axis=1, keepdims=True)
    X1 = S1 - mu1
    X2 = S2 - mu2

    # 2. Compute variance of X1 used for scale.
    var1 = np.sum(X1**2)

    # 3. The outer product of X1 and X2.
    K = X1.dot(X2.T)

    # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are
    # singular vectors of K.
    U, s, Vh = np.linalg.svd(K)
    V = Vh.T
    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = np.eye(U.shape[0])
    Z[-1, -1] *= np.sign(np.linalg.det(U.dot(V.T)))
    # Construct R.
    R = V.dot(Z.dot(U.T))

    # 5. Recover scale.
    scale = np.trace(R.dot(K)) / var1

    # 6. Recover translation.
    t = mu2 - scale*(R.dot(mu1))

    # 7. Error:
    S1_hat = scale*R.dot(S1) + t

    if transposed:
        S1_hat = S1_hat.T

    return S1_hat

def compute_similarity_transform_batch(S1, S2):
    """Batched version of compute_similarity_transform."""
    S1_hat = np.zeros_like(S1)
    for i in range(S1.shape[0]):
        S1_hat[i] = compute_similarity_transform(S1[i], S2[i])
    return S1_hat

def reconstruction_error(S1, S2, reduction='mean'):
    """Do Procrustes alignment and compute reconstruction error."""
    S1_hat = compute_similarity_transform_batch(S1, S2)
    re = np.sqrt( ((S1_hat - S2)** 2).sum(axis=-1)).mean(axis=-1)
    if reduction == 'mean':
        re = re.mean()
    elif reduction == 'sum':
        re = re.sum()
    return re

def compute_accel(joints):
    """
    Computes acceleration of 3D joints.
    Args:
        joints (Nx25x3).
    Returns:
        Accelerations (N-2).
    """
    velocities = joints[1:] - joints[:-1]
    acceleration = velocities[1:] - velocities[:-1]
    acceleration_normed = np.linalg.norm(acceleration, axis=2)
    return np.mean(acceleration_normed, axis=1)


def compute_error_accel(joints_gt, joints_pred, vis=None):
    """
    Computes acceleration error:
        1/(n-2) \sum_{i=1}^{n-1} X_{i-1} - 2X_i + X_{i+1}
    Note that for each frame that is not visible, three entries in the
    acceleration error should be zero'd out.
    Args:
        joints_gt (Nx14x3).
        joints_pred (Nx14x3).
        vis (N).
    Returns:
        error_accel (N-2).
    """
    # (N-2)x14x3
    accel_gt = joints_gt[:-2] - 2 * joints_gt[1:-1] + joints_gt[2:]
    accel_pred = joints_pred[:-2] - 2 * joints_pred[1:-1] + joints_pred[2:]

    normed = np.linalg.norm(accel_pred - accel_gt, axis=2)

    if vis is None:
        new_vis = np.ones(len(normed), dtype=bool)
    else:
        invis = np.logical_not(vis)
        invis1 = np.roll(invis, -1)
        invis2 = np.roll(invis, -2)
        new_invis = np.logical_or(invis, np.logical_or(invis1, invis2))[:-2]
        new_vis = np.logical_not(new_invis)

    return np.mean(normed[new_vis], axis=1)

utils/pose_utils.py
sha2nkt/ipman-r
import cv2
import torch
import numpy as np
import matplotlib.pyplot as plt
import os
from .rotation_converter import batch_euler2matrix


def plot_3d_andCam(p3d, camWorld):
    from .kp_utils import get_smpl_joint_names, get_spin_joint_names
    from mpl_toolkits.mplot3d import Axes3D
    jn = get_spin_joint_names()[25:]
    import matplotlib.cm as cm
    colors = cm.tab20c(np.linspace(0, 1, p3d.shape[0]))
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    for j in range(p3d.shape[0]):
        ax.scatter(p3d[j, 0], p3d[j, 1], p3d[j, 2], c=colors[j], label=j)
        label = '(%d, %d, %d), %s' % (p3d[j, 0], p3d[j, 1], p3d[j, 2], j)
        # ax.text(p3d[j, 0], p3d[j, 1], p3d[j, 2], jn[j])
    ax.set_xlim(-3, 3)
    ax.set_ylim(-3, 3)
    ax.set_zlim(0, 2)
    # ax.scatter(p3d[:, 0], p3d[:, 1], p3d[:, 2])


    # ax.scatter(camWorld[0,0], camWorld[0,1], camWorld[0,2], c='black', label='Camera')
    # ax.text(camWorld[0,0], camWorld[0,1], camWorld[0,2], 'Camera')
    # plt.legend()


    max_range = np.max(np.array([np.max(p3d[j, 0]) + np.min(p3d[j, 0]), np.max(p3d[j, 1]) + np.min(p3d[j, 1]), np.max(p3d[j, 2]) + np.min(p3d[j, 2])]))
    Xb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2, -1:2:2][0].flatten() + 0.5 * (np.max(p3d[j, 0]) + np.min(p3d[j, 0]))
    Yb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2, -1:2:2][1].flatten() + 0.5 * (np.max(p3d[j, 1]) + np.min(p3d[j, 1]))
    Zb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2, -1:2:2][2].flatten() + 0.5 * (np.max(p3d[j, 2]) + np.min(p3d[j, 2]))
    # Comment or uncomment following both lines to test the fake bounding box:
    # for xb, yb, zb in zip(Xb, Yb, Zb):
    #     ax.plot([xb], [yb], [zb], 'w')


    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')

    plt.show()
    print('wop')


def focalLength_mm2px(focalLength,dslr_sens,  focalPoint):
    focal_pixel = (focalLength / dslr_sens) * focalPoint *2
    return focal_pixel


def toCamCoords(j3d, camPosWorld):
    # transform gt to camera coordinate frame
    j3d = j3d - camPosWorld
    return  j3d


def unreal2cv2(points):
    '''

    :param points: Nx3
    :return:
    '''
    # x --> y, y --> z, z --> x
    points = np.roll(points, 2, 1)
    # change direction of y
    points = points * np.array([1,-1,1])
    return points


# def smpl2opencv(j3d):
#     # swap axis 1 and axis 2
#     j3d[:, [1, 2]] = j3d[:, [2, 1]]
#     j3d = j3d * np.array([1,1,-1])
#     return j3d


def smpl2opencv(j3d):
    # change sign of axis 1 and axis 2
    j3d = j3d * np.array([1,-1,-1])
    return j3d


def project_point(joint, RT, KKK):

    P = np.dot(KKK,RT)
    joints_2d = np.dot(P, joint)
    joints_2d = joints_2d[0:2] / joints_2d[2]

    return joints_2d


def project_2d(args, df, i, pNum, joints3d, imgPath, imgHeight, imgWidth, constants, meanPose=False):
    if args.scene3d:
        camPosWorld = [df.iloc[i]['camX'], df.iloc[i]['camY'], df.iloc[i]['camZ']]
        camYaw = df.iloc[i]['camYaw']
    else:
        camPosWorld = constants['camPosWorld']
        camYaw = constants['camPitch']

    if meanPose:
        transl = np.array([0,0,0])
        yawSMPL = 0
        trans3d = [0,0,0]
    else:
        transl = df.iloc[i]['gt'][pNum]['transl']
        yawSMPL = df.iloc[i]['YawLocal'][pNum]
        trans3d = [df.iloc[i]['XLocal'][pNum],
                   df.iloc[i]['YLocal'][pNum],
                   df.iloc[i]['ZLocal'][pNum]]

    if not (args.scene3d):
        gt2d, gt3d_camCoord, valid, valJoints = project2d_hdri(joints3d, focalLength=constants['focalLength'],
                                                           trans3d=trans3d,
                                                           dslr_sens_width=constants['dslr_sens_width'],
                                                           dslr_sens_height=constants['dslr_sens_height'],
                                                           camPosWorld=camPosWorld,
                                                           cy=imgHeight / 2,
                                                           cx=imgWidth / 2,
                                                           imgPath=imgPath,
                                                           yawSMPL=yawSMPL,
                                                           transl=transl,
                                                           debug=False,
                                                           checkValid=True,
                                                           constants=constants,
                                                           meanPose=meanPose)
    else:
        gt2d, gt3d_camCoord, valid, valJoints = project2d_world( joints3d, focalLength=constants['focalLength'],
                                                                 trans3d=trans3d,
                                                                 dslr_sens_width=constants['dslr_sens_width'],
                                                                 dslr_sens_height=constants['dslr_sens_height'],
                                                                 camPosWorld=camPosWorld,
                                                                 cy=imgHeight / 2,
                                                                 cx=imgWidth / 2,
                                                                 imgPath=imgPath,
                                                                 yawSMPL=yawSMPL,
                                                                 transl=transl,
                                                                 debug=False,
                                                                 checkValid=True,
                                                                 constants=constants,
                                                                 camYaw=camYaw,
                                                                 meanPose=meanPose)

    # gt2d are 2d projections of gt3d_camCoord
    return gt2d, gt3d_camCoord, valid, valJoints


def project2d_hdri(j3d, trans3d, focalLength, dslr_sens_width, dslr_sens_height, cy, cx, camPosWorld, yawSMPL,
                   imgPath=None, transl=None, debug=False, checkValid=False, constants=None, meanPose=False):
    """

    :param j3d:
    :param trans3d:
    :param focalLength:
    :param dslr_sens_width:
    :param dslr_sens_height:
    :param cy:
    :param cx:
    :param camPosWorld:
    :param yawSMPL:
    :param imgPath:
    :param transl:
    :return: 2d projection and 3d joints in camera coordinate system
    """

    focalLength_x = focalLength_mm2px(focalLength, dslr_sens_width, cx)
    focalLength_y = focalLength_mm2px(focalLength, dslr_sens_height, cy)

    camMat =  np.array( [[focalLength_x, 0, cx],
                        [0, focalLength_y,  cy],
                        [0 ,0,            1]])


    #camPosWorld and trans3d are in cm. Transform to meter
    trans3d = np.array(trans3d) / 100
    trans3d = unreal2cv2(np.reshape(trans3d,(1,3)))
    camPosWorld = np.array(camPosWorld) / 100
    camPosWorld = unreal2cv2(np.reshape(camPosWorld, (1,3)))

    # only when doing that with gt (other not necessary)
    # j3d = toCamCoords(j3d, camPosWorld)


    #get points in camera coordinate system
    j3d = smpl2opencv(j3d)

    #scans have a 90deg rotation, but for mean pose from vposer there is no such rotation
    if meanPose:
        rotMat, _ = cv2.Rodrigues(np.array([[0, (yawSMPL) / 180 * np.pi, 0]], dtype=float))
    else:
        rotMat, _ = cv2.Rodrigues(np.array([[0, (yawSMPL-90)/180*np.pi, 0]], dtype=float))

    j3d = np.matmul(rotMat, j3d.T).T

    j3d = j3d + trans3d
    #j3d = j3d - np.array([0, 0.8, 0])

    # plot_3d_andCam(j3d[:22,:], camPosWorld)

    camera_rotationMatrix, _ = cv2.Rodrigues(np.array([constants['camPitch']/180 * np.pi, 0., 0.]).reshape(3,1))
    j3d_new = np.matmul(camera_rotationMatrix, j3d.T - camPosWorld.T).T
    camPos_new = np.matmul(camera_rotationMatrix, camPosWorld.T - camPosWorld.T).T

    # plot_3d_andCam(j3d_new[:22, :], camPos_new)


    #do projection
    #rt is cannonical, since coordinate system was already changed to cam coords
    #
    RT = np.concatenate((np.diag([1., 1., 1.]), np.zeros((3, 1))), axis=1)
    j2d = np.zeros((j3d_new.shape[0], 2))
    for i in range(j3d_new.shape[0]):
        j2d[i,:] = project_point(np.concatenate([j3d_new[i,:], np.array([1])]), RT, camMat)



    #prepare imgpath
    if not(imgPath is None):

        # imgPath = imgPath.split('/')
        # imgPath[-2] = 'images' #hardcoded
        # imgName = imgPath[-1].split('_')[:2]
        # imgPath = os.path.join(*imgPath[:-1])
        #
        # imgName = imgName[0] + '_' + imgName[1] + '.png'
        # imgPath = os.path.join(imgPath,imgName)
        #imgPath = os.path.join(*imgPath)

        img = cv2.imread('/'+imgPath)
        img = img[:,:,::-1]


    if debug:
        import matplotlib.cm as cm
        colors = cm.tab20c(np.linspace(0, 1, 25))
        fig = plt.figure(dpi=300)
        ax = fig.add_subplot(111)
        if not (imgPath is None):
            ax.imshow(img)
        for i in range(22):
            ax.scatter(j2d[i,0], j2d[i,1], c=colors[i], s=10)
            # ax.text(j2d[i,0], j2d[i,1], str(i))
        # plt.show()

        if not (imgPath is None):
            savename = imgPath.split('/')[-1]
            savename=savename.replace('.pkl','.jpg')
            plt.savefig('./figs/tmp'+savename)
            plt.close('all')


    #check if person is valid ground truth (50% of joints visible)
    if checkValid:
        #S: valid is a boolean scalar, valJoints is a boolean flag for valid joints
        valid, valJoints = checkValidGT(j2d, cx*2, cy*2, return_val_joints=True)
        return j2d, j3d_new, valid, valJoints

    return j2d, j3d_new


def translation_unreal2cv(translation):
    # translation are in cm. Transform to meter
    translation = np.array(translation) / 100
    return unreal2cv2(np.reshape(translation, (1, 3)))


def rotation_unreal2cv(pitch, yaw, roll):
    # Unreal camera parameters are euler angles not axis angles.
    rotmat = batch_euler2matrix(
        torch.tensor(
            [[-pitch / 180 * np.pi, -yaw / 180 * np.pi, -roll / 180 * np.pi]]
        )
    )[0].numpy()
    return rotmat


def project2d_world(j3d, smpl_trans, smpl_rotmat, cam_pos_world, cam_intrinsics, cam_rotation_mat,
                    img_path=None, debug=False, check_valid=False, mean_pose=False,
                    img_w=1280, img_h=720, debug_file_name=''):

    # only when doing that with gt (other not necessary)
    # j3d = toCamCoords(j3d, cam_pos_world)

    # get points in camera coordinate system
    j3d = smpl2opencv(j3d)

    j3d = np.matmul(smpl_rotmat, j3d.T).T

    j3d = j3d + smpl_trans

    # plot_3d_andCam(j3d[:22,:], cam_pos_world)

    j3d_new = np.matmul(cam_rotation_mat, j3d.T - cam_pos_world.T).T

    # do projection
    # rt is cannonical, since coordinate system was already changed to cam coords
    RT = np.concatenate((np.diag([1., 1., 1.]), np.zeros((3, 1))), axis=1)
    j2d = np.zeros((j3d_new.shape[0], 2))
    for i in range(j3d_new.shape[0]):
        j2d[i,:] = project_point(np.concatenate([j3d_new[i,:], np.array([1])]), RT, cam_intrinsics)


    if debug:
        from .kp_utils import get_smpl_joint_names, get_spin_joint_names
        import matplotlib.cm as cm

        img = cv2.imread(img_path)
        img = img[:, :, ::-1]

        cam_pos_new = np.matmul(cam_rotation_mat, cam_pos_world.T - cam_pos_world.T).T
        # plot_3d_andCam(j3d_new[25:, :], cam_pos_new)

        j2d[:, 0] = np.clip(j2d[:, 0], 0, img_w)
        j2d[:, 1] = np.clip(j2d[:, 1], 0, img_h)
        # breakpoint()
        jn = get_spin_joint_names() # get_smpl_joint_names()
        colors = cm.tab20c(np.linspace(0, 1, len(j2d[:,0])))
        fig = plt.figure(figsize=(19.2, 10.8))
        ax = fig.add_subplot(111)
        if not (img_path is None):
            ax.imshow(img)
        for i in range(len(j2d[:,0])):
            ax.scatter(j2d[i,0], j2d[i,1], s=10) # , c=colors[i])
            ax.text(j2d[i,0], j2d[i,1], jn[i], fontsize=5)

        if debug_file_name != '':
            plt.savefig(debug_file_name, bbox_inches='tight')
        else:
            plt.show()

        plt.close()


    # check if person is valid ground truth (50% of joints visible)
    # if check_valid:
    #     valid, valJoints = checkValidGT(j2d, cx*2, cy*2, return_val_joints=True)
    #     return j2d, j3d_new, valid, valJoints

    return j2d, j3d_new


def get_head_bbox_size(j2d):
    head2d = np.concatenate([j2d[55:60, :], j2d[76:,:]], axis=0)

    minPt = np.min(head2d, axis=0)
    maxPt = np.max(head2d, axis=0)
    bbxDiagSize = np.linalg.norm(maxPt-minPt,ord=2)

    return bbxDiagSize


def checkValidGT(j2d, resx, resy, threshold = 0.5, return_val_joints=False, numJointsToCheck=22):
    #todo: function not tested
    valX = np.logical_and(j2d[:, 0] > 0,  j2d[:,0] < resx )
    valY = np.logical_and(j2d[:, 1] > 0, j2d[:, 1] < resy)

    valJoint = np.logical_and(valX, valY)

    if np.sum(valJoint[:numJointsToCheck]) / numJointsToCheck < threshold:
        if return_val_joints:
            return False, valJoint
        else:
            return False
    else:
        if return_val_joints:
            return True, valJoint
        else:
            return True
utils/projection.py
sha2nkt/ipman-r
import os
os.environ['PYOPENGL_PLATFORM'] = 'egl'
# # hack to get the correct gpu device id on cluster
# #os.environ['PYOPENGL_PLATFORM'] = 'egl'
# #os.environ['EGL_DEVICE_ID'] = os.environ['GPU_DEVICE_ORDINAL'].split(',')[0]
import torch
from torchvision.utils import make_grid
import cv2
import numpy as np
import pyrender
import matplotlib as mpl
import trimesh
from vis_utils.mesh_utils import get_checkerboard_plane, GMoF_unscaled
from utils.geometry import get_view_matrix
import PIL.Image as pil_img


class Renderer:
    """
    Renderer used for visualizing the SMPL model
    Code adapted from https://github.com/vchoutas/smplify-x
    """
    def __init__(self, focal_length=5000, img_res=224, faces=None):

        self.focal_length = focal_length

        if isinstance(img_res, tuple):
            self.camera_center = [img_res[1] // 2, img_res[0] // 2]
            self.renderer = pyrender.OffscreenRenderer(viewport_width=img_res[1],
                                                       viewport_height=img_res[0],
                                                       point_size=1.0)
        else:
            #square image
            self.camera_center = [img_res // 2, img_res // 2]
            self.renderer = pyrender.OffscreenRenderer(viewport_width=img_res,
                                                       viewport_height=img_res,
                                                       point_size=1.0)
        self.faces = faces

    def visualize_tb(self, vertices, camera_translation, images):
        vertices = vertices.cpu().numpy()
        camera_translation = camera_translation.cpu().numpy()
        images = images.cpu()
        images_np = np.transpose(images.numpy(), (0,2,3,1))
        rend_imgs = []
        for i in range(vertices.shape[0]):
            rend_img = torch.from_numpy(np.transpose(self.__call__(vertices[i], camera_translation[i], images_np[i]), (2,0,1))).float()
            rend_imgs.append(images[i])
            rend_imgs.append(rend_img)
        rend_imgs = make_grid(rend_imgs, nrow=2)
        return rend_imgs

    def __call__(self, vertices, camera_translation, image):
        material = pyrender.MetallicRoughnessMaterial(
            metallicFactor=0.2,
            alphaMode='OPAQUE',
            baseColorFactor=(0.8, 0.3, 0.3, 1.0))

        camera_translation[0] *= -1.

        mesh = trimesh.Trimesh(vertices, self.faces)
        rot = trimesh.transformations.rotation_matrix(
            np.radians(180), [1, 0, 0])
        mesh.apply_transform(rot)
        mesh = pyrender.Mesh.from_trimesh(mesh, material=material)

        scene = pyrender.Scene(ambient_light=(0.5, 0.5, 0.5))
        scene.add(mesh, 'mesh')

        camera_pose = np.eye(4)
        camera_pose[:3, 3] = camera_translation
        camera = pyrender.IntrinsicsCamera(fx=self.focal_length, fy=self.focal_length,
                                           cx=self.camera_center[0], cy=self.camera_center[1])
        scene.add(camera, pose=camera_pose)


        light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
        light_pose = np.eye(4)

        light_pose[:3, 3] = np.array([0, -1, 1])
        scene.add(light, pose=light_pose)

        light_pose[:3, 3] = np.array([0, 1, 1])
        scene.add(light, pose=light_pose)

        light_pose[:3, 3] = np.array([1, 1, 2])
        scene.add(light, pose=light_pose)

        color, rend_depth = self.renderer.render(scene, flags=pyrender.RenderFlags.RGBA)
        color = color.astype(np.float32) / 255.0
        valid_mask = (rend_depth > 0)[:,:,None]
        output_img = (color[:, :, :3] * valid_mask +
                  (1 - valid_mask) * image)
        return output_img

def vis_heatmap(heatmap):
    imgpath = './data/checkerboard.png'
    image = cv2.imread(imgpath).astype(np.uint8)
    heatmap = np.clip(heatmap*255, 0, 255).astype(np.uint8)
    resized_heatmap = cv2.resize(heatmap, (int(image.shape[1]), int(image.shape[0])))
    # idx = resized_heatmap.argmax()

    # center_y = idx // image.shape[1]
    # center_x = idx % image.shape[1]

    colored_heatmap = cv2.applyColorMap(resized_heatmap.astype(np.uint8), cv2.COLORMAP_JET)
    # colored_heatmap = cv2.applyColorMap(heatmap.astype(np.uint8), cv2.COLORMAP_JET)
    masked_image = colored_heatmap*0.8 + image*0.2

    # cv2.circle(masked_image, (center_x, center_y), 3, [0, 0, 255], -1)
    # cv2.circle(masked_image, (center_x, center_y), 3, [0, 0, 255], -1)
    return masked_image.astype(np.uint8)

def  overlay_mesh(verts, faces, camera_transl, focal_length_x, focal_length_y, camera_center,
        H, W, img, camera_rotation=None, rotaround=None, topview=False, bottomview=False, viewer=False, gmof_rho= 0.01,
                 draw_ground_plane=True, draw_contact=True, draw_support_plane=False, com=None, xcom=None):

    material = pyrender.MetallicRoughnessMaterial(
                metallicFactor=0.0,
                alphaMode='OPAQUE',
                baseColorFactor=(1.0, 1.0, 0.9, 1.0))

    support_plane_material = pyrender.MetallicRoughnessMaterial(
        metallicFactor=0.2,
        alphaMode='BLEND',
        baseColorFactor=(1.0, 1.0, 1.0, 0.6))

    out_mesh = trimesh.Trimesh(verts, faces, process=False)

    # Set mesh vertex color based on distance from the ground plane
    ground_plane_height = 0.0 # obtained by visualization on the presented pose
    view_offset = 0.0 # groundplane at 0 not visible, lower mesh and ground plane just for viewing

    # Calculate contact vertices
    vertex_height = (out_mesh.vertices[:, 1] - ground_plane_height)
    vertex_height_robustified = GMoF_unscaled(rho=gmof_rho)(vertex_height)


    if draw_contact:
        cmap = mpl.cm.get_cmap('jet')
        vertex_colors = [cmap(1 - x) for x in vertex_height_robustified]
        out_mesh.visual.vertex_colors = vertex_colors
    else:
        color = np.array([1.00, 0.55, 0.41]) # salmon1
        out_mesh.visual.vertex_colors = color

    # draw ground_plane
    if draw_ground_plane:
        ground_mesh = get_checkerboard_plane()
        ground_pose = np.eye(4)
        if rotaround is not None:
            if topview or bottomview:
                ground_pose = trimesh.transformations.rotation_matrix(np.radians(90), [1, 0, 0]) # make groundplane frontal
            else:
                ground_pose = trimesh.transformations.rotation_matrix(np.radians(0), [1, 0, 0])
            # ground_pose[:3, 3] = np.array([0, ground_plane_height+view_offset, 0])
            # ground_mesh = [mesh.apply_transform(ground_pose) for mesh in ground_mesh]

    # draw support plane that is perpendicular to the ground plane and passes through the centroid of contact vertices
    if draw_support_plane:
        support_mesh = get_checkerboard_plane()
        # Align with y-z plane
        support_pose = trimesh.transformations.rotation_matrix(np.radians(90), [0, 1, 0])
        # get z offset
        vertex_idx = vertex_height_robustified < 0.2
        centroid = np.mean(out_mesh.vertices[vertex_idx, :], axis=0)
        support_plane_offset = centroid[0] # ToDo: fit to the contact vertices to minimize distance instead
        support_pose[:3, 3] = np.array([support_plane_offset, 0, 0])


    if rotaround is not None:
        if topview:
            rot = trimesh.transformations.rotation_matrix(
                np.radians(90), [1, 0, 0], [0,0,0])
        elif bottomview:
            rot = trimesh.transformations.rotation_matrix(
                np.radians(270), [1, 0, 0], [0, 0, 0])
        else:
            rot = trimesh.transformations.rotation_matrix(
                np.radians(rotaround), [0, 1, 0], out_mesh.vertices[4297])
        out_mesh.apply_transform(rot)
        if draw_support_plane:
            support_pose = np.dot(rot, support_pose)
        view_matrix = get_view_matrix(
            pitch=180,yaw=-0.5,roll=0, tx=camera_transl[0], ty=camera_transl[1], tz=camera_transl[2]
        )
    else:
        camera_pose_gl = get_view_matrix(
            pitch=180+camera_rotation[0],yaw=-camera_rotation[1],roll=-camera_rotation[2],
            tx=camera_transl[0], ty=-camera_transl[1], tz=-camera_transl[2]
        ) # note coversion from camera_orientation to global orientation
        out_mesh.vertices = camera_pose_gl[:3, 3] + np.einsum('mn,jn->jm', camera_pose_gl[:3, :3], verts)

        if draw_ground_plane:
            ground_mesh = [mesh.apply_transform(camera_pose_gl) for mesh in ground_mesh]
        if draw_support_plane:
            support_pose = np.dot(camera_pose_gl, support_pose)

        view_matrix = np.eye(4)

    # add origin sphere
    origin_sphere = trimesh.creation.uv_sphere(0.05)
    yellow = [255, 255, 0, 255]
    origin_sphere.visual.face_colors = yellow
    origin_sphere_mesh = pyrender.Mesh.from_trimesh(origin_sphere, smooth=False)
    origin_pose = np.eye(4)

    # add com sphere
    if com is not None:
        com_sphere = trimesh.creation.uv_sphere(0.05)
        green = [0, 255, 0, 255]
        com_sphere.visual.face_colors = green
        com_sphere_mesh = pyrender.Mesh.from_trimesh(com_sphere, smooth=False)
        com_pose = np.eye(4)
        if rotaround is not None:
            origin_pose[:3, 3] = np.array([0, view_offset, 0])  # just for viewing, lower the mesh with the ground plane
            if bottomview:
                com_pose[:3, 3] = np.array([com[0], com[2], 0])
            else:
                com_pose[:3, 3] = np.array([com[0], 0, com[2]])  # visualize com projection on the ground plane

    # add com sphere
    if xcom is not None:
        xcom_sphere = trimesh.creation.uv_sphere(0.05)
        orange = [255, 99, 71, 255]
        xcom_sphere.visual.face_colors = orange
        xcom_sphere_mesh = pyrender.Mesh.from_trimesh(xcom_sphere, smooth=False)
        xcom_pose = np.eye(4)
        if rotaround is not None:
            origin_pose[:3, 3] = np.array([0, view_offset, 0])  # just for viewing, lower the mesh with the ground plane
            if bottomview:
                xcom_pose[:3, 3] = np.array([xcom[0], xcom[2], 0])
            else:
                xcom_pose[:3, 3] = np.array([xcom[0], 0, xcom[2]])  # visualize xcom projection on the ground plane

    # add mesh to scene
    mesh = pyrender.Mesh.from_trimesh(
        out_mesh,
        material=material)
    if img is not None:
        scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0],
                               ambient_light=(0.3, 0.3, 0.3))
    else:
        scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0],
                               ambient_light=(0.3, 0.3, 0.3))
    mesh_pose = np.eye(4)
    if rotaround is not None:
        mesh_pose[:3, 3] = np.array([0, view_offset, 0]) # just for viewing, lower the mesh with the ground plane
    scene.add(mesh, pose=mesh_pose, name='mesh')

    if draw_ground_plane:
        ground_mesh = pyrender.Mesh.from_trimesh(
            ground_mesh,
            smooth=False)
        # scene.add(ground_mesh, pose=ground_pose, name='ground_plane')
        scene.add(ground_mesh, name='ground_plane')

    if draw_support_plane:
        support_mesh = [mesh.apply_transform(support_pose) for mesh in support_mesh]
        support_mesh = pyrender.Mesh.from_trimesh(
            support_mesh,
            smooth=False,
            material=support_plane_material
        )
        scene.add(support_mesh, name='support_plane')

    scene.add(origin_sphere_mesh, pose=origin_pose, name='origin_sphere')
    if com is not None:
        scene.add(com_sphere_mesh, pose=com_pose, name='com_sphere')
    if xcom is not None:
        scene.add(xcom_sphere_mesh, pose=xcom_pose, name='xcom_sphere')

    # add mesh for camera
    if img is None and topview is False:
        # to lower the ground-plane so person doesn't get cut
        cam_y_offset = 150
    elif topview:
        cam_y_offset = -100
    else:
        cam_y_offset = 0
    pyrencamera = pyrender.camera.IntrinsicsCamera(
        fx=focal_length_x, fy=focal_length_y,
        cx=camera_center[0], cy=camera_center[1] + cam_y_offset)
    scene.add(pyrencamera, pose=view_matrix)

    # create and add light
    light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=1)
    light_pose = np.eye(4)
    for lp in [[1,1,1], [-1,1,1],[1,-1,1],[-1,-1,1]]:
        light_pose[:3, 3] = out_mesh.vertices.mean(0) + np.array(lp)
        #out_mesh.vertices.mean(0) + np.array(lp)
        scene.add(light, pose=light_pose)

    if viewer:
        pyrender.Viewer(scene, use_raymond_lighting=True)
        return 0
    else:
        r = pyrender.OffscreenRenderer(viewport_width=W,
                                       viewport_height=H,
                                       point_size=1.0)
        color, _ = r.render(scene, flags=pyrender.RenderFlags.RGBA)
        color = color.astype(np.float32) / 255.0

        if img is not None:
            valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
            input_img = img.detach().cpu().numpy()
            output_img = (color[:, :, :-1] * valid_mask +
                          (1 - valid_mask) * input_img)
        else:
            output_img = color

        output_img = pil_img.fromarray((output_img * 255).astype(np.uint8))
        output_img= np.asarray(output_img)[:,:,:3]

        return output_img

utils/renderer.py
sha2nkt/ipman-r
import torch
import torch.nn.functional as F

''' Rotation Converter
    This function is borrowed from https://github.com/kornia/kornia

ref: https://kornia.readthedocs.io/en/v0.1.2/_modules/torchgeometry/core/conversions.html#
Repre: euler angle(3), axis angle(3), rotation matrix(3x3), quaternion(4), continuous rotation representation (6)
batch_rodrigues: axis angle -> matrix 
'''
pi = torch.Tensor([3.14159265358979323846])

def rad2deg(tensor):
    """Function that converts angles from radians to degrees.

    See :class:`~torchgeometry.RadToDeg` for details.

    Args:
        tensor (Tensor): Tensor of arbitrary shape.

    Returns:
        Tensor: Tensor with same shape as input.

    Example:
        >>> input = tgm.pi * torch.rand(1, 3, 3)
        >>> output = tgm.rad2deg(input)
    """
    if not torch.is_tensor(tensor):
        raise TypeError("Input type is not a torch.Tensor. Got {}"
                        .format(type(tensor)))

    return 180. * tensor / pi.to(tensor.device).type(tensor.dtype)

def deg2rad(tensor):
    """Function that converts angles from degrees to radians.

    See :class:`~torchgeometry.DegToRad` for details.

    Args:
        tensor (Tensor): Tensor of arbitrary shape.

    Returns:
        Tensor: Tensor with same shape as input.

    Examples::

        >>> input = 360. * torch.rand(1, 3, 3)
        >>> output = tgm.deg2rad(input)
    """
    if not torch.is_tensor(tensor):
        raise TypeError("Input type is not a torch.Tensor. Got {}"
                        .format(type(tensor)))

    return tensor * pi.to(tensor.device).type(tensor.dtype) / 180.

######### to quaternion
def euler_to_quaternion(r):
    x = r[..., 0]
    y = r[..., 1]
    z = r[..., 2]

    z = z/2.0
    y = y/2.0
    x = x/2.0
    cz = torch.cos(z)
    sz = torch.sin(z)
    cy = torch.cos(y)
    sy = torch.sin(y)
    cx = torch.cos(x)
    sx = torch.sin(x)
    quaternion = torch.zeros_like(r.repeat(1,2))[..., :4].to(r.device)
    quaternion[..., 0] += cx*cy*cz - sx*sy*sz
    quaternion[..., 1] += cx*sy*sz + cy*cz*sx
    quaternion[..., 2] += cx*cz*sy - sx*cy*sz
    quaternion[..., 3] += cx*cy*sz + sx*cz*sy
    return quaternion

def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):
    """Convert 3x4 rotation matrix to 4d quaternion vector

    This algorithm is based on algorithm described in
    https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201

    Args:
        rotation_matrix (Tensor): the rotation matrix to convert.

    Return:
        Tensor: the rotation in quaternion

    Shape:
        - Input: :math:`(N, 3, 4)`
        - Output: :math:`(N, 4)`

    Example:
        >>> input = torch.rand(4, 3, 4)  # Nx3x4
        >>> output = tgm.rotation_matrix_to_quaternion(input)  # Nx4
    """
    if not torch.is_tensor(rotation_matrix):
        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
            type(rotation_matrix)))

    if len(rotation_matrix.shape) > 3:
        raise ValueError(
            "Input size must be a three dimensional tensor. Got {}".format(
                rotation_matrix.shape))
    # if not rotation_matrix.shape[-2:] == (3, 4):
    #     raise ValueError(
    #         "Input size must be a N x 3 x 4  tensor. Got {}".format(
    #             rotation_matrix.shape))

    rmat_t = torch.transpose(rotation_matrix, 1, 2)

    mask_d2 = rmat_t[:, 2, 2] < eps

    mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1]
    mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1]

    t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
    q0 = torch.stack([rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
                      t0, rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
                      rmat_t[:, 2, 0] + rmat_t[:, 0, 2]], -1)
    t0_rep = t0.repeat(4, 1).t()

    t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
    q1 = torch.stack([rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
                      rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
                      t1, rmat_t[:, 1, 2] + rmat_t[:, 2, 1]], -1)
    t1_rep = t1.repeat(4, 1).t()

    t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
    q2 = torch.stack([rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
                      rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
                      rmat_t[:, 1, 2] + rmat_t[:, 2, 1], t2], -1)
    t2_rep = t2.repeat(4, 1).t()

    t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
    q3 = torch.stack([t3, rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
                      rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
                      rmat_t[:, 0, 1] - rmat_t[:, 1, 0]], -1)
    t3_rep = t3.repeat(4, 1).t()

    mask_c0 = mask_d2 * mask_d0_d1.float()
    mask_c1 = mask_d2 * (1 - mask_d0_d1.float())
    mask_c2 = (1 - mask_d2.float()) * mask_d0_nd1
    mask_c3 = (1 - mask_d2.float()) * (1 - mask_d0_nd1.float())
    mask_c0 = mask_c0.view(-1, 1).type_as(q0)
    mask_c1 = mask_c1.view(-1, 1).type_as(q1)
    mask_c2 = mask_c2.view(-1, 1).type_as(q2)
    mask_c3 = mask_c3.view(-1, 1).type_as(q3)

    q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3
    q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 +  # noqa
                    t2_rep * mask_c2 + t3_rep * mask_c3)  # noqa
    q *= 0.5
    return q

def angle_axis_to_quaternion(angle_axis: torch.Tensor) -> torch.Tensor:
    """Convert an angle axis to a quaternion.

    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h

    Args:
        angle_axis (torch.Tensor): tensor with angle axis.

    Return:
        torch.Tensor: tensor with quaternion.

    Shape:
        - Input: :math:`(*, 3)` where `*` means, any number of dimensions
        - Output: :math:`(*, 4)`

    Example:
        >>> angle_axis = torch.rand(2, 4)  # Nx4
        >>> quaternion = tgm.angle_axis_to_quaternion(angle_axis)  # Nx3
    """
    if not torch.is_tensor(angle_axis):
        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
            type(angle_axis)))

    if not angle_axis.shape[-1] == 3:
        raise ValueError("Input must be a tensor of shape Nx3 or 3. Got {}"
                         .format(angle_axis.shape))
    # unpack input and compute conversion
    a0: torch.Tensor = angle_axis[..., 0:1]
    a1: torch.Tensor = angle_axis[..., 1:2]
    a2: torch.Tensor = angle_axis[..., 2:3]
    theta_squared: torch.Tensor = a0 * a0 + a1 * a1 + a2 * a2

    theta: torch.Tensor = torch.sqrt(theta_squared)
    half_theta: torch.Tensor = theta * 0.5

    mask: torch.Tensor = theta_squared > 0.0
    ones: torch.Tensor = torch.ones_like(half_theta)

    k_neg: torch.Tensor = 0.5 * ones
    k_pos: torch.Tensor = torch.sin(half_theta) / theta
    k: torch.Tensor = torch.where(mask, k_pos, k_neg)
    w: torch.Tensor = torch.where(mask, torch.cos(half_theta), ones)

    quaternion: torch.Tensor = torch.zeros_like(angle_axis)
    quaternion[..., 0:1] += a0 * k
    quaternion[..., 1:2] += a1 * k
    quaternion[..., 2:3] += a2 * k
    return torch.cat([w, quaternion], dim=-1)

#### quaternion to
def quaternion_to_rotation_matrix(quat):
    """Convert quaternion coefficients to rotation matrix.
    Args:
        quat: size = [B, 4] 4 <===>(w, x, y, z)
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    norm_quat = quat
    norm_quat = norm_quat / norm_quat.norm(p=2, dim=1, keepdim=True)
    w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:, 2], norm_quat[:, 3]

    B = quat.size(0)

    w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
    wx, wy, wz = w * x, w * y, w * z
    xy, xz, yz = x * y, x * z, y * z

    rotMat = torch.stack([w2 + x2 - y2 - z2, 2 * xy - 2 * wz, 2 * wy + 2 * xz,
                          2 * wz + 2 * xy, w2 - x2 + y2 - z2, 2 * yz - 2 * wx,
                          2 * xz - 2 * wy, 2 * wx + 2 * yz, w2 - x2 - y2 + z2], dim=1).view(B, 3, 3)
    return rotMat

def quaternion_to_angle_axis(quaternion: torch.Tensor):
    """Convert quaternion vector to angle axis of rotation. TODO: CORRECT

    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h

    Args:
        quaternion (torch.Tensor): tensor with quaternions.

    Return:
        torch.Tensor: tensor with angle axis of rotation.

    Shape:
        - Input: :math:`(*, 4)` where `*` means, any number of dimensions
        - Output: :math:`(*, 3)`

    Example:
        >>> quaternion = torch.rand(2, 4)  # Nx4
        >>> angle_axis = tgm.quaternion_to_angle_axis(quaternion)  # Nx3
    """
    if not torch.is_tensor(quaternion):
        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
            type(quaternion)))

    if not quaternion.shape[-1] == 4:
        raise ValueError("Input must be a tensor of shape Nx4 or 4. Got {}"
                         .format(quaternion.shape))
    # unpack input and compute conversion
    q1: torch.Tensor = quaternion[..., 1]
    q2: torch.Tensor = quaternion[..., 2]
    q3: torch.Tensor = quaternion[..., 3]
    sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3

    sin_theta: torch.Tensor = torch.sqrt(sin_squared_theta)
    cos_theta: torch.Tensor = quaternion[..., 0]
    two_theta: torch.Tensor = 2.0 * torch.where(
        cos_theta < 0.0,
        torch.atan2(-sin_theta, -cos_theta),
        torch.atan2(sin_theta, cos_theta))

    k_pos: torch.Tensor = two_theta / sin_theta
    k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta).to(quaternion.device)
    k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg)

    angle_axis: torch.Tensor = torch.zeros_like(quaternion).to(quaternion.device)[..., :3]
    angle_axis[..., 0] += q1 * k
    angle_axis[..., 1] += q2 * k
    angle_axis[..., 2] += q3 * k
    return angle_axis

#### batch converter
def batch_euler2axis(r):
    return quaternion_to_angle_axis(euler_to_quaternion(r))

def batch_euler2matrix(r):
    return quaternion_to_rotation_matrix(euler_to_quaternion(r))

def batch_matrix2euler(rot_mats):
    # Calculates rotation matrix to euler angles
    # Careful for extreme cases of eular angles like [0.0, pi, 0.0]
    ### only y biw
    # TODO: add x, z
    sy = torch.sqrt(rot_mats[:, 0, 0] * rot_mats[:, 0, 0] +
                    rot_mats[:, 1, 0] * rot_mats[:, 1, 0])
    return torch.atan2(-rot_mats[:, 2, 0], sy)

def rotmat2euler(R):
    """
    Converts a rotation matrix to Euler angles
    Matlab port to python for evaluation purposes
    https://github.com/asheshjain399/RNNexp/blob/srnn/structural_rnn/CRFProblems/H3.6m/mhmublv/Motion/RotMat2Euler.m#L1
    Args
    R: a Bx3x3 rotation matrix
    Returns
    eul: a 3x1 Euler angle representation of R
    """
    if R[0,2] == 1 or R[0,2] == -1:
    # special case
        E3   = 0 # set arbitrarily
        dlta = np.arctan2( R[0,1], R[0,2] )
        if R[0,2] == -1:
            E2 = np.pi/2
            E1 = E3 + dlta
        else:
            E2 = -np.pi/2
            E1 = -E3 + dlta
    else:
        E2 = -np.arcsin( R[0,2] )
        E1 = np.arctan2( R[1,2]/np.cos(E2), R[2,2]/np.cos(E2) )
        E3 = np.arctan2( R[0,1]/np.cos(E2), R[0,0]/np.cos(E2) )
    eul = np.array([E1, E2, E3])
    return eul

def batch_matrix2axis(rot_mats):
    return quaternion_to_angle_axis(rotation_matrix_to_quaternion(rot_mats))

def batch_axis2matrix(theta):
    # angle axis to rotation matrix
    # theta N x 3
    # return quat2mat(quat)
    # batch_rodrigues
    return quaternion_to_rotation_matrix(angle_axis_to_quaternion(theta))

def batch_axis2euler(theta):
    return batch_matrix2euler(batch_axis2matrix(theta))

def batch_axis2euler(r):
    return rot_mat_to_euler(batch_rodrigues(r))

def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
    '''  same as batch_matrix2axis
    Calculates the rotation matrices for a batch of rotation vectors
        Parameters
        ----------
        rot_vecs: torch.tensor Nx3
            array of N axis-angle vectors
        Returns
        -------
        R: torch.tensor Nx3x3
            The rotation matrices for the given axis-angle parameters
    Code from smplx/flame, what PS people often use
    '''

    batch_size = rot_vecs.shape[0]
    device = rot_vecs.device

    angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
    rot_dir = rot_vecs / angle

    cos = torch.unsqueeze(torch.cos(angle), dim=1)
    sin = torch.unsqueeze(torch.sin(angle), dim=1)

    # Bx1 arrays
    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
    K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)

    zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))

    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
    return rot_mat

def batch_cont2matrix(module_input):
    ''' Decoder for transforming a latent representation to rotation matrices

        Implements the decoding method described in:
        "On the Continuity of Rotation Representations in Neural Networks"
        Code from https://github.com/vchoutas/expose
    '''
    batch_size = module_input.shape[0]
    reshaped_input = module_input.reshape(-1, 3, 2)

    # Normalize the first vector
    b1 = F.normalize(reshaped_input[:, :, 0].clone(), dim=1)

    dot_prod = torch.sum(
        b1 * reshaped_input[:, :, 1].clone(), dim=1, keepdim=True)
    # Compute the second vector by finding the orthogonal complement to it
    b2 = F.normalize(reshaped_input[:, :, 1] - dot_prod * b1, dim=1)
    # Finish building the basis by taking the cross product
    b3 = torch.cross(b1, b2, dim=1)
    rot_mats = torch.stack([b1, b2, b3], dim=-1)

    return rot_mats.view(batch_size, -1, 3, 3)

utils/rotation_converter.py
sha2nkt/ipman-r
from __future__ import division
import os
import datetime

import torch

class CheckpointSaver():
    """Class that handles saving and loading checkpoints during training."""
    def __init__(self, save_dir, save_steps=1000):
        self.save_dir = os.path.abspath(save_dir)
        self.save_steps = save_steps
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir)
        self.get_latest_checkpoint()
        return

    def exists_checkpoint(self, checkpoint_file=None):
        """Check if a checkpoint exists in the current directory."""
        if checkpoint_file is None:
            return False if self.latest_checkpoint is None else True
        else:
            return os.path.isfile(checkpoint_file)
    
    def save_checkpoint(self, models, optimizers, epoch, batch_idx, batch_size, dataset_perm, total_step_count):
        """Save checkpoint."""
        timestamp = datetime.datetime.now()
        checkpoint_filename = os.path.abspath(os.path.join(self.save_dir, timestamp.strftime('%Y_%m_%d-%H_%M_%S') + f'_epoch{epoch}_stepcount_{total_step_count}.pt'))
        checkpoint_filename = os.path.abspath(os.path.join(self.save_dir, timestamp.strftime('%Y_%m_%d-%H_%M_%S') + f'_epoch{epoch}_stepcount_{total_step_count}.pt'))
        # alt_dir = '/ps/project/common/renderpeople_initialfit/4Priyanka/spin_ft_on_agora/logs/agora_train_full_spinfix/checkpoints'
        # alt_checkpoint_filename = os.path.abspath(os.path.join(alt_dir, timestamp.strftime('%Y_%m_%d-%H_%M_%S') + f'_epoch{epoch}.pt'))
        checkpoint = {}
        for model in models:
            checkpoint[model] = models[model].state_dict()
        for optimizer in optimizers:
            checkpoint[optimizer] = optimizers[optimizer].state_dict()
        checkpoint['epoch'] = epoch
        checkpoint['batch_idx'] = batch_idx
        checkpoint['batch_size'] = batch_size
        checkpoint['dataset_perm'] = dataset_perm
        checkpoint['total_step_count'] = total_step_count
        print(timestamp, 'Epoch:', epoch, 'Iteration:', batch_idx)
        print('Saving checkpoint file [' + checkpoint_filename + ']')
        torch.save(checkpoint, checkpoint_filename)
        # torch.save(checkpoint, alt_checkpoint_filename)
        return

    def load_checkpoint(self, models, optimizers, checkpoint_file=None):
        """Load a checkpoint."""
        if checkpoint_file is None:
            print('Loading latest checkpoint [' + self.latest_checkpoint + ']')
            checkpoint_file = self.latest_checkpoint
        checkpoint = torch.load(checkpoint_file)
        for model in models:
            if model in checkpoint:
                models[model].load_state_dict(checkpoint[model])
        for optimizer in optimizers:
            if optimizer in checkpoint:
                optimizers[optimizer].load_state_dict(checkpoint[optimizer])
        return {'epoch': checkpoint['epoch'],
                'batch_idx': checkpoint['batch_idx'],
                'batch_size': checkpoint['batch_size'],
                'dataset_perm': checkpoint['dataset_perm'],
                'total_step_count': checkpoint['total_step_count']}

    def get_latest_checkpoint(self):
        """Get filename of latest checkpoint if it exists."""
        checkpoint_list = [] 
        for dirpath, dirnames, filenames in os.walk(self.save_dir):
            for filename in filenames:
                if filename.endswith('.pt'):
                    checkpoint_list.append(os.path.abspath(os.path.join(dirpath, filename)))
        checkpoint_list = sorted(checkpoint_list)
        self.latest_checkpoint =  None if (len(checkpoint_list) is 0) else checkpoint_list[-1]
        return

utils/saver.py
sha2nkt/ipman-r
import os
import json
import argparse
import numpy as np
from collections import namedtuple

class TrainOptions():

    def __init__(self):
        self.parser = argparse.ArgumentParser()

        req = self.parser.add_argument_group('Required')
        req.add_argument('--name', required=True, help='Name of the experiment')

        gen = self.parser.add_argument_group('General')
        gen.add_argument('--time_to_run', type=int, default=np.inf, help='Total time to run in seconds. Used for training in environments with timing constraints')
        gen.add_argument('--resume', dest='resume', default=False, action='store_true', help='Resume from checkpoint (Use latest checkpoint by default')
        gen.add_argument('--num_workers', type=int, default=8, help='Number of processes used for data loading')
        pin = gen.add_mutually_exclusive_group()
        pin.add_argument('--pin_memory', dest='pin_memory', action='store_true')
        pin.add_argument('--no_pin_memory', dest='pin_memory', action='store_false')
        gen.set_defaults(pin_memory=True)

        io = self.parser.add_argument_group('io')
        io.add_argument('--log_dir', default='logs', help='Directory to store logs')
        io.add_argument('--checkpoint', default=None, help='Path to checkpoint')
        io.add_argument('--from_json', default=None, help='Load options from json file instead of the command line')
        io.add_argument('--pretrained_checkpoint', default=None, help='Load a pretrained checkpoint at the beginning training') 

        train = self.parser.add_argument_group('Training Options')
        train.add_argument('--num_epochs', type=int, default=200, help='Total number of training epochs')
        train.add_argument("--lr", type=float, default=5e-5, help="Learning rate")
        train.add_argument("--cop_w", type=float, default=10, help="pressure term for outside vertices in center of support computation")
        train.add_argument("--cop_k", type=float, default=100, help="pressure term/spring constant for inside vertices in center of support computation")
        train.add_argument("--in_alpha1", type=float, default=1,
                           help="alpha1 constant for the inside push term")
        train.add_argument("--in_alpha2", type=float, default=0.5,
                           help="alpha2 constant for the inside push term")
        train.add_argument("--out_alpha1", type=float, default=1,
                           help="alpha1 constant for the outside pull term")
        train.add_argument("--out_alpha2", type=float, default=0.15,
                           help="alpha2 constant for the outside pull term")
        train.add_argument("--contact_thresh", type=float, default=0.1,
                           help="threshold for contact")
        train.add_argument('--batch_size', type=int, default=64, help='Batch size')
        train.add_argument('--summary_steps', type=int, default=100, help='Summary saving frequency')
        train.add_argument('--test_steps', type=int, default=1000, help='Testing frequency during training')
        train.add_argument('--run_eval', default=False, action='store_true', help='Runs everytime model checkpoint is saved on disk')
        train.add_argument('--checkpoint_steps', type=int, default=10000, help='Checkpoint saving frequency')
        train.add_argument('--img_res', type=int, default=224, help='Rescale bounding boxes to size [img_res, img_res] before feeding them in the network') 
        train.add_argument('--rot_factor', type=float, default=30, help='Random rotation in the range [-rot_factor, rot_factor]') 
        train.add_argument('--noise_factor', type=float, default=0.4, help='Randomly multiply pixel values with factor in the range [1-noise_factor, 1+noise_factor]') 
        train.add_argument('--scale_factor', type=float, default=0.25, help='Rescale bounding boxes by a factor of [1-scale_factor,1+scale_factor]') 
        train.add_argument('--ignore_3d', default=False, action='store_true', help='Ignore GT 3D data (for unpaired experiments') 
        train.add_argument('--is_agora', default=False, action='store_true', help='Enable AGORA supervised training')
        train.add_argument('--is_others', default=False, action='store_true', help='Enable SMPLIFY training with other SPIN dataset except for H36M')
        train.add_argument('--shape_loss_weight', default=0, type=float, help='Weight of per-vertex loss')
        train.add_argument('--keypoint_loss_weight', default=5., type=float, help='Weight of 2D and 3D keypoint loss') 
        train.add_argument('--pose_loss_weight', default=1., type=float, help='Weight of SMPL pose loss') 
        train.add_argument('--beta_loss_weight', default=0.001, type=float, help='Weight of SMPL betas loss')
        train.add_argument('--stability_loss_weight', default=1., type=float, help='Weight for the stability loss')
        train.add_argument('--inside_push_loss_weight', default=1., type=float, help='Weight for the push term in ground intersection loss')
        train.add_argument('--outside_pull_loss_weight', default=1., type=float, help='Weight for the pull term in ground intersection loss')
        train.add_argument('--openpose_train_weight', default=0., help='Weight for OpenPose keypoints during training')
        train.add_argument('--gt_train_weight', default=1., help='Weight for GT keypoints during training') 
        train.add_argument('--run_smplify', default=False, action='store_true', help='Run SMPLify during training') 
        train.add_argument('--smplify_threshold', type=float, default=100., help='Threshold for ignoring SMPLify fits during training') 
        train.add_argument('--num_smplify_iters', default=100, type=int, help='Number of SMPLify iterations') 

        shuffle_train = train.add_mutually_exclusive_group()
        shuffle_train.add_argument('--shuffle_train', dest='shuffle_train', action='store_true', help='Shuffle training data')
        shuffle_train.add_argument('--no_shuffle_train', dest='shuffle_train', action='store_false', help='Don\'t shuffle training data')
        shuffle_train.set_defaults(shuffle_train=True)
        return 

    def parse_args(self):
        """Parse input arguments."""
        self.args = self.parser.parse_args()
        # If config file is passed, override all arguments with the values from the config file
        if self.args.from_json is not None:
            path_to_json = os.path.abspath(self.args.from_json)
            with open(path_to_json, "r") as f:
                json_args = json.load(f)
                json_args = namedtuple("json_args", json_args.keys())(**json_args)
                return json_args
        else:
            self.args.log_dir = os.path.join(os.path.abspath(self.args.log_dir), self.args.name)
            self.args.summary_dir = os.path.join(self.args.log_dir, 'tensorboard')
            if not os.path.exists(self.args.log_dir):
                os.makedirs(self.args.log_dir)
            self.args.checkpoint_dir = os.path.join(self.args.log_dir, 'checkpoints')
            if not os.path.exists(self.args.checkpoint_dir):
                os.makedirs(self.args.checkpoint_dir)
            self.save_dump()
            return self.args

    def save_dump(self):
        """Store all argument values to a json file.
        The default location is logs/expname/config.json.
        """
        if not os.path.exists(self.args.log_dir):
            os.makedirs(self.args.log_dir)
        with open(os.path.join(self.args.log_dir, "config.json"), "w") as f:
            json.dump(vars(self.args), f, indent=4)
        return

utils/train_options.py
sha2nkt/ipman-r
ESSENTIALS_DIR = 'data/essentials'
SMPLX_VERT_SEG = 'data/meshscripts/smplx_vert_segmentation.json'
SMPL_VERT_SEG = 'data/meshscripts/smpl_vert_segmentation.json'
SEGMENTS_DIR = 'data/essentials/yogi_segments'

frame_select_dict_combined = {
    'Akarna_Dhanurasana-a': 1226,
    'Akarna_Dhanurasana-b': 1126,
    'Akarna_Dhanurasana-c': 730,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-a': 872,
    'Boat_Pose_or_Paripurna_Navasana_-a': 486,
    'Boat_Pose_or_Paripurna_Navasana_-b': 392,
    'Boat_Pose_or_Paripurna_Navasana_-c': 340,
    'Boat_Pose_or_Paripurna_Navasana_-d': 300,
    'Boat_Pose_or_Paripurna_Navasana_-e': 314,
    'Boat_Pose_or_Paripurna_Navasana_-f': 280,
    'Bound_Angle_Pose_or_Baddha_Konasana_-a': 622,
    'Bound_Angle_Pose_or_Baddha_Konasana_-b': 510,
    'Bound_Angle_Pose_or_Baddha_Konasana_-c': 462,
    'Bound_Angle_Pose_or_Baddha_Konasana_-f': 390,
    'Bound_Angle_Pose_or_Baddha_Konasana_-g': 274,
    'Camel_Pose_or_Ustrasana_-c': 430,
    'Camel_Pose_or_Ustrasana_-d': 388,
    'Cat_Cow_Pose_or_Marjaryasana_-a': 326,
    'Cat_Cow_Pose_or_Marjaryasana_-b': 234,
    'Chair_Pose_or_Utkatasana_-a': 266,
    'Chair_Pose_or_Utkatasana_-b': 272,
    'Chair_Pose_or_Utkatasana_-c': 368,
    'Child_Pose_or_Balasana_-a': 468,
    'Cobra_Pose_or_Bhujangasana_-a': 594,
    'Cobra_Pose_or_Bhujangasana_-b': 470,
    'Cobra_Pose_or_Bhujangasana_-c': 444,
    'Cobra_Pose_or_Bhujangasana_-d': 476,
    'Corpse_Pose_or_Savasana_-a': 376,
    'Cow_Face_Pose_or_Gomukhasana_-a': 990,
    'Cow_Face_Pose_or_Gomukhasana_-b': 650,
    'Cow_Face_Pose_or_Gomukhasana_-c': 628,
    'Cow_Face_Pose_or_Gomukhasana_-e': 578,
    'Crane_(Crow)_Pose_or_Bakasana_-a': 848,
    'Crane_(Crow)_Pose_or_Bakasana_-b': 574,
    'Dolphin_Plank_Pose_or_Makara_Adho_Mukha_Svanasana_-a': 468,
    'Dolphin_Pose_or_Ardha_Pincha_Mayurasana_-a': 416,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-a': 432,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-b': 380,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-c': 340,
    'Eagle_Pose_or_Garudasana_-a': 580,
    'Eight-Angle_Pose_or_Astavakrasana_-a': 686,
    'Extended_Puppy_Pose_or_Uttana_Shishosana_-a': 464,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-a': 590,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-b': 538,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-d': 634,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-e': 510,
    'Extended_Revolved_Triangle_Pose_or_Utthita_Trikonasana_-a': 496,
    'Extended_Revolved_Triangle_Pose_or_Utthita_Trikonasana_-b': 572,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-a': 1216,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-b': 752,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-c': 1546,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-d': 766,
    'Firefly_Pose_or_Tittibhasana_-a': 644,
    'Firefly_Pose_or_Tittibhasana_-b': 496,
    'Fish_Pose_or_Matsyasana_-a': 1286,
    'Fish_Pose_or_Matsyasana_-c': 1160,
    'Fish_Pose_or_Matsyasana_-d': 808,
    'Four-Limbed_Staff_Pose_or_Chaturanga_Dandasana_-a': 388,
    'Four-Limbed_Staff_Pose_or_Chaturanga_Dandasana_-b': 410,
    'Frog_Pose_or_Bhekasana-a': 886,
    'Garland_Pose_or_Malasana_-a': 334,
    'Gate_Pose_or_Parighasana_-a': 540,
    'Half_Lord_of_the_Fishes_Pose_or_Ardha_Matsyendrasana_-a': 568,
    'Half_Moon_Pose_or_Ardha_Chandrasana_-a': 750,
    'Half_Moon_Pose_or_Ardha_Chandrasana_-b': 688,
    'Happy_Baby_Pose-a': 540,
    'Head-to-Knee_Forward_Bend_pose_or_Janu_Sirsasana_-a': 680,
    'Heron_Pose_or_Krounchasana_-a': 658,
    'Intense_Side_Stretch_Pose_or_Parsvottanasana_-a': 592,
    'Legs-Up-the-Wall_Pose_or_Viparita_Karani_-a': 456,
    'Locust_Pose_or_Salabhasana_-a': 458,
    'Locust_Pose_or_Salabhasana_-b': 446,
    'Low_Lunge_pose_or_Anjaneyasana_-a': 848,
    'Low_Lunge_pose_or_Anjaneyasana_-b': 526,
    'Low_Lunge_pose_or_Anjaneyasana_-c': 600,
    'Low_Lunge_pose_or_Anjaneyasana_-d': 684,
    'Noose_Pose_or_Pasasana_-b': 544,
    'Peacock_Pose_or_Mayurasana_-a': 732,
    'Peacock_Pose_or_Mayurasana_-b': 1296,
    'Plank_Pose_or_Kumbhakasana_-a': 398,
    'Plow_Pose_or_Halasana_-a': 742,
    'Plow_Pose_or_Halasana_-b': 720,
    'Pose_Dedicated_to_the_Sage_Koundinya_or_Eka_Pada_Koundinyanasana_I_and_II-a': 484,
    'Pose_Dedicated_to_the_Sage_Koundinya_or_Eka_Pada_Koundinyanasana_I_and_II-b': 514,
    'Revolved_Head-to-Knee_Pose_or_Parivrtta_Janu_Sirsasana_-b': 666,
    'Scale_Pose_or_Tolasana_-a': 914,
    'Scorpion_pose_or_vrischikasana-a': 1400,
    'Scorpion_pose_or_vrischikasana-c': 1082,
    'Shoulder-Pressing_Pose_or_Bhujapidasana_-a': 580,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-a': 746,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-b': 738,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-c': 1162,
    'Side_Plank_Pose_or_Vasisthasana_-a': 500,
    'Side_Plank_Pose_or_Vasisthasana_-b': 602,
    'Side_Plank_Pose_or_Vasisthasana_-c': 888,
    'Side_Plank_Pose_or_Vasisthasana_-d': 500,
    'Side_Plank_Pose_or_Vasisthasana_-e': 554,
    'Sitting_pose_1_(normal)-a': 630,
    'Sitting_pose_1_(normal)-b': 358,
    'Sitting_pose_1_(normal)-c': 440,
    'Split_pose-a': 938,
    'Split_pose-b': 768,
    'Split_pose-c': 462,
    'Staff_Pose_or_Dandasana_-a': 514,
    'Staff_Pose_or_Dandasana_-b': 458,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-a': 456,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-b': 506,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-c': 396,
    'Standing_Forward_Bend_pose_or_Uttanasana_-a': 442,
    'Standing_Forward_Bend_pose_or_Uttanasana_-b': 382,
    'Standing_Split_pose_or_Urdhva_Prasarita_Eka_Padasana_-a': 380,
    'Standing_Split_pose_or_Urdhva_Prasarita_Eka_Padasana_-b': 370,
    'Supported_Headstand_pose_or_Salamba_Sirsasana_-b': 826,
    'Supported_Shoulderstand_pose_or_Salamba_Sarvangasana_-a': 640,
    'Supta_Baddha_Konasana_-b': 754,
    'Supta_Virasana_Vajrasana-c': 966,
    'Tree_Pose_or_Vrksasana_-a': 446,
    'Tree_Pose_or_Vrksasana_-b': 360,
    'Upward_Bow_(Wheel)_Pose_or_Urdhva_Dhanurasana_-a': 612,
    'Upward_Bow_(Wheel)_Pose_or_Urdhva_Dhanurasana_-b': 788,
    'Upward_Plank_Pose_or_Purvottanasana_-a': 592,
    'Upward_Plank_Pose_or_Purvottanasana_-b': 550,
    'viparita_virabhadrasana_or_reverse_warrior_pose-a': 538,
    'viparita_virabhadrasana_or_reverse_warrior_pose-b': 436,
    'Virasana_or_Vajrasana-a': 438,
    'Virasana_or_Vajrasana-c': 428,
    'Warrior_I_Pose_or_Virabhadrasana_I_-a': 518,
    'Warrior_II_Pose_or_Virabhadrasana_II_-a': 446,
    'Warrior_III_Pose_or_Virabhadrasana_III_-a': 630,
    'Warrior_III_Pose_or_Virabhadrasana_III_-b': 400,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-a': 1172,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-b': 682,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-c': 562,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-d': 610,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-a': 872,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-c': 826,
    'Wild_Thing_pose_or_Camatkarasana_-a': 692,
    'Handstand_pose_or_Adho_Mukha_Vrksasana_-a': 622,
    'Scorpion_pose_or_vrischikasana-b': 1946,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-b': 720,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-c': 540,
    'Bound_Angle_Pose_or_Baddha_Konasana_-d': 1090,
    'Bound_Angle_Pose_or_Baddha_Konasana_-e': 788,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-a': 420,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-b': 600,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-c': 540,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-d': 550,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-e': 582,
    'Camel_Pose_or_Ustrasana_-a': 684,
    'Camel_Pose_or_Ustrasana_-b': 360,
    'Child_Pose_or_Balasana_-b': 360,
    'Cow_Face_Pose_or_Gomukhasana_-d': 780,
    'Fish_Pose_or_Matsyasana_-b': 850,
    'Frog_Pose_or_Bhekasana-b': 580,
    'Frog_Pose_or_Bhekasana-c': 600,
    'Garland_Pose_or_Malasana_-b': 780,
    'Garland_Pose_or_Malasana_-c': 360,
    'Gate_Pose_or_Parighasana_-b': 480,
    'Heron_Pose_or_Krounchasana_-b': 1020,
    'Heron_Pose_or_Krounchasana_-c': 600,
    'Intense_Side_Stretch_Pose_or_Parsvottanasana_-b': 1080,
    'Locust_Pose_or_Salabhasana_-c': 720,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-a': 480,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-c': 1620,
    'Noose_Pose_or_Pasasana_-a': 780,
    'Rajakapotasana-a': 1536,
    'Rajakapotasana-b': 1500,
    'Rajakapotasana-c': 1260,
    'Rajakapotasana-d': 1500,
    'Reclining_Hand-to-Big-Toe_Pose_or_Supta_Padangusthasana_-a': 420,
    'Reclining_Hand-to-Big-Toe_Pose_or_Supta_Padangusthasana_-b': 600,
    'Revolved_Head-to-Knee_Pose_or_Parivrtta_Janu_Sirsasana_-a': 900,
    'Bow_Pose_or_Dhanurasana_-a': 660,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-c': 360,
    'Head-to-Knee_Forward_Bend_pose_or_Janu_Sirsasana_-b': 660,
    'Seated_Forward_Bend_pose_or_Paschimottanasana_-a': 600,
    'Side-Reclining_Leg_Lift_pose_or_Anantasana_-a': 780,
    'Supported_Headstand_pose_or_Salamba_Sirsasana_-a': 754,
    'Supta_Baddha_Konasana_-a': 620,
    'Supta_Virasana_Vajrasana-a': 800,
    'Supta_Virasana_Vajrasana-b': 1100,
    'Supta_Virasana_Vajrasana-d': 780,
    'Tortoise_Pose-a': 960,
    'Upward_Facing_Two-Foot_Staff_Pose_or_Dwi_Pada_Viparita_Dandasana_-a': 1080,
    'Virasana_or_Vajrasana-b': 480,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-b': 840,
    'Wind_Relieving_pose_or_Pawanmuktasana-a': 320,
    'Wind_Relieving_pose_or_Pawanmuktasana-b': 360,
    'Yogic_sleep_pose-a': 1624,
    'Yogic_sleep_pose-b': 1080
}

frame_select_dict_first_session = {
    'Akarna_Dhanurasana-a': 1226,
    'Akarna_Dhanurasana-b': 1126,
    'Akarna_Dhanurasana-c': 730,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-a': 872,
    'Boat_Pose_or_Paripurna_Navasana_-a': 486,
    'Boat_Pose_or_Paripurna_Navasana_-b': 392,
    'Boat_Pose_or_Paripurna_Navasana_-c': 340,
    'Boat_Pose_or_Paripurna_Navasana_-d': 300,
    'Boat_Pose_or_Paripurna_Navasana_-e': 314,
    'Boat_Pose_or_Paripurna_Navasana_-f': 280,
    'Bound_Angle_Pose_or_Baddha_Konasana_-a': 622,
    'Bound_Angle_Pose_or_Baddha_Konasana_-b': 510,
    'Bound_Angle_Pose_or_Baddha_Konasana_-c': 462,
    'Bound_Angle_Pose_or_Baddha_Konasana_-f': 390,
    'Bound_Angle_Pose_or_Baddha_Konasana_-g': 274,
    'Camel_Pose_or_Ustrasana_-c': 430,
    'Camel_Pose_or_Ustrasana_-d': 388,
    'Cat_Cow_Pose_or_Marjaryasana_-a': 326,
    'Cat_Cow_Pose_or_Marjaryasana_-b': 234,
    'Chair_Pose_or_Utkatasana_-a': 266,
    'Chair_Pose_or_Utkatasana_-b': 272,
    'Chair_Pose_or_Utkatasana_-c': 368,
    'Child_Pose_or_Balasana_-a': 468,
    'Cobra_Pose_or_Bhujangasana_-a': 594,
    'Cobra_Pose_or_Bhujangasana_-b': 470,
    'Cobra_Pose_or_Bhujangasana_-c': 444,
    'Cobra_Pose_or_Bhujangasana_-d': 476,
    'Corpse_Pose_or_Savasana_-a': 376,
    'Cow_Face_Pose_or_Gomukhasana_-a': 990,
    'Cow_Face_Pose_or_Gomukhasana_-b': 650,
    'Cow_Face_Pose_or_Gomukhasana_-c': 628,
    'Cow_Face_Pose_or_Gomukhasana_-e': 578,
    'Crane_(Crow)_Pose_or_Bakasana_-a': 848,
    'Crane_(Crow)_Pose_or_Bakasana_-b': 574,
    'Dolphin_Plank_Pose_or_Makara_Adho_Mukha_Svanasana_-a': 468,
    'Dolphin_Pose_or_Ardha_Pincha_Mayurasana_-a': 416,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-a': 432,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-b': 380,
    'Downward-Facing_Dog_pose_or_Adho_Mukha_Svanasana_-c': 340,
    'Eagle_Pose_or_Garudasana_-a': 580,
    'Eight-Angle_Pose_or_Astavakrasana_-a': 686,
    'Extended_Puppy_Pose_or_Uttana_Shishosana_-a': 464,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-a': 590,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-b': 538,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-d': 634,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-e': 510,
    'Extended_Revolved_Triangle_Pose_or_Utthita_Trikonasana_-a': 496,
    'Extended_Revolved_Triangle_Pose_or_Utthita_Trikonasana_-b': 572,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-a': 1216,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-b': 752,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-c': 1546,
    'Feathered_Peacock_Pose_or_Pincha_Mayurasana_-d': 766,
    'Firefly_Pose_or_Tittibhasana_-a': 644,
    'Firefly_Pose_or_Tittibhasana_-b': 496,
    'Fish_Pose_or_Matsyasana_-a': 1286,
    'Fish_Pose_or_Matsyasana_-c': 1160,
    'Fish_Pose_or_Matsyasana_-d': 808,
    'Four-Limbed_Staff_Pose_or_Chaturanga_Dandasana_-a': 388,
    'Four-Limbed_Staff_Pose_or_Chaturanga_Dandasana_-b': 410,
    'Frog_Pose_or_Bhekasana-a': 886,
    'Garland_Pose_or_Malasana_-a': 334,
    'Gate_Pose_or_Parighasana_-a': 540,
    'Half_Lord_of_the_Fishes_Pose_or_Ardha_Matsyendrasana_-a': 568,
    'Half_Moon_Pose_or_Ardha_Chandrasana_-a': 750,
    'Half_Moon_Pose_or_Ardha_Chandrasana_-b': 688,
    'Happy_Baby_Pose-a': 540,
    'Head-to-Knee_Forward_Bend_pose_or_Janu_Sirsasana_-a': 680,
    'Heron_Pose_or_Krounchasana_-a': 658,
    'Intense_Side_Stretch_Pose_or_Parsvottanasana_-a': 592,
    'Legs-Up-the-Wall_Pose_or_Viparita_Karani_-a': 456,
    'Locust_Pose_or_Salabhasana_-a': 458,
    'Locust_Pose_or_Salabhasana_-b': 446,
    'Low_Lunge_pose_or_Anjaneyasana_-a': 848,
    'Low_Lunge_pose_or_Anjaneyasana_-b': 526,
    'Low_Lunge_pose_or_Anjaneyasana_-c': 600,
    'Low_Lunge_pose_or_Anjaneyasana_-d': 684,
    'Noose_Pose_or_Pasasana_-b': 544,
    'Peacock_Pose_or_Mayurasana_-a': 732,
    'Peacock_Pose_or_Mayurasana_-b': 1296,
    'Plank_Pose_or_Kumbhakasana_-a': 398,
    'Plow_Pose_or_Halasana_-a': 742,
    'Pose_Dedicated_to_the_Sage_Koundinya_or_Eka_Pada_Koundinyanasana_I_and_II-a': 484,
    'Pose_Dedicated_to_the_Sage_Koundinya_or_Eka_Pada_Koundinyanasana_I_and_II-b': 514,
    'Revolved_Head-to-Knee_Pose_or_Parivrtta_Janu_Sirsasana_-b': 666,
    'Scale_Pose_or_Tolasana_-a': 914,
    'Scorpion_pose_or_vrischikasana-a': 1400,
    'Scorpion_pose_or_vrischikasana-c': 1082,
    'Shoulder-Pressing_Pose_or_Bhujapidasana_-a': 580,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-a': 746,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-b': 738,
    'Side_Crane_(Crow)_Pose_or_Parsva_Bakasana_-c': 1162,
    'Side_Plank_Pose_or_Vasisthasana_-a': 500,
    'Side_Plank_Pose_or_Vasisthasana_-b': 602,
    'Side_Plank_Pose_or_Vasisthasana_-c': 888,
    'Side_Plank_Pose_or_Vasisthasana_-d': 500,
    'Side_Plank_Pose_or_Vasisthasana_-e': 554,
    'Sitting_pose_1_(normal)-a': 630,
    'Sitting_pose_1_(normal)-b': 358,
    'Sitting_pose_1_(normal)-c': 440,
    'Split_pose-a': 938,
    'Split_pose-b': 768,
    'Split_pose-c': 462,
    'Staff_Pose_or_Dandasana_-a': 514,
    'Staff_Pose_or_Dandasana_-b': 458,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-a': 456,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-b': 506,
    'Standing_big_toe_hold_pose_or_Utthita_Padangusthasana-c': 396,
    'Standing_Forward_Bend_pose_or_Uttanasana_-a': 442,
    'Standing_Forward_Bend_pose_or_Uttanasana_-b': 382,
    'Standing_Split_pose_or_Urdhva_Prasarita_Eka_Padasana_-a': 380,
    'Standing_Split_pose_or_Urdhva_Prasarita_Eka_Padasana_-b': 370,
    'Supported_Headstand_pose_or_Salamba_Sirsasana_-b': 826,
    'Supported_Shoulderstand_pose_or_Salamba_Sarvangasana_-a': 640,
    'Supta_Baddha_Konasana_-b': 754,
    'Supta_Virasana_Vajrasana-c': 966,
    'Tree_Pose_or_Vrksasana_-a': 446,
    'Tree_Pose_or_Vrksasana_-b': 360,
    'Upward_Bow_(Wheel)_Pose_or_Urdhva_Dhanurasana_-a': 612,
    'Upward_Bow_(Wheel)_Pose_or_Urdhva_Dhanurasana_-b': 788,
    'Upward_Plank_Pose_or_Purvottanasana_-a': 592,
    'Upward_Plank_Pose_or_Purvottanasana_-b': 550,
    'viparita_virabhadrasana_or_reverse_warrior_pose-a': 538,
    'viparita_virabhadrasana_or_reverse_warrior_pose-b': 436,
    'Virasana_or_Vajrasana-a': 438,
    'Virasana_or_Vajrasana-c': 428,
    'Warrior_I_Pose_or_Virabhadrasana_I_-a': 518,
    'Warrior_II_Pose_or_Virabhadrasana_II_-a': 446,
    'Warrior_III_Pose_or_Virabhadrasana_III_-a': 630,
    'Warrior_III_Pose_or_Virabhadrasana_III_-b': 400,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-a': 1172,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-b': 682,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-c': 562,
    'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-d': 610,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-a': 872,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-c': 826,
    'Wild_Thing_pose_or_Camatkarasana_-a': 692
}

frame_select_dict_second_session = {
    'Handstand_pose_or_Adho_Mukha_Vrksasana_-a': 622,
    'Scorpion_pose_or_vrischikasana-b': 1946,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-b': 720,
    'Bharadvajas_Twist_pose_or_Bharadvajasana_I_-c': 540,
    'Bound_Angle_Pose_or_Baddha_Konasana_-d': 1090,
    'Bound_Angle_Pose_or_Baddha_Konasana_-e': 788,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-a': 420,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-b': 600,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-c': 540,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-d': 550,
    'Bridge_Pose_or_Setu_Bandha_Sarvangasana_-e': 582,
    'Camel_Pose_or_Ustrasana_-a': 684,
    'Camel_Pose_or_Ustrasana_-b': 360,
    'Child_Pose_or_Balasana_-b': 360,
    'Cow_Face_Pose_or_Gomukhasana_-d': 780,
    'Fish_Pose_or_Matsyasana_-b': 850,
    'Frog_Pose_or_Bhekasana-b': 580,
    'Frog_Pose_or_Bhekasana-c': 600,
    'Garland_Pose_or_Malasana_-b': 780,
    'Garland_Pose_or_Malasana_-c': 360,
    'Gate_Pose_or_Parighasana_-b': 480,
    'Heron_Pose_or_Krounchasana_-b': 1020,
    'Heron_Pose_or_Krounchasana_-c': 600,
    'Intense_Side_Stretch_Pose_or_Parsvottanasana_-b': 1080,
    'Locust_Pose_or_Salabhasana_-c': 720,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-a': 480,
    'Lord_of_the_Dance_Pose_or_Natarajasana_-c': 1620,
    'Noose_Pose_or_Pasasana_-a': 780,
    'Plow_Pose_or_Halasana_-b': 720,
    'Rajakapotasana-a': 1536,
    'Rajakapotasana-b': 1500,
    'Rajakapotasana-c': 1260,
    'Rajakapotasana-d': 1500,
    'Reclining_Hand-to-Big-Toe_Pose_or_Supta_Padangusthasana_-a': 420,
    'Reclining_Hand-to-Big-Toe_Pose_or_Supta_Padangusthasana_-b': 600,
    'Revolved_Head-to-Knee_Pose_or_Parivrtta_Janu_Sirsasana_-a': 900,
    'Bow_Pose_or_Dhanurasana_-a': 660,
    'Extended_Revolved_Side_Angle_Pose_or_Utthita_Parsvakonasana_-c': 360,
    'Head-to-Knee_Forward_Bend_pose_or_Janu_Sirsasana_-b': 660,
    'Seated_Forward_Bend_pose_or_Paschimottanasana_-a': 600,
    'Side-Reclining_Leg_Lift_pose_or_Anantasana_-a': 780,
    'Supported_Headstand_pose_or_Salamba_Sirsasana_-a': 754,
    'Supta_Baddha_Konasana_-a': 620,
    'Supta_Virasana_Vajrasana-a': 800,
    'Supta_Virasana_Vajrasana-b': 1100,
    'Supta_Virasana_Vajrasana-d': 780,
    'Tortoise_Pose-a': 960,
    'Upward_Facing_Two-Foot_Staff_Pose_or_Dwi_Pada_Viparita_Dandasana_-a': 1080,
    'Virasana_or_Vajrasana-b': 480,
    'Wide-Legged_Forward_Bend_pose_or_Prasarita_Padottanasana_-b': 840,
    'Wind_Relieving_pose_or_Pawanmuktasana-a': 320,
    'Wind_Relieving_pose_or_Pawanmuktasana-b': 360,
    'Yogic_sleep_pose-a': 1624,
    'Yogic_sleep_pose-b': 1080
}

# list of mismatch names: vicon_name:ioi_name
name_changes_first_session = {
    'Akarna_Dhanurasana-a': 'Akarna_Dhanurasana-b_1',
    'Akarna_Dhanurasana-b': 'Akarna_Dhanurasana-b_2',
    'Side-Reclining_Leg_Lift_pose_or_Anantasana_-a': 'Wide-Angle_Seated_Forward_Bend_pose_or_Upavistha_Konasana_-a',
    'Rajakapotasana-c': 'check_name',
    'Side-Reclining_Leg_Lift_pose_or_Anantasana_-a': 'CoreView2',
    'Pose_Dedicated_to_the_Sage_Koundinya_or_Eka_Pada_Koundinyanasana_I_and_II-a': 'Pose_Dedicated_to_the_Sage_Ko_or_Eka_Pada_Kound_I_and_II-a'

}

# list of rejected names with frame_nums useful in matching IOI with Vicon
reject_names_first_session = []

vis_utils/constants.py
sha2nkt/ipman-r
# borrowed from https://github.com/CZ-Wu/GPNet/blob/d8df3b1489800626d4f503f62d2ae8daffe8b686/tools/visualization.py

import trimesh
import os.path as osp
import numpy as np
import torch
import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), '../../..'))
from vis_utils.constants import ESSENTIALS_DIR
import torch.nn as nn

def batch_face_normals(triangles):
    # Calculate the edges of the triangles
    # Size: BxFx3
    edge0 = triangles[:, :, 1] - triangles[:, :, 0]
    edge1 = triangles[:, :, 2] - triangles[:, :, 0]
    # Compute the cross product of the edges to find the normal vector of
    # the triangle
    aCrossb = torch.cross(edge0, edge1, dim=2)
    # Normalize the result to get a unit vector
    normals = aCrossb / torch.norm(aCrossb, 2, dim=2, keepdim=True)

    return normals

def sparse_batch_mm(m1, m2):
    """
    https://github.com/pytorch/pytorch/issues/14489

    m1: sparse matrix of size N x M
    m2: dense matrix of size B x M x K
    returns m1@m2 matrix of size B x N x K
    """

    batch_size = m2.shape[0]
    # stack m2 into columns: (B x N x K) -> (N, B, K) -> (N, B * K)
    m2_stack = m2.transpose(0, 1).reshape(m1.shape[1], -1)
    result = m1.mm(m2_stack).reshape(m1.shape[0], batch_size, -1) \
               .transpose(1, 0)
    return result

class HDfier():
    def __init__(self, model_type='smpl', device='cuda'):
        hd_operator_path = osp.join(ESSENTIALS_DIR, 'hd_model', model_type,
                                    f'{model_type}_neutral_hd_vert_regressor_sparse.npz')
        hd_operator = np.load(hd_operator_path)
        self.hd_operator = torch.sparse.FloatTensor(
            torch.tensor(hd_operator['index_row_col']),
            torch.tensor(hd_operator['values']),
            torch.Size(hd_operator['size']))
        self.model_type = model_type
        self.device = device

    def hdfy_mesh(self, vertices):
        """
        Applies a regressor that maps SMPL vertices to uniformly distributed vertices
        """
        # device = body.vertices.device
        # check if vertices ndim are 3, if not , add a new axis
        if vertices.dim() != 3:
            # batchify the vertices
            vertices = vertices[None, :, :]

        # check if vertices are an ndarry, if yes, make pytorch tensor
        if isinstance(vertices, np.ndarray):
            vertices = torch.from_numpy(vertices).to(self.device)

        vertices = vertices.to(torch.double)

        if self.hd_operator.device != vertices.device:
            self.hd_operator = self.hd_operator.to(vertices.device)
        hd_verts = sparse_batch_mm(self.hd_operator, vertices).to(torch.float)
        return hd_verts

def get_submesh(verts, faces, verts_retained=None, faces_retained=None, min_vert_in_face=2):
    '''
        Given a mesh, create a (smaller) submesh
        indicate faces or verts to retain as indices or boolean

        @return new_verts: the new array of 3D vertices
                new_faces: the new array of faces
                bool_faces: the faces indices wrt the input mesh
                vetex_ids: the vertex_ids wrt the input mesh
        '''

    if verts_retained is not None:
        # Transform indices into bool array
        if verts_retained.dtype != 'bool':
            vert_mask = np.zeros(len(verts), dtype=bool)
            vert_mask[verts_retained] = True
        else:
            vert_mask = verts_retained

        # Faces with at least min_vert_in_face vertices
        bool_faces = np.sum(vert_mask[faces.ravel()].reshape(-1, 3), axis=1) > min_vert_in_face

    elif faces_retained is not None:
        # Transform indices into bool array
        if faces_retained.dtype != 'bool':
            bool_faces = np.zeros(len(faces_retained), dtype=bool)
        else:
            bool_faces = faces_retained

    new_faces = faces[bool_faces]
    # just in case additional vertices are added
    vertex_ids = list(set(new_faces.ravel()))

    oldtonew = -1 * np.ones([len(verts)])
    oldtonew[vertex_ids] = range(0, len(vertex_ids))

    new_verts = verts[vertex_ids]
    new_faces = oldtonew[new_faces].astype('int32')

    return (new_verts, new_faces, bool_faces, vertex_ids)

def get_world_mesh_list(planeWidth=4, axisHeight=0.7, axisRadius=0.02, add_plane=True):
    groundColor = [220, 220, 220, 255]    # face_colors: [R, G, B, transparency]
    xColor = [255, 0, 0, 128]
    yColor = [0, 255, 0, 128]
    zColor = [0, 0, 255, 128]

    if add_plane:
        ground = trimesh.primitives.Box(
            center=[0, 0, -0.0001],
            extents=[planeWidth, planeWidth, 0.0002]
        )
        ground.visual.face_colors = groundColor

    xAxis = trimesh.primitives.Cylinder(
        radius=axisRadius,
        height=axisHeight,
    )
    xAxis.apply_transform(matrix=np.mat(
        ((  0,  0, 1, axisHeight / 2),
        (   0,  1, 0, 0),
        (   -1, 0, 0, 0),
        (   0,  0, 0, 1))
    ))
    xAxis.visual.face_colors = xColor
    yAxis = trimesh.primitives.Cylinder(
        radius=axisRadius,
        height=axisHeight,
    )
    yAxis.apply_transform(matrix=np.mat(
        ((  1, 0, 0, 0),
        (   0, 0, -1, axisHeight / 2),
        (   0, 1, 0, 0),
        (   0, 0, 0, 1))
    ))
    yAxis.visual.face_colors = yColor
    zAxis = trimesh.primitives.Cylinder(
        radius=axisRadius,
        height=axisHeight,
    )
    zAxis.apply_transform(matrix=np.mat(
        ((1, 0, 0, 0),
         (0, 1, 0, 0),
         (0, 0, 1, axisHeight / 2),
         (0, 0, 0, 1))
    ))
    zAxis.visual.face_colors = zColor
    xBox = trimesh.primitives.Box(
        extents=[axisRadius * 3, axisRadius * 3, axisRadius * 3]
    )
    xBox.apply_translation((axisHeight, 0, 0))
    xBox.visual.face_colors = xColor
    yBox = trimesh.primitives.Box(
        extents=[axisRadius * 3, axisRadius * 3, axisRadius * 3]
    )
    yBox.apply_translation((0, axisHeight, 0))
    yBox.visual.face_colors = yColor
    zBox = trimesh.primitives.Box(
        extents=[axisRadius * 3, axisRadius * 3, axisRadius * 3]
    )
    zBox.apply_translation((0, 0, axisHeight))
    zBox.visual.face_colors = zColor
    if add_plane:
        worldMeshList = [ground, xAxis, yAxis, zAxis, xBox, yBox, zBox]
    else:
        worldMeshList = [xAxis, yAxis, zAxis, xBox, yBox, zBox]
    return worldMeshList


def get_checkerboard_plane(plane_width=5, num_boxes=9, center=True):

    pw = plane_width/num_boxes
    white = [220, 220, 220, 255]
    black = [35, 35, 35, 255]

    meshes = []
    for i in range(num_boxes):
        for j in range(num_boxes):
            c = i * pw, j * pw
            ground = trimesh.primitives.Box(
                center=[0, 0, -0.0001],
                extents=[pw, pw, 0.0002]
            )
            if center:
                c = c[0]+(pw/2)-(plane_width/2), c[1]+(pw/2)-(plane_width/2)
            # trans = trimesh.transformations.scale_and_translate(scale=1, translate=[c[0], c[1], 0])
            ground.apply_translation([c[0], c[1], 0])
            # ground.apply_transform(trimesh.transformations.rotation_matrix(np.rad2deg(-120), direction=[1,0,0]))
            ground.visual.face_colors = black if ((i+j) % 2) == 0 else white
            meshes.append(ground)

    # orient mesh be ground in pyrender axis system
    pose = trimesh.transformations.rotation_matrix(np.radians(90), [1, 0, 0])
    meshes = [mesh.apply_transform(pose) for mesh in meshes]
    return meshes


def get_meshed_plane(num_boxes=9):
    def makeVertexGrid(n, m):
        #vertices are incremented in X and Y but remain at zero in Z
        x = np.repeat(np.arange(-n/2, n/2), m)
        y = np.tile(np.arange(-m/2, m/2), n)
        z = np.zeros(n * m)
        grid = np.column_stack((x, y, z))
        return grid

    def makeFaceGrid(n, m):
        # Each face references three vertices by their index numbers, in counterclockwise order. for every vertex i not
        # in the right column or the top row, triangles [i, i+m,i+m+1] and [i, i+m+1, i+1] are created, where m is the
        # number of rows in the array.
        grid = np.zeros((2 * (m - 1) * (n - 1), 3))
        counter = 0
        for i in range((n - 1) * m):
            if i % m == m - 1:
                pass
            else:
                grid[counter] = [i, i + m, i + m + 1]
                grid[counter + 1] = [i, i + m + 1, i + 1]
                counter += 2
        return grid

    vertices = makeVertexGrid(num_boxes, num_boxes)
    faces = makeFaceGrid(num_boxes, num_boxes)

    # Add faces at the top of the mesh to avoid backface culling during visualization
    faces = np.vstack((faces, np.fliplr(faces)))

    white = [220, 220, 220, 255]
    black = [35, 35, 35, 255]
    face_colors = [black if i % 2 == 0 else white for i, _ in enumerate(faces)]
    ground = trimesh.Trimesh(vertices=vertices, faces=faces, face_colors=face_colors)
    # ground.visual.face_colors = white
    return ground

class GMoF(nn.Module):
    def __init__(self, rho=1):
        super(GMoF, self).__init__()
        self.rho = rho

    def extra_repr(self):
        return 'rho = {}'.format(self.rho)

    def forward(self, residual):
        squared_res = residual ** 2
        dist = torch.div(squared_res, squared_res + self.rho ** 2)
        return self.rho ** 2 * dist

class GMoF_unscaled(nn.Module):
    def __init__(self, rho=1):
        super(GMoF_unscaled, self).__init__()
        self.rho = rho

    def extra_repr(self):
        return 'rho = {}'.format(self.rho)

    def forward(self, residual):
        squared_res = residual ** 2
        if type(residual) is torch.Tensor:
            dist = torch.div(squared_res, squared_res + self.rho ** 2)
        else:
            dist = squared_res/(squared_res + self.rho ** 2)
        return dist



def getNewCoordinate(axisHeight=0.05, axisRadius=0.001):
    xColor = [200, 50, 0, 128]
    yColor = [0, 200, 50, 128]
    zColor = [50, 0, 200, 128]

    xAxis2 = trimesh.primitives.Cylinder(
        radius=axisRadius,
        height=axisHeight,
    )
    xAxis2.apply_transform(matrix=np.mat(
        ((  0,  0, 1, axisHeight / 2),
        (   0,  1, 0, 0),
        (   -1, 0, 0, 0),
        (   0,  0, 0, 1))
    ))
    xAxis2.visual.face_colors = xColor
    yAxis2 = trimesh.primitives.Cylinder(
        radius=axisRadius,
        height=axisHeight,
    )
    yAxis2.apply_transform(matrix=np.mat(
        ((  1, 0, 0, 0),
        (   0, 0, -1, axisHeight / 2),
        (   0, 1, 0, 0),
        (   0, 0, 0, 1))
    ))
    yAxis2.visual.face_colors = yColor
    zAxis2 = trimesh.primitives.Cylinder(
        radius=axisRadius,
        height=axisHeight,
    )
    zAxis2.apply_transform(matrix=np.mat(
        ((1, 0, 0, 0),
         (0, 1, 0, 0),
         (0, 0, 1, axisHeight / 2),
         (0, 0, 0, 1))
    ))
    zAxis2.visual.face_colors = zColor
    xBox2 = trimesh.primitives.Box(
        extents=[axisRadius * 3, axisRadius * 3, axisRadius * 3]
    )
    xBox2.apply_translation((axisHeight, 0, 0))
    xBox2.visual.face_colors = xColor
    yBox2 = trimesh.primitives.Box(
        extents=[axisRadius * 3, axisRadius * 3, axisRadius * 3]
    )
    yBox2.apply_translation((0, axisHeight, 0))
    yBox2.visual.face_colors = yColor
    zBox2 = trimesh.primitives.Box(
        extents=[axisRadius * 3, axisRadius * 3, axisRadius * 3]
    )
    zBox2.apply_translation((0, 0, axisHeight))
    zBox2.visual.face_colors = zColor

    return 1


def meshVisualization(mesh):
    worldMeshList = get_world_mesh_list(planeWidth=0.2, axisHeight=0.05, axisRadius=0.001)
    mesh.visual.face_colors = [255, 128, 255, 200]
    worldMeshList.append(mesh)
    scene = trimesh.Scene(worldMeshList)
    scene.show()


def meshPairVisualization(mesh1, mesh2):
    worldMeshList = get_world_mesh_list(planeWidth=0.2, axisHeight=0.05, axisRadius=0.001)
    mesh1.visual.face_colors = [255, 128, 255, 200]
    mesh2.visual.face_colors = [255, 255, 128, 200]

    worldMeshList.append((mesh1, mesh2))
    scene = trimesh.Scene(worldMeshList)
    scene.show()


if __name__ == '__main__':
    # meshVisualization(trimesh.convex.convex_hull(np.array([[0.1, 0, 0], [0, 0.1, 0], [0, 0, 0.1], [0, 0, 0]])))
    scene = trimesh.Scene(get_checkerboard_plane())
    scene.add_geometry(get_world_mesh_list(add_plane=False))
    scene.show()
vis_utils/mesh_utils.py
sha2nkt/ipman-r
import os
import os.path as osp
import torch
import trimesh
import pyrender
import numpy as np
import matplotlib as mpl
import cv2
import pickle as pkl

import PIL.Image as pil_img
from PIL import ImageFont
from PIL import ImageDraw
from constants import CONTACT_THRESH, GROUND_OFFSETS
from vis_utils.mesh_utils import get_checkerboard_plane, get_meshed_plane, GMoF_unscaled, HDfier
# from lib.core.part_volumes import PartVolume
from utils.geometry import ea2rm
from stability.stability_metrics import StabilityMetrics
# from psbody.mesh.colors import name_to_rgb
os.environ['PYOPENGL_PLATFORM'] = 'egl'

from smplx import SMPL
SMPL_MODEL_DIR = '/ps/project/alignment/models/smplx_internal/smpl/'
SMPL_BATCH_SIZE = 500
smpl_obj = SMPL(
            SMPL_MODEL_DIR,
            batch_size=SMPL_BATCH_SIZE,
            create_transl=True
        )

STABILITY_METRICS = StabilityMetrics()

def vis_vert_with_ground(vertices, trans, seq_name, vis_path='dummy', device='cuda',
                         start_idx=0, sub_sample=5, ground_offset=GROUND_OFFSETS['h36m'], contact_thresh=CONTACT_THRESH, imgnames=[],
                         smpl_batch_size=SMPL_BATCH_SIZE, visualize=True):
    global smpl_obj
    smpl_obj = smpl_obj.to(device)
    with torch.no_grad():
        vertices += trans
        cam_rot = trimesh.transformations.rotation_matrix(np.radians(-30), [1, 0, 0])
        img_count = start_idx
        out_dir = os.path.join(vis_path, seq_name)
        os.makedirs(out_dir, exist_ok=True)
        faces = torch.tensor(smpl_obj.faces.astype(np.int64), dtype=torch.long).to(device)
        # rotate the vertices to align with the pyrender coordinates, no aligned with y axis
        R1 = ea2rm(torch.tensor([[np.radians(270)]]), torch.tensor([[np.radians(0)]]),
                   torch.tensor([[np.radians(0)]])).float().to(device).expand(vertices.shape[0], 3, 3)
        vertices = torch.einsum('bki,bji->bjk', [R1, vertices])

        vertices[:, :, 1] += -ground_offset
        in_bos_label, com, contact_metric, contact_mask = STABILITY_METRICS.check_bos_interior(vertices, faces,
                                                                                 contact_thresh=contact_thresh)

        for fnum in range(min(smpl_batch_size // sub_sample, len(vertices))):
            # check stability (com inside bos)
            if visualize:
                # camera values come from visual inspection in yogi/smplify-xmc-support/rendering
                view1 = overlay_mesh(vertices[fnum].cpu().numpy(),
                                     smpl_obj.faces,
                                     camera_transl=np.array([-1.7, 2, 7]),
                                     # camera_transl=np.array([-0, 2, 6]), # for last two sequences in rich test
                                     focal_length_x=1000,
                                     focal_length_y=1000,
                                     camera_center=[150, 250],
                                     H=1000,
                                     W=1000,
                                     img=None,
                                     camera_rotation=cam_rot,
                                     rotaround=0,
                                     com=com.cpu().numpy(),
                                     gmof_rho=contact_thresh)
                # if img_count == 396:
                view2 = overlay_mesh(vertices[fnum].cpu().numpy(),
                                   smpl_obj.faces,
                                   camera_transl=np.array([-1.7, 2, 5]),
                                   # camera_transl=np.array([-0, 2, 6]), # for last two sequences in rich test
                                   focal_length_x=1000,
                                   focal_length_y=1000,
                                   camera_center=[150, 250],
                                   H=1000,
                                   W=1000,
                                   img=None,
                                   camera_rotation=cam_rot,
                                   rotaround=45,
                                   com=com.cpu().numpy(),
                                     gmof_rho=contact_thresh)
                view3 = overlay_mesh(vertices[fnum].cpu().numpy(),
                                   smpl_obj.faces,
                                     camera_transl=np.array([-1.7, 2, 5]),
                                     # camera_transl=np.array([-0, 2, 6]), # for last two sequences in rich test
                                   focal_length_x=1000,
                                   focal_length_y=1000,
                                   camera_center=[150, 250],
                                   H=1000,
                                   W=1000,
                                   img=None,
                                   camera_rotation=cam_rot,
                                   rotaround=90,
                                   com=com.cpu().numpy(),
                                     gmof_rho=contact_thresh)
                img_bottom = overlay_mesh(vertices[fnum].cpu().numpy(),
                                          smpl_obj.faces,
                                          camera_transl=np.array([-1.7, 2, 7]),
                                          # camera_transl=np.array([-0, 2, 7]), # for last two sequences in rich test
                                          focal_length_x=1000,
                                          focal_length_y=1000,
                                          camera_center=[150, 250],
                                          H=1000,
                                          W=1000,
                                          img=None,
                                          camera_rotation=cam_rot,
                                          rotaround=0,
                                          com=com.cpu().numpy(),
                                          bottomview=True,
                                          draw_ground_plane=False,
                                          gmof_rho=contact_thresh)

                # save image
                if len(imgnames) == 0:
                    filename = os.path.join(out_dir, f'{img_count:04d}.png')
                else:
                    new_img_name = '_'.join(imgnames[fnum].split('/')[-6:])
                    filename = os.path.join(out_dir, new_img_name)
                print(f'Saved filename: {filename}')

                IMG = np.hstack((view1, view2, view3, img_bottom))
                IMG = pil_img.fromarray(IMG)
                # draw = ImageDraw.Draw(IMG)
                # font = ImageFont.truetype("arial.ttf", 100)
                # if in_bos_label[fnum] == 1:
                #     draw.text((0, 0), "Stable", (0, 255, 0), font=font)
                # elif in_bos_label[fnum] == -1:
                #     draw.text((0, 0), "Unstable: No contact", (255, 69, 0), font=font)
                # else:
                #     draw.text((0, 0), "Unstable: In contact", (255, 0, 0), font=font)
                IMG.save(filename)
                img_count += 1
    return in_bos_label.squeeze(), contact_metric.squeeze(), contact_mask.squeeze()

def vis_smpl_with_ground(thetas, transl, betas, seq_name, vis_path, device='cuda',
                         start_idx=0, sub_sample=5, ground_offset=GROUND_OFFSETS['h36m'], contact_thresh=CONTACT_THRESH,
                         smpl_batch_size=SMPL_BATCH_SIZE, visualize=True):
    global smpl_obj
    smpl_obj = smpl_obj.to(device)
    with torch.no_grad():
        mm2m = 1000
        if isinstance(betas, np.ndarray):
            betas = torch.FloatTensor(betas).to(device)
        if isinstance(thetas, np.ndarray):
            thetas = torch.FloatTensor(thetas).to(device)
        if isinstance(transl, np.ndarray):
            transl = torch.FloatTensor(transl).to(device)

        faces = torch.tensor(smpl_obj.faces.astype(np.int64), dtype=torch.long).to(device)

        cam_rot = trimesh.transformations.rotation_matrix(np.radians(-30), [1, 0, 0]) # for world
        # cam_rot = trimesh.transformations.rotation_matrix(np.radians(0), [1, 0, 0]) # for cam 1
        img_count = start_idx
        out_dir = os.path.join(vis_path, seq_name)
        os.makedirs(out_dir, exist_ok=True)
        in_bos_label_accumulator = []
        contact_metric_accumulator = []
        contact_mask_accumulator = []
        for i in range(0, len(thetas), smpl_batch_size):
            gt_pose = thetas[i:i+smpl_batch_size:sub_sample, :]
            gt_beta = betas[i:i+smpl_batch_size:sub_sample, :]
            gt_transl = transl[i:i+smpl_batch_size:sub_sample, :]
            # add ground plane offset to smpl mesh along z
            gt_transl[:, -1] += -ground_offset
            smpl_output = smpl_obj(
                betas=gt_beta,
                body_pose=gt_pose[:, 3:],
                global_orient=gt_pose[:, :3],
                transl=gt_transl,
            )
            vertices = smpl_output.vertices
            # rotate the vertices to align with the pyrender coordinates, no aligned with y axis
            R1 = ea2rm(torch.tensor([[np.radians(270)]]), torch.tensor([[np.radians(0)]]),
                       torch.tensor([[np.radians(0)]])).float().to(device).expand(vertices.shape[0], 3, 3) # for world
            # R1 = ea2rm(torch.tensor([[np.radians(180)]]), torch.tensor([[np.radians(0)]]),
            #            torch.tensor([[np.radians(0)]])).float().to(device).expand(vertices.shape[0], 3, 3) # for camera
            vertices = torch.einsum('bki,bji->bjk', [R1, vertices])

            # check stability (com inside bos)
            in_bos_label, com, contact_metric, contact_mask = STABILITY_METRICS.check_bos_interior(vertices, faces,
                                                                                     contact_thresh=contact_thresh)
            in_bos_label_accumulator.append(in_bos_label)
            contact_metric_accumulator.append(contact_metric)
            contact_mask_accumulator.append(contact_mask)

            for fnum in range(min(smpl_batch_size//sub_sample, len(vertices))):
                if visualize:
                    # camera values come from visual inspection in yogi/smplify-xmc-support/rendering
                    view1 = overlay_mesh(vertices[fnum].cpu().numpy(),
                                 smpl_obj.faces,
                                 camera_transl=np.array([-1.7, 2, 7]), # for world
                                 # camera_transl=np.array([-1.2, 1.0, 9]), # for camera
                                 focal_length_x=1000,
                                 focal_length_y=1000,
                                 camera_center=[150, 250],
                                 H=1000,
                                 W=1000,
                                 img=None,
                                 camera_rotation=cam_rot,
                                 rotaround=0,
                                 com=com.cpu().numpy(),
                                         gmof_rho=contact_thresh)
                    # view2 = overlay_mesh(vertices[fnum].cpu().numpy(),
                    #                    smpl_obj.faces,
                    #                    camera_transl=np.array([-1.7, 2, 5]),
                    #                    focal_length_x=1000,
                    #                    focal_length_y=1000,
                    #                    camera_center=[150, 250],
                    #                    H=1000,
                    #                    W=1000,
                    #                    img=None,
                    #                    camera_rotation=cam_rot,
                    #                    rotaround=45,
                    #                    com=com.cpu().numpy(), )
                    img_bottom = overlay_mesh(vertices[fnum].cpu().numpy(),
                                       smpl_obj.faces,
                                       camera_transl=np.array([-1.7, 2, 7]),
                                       focal_length_x=1000,
                                       focal_length_y=1000,
                                       camera_center=[150, 250],
                                       H=1000,
                                       W=1000,
                                       img=None,
                                       camera_rotation=cam_rot,
                                       rotaround=0,
                                       com=com.cpu().numpy(),
                                              bottomview=True,
                                              draw_ground_plane=False,
                                              gmof_rho=contact_thresh)

                    # save image
                    filename = os.path.join(out_dir, f'{img_count:04d}.png')
                    print(f'Saved filename: {filename}')
                    img_count += 1
                    IMG = np.hstack((view1, img_bottom))
                    IMG = pil_img.fromarray(IMG)
                    draw = ImageDraw.Draw(IMG)
                    font = ImageFont.truetype("arial.ttf", 100)
                    if in_bos_label[fnum] == 1:
                        draw.text((0, 0), "Stable", (0, 255, 0), font=font)
                    elif in_bos_label[fnum] == -1:
                        draw.text((0, 0), "Unstable: No contact", (255, 69, 0), font=font)
                    else:
                        draw.text((0, 0), "Unstable: In contact", (255, 0, 0), font=font)
                    IMG.save(filename)
                    # md_dict = {'beta': gt_beta[fnum],
                    #            'pose': gt_pose[fnum],
                    #            'transl': gt_transl[fnum]}
                    #
                    # # save md dict as pickle in out_dir
                    # md_name = os.path.join(out_dir, f'{img_count:04d}.pkl')
                    # with open(md_name, 'wb') as fp:
                    #     pkl.dump(md_dict, fp)
    in_bos_label_accumulator = torch.cat(in_bos_label_accumulator, dim=0)
    contact_metric_accumulator = torch.cat(contact_metric_accumulator, dim=0)
    contact_mask_accumulator = torch.cat(contact_mask_accumulator, dim=0)
    return in_bos_label_accumulator.squeeze(), contact_metric_accumulator.squeeze(), contact_mask_accumulator.squeeze()

def get_view_matrix(pitch=0, yaw=0, roll=0, tx=0, ty=0, tz=0):
    camera_rot = create_rotmat(pitch, yaw, roll)

    view_matrix = np.eye(4)
    view_matrix[:3, :3] = camera_rot
    view_matrix[0, 3] = tx
    view_matrix[1, 3] = ty
    view_matrix[2, 3] = tz

    return view_matrix


def create_rotmat(pitch, yaw, roll):
    pitch = (180 + pitch) * np.pi / 180
    yaw = yaw * np.pi / 180
    roll = roll * np.pi / 180

    cos_p, sin_p = np.cos(pitch), np.sin(pitch)
    cos_y, sin_y = np.cos(yaw), np.sin(yaw)
    cos_r, sin_r = np.cos(roll), np.sin(roll)

    R = np.array(
        [[cos_y * cos_r, sin_p * sin_y * cos_r - cos_p * sin_r, cos_p * sin_y * cos_r + sin_p * sin_r],
         [cos_y * sin_r, sin_p * sin_y * sin_r + cos_p * cos_r, cos_p * sin_y * sin_r - sin_p * cos_r],
         [-sin_y, sin_p * cos_y, cos_p * cos_y]]
    )

    return R

def vis_heatmap(heatmap):
    imgpath = './data/checkerboard.png'
    image = cv2.imread(imgpath).astype(np.uint8)
    heatmap = np.clip(heatmap*255, 0, 255).astype(np.uint8)
    resized_heatmap = cv2.resize(heatmap, (int(image.shape[1]), int(image.shape[0])))
    # idx = resized_heatmap.argmax()

    # center_y = idx // image.shape[1]
    # center_x = idx % image.shape[1]

    colored_heatmap = cv2.applyColorMap(resized_heatmap.astype(np.uint8), cv2.COLORMAP_JET)
    # colored_heatmap = cv2.applyColorMap(heatmap.astype(np.uint8), cv2.COLORMAP_JET)
    masked_image = colored_heatmap*0.8 + image*0.2

    # cv2.circle(masked_image, (center_x, center_y), 3, [0, 0, 255], -1)
    # cv2.circle(masked_image, (center_x, center_y), 3, [0, 0, 255], -1)
    return masked_image.astype(np.uint8)

def  overlay_mesh(verts, faces, camera_transl, focal_length_x, focal_length_y, camera_center,
        H, W, img, camera_rotation=None, rotaround=None, topview=False, bottomview=False, viewer=False, gmof_rho= 0.01,
                 draw_ground_plane=True, draw_contact=True, draw_support_plane=False, com=None):


    material = pyrender.MetallicRoughnessMaterial(
                metallicFactor=0.0,
                alphaMode='OPAQUE',
                baseColorFactor=(1.0, 1.0, 0.9, 1.0))

    support_plane_material = pyrender.MetallicRoughnessMaterial(
        metallicFactor=0.2,
        alphaMode='BLEND',
        baseColorFactor=(1.0, 1.0, 1.0, 0.6))

    out_mesh = trimesh.Trimesh(verts, faces, process=False)

    # Set mesh vertex color based on distance from the ground plane
    ground_plane_height = 0.0 # obtained by visualization on the presented pose
    view_offset = 0.0 # groundplane at 0 not visible, lower mesh and ground plane just for viewing

    # Calculate contact vertices
    vertex_height = (out_mesh.vertices[:, 1] - ground_plane_height)
    vertex_height_robustified = GMoF_unscaled(rho=gmof_rho)(vertex_height)


    if draw_contact:
        cmap = mpl.cm.get_cmap('jet')
        vertex_colors = [cmap(1 - x) for x in vertex_height_robustified]
        out_mesh.visual.vertex_colors = vertex_colors
    else:
        color = np.array([1.00, 0.55, 0.41]) # salmon1
        out_mesh.visual.vertex_colors = color

    # draw ground_plane
    if draw_ground_plane:
        ground_mesh = get_checkerboard_plane()
        ground_pose = np.eye(4)
        if rotaround is not None:
            if topview or bottomview:
                ground_pose = trimesh.transformations.rotation_matrix(np.radians(90), [1, 0, 0]) # make groundplane frontal
            else:
                ground_pose = trimesh.transformations.rotation_matrix(np.radians(0), [1, 0, 0])
            # ground_pose[:3, 3] = np.array([0, ground_plane_height+view_offset, 0])
            # ground_mesh = [mesh.apply_transform(ground_pose) for mesh in ground_mesh]

    # draw support plane that is perpendicular to the ground plane and passes through the centroid of contact vertices
    if draw_support_plane:
        support_mesh = get_checkerboard_plane()
        # Align with y-z plane
        support_pose = trimesh.transformations.rotation_matrix(np.radians(90), [0, 1, 0])
        # get z offset
        vertex_idx = vertex_height_robustified < 0.2
        centroid = np.mean(out_mesh.vertices[vertex_idx, :], axis=0)
        support_plane_offset = centroid[0] # ToDo: fit to the contact vertices to minimize distance instead
        support_pose[:3, 3] = np.array([support_plane_offset, 0, 0])


    if rotaround is not None:
        if topview:
            rot = trimesh.transformations.rotation_matrix(
                np.radians(90), [1, 0, 0], [0,0,0])
        elif bottomview:
            rot = trimesh.transformations.rotation_matrix(
                np.radians(270), [1, 0, 0], [0, 0, 0])
        else:
            rot = trimesh.transformations.rotation_matrix(
                np.radians(rotaround), [0, 1, 0], out_mesh.vertices[4297])
        out_mesh.apply_transform(rot)
        if draw_support_plane:
            support_pose = np.dot(rot, support_pose)
        view_matrix = get_view_matrix(
            pitch=180,yaw=-0.5,roll=0, tx=camera_transl[0], ty=camera_transl[1], tz=camera_transl[2]
        )
    else:
        camera_pose_gl = get_view_matrix(
            pitch=180+camera_rotation[0],yaw=-camera_rotation[1],roll=-camera_rotation[2],
            tx=camera_transl[0], ty=-camera_transl[1], tz=-camera_transl[2]
        ) # note coversion from camera_orientation to global orientation
        out_mesh.vertices = camera_pose_gl[:3, 3] + np.einsum('mn,jn->jm', camera_pose_gl[:3, :3], verts)

        if draw_ground_plane:
            ground_mesh = [mesh.apply_transform(camera_pose_gl) for mesh in ground_mesh]
        if draw_support_plane:
            support_pose = np.dot(camera_pose_gl, support_pose)

        view_matrix = np.eye(4)

    # add origin sphere
    origin_sphere = trimesh.creation.uv_sphere(0.05)
    yellow = [255, 255, 0, 255]
    origin_sphere.visual.face_colors = yellow
    origin_sphere_mesh = pyrender.Mesh.from_trimesh(origin_sphere, smooth=False)
    origin_pose = np.eye(4)

    # add com sphere
    com_sphere = trimesh.creation.uv_sphere(0.05)
    green = [0, 255, 0, 255]
    com_sphere.visual.face_colors = green
    com_sphere_mesh = pyrender.Mesh.from_trimesh(com_sphere, smooth=False)
    com_pose = np.eye(4)
    if rotaround is not None:
        origin_pose[:3, 3] = np.array([0, view_offset, 0])  # just for viewing, lower the mesh with the ground plane
        if bottomview:
            com_pose[:3, 3] = np.array([com[0, 0], com[0, 2], 0])
        else:
            com_pose[:3, 3] = np.array([com[0, 0], 0, com[0, 2]])  # visualize com projection on the ground plane

    # add mesh to scene
    mesh = pyrender.Mesh.from_trimesh(
        out_mesh,
        material=material)
    if img is not None:
        scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0],
                               ambient_light=(0.3, 0.3, 0.3))
    else:
        scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0],
                               ambient_light=(0.3, 0.3, 0.3))
    mesh_pose = np.eye(4)
    if rotaround is not None:
        mesh_pose[:3, 3] = np.array([0, view_offset, 0]) # just for viewing, lower the mesh with the ground plane
    scene.add(mesh, pose=mesh_pose, name='mesh')

    if draw_ground_plane:
        ground_mesh = pyrender.Mesh.from_trimesh(
            ground_mesh,
            smooth=False)
        # scene.add(ground_mesh, pose=ground_pose, name='ground_plane')
        scene.add(ground_mesh, name='ground_plane')

    if draw_support_plane:
        support_mesh = [mesh.apply_transform(support_pose) for mesh in support_mesh]
        support_mesh = pyrender.Mesh.from_trimesh(
            support_mesh,
            smooth=False,
            material=support_plane_material
        )
        scene.add(support_mesh, name='support_plane')

    scene.add(origin_sphere_mesh, pose=origin_pose, name='origin_sphere')
    scene.add(com_sphere_mesh, pose=com_pose, name='com_sphere')

    # add mesh for camera
    if img is None and topview is False:
        # to lower the ground-plane so person doesn't get cut
        cam_y_offset = 150
    elif topview:
        cam_y_offset = -100
    else:
        cam_y_offset = 0
    pyrencamera = pyrender.camera.IntrinsicsCamera(
        fx=focal_length_x, fy=focal_length_y,
        cx=camera_center[0], cy=camera_center[1] + cam_y_offset)
    scene.add(pyrencamera, pose=view_matrix)

    # create and add light
    light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=1)
    light_pose = np.eye(4)
    for lp in [[1,1,1], [-1,1,1],[1,-1,1],[-1,-1,1]]:
        light_pose[:3, 3] = out_mesh.vertices.mean(0) + np.array(lp)
        #out_mesh.vertices.mean(0) + np.array(lp)
        scene.add(light, pose=light_pose)

    if viewer:
        pyrender.Viewer(scene, use_raymond_lighting=True)
        return 0
    else:
        r = pyrender.OffscreenRenderer(viewport_width=W,
                                       viewport_height=H,
                                       point_size=1.0)
        color, _ = r.render(scene, flags=pyrender.RenderFlags.RGBA)
        color = color.astype(np.float32) / 255.0

        if img is not None:
            valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
            input_img = img.detach().cpu().numpy()
            output_img = (color[:, :, :-1] * valid_mask +
                          (1 - valid_mask) * input_img)
        else:
            output_img = color

        output_img = pil_img.fromarray((output_img * 255).astype(np.uint8))
        output_img= np.asarray(output_img)[:,:,:3]

        return output_img
vis_utils/world_vis.py
sha2nkt/ipman-r