added missing files

utils/cluster.py

@@ -0,0 +1,99 @@
+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),
+}
+
+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,
+        condor_dir,
+        cfg_file,
+        num_exp=1,
+        exp_opts=None,
+        bid_amount=10,
+        num_workers=2,
+        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.
+    # this makes sure the correct config is copied even if you change the config file
+    # after starting the experiment and before the first job is submitted
+    temp_config_dir = os.path.join(os.path.dirname(condor_dir), 'temp_configs', exp_name)
+    os.makedirs(temp_config_dir, exist_ok=True)
+    new_cfg_file = os.path.join(temp_config_dir, '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])
+
+    condor_log_dir = os.path.join(condor_dir, 'condorlog', exp_name)
+    os.makedirs(condor_log_dir, exist_ok=True)
+    submission = f'executable = {condor_log_dir}/{exp_name}_run.sh\n' \
+                 'arguments = $(Process) $(Cluster)\n' \
+                 f'error = {condor_log_dir}/{exp_name}_$(Cluster).$(Process).err\n' \
+                 f'output = {condor_log_dir}/{exp_name}_$(Cluster).$(Process).out\n' \
+                 f'log = {condor_log_dir}/{exp_name}_$(Cluster).$(Process).log\n' \
+                 f'request_memory = {memory}\n' \
+                 f'request_cpus={int(num_workers)}\n' \
+                 f'request_gpus={num_gpus}\n' \
+                 f'requirements={gpus}\n' \
+                 f'+MaxRunningPrice = 500\n' \
+                 f'queue {num_exp}'
+                 # f'request_cpus={int(num_workers/2)}\n' \
+                 # f'+RunningPriceExceededAction = \"kill\"\n' \
+    print('<<< Condor Submission >>> ')
+    print(submission)
+
+    with open(f'{condor_log_dir}/{exp_name}_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_log_dir}')
+    logger.info(f'The outputs for this experiments can be found under: {output_dir}')
+    ## This is the trick. Notice there is no --cluster here
+    bash = 'export PYTHONBUFFERED=1\n export PATH=$PATH\n ' \
+           f'{sys.executable} {script} --cfg {new_cfg_file} --cfg_id $1'
+
+    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_log_dir}/{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}', f'{condor_log_dir}/{exp_name}_submit.sub']
+    logger.info('Executing ' + ' '.join(cmd))
+    subprocess.call(cmd)

utils/colorwheel.py

@@ -0,0 +1,22 @@
+import cv2
+import numpy as np
+
+
+def make_color_wheel_image(img_width, img_height):
+    """
+    Creates a color wheel based image of given width and height
+    Args:
+        img_width (int):
+        img_height (int):
+
+    Returns:
+        opencv image (numpy array): color wheel based image
+    """
+    hue = np.fromfunction(lambda i, j: (np.arctan2(i-img_height/2, img_width/2-j) + np.pi)*(180/np.pi)/2,
+                          (img_height, img_width), dtype=np.float)
+    saturation = np.ones((img_height, img_width)) * 255
+    value = np.ones((img_height, img_width)) * 255
+    hsl = np.dstack((hue, saturation, value))
+    color_map = cv2.cvtColor(np.array(hsl, dtype=np.uint8), cv2.COLOR_HSV2BGR)
+    return color_map
+

utils/config.py

@@ -0,0 +1,196 @@
+import itertools
+import operator
+import os
+import shutil
+import time
+from functools import reduce
+from typing import List, Union
+
+import configargparse
+import yaml
+from flatten_dict import flatten, unflatten
+from loguru import logger
+from yacs.config import CfgNode as CN
+
+from utils.cluster import execute_task_on_cluster
+from utils.default_hparams import hparams
+
+
+def parse_args():
+    def add_common_cmdline_args(parser):
+        # for cluster runs
+        parser.add_argument('--cfg', required=True, 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=10, help='amount of bid for cluster')
+        parser.add_argument('--memory', type=int, default=64000, help='memory amount for cluster')
+        parser.add_argument('--gpu_min_mem', type=int, default=12000, 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')
+        return parser
+
+    # For Blender main parser
+    arg_formatter = configargparse.ArgumentDefaultsHelpFormatter
+    cfg_parser = configargparse.YAMLConfigFileParser
+    description = 'PyTorch implementation of DECO'
+
+    parser = configargparse.ArgumentParser(formatter_class=arg_formatter,
+                                      config_file_parser_class=cfg_parser,
+                                      description=description,
+                                      prog='deco')
+
+    parser = add_common_cmdline_args(parser)
+
+    args = parser.parse_args()
+    print(args, end='\n\n')
+
+    return args
+
+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 get_from_dict(dict, keys):
+    return reduce(operator.getitem, keys, dict)
+
+def save_dict_to_yaml(obj, filename, mode='w'):
+    with open(filename, mode) as f:
+        yaml.dump(obj, f, default_flow_style=False)
+
+def run_grid_search_experiments(
+        args,
+        script='train.py',
+        change_wt_name=True
+):
+    cfg = yaml.safe_load(open(args.cfg))
+    # parse config file to split into a list of configs with tuning hyperparameters separated
+    # Also return the names of tuned hyperparameters hyperparameters
+    different_configs, hyperparams = get_grid_search_configs(
+        cfg,
+        excluded_keys=['TRAINING/DATASETS', 'TRAINING/DATASET_MIX_PDF', 'VALIDATION/DATASETS'],
+    )
+    logger.info(f'Grid search hparams: \n {hyperparams}')
+
+    # The config file may be missing some default values, so we need to add them
+    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[args.cfg_id])
+
+    if args.cluster:
+        execute_task_on_cluster(
+            script=script,
+            exp_name=config_to_run.EXP_NAME,
+            output_dir=config_to_run.OUTPUT_DIR,
+            condor_dir=config_to_run.CONDOR_DIR,
+            cfg_file=args.cfg,
+            num_exp=len(different_configs),
+            bid_amount=args.bid,
+            num_workers=config_to_run.DATASET.NUM_WORKERS,
+            memory=args.memory,
+            exp_opts=args.opts,
+            gpu_min_mem=args.gpu_min_mem,
+            gpu_arch=args.gpu_arch,
+        )
+        exit()
+
+    # ==== create logdir using hyperparam settings
+    logtime = time.strftime('%d-%m-%Y_%H-%M-%S')
+    logdir = f'{logtime}_{config_to_run.EXP_NAME}'
+    wt_file = config_to_run.EXP_NAME + '_'
+    for hp in hyperparams:
+        v = get_from_dict(different_configs[args.cfg_id], hp.split('/'))
+        logdir += f'_{hp.replace("/", ".").replace("_", "").lower()}-{v}'
+        wt_file += f'{hp.replace("/", ".").replace("_", "").lower()}-{v}_'
+    logdir = os.path.join(config_to_run.OUTPUT_DIR, logdir)
+    os.makedirs(logdir, exist_ok=True)
+    config_to_run.LOGDIR = logdir
+
+    wt_file += 'best.pth'
+    wt_path = os.path.join(os.path.dirname(config_to_run.TRAINING.BEST_MODEL_PATH), wt_file)
+    if change_wt_name: config_to_run.TRAINING.BEST_MODEL_PATH = wt_path
+
+    shutil.copy(src=args.cfg, dst=os.path.join(logdir, 'config.yaml'))
+
+    # save config
+    save_dict_to_yaml(
+        unflatten(flatten(config_to_run)),
+        os.path.join(config_to_run.LOGDIR, 'config_to_run.yaml')
+    )
+
+    return config_to_run

utils/default_hparams.py

@@ -0,0 +1,45 @@
+from yacs.config import CfgNode as CN
+
+# Set default hparams to construct new default config
+# Make sure the defaults are same as in parser
+hparams = CN()
+
+# General settings
+hparams.EXP_NAME = 'default'
+hparams.PROJECT_NAME = 'default'
+hparams.OUTPUT_DIR = 'deco_results/'
+hparams.CONDOR_DIR = '/is/cluster/work/achatterjee/condor/rich/'
+hparams.LOGDIR = ''
+
+# Dataset hparams
+hparams.DATASET = CN()
+hparams.DATASET.BATCH_SIZE = 64
+hparams.DATASET.NUM_WORKERS = 4
+hparams.DATASET.NORMALIZE_IMAGES = True
+
+# Optimizer hparams
+hparams.OPTIMIZER = CN()
+hparams.OPTIMIZER.TYPE = 'adam'
+hparams.OPTIMIZER.LR = 5e-5
+hparams.OPTIMIZER.NUM_UPDATE_LR = 10
+
+# Training hparams
+hparams.TRAINING = CN()
+hparams.TRAINING.ENCODER = 'hrnet'
+hparams.TRAINING.CONTEXT = True
+hparams.TRAINING.NUM_EPOCHS = 50
+hparams.TRAINING.SUMMARY_STEPS = 100
+hparams.TRAINING.CHECKPOINT_EPOCHS = 5
+hparams.TRAINING.NUM_EARLY_STOP = 10
+hparams.TRAINING.DATASETS = ['rich']
+hparams.TRAINING.DATASET_MIX_PDF = ['1.']
+hparams.TRAINING.DATASET_ROOT_PATH = '/is/cluster/work/achatterjee/rich/npzs'
+hparams.TRAINING.BEST_MODEL_PATH = '/is/cluster/work/achatterjee/weights/rich/exp/rich_exp.pth'
+hparams.TRAINING.LOSS_WEIGHTS = 1.
+hparams.TRAINING.PAL_LOSS_WEIGHTS = 1.
+
+# Training hparams
+hparams.VALIDATION = CN()
+hparams.VALIDATION.SUMMARY_STEPS = 100
+hparams.VALIDATION.DATASETS = ['rich']
+hparams.VALIDATION.MAIN_DATASET = 'rich'

utils/diff_renderer.py

@@ -0,0 +1,287 @@
+# from https://gitlab.tuebingen.mpg.de/mkocabas/projects/-/blob/master/pare/pare/utils/diff_renderer.py
+
+import torch
+import numpy as np
+import torch.nn as nn
+
+from pytorch3d.renderer import (
+        PerspectiveCameras,
+        RasterizationSettings,
+        DirectionalLights,
+        BlendParams,
+        HardFlatShader,
+        MeshRasterizer,
+        TexturesVertex,
+        TexturesAtlas
+    )
+from pytorch3d.structures import Meshes
+
+from .image_utils import get_default_camera
+from .smpl_uv import get_tenet_texture
+
+
+class MeshRendererWithDepth(nn.Module):
+    """
+    A class for rendering a batch of heterogeneous meshes. The class should
+    be initialized with a rasterizer and shader class which each have a forward
+    function.
+    """
+
+    def __init__(self, rasterizer, shader):
+        super().__init__()
+        self.rasterizer = rasterizer
+        self.shader = shader
+
+    def forward(self, meshes_world, **kwargs) -> torch.Tensor:
+        """
+        Render a batch of images from a batch of meshes by rasterizing and then
+        shading.
+
+        NOTE: If the blur radius for rasterization is > 0.0, some pixels can
+        have one or more barycentric coordinates lying outside the range [0, 1].
+        For a pixel with out of bounds barycentric coordinates with respect to a
+        face f, clipping is required before interpolating the texture uv
+        coordinates and z buffer so that the colors and depths are limited to
+        the range for the corresponding face.
+        """
+        fragments = self.rasterizer(meshes_world, **kwargs)
+        images = self.shader(fragments, meshes_world, **kwargs)
+
+        mask = (fragments.zbuf > -1).float()
+
+        zbuf = fragments.zbuf.view(images.shape[0], -1)
+        # print(images.shape, zbuf.shape)
+        depth = (zbuf - zbuf.min(-1, keepdims=True).values) / \
+                (zbuf.max(-1, keepdims=True).values - zbuf.min(-1, keepdims=True).values)
+        depth = depth.reshape(*images.shape[:3] + (1,))
+
+        images = torch.cat([images[:, :, :, :3], mask, depth], dim=-1)
+        return images
+
+
+class DifferentiableRenderer(nn.Module):
+    def __init__(
+            self,
+            img_h,
+            img_w,
+            focal_length,
+            device='cuda',
+            background_color=(0.0, 0.0, 0.0),
+            texture_mode='smplpix',
+            vertex_colors=None,
+            face_textures=None,
+            smpl_faces=None,
+            is_train=False,
+            is_cam_batch=False,
+    ):
+        super(DifferentiableRenderer, self).__init__()
+        self.x = 'a'
+        self.img_h = img_h
+        self.img_w = img_w
+        self.device = device
+        self.focal_length = focal_length
+        K, R = get_default_camera(focal_length, img_h, img_w, is_cam_batch=is_cam_batch)
+        K, R = K.to(device), R.to(device)
+
+        # T = torch.tensor([[0, 0, 2.5 * self.focal_length / max(self.img_h, self.img_w)]]).to(device)
+        if is_cam_batch:
+            T = torch.zeros((K.shape[0], 3)).to(device)
+        else:
+            T = torch.tensor([[0.0, 0.0, 0.0]]).to(device)
+        self.background_color = background_color
+        self.renderer = None
+        smpl_faces = smpl_faces
+
+        if texture_mode == 'smplpix':
+            face_colors = get_tenet_texture(mode=texture_mode).to(device).float()
+            vertex_colors = torch.from_numpy(
+                np.load(f'data/smpl/{texture_mode}_vertex_colors.npy')[:,:3]
+            ).unsqueeze(0).to(device).float()
+        if texture_mode == 'partseg':
+            vertex_colors = vertex_colors[..., :3].unsqueeze(0).to(device)
+            face_colors = face_textures.to(device)
+        if texture_mode == 'deco':
+            vertex_colors = vertex_colors[..., :3].to(device)
+            face_colors = face_textures.to(device)
+
+        self.register_buffer('K', K)
+        self.register_buffer('R', R)
+        self.register_buffer('T', T)
+        self.register_buffer('face_colors', face_colors)
+        self.register_buffer('vertex_colors', vertex_colors)
+        self.register_buffer('smpl_faces', smpl_faces)
+
+        self.set_requires_grad(is_train)
+
+    def set_requires_grad(self, val=False):
+        self.K.requires_grad_(val)
+        self.R.requires_grad_(val)
+        self.T.requires_grad_(val)
+        self.face_colors.requires_grad_(val)
+        self.vertex_colors.requires_grad_(val)
+        # check if smpl_faces is a FloatTensor as requires_grad_ is not defined for LongTensor
+        if isinstance(self.smpl_faces, torch.FloatTensor):
+            self.smpl_faces.requires_grad_(val)
+
+    def forward(self, vertices, faces=None, R=None, T=None):
+        raise NotImplementedError
+
+
+class Pytorch3D(DifferentiableRenderer):
+    def __init__(
+            self,
+            img_h,
+            img_w,
+            focal_length,
+            device='cuda',
+            background_color=(0.0, 0.0, 0.0),
+            texture_mode='smplpix',
+            vertex_colors=None,
+            face_textures=None,
+            smpl_faces=None,
+            model_type='smpl',
+            is_train=False,
+            is_cam_batch=False,
+    ):
+        super(Pytorch3D, self).__init__(
+            img_h,
+            img_w,
+            focal_length,
+            device=device,
+            background_color=background_color,
+            texture_mode=texture_mode,
+            vertex_colors=vertex_colors,
+            face_textures=face_textures,
+            smpl_faces=smpl_faces,
+            is_train=is_train,
+            is_cam_batch=is_cam_batch,
+        )
+
+        # this R converts the camera from pyrender NDC to
+        # OpenGL coordinate frame. It is basicall R(180, X) x R(180, Y)
+        # I manually defined it here for convenience
+        self.R = self.R @ torch.tensor(
+            [[[ -1.0,  0.0, 0.0],
+              [  0.0, -1.0, 0.0],
+              [  0.0,  0.0, 1.0]]],
+            dtype=self.R.dtype, device=self.R.device,
+        )
+
+        if is_cam_batch:
+            focal_length = self.focal_length
+        else:
+            focal_length = self.focal_length[None, :]
+
+        principal_point = ((self.img_w // 2, self.img_h // 2),)
+        image_size = ((self.img_h, self.img_w),)
+
+        cameras = PerspectiveCameras(
+            device=self.device,
+            focal_length=focal_length,
+            principal_point=principal_point,
+            R=self.R,
+            T=self.T,
+            in_ndc=False,
+            image_size=image_size,
+        )
+
+        for param in cameras.parameters():
+            param.requires_grad_(False)
+
+        raster_settings = RasterizationSettings(
+            image_size=(self.img_h, self.img_w),
+            blur_radius=0.0,
+            max_faces_per_bin=20000,
+            faces_per_pixel=1,
+        )
+
+        lights = DirectionalLights(
+            device=self.device,
+            ambient_color=((1.0, 1.0, 1.0),),
+            diffuse_color=((0.0, 0.0, 0.0),),
+            specular_color=((0.0, 0.0, 0.0),),
+            direction=((0, 1, 0),),
+        )
+
+        blend_params = BlendParams(background_color=self.background_color)
+
+        shader = HardFlatShader(device=self.device,
+                                cameras=cameras,
+                                blend_params=blend_params,
+                                lights=lights)
+
+        self.textures = TexturesVertex(verts_features=self.vertex_colors)
+
+        self.renderer = MeshRendererWithDepth(
+            rasterizer=MeshRasterizer(
+                cameras=cameras,
+                raster_settings=raster_settings
+            ),
+            shader=shader,
+        )
+
+    def forward(self, vertices, faces=None, R=None, T=None, face_atlas=None):
+        batch_size = vertices.shape[0]
+        if faces is None:
+            faces = self.smpl_faces.expand(batch_size, -1, -1)
+
+        if R is None:
+            R = self.R.expand(batch_size, -1, -1)
+
+        if T is None:
+            T = self.T.expand(batch_size, -1)
+
+        # convert camera translation to pytorch3d coordinate frame
+        T = torch.bmm(R, T.unsqueeze(-1)).squeeze(-1)
+
+        vertex_textures = TexturesVertex(
+            verts_features=self.vertex_colors.expand(batch_size, -1, -1)
+        )
+
+        # face_textures needed because vertex_texture cause interpolation at boundaries
+        if face_atlas:
+            face_textures = TexturesAtlas(atlas=face_atlas)
+        else:
+            face_textures = TexturesAtlas(atlas=self.face_colors)
+
+        # we may need to rotate the mesh
+        meshes = Meshes(verts=vertices, faces=faces, textures=face_textures)
+        images = self.renderer(meshes, R=R, T=T)
+        images = images.permute(0, 3, 1, 2)
+        return images
+
+
+class NeuralMeshRenderer(DifferentiableRenderer):
+    def __init__(self, *args, **kwargs):
+        import neural_renderer as nr
+
+        super(NeuralMeshRenderer, self).__init__(*args, **kwargs)
+
+        self.neural_renderer = nr.Renderer(
+            dist_coeffs=None,
+            orig_size=self.img_size,
+            image_size=self.img_size,
+            light_intensity_ambient=1,
+            light_intensity_directional=0,
+            anti_aliasing=False,
+        )
+
+    def forward(self, vertices, faces=None, R=None, T=None):
+        batch_size = vertices.shape[0]
+        if faces is None:
+            faces = self.smpl_faces.expand(batch_size, -1, -1)
+
+        if R is None:
+            R = self.R.expand(batch_size, -1, -1)
+
+        if T is None:
+            T = self.T.expand(batch_size, -1)
+        rgb, depth, mask = self.neural_renderer(
+            vertices,
+            faces,
+            textures=self.face_colors.expand(batch_size, -1, -1, -1, -1, -1),
+            K=self.K.expand(batch_size, -1, -1),
+            R=R,
+            t=T.unsqueeze(1),
+        )
+        return torch.cat([rgb, depth.unsqueeze(1), mask.unsqueeze(1)], dim=1)

utils/get_cfg.py

@@ -0,0 +1,17 @@
+from yacs.config import CfgNode
+
+_VALID_TYPES = {tuple, list, str, int, float, bool}
+
+
+def convert_to_dict(cfg_node, key_list=[]):
+    """ Convert a config node to dictionary """
+    if not isinstance(cfg_node, CfgNode):
+        if type(cfg_node) not in _VALID_TYPES:
+            print("Key {} with value {} is not a valid type; valid types: {}".format(
+                ".".join(key_list), type(cfg_node), _VALID_TYPES), )
+        return cfg_node
+    else:
+        cfg_dict = dict(cfg_node)
+        for k, v in cfg_dict.items():
+            cfg_dict[k] = convert_to_dict(v, key_list + [k])
+        return cfg_dict

utils/hrnet.py

@@ -0,0 +1,625 @@
+import os
+
+import torch
+import torch.nn as nn
+from loguru import logger
+import torch.nn.functional as F
+from yacs.config import CfgNode as CN
+
+models = [
+    'hrnet_w32',
+    'hrnet_w48',
+]
+
+BN_MOMENTUM = 0.1
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
+                               bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion,
+                                  momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class HighResolutionModule(nn.Module):
+    def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
+                 num_channels, fuse_method, multi_scale_output=True):
+        super(HighResolutionModule, self).__init__()
+        self._check_branches(
+            num_branches, blocks, num_blocks, num_inchannels, num_channels)
+
+        self.num_inchannels = num_inchannels
+        self.fuse_method = fuse_method
+        self.num_branches = num_branches
+
+        self.multi_scale_output = multi_scale_output
+
+        self.branches = self._make_branches(
+            num_branches, blocks, num_blocks, num_channels)
+        self.fuse_layers = self._make_fuse_layers()
+        self.relu = nn.ReLU(True)
+
+    def _check_branches(self, num_branches, blocks, num_blocks,
+                        num_inchannels, num_channels):
+        if num_branches != len(num_blocks):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(
+                num_branches, len(num_blocks))
+            logger.error(error_msg)
+            raise ValueError(error_msg)
+
+        if num_branches != len(num_channels):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(
+                num_branches, len(num_channels))
+            logger.error(error_msg)
+            raise ValueError(error_msg)
+
+        if num_branches != len(num_inchannels):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(
+                num_branches, len(num_inchannels))
+            logger.error(error_msg)
+            raise ValueError(error_msg)
+
+    def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
+                         stride=1):
+        downsample = None
+        if stride != 1 or \
+           self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.num_inchannels[branch_index],
+                    num_channels[branch_index] * block.expansion,
+                    kernel_size=1, stride=stride, bias=False
+                ),
+                nn.BatchNorm2d(
+                    num_channels[branch_index] * block.expansion,
+                    momentum=BN_MOMENTUM
+                ),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.num_inchannels[branch_index],
+                num_channels[branch_index],
+                stride,
+                downsample
+            )
+        )
+        self.num_inchannels[branch_index] = \
+            num_channels[branch_index] * block.expansion
+        for i in range(1, num_blocks[branch_index]):
+            layers.append(
+                block(
+                    self.num_inchannels[branch_index],
+                    num_channels[branch_index]
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def _make_branches(self, num_branches, block, num_blocks, num_channels):
+        branches = []
+
+        for i in range(num_branches):
+            branches.append(
+                self._make_one_branch(i, block, num_blocks, num_channels)
+            )
+
+        return nn.ModuleList(branches)
+
+    def _make_fuse_layers(self):
+        if self.num_branches == 1:
+            return None
+
+        num_branches = self.num_branches
+        num_inchannels = self.num_inchannels
+        fuse_layers = []
+        for i in range(num_branches if self.multi_scale_output else 1):
+            fuse_layer = []
+            for j in range(num_branches):
+                if j > i:
+                    fuse_layer.append(
+                        nn.Sequential(
+                            nn.Conv2d(
+                                num_inchannels[j],
+                                num_inchannels[i],
+                                1, 1, 0, bias=False
+                            ),
+                            nn.BatchNorm2d(num_inchannels[i]),
+                            nn.Upsample(scale_factor=2**(j-i), mode='nearest')
+                        )
+                    )
+                elif j == i:
+                    fuse_layer.append(None)
+                else:
+                    conv3x3s = []
+                    for k in range(i-j):
+                        if k == i - j - 1:
+                            num_outchannels_conv3x3 = num_inchannels[i]
+                            conv3x3s.append(
+                                nn.Sequential(
+                                    nn.Conv2d(
+                                        num_inchannels[j],
+                                        num_outchannels_conv3x3,
+                                        3, 2, 1, bias=False
+                                    ),
+                                    nn.BatchNorm2d(num_outchannels_conv3x3)
+                                )
+                            )
+                        else:
+                            num_outchannels_conv3x3 = num_inchannels[j]
+                            conv3x3s.append(
+                                nn.Sequential(
+                                    nn.Conv2d(
+                                        num_inchannels[j],
+                                        num_outchannels_conv3x3,
+                                        3, 2, 1, bias=False
+                                    ),
+                                    nn.BatchNorm2d(num_outchannels_conv3x3),
+                                    nn.ReLU(True)
+                                )
+                            )
+                    fuse_layer.append(nn.Sequential(*conv3x3s))
+            fuse_layers.append(nn.ModuleList(fuse_layer))
+
+        return nn.ModuleList(fuse_layers)
+
+    def get_num_inchannels(self):
+        return self.num_inchannels
+
+    def forward(self, x):
+        if self.num_branches == 1:
+            return [self.branches[0](x[0])]
+
+        for i in range(self.num_branches):
+            x[i] = self.branches[i](x[i])
+
+        x_fuse = []
+
+        for i in range(len(self.fuse_layers)):
+            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
+            for j in range(1, self.num_branches):
+                if i == j:
+                    y = y + x[j]
+                else:
+                    y = y + self.fuse_layers[i][j](x[j])
+            x_fuse.append(self.relu(y))
+
+        return x_fuse
+
+
+blocks_dict = {
+    'BASIC': BasicBlock,
+    'BOTTLENECK': Bottleneck
+}
+
+
+class PoseHighResolutionNet(nn.Module):
+
+    def __init__(self, cfg):
+        self.inplanes = 64
+        extra = cfg['MODEL']['EXTRA']
+        super(PoseHighResolutionNet, self).__init__()
+
+        self.cfg = extra
+
+        # stem net
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
+                               bias=False)
+        self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.layer1 = self._make_layer(Bottleneck, 64, 4)
+
+        self.stage2_cfg = extra['STAGE2']
+        num_channels = self.stage2_cfg['NUM_CHANNELS']
+        block = blocks_dict[self.stage2_cfg['BLOCK']]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))
+        ]
+        self.transition1 = self._make_transition_layer([256], num_channels)
+        self.stage2, pre_stage_channels = self._make_stage(
+            self.stage2_cfg, num_channels)
+
+        self.stage3_cfg = extra['STAGE3']
+        num_channels = self.stage3_cfg['NUM_CHANNELS']
+        block = blocks_dict[self.stage3_cfg['BLOCK']]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))
+        ]
+        self.transition2 = self._make_transition_layer(
+            pre_stage_channels, num_channels)
+        self.stage3, pre_stage_channels = self._make_stage(
+            self.stage3_cfg, num_channels)
+
+        self.stage4_cfg = extra['STAGE4']
+        num_channels = self.stage4_cfg['NUM_CHANNELS']
+        block = blocks_dict[self.stage4_cfg['BLOCK']]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))
+        ]
+        self.transition3 = self._make_transition_layer(
+            pre_stage_channels, num_channels)
+        self.stage4, pre_stage_channels = self._make_stage(
+            self.stage4_cfg, num_channels, multi_scale_output=True)
+
+        self.final_layer = nn.Conv2d(
+            in_channels=pre_stage_channels[0],
+            out_channels=cfg['MODEL']['NUM_JOINTS'],
+            kernel_size=extra['FINAL_CONV_KERNEL'],
+            stride=1,
+            padding=1 if extra['FINAL_CONV_KERNEL'] == 3 else 0
+        )
+
+        self.pretrained_layers = extra['PRETRAINED_LAYERS']
+
+        if extra.DOWNSAMPLE and extra.USE_CONV:
+            self.downsample_stage_1 = self._make_downsample_layer(3, num_channel=self.stage2_cfg['NUM_CHANNELS'][0])
+            self.downsample_stage_2 = self._make_downsample_layer(2, num_channel=self.stage2_cfg['NUM_CHANNELS'][-1])
+            self.downsample_stage_3 = self._make_downsample_layer(1, num_channel=self.stage3_cfg['NUM_CHANNELS'][-1])
+        elif not extra.DOWNSAMPLE and extra.USE_CONV:
+            self.upsample_stage_2 = self._make_upsample_layer(1, num_channel=self.stage2_cfg['NUM_CHANNELS'][-1])
+            self.upsample_stage_3 = self._make_upsample_layer(2, num_channel=self.stage3_cfg['NUM_CHANNELS'][-1])
+            self.upsample_stage_4 = self._make_upsample_layer(3, num_channel=self.stage4_cfg['NUM_CHANNELS'][-1])
+
+    def _make_transition_layer(
+            self, num_channels_pre_layer, num_channels_cur_layer):
+        num_branches_cur = len(num_channels_cur_layer)
+        num_branches_pre = len(num_channels_pre_layer)
+
+        transition_layers = []
+        for i in range(num_branches_cur):
+            if i < num_branches_pre:
+                if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
+                    transition_layers.append(
+                        nn.Sequential(
+                            nn.Conv2d(
+                                num_channels_pre_layer[i],
+                                num_channels_cur_layer[i],
+                                3, 1, 1, bias=False
+                            ),
+                            nn.BatchNorm2d(num_channels_cur_layer[i]),
+                            nn.ReLU(inplace=True)
+                        )
+                    )
+                else:
+                    transition_layers.append(None)
+            else:
+                conv3x3s = []
+                for j in range(i+1-num_branches_pre):
+                    inchannels = num_channels_pre_layer[-1]
+                    outchannels = num_channels_cur_layer[i] \
+                        if j == i-num_branches_pre else inchannels
+                    conv3x3s.append(
+                        nn.Sequential(
+                            nn.Conv2d(
+                                inchannels, outchannels, 3, 2, 1, bias=False
+                            ),
+                            nn.BatchNorm2d(outchannels),
+                            nn.ReLU(inplace=True)
+                        )
+                    )
+                transition_layers.append(nn.Sequential(*conv3x3s))
+
+        return nn.ModuleList(transition_layers)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.inplanes, planes * block.expansion,
+                    kernel_size=1, stride=stride, bias=False
+                ),
+                nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def _make_stage(self, layer_config, num_inchannels,
+                    multi_scale_output=True):
+        num_modules = layer_config['NUM_MODULES']
+        num_branches = layer_config['NUM_BRANCHES']
+        num_blocks = layer_config['NUM_BLOCKS']
+        num_channels = layer_config['NUM_CHANNELS']
+        block = blocks_dict[layer_config['BLOCK']]
+        fuse_method = layer_config['FUSE_METHOD']
+
+        modules = []
+        for i in range(num_modules):
+            # multi_scale_output is only used last module
+            if not multi_scale_output and i == num_modules - 1:
+                reset_multi_scale_output = False
+            else:
+                reset_multi_scale_output = True
+
+            modules.append(
+                HighResolutionModule(
+                    num_branches,
+                    block,
+                    num_blocks,
+                    num_inchannels,
+                    num_channels,
+                    fuse_method,
+                    reset_multi_scale_output
+                )
+            )
+            num_inchannels = modules[-1].get_num_inchannels()
+
+        return nn.Sequential(*modules), num_inchannels
+
+    def _make_upsample_layer(self, num_layers, num_channel, kernel_size=3):
+        layers = []
+        for i in range(num_layers):
+            layers.append(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True))
+            layers.append(
+                nn.Conv2d(
+                    in_channels=num_channel, out_channels=num_channel,
+                    kernel_size=kernel_size, stride=1, padding=1, bias=False,
+                )
+            )
+            layers.append(nn.BatchNorm2d(num_channel, momentum=BN_MOMENTUM))
+            layers.append(nn.ReLU(inplace=True))
+
+        return nn.Sequential(*layers)
+
+    def _make_downsample_layer(self, num_layers, num_channel, kernel_size=3):
+        layers = []
+        for i in range(num_layers):
+            layers.append(
+                nn.Conv2d(
+                    in_channels=num_channel, out_channels=num_channel,
+                    kernel_size=kernel_size, stride=2, padding=1, bias=False,
+                )
+            )
+            layers.append(nn.BatchNorm2d(num_channel, momentum=BN_MOMENTUM))
+            layers.append(nn.ReLU(inplace=True))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+        x = self.layer1(x)
+
+        x_list = []
+        for i in range(self.stage2_cfg['NUM_BRANCHES']):
+            if self.transition1[i] is not None:
+                x_list.append(self.transition1[i](x))
+            else:
+                x_list.append(x)
+        y_list = self.stage2(x_list)
+
+        x_list = []
+        for i in range(self.stage3_cfg['NUM_BRANCHES']):
+            if self.transition2[i] is not None:
+                x_list.append(self.transition2[i](y_list[-1]))
+            else:
+                x_list.append(y_list[i])
+        y_list = self.stage3(x_list)
+
+        x_list = []
+        for i in range(self.stage4_cfg['NUM_BRANCHES']):
+            if self.transition3[i] is not None:
+                x_list.append(self.transition3[i](y_list[-1]))
+            else:
+                x_list.append(y_list[i])
+        x = self.stage4(x_list)
+
+        if self.cfg.DOWNSAMPLE:
+            if self.cfg.USE_CONV:
+                # Downsampling with strided convolutions
+                x1 = self.downsample_stage_1(x[0])
+                x2 = self.downsample_stage_2(x[1])
+                x3 = self.downsample_stage_3(x[2])
+                x = torch.cat([x1, x2, x3, x[3]], 1)
+            else:
+                # Downsampling with interpolation
+                x0_h, x0_w = x[3].size(2), x[3].size(3)
+                x1 = F.interpolate(x[0], size=(x0_h, x0_w), mode='bilinear', align_corners=True)
+                x2 = F.interpolate(x[1], size=(x0_h, x0_w), mode='bilinear', align_corners=True)
+                x3 = F.interpolate(x[2], size=(x0_h, x0_w), mode='bilinear', align_corners=True)
+                x = torch.cat([x1, x2, x3, x[3]], 1)
+        else:
+            if self.cfg.USE_CONV:
+                # Upsampling with interpolations + convolutions
+                x1 = self.upsample_stage_2(x[1])
+                x2 = self.upsample_stage_3(x[2])
+                x3 = self.upsample_stage_4(x[3])
+                x = torch.cat([x[0], x1, x2, x3], 1)
+            else:
+                # Upsampling with interpolation
+                x0_h, x0_w = x[0].size(2), x[0].size(3)
+                x1 = F.interpolate(x[1], size=(x0_h, x0_w), mode='bilinear', align_corners=True)
+                x2 = F.interpolate(x[2], size=(x0_h, x0_w), mode='bilinear', align_corners=True)
+                x3 = F.interpolate(x[3], size=(x0_h, x0_w), mode='bilinear', align_corners=True)
+                x = torch.cat([x[0], x1, x2, x3], 1)
+
+        return x
+
+    def init_weights(self, pretrained=''):
+        logger.info('=> init weights from normal distribution')
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ['bias']:
+                        nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ['bias']:
+                        nn.init.constant_(m.bias, 0)
+
+        if os.path.isfile(pretrained):
+            pretrained_state_dict = torch.load(pretrained)
+            logger.info('=> loading pretrained model {}'.format(pretrained))
+
+            need_init_state_dict = {}
+            for name, m in pretrained_state_dict.items():
+                if name.split('.')[0] in self.pretrained_layers \
+                   or self.pretrained_layers[0] is '*':
+                    need_init_state_dict[name] = m
+            self.load_state_dict(need_init_state_dict, strict=False)
+        elif pretrained:
+            logger.warning('IMPORTANT WARNING!! Please download pre-trained models if you are in TRAINING mode!')
+            # raise ValueError('{} is not exist!'.format(pretrained))
+
+
+def get_pose_net(cfg, is_train):
+    model = PoseHighResolutionNet(cfg)
+
+    if is_train and cfg['MODEL']['INIT_WEIGHTS']:
+        model.init_weights(cfg['MODEL']['PRETRAINED'])
+
+    return model
+
+
+def get_cfg_defaults(pretrained, width=32, downsample=False, use_conv=False):
+    # pose_multi_resoluton_net related params
+    HRNET = CN()
+    HRNET.PRETRAINED_LAYERS = [
+        'conv1', 'bn1', 'conv2', 'bn2', 'layer1', 'transition1',
+        'stage2', 'transition2', 'stage3', 'transition3', 'stage4',
+    ]
+    HRNET.STEM_INPLANES = 64
+    HRNET.FINAL_CONV_KERNEL = 1
+    HRNET.STAGE2 = CN()
+    HRNET.STAGE2.NUM_MODULES = 1
+    HRNET.STAGE2.NUM_BRANCHES = 2
+    HRNET.STAGE2.NUM_BLOCKS = [4, 4]
+    HRNET.STAGE2.NUM_CHANNELS = [width, width*2]
+    HRNET.STAGE2.BLOCK = 'BASIC'
+    HRNET.STAGE2.FUSE_METHOD = 'SUM'
+    HRNET.STAGE3 = CN()
+    HRNET.STAGE3.NUM_MODULES = 4
+    HRNET.STAGE3.NUM_BRANCHES = 3
+    HRNET.STAGE3.NUM_BLOCKS = [4, 4, 4]
+    HRNET.STAGE3.NUM_CHANNELS = [width, width*2, width*4]
+    HRNET.STAGE3.BLOCK = 'BASIC'
+    HRNET.STAGE3.FUSE_METHOD = 'SUM'
+    HRNET.STAGE4 = CN()
+    HRNET.STAGE4.NUM_MODULES = 3
+    HRNET.STAGE4.NUM_BRANCHES = 4
+    HRNET.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]
+    HRNET.STAGE4.NUM_CHANNELS = [width, width*2, width*4, width*8]
+    HRNET.STAGE4.BLOCK = 'BASIC'
+    HRNET.STAGE4.FUSE_METHOD = 'SUM'
+    HRNET.DOWNSAMPLE = downsample
+    HRNET.USE_CONV = use_conv
+
+    cfg = CN()
+    cfg.MODEL = CN()
+    cfg.MODEL.INIT_WEIGHTS = True
+    cfg.MODEL.PRETRAINED = pretrained  # 'data/pretrained_models/hrnet_w32-36af842e.pth'
+    cfg.MODEL.EXTRA = HRNET
+    cfg.MODEL.NUM_JOINTS = 24
+    return cfg
+
+
+def hrnet_w32(
+        pretrained=True,
+        pretrained_ckpt='data/weights/pose_hrnet_w32_256x192.pth',
+        downsample=False,
+        use_conv=False,
+):
+    cfg = get_cfg_defaults(pretrained_ckpt, width=32, downsample=downsample, use_conv=use_conv)
+    return get_pose_net(cfg, is_train=True)
+
+
+def hrnet_w48(
+        pretrained=True,
+        pretrained_ckpt='data/weights/pose_hrnet_w48_256x192.pth',
+        downsample=False,
+        use_conv=False,
+):
+    cfg = get_cfg_defaults(pretrained_ckpt, width=48, downsample=downsample, use_conv=use_conv)
+    return get_pose_net(cfg, is_train=True)

utils/image_utils.py

@@ -0,0 +1,444 @@
+"""
+This file contains functions that are used to perform data augmentation.
+"""
+import cv2
+import torch
+import json
+from skimage.transform import rotate, resize
+import numpy as np
+import jpeg4py as jpeg
+from trimesh.visual import color
+
+# from ..core import constants
+# from .vibe_image_utils import gen_trans_from_patch_cv
+from .kp_utils import map_smpl_to_common, get_smpl_joint_names
+
+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(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 = rotate(new_img, rot) # scipy.misc.imrotate(new_img, rot)
+        new_img = new_img[pad:-pad, pad:-pad]
+
+    # resize image
+    new_img = resize(new_img, res) # scipy.misc.imresize(new_img, res)
+    return new_img
+
+
+def crop_cv2(img, center, scale, res, rot=0):
+    c_x, c_y = center
+    c_x, c_y = int(round(c_x)), int(round(c_y))
+    patch_width, patch_height = int(round(res[0])), int(round(res[1]))
+    bb_width = bb_height = int(round(scale * 200.))
+
+    trans = gen_trans_from_patch_cv(
+        c_x, c_y, bb_width, bb_height,
+        patch_width, patch_height,
+        scale=1.0, rot=rot, inv=False,
+    )
+
+    crop_img = cv2.warpAffine(
+        img, trans, (int(patch_width), int(patch_height)),
+        flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT
+    )
+
+    return crop_img
+
+
+def get_random_crop_coords(height, width, crop_height, crop_width, h_start, w_start):
+    y1 = int((height - crop_height) * h_start)
+    y2 = y1 + crop_height
+    x1 = int((width - crop_width) * w_start)
+    x2 = x1 + crop_width
+    return x1, y1, x2, y2
+
+
+def random_crop(center, scale, crop_scale_factor, axis='all'):
+    '''
+    center: bbox center [x,y]
+    scale: bbox height / 200
+    crop_scale_factor: amount of cropping to be applied
+    axis: axis which cropping will be applied
+        "x": center the y axis and get random crops in x
+        "y": center the x axis and get random crops in y
+        "all": randomly crop from all locations
+    '''
+    orig_size = int(scale * 200.)
+    ul = (center - (orig_size / 2.)).astype(int)
+
+    crop_size = int(orig_size * crop_scale_factor)
+
+    if axis == 'all':
+        h_start = np.random.rand()
+        w_start = np.random.rand()
+    elif axis == 'x':
+        h_start = np.random.rand()
+        w_start = 0.5
+    elif axis == 'y':
+        h_start = 0.5
+        w_start = np.random.rand()
+    else:
+        raise ValueError(f'axis {axis} is undefined!')
+
+    x1, y1, x2, y2 = get_random_crop_coords(
+        height=orig_size,
+        width=orig_size,
+        crop_height=crop_size,
+        crop_width=crop_size,
+        h_start=h_start,
+        w_start=w_start,
+    )
+    scale = (y2 - y1) / 200.
+    center = ul + np.array([(y1 + y2) / 2, (x1 + x2) / 2])
+    return center, scale
+
+
+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 = resize(img, crop_shape) #, interp='nearest') # 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 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
+
+
+def denormalize_images(images):
+    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)
+    return images
+
+
+def read_img(img_fn):
+    #  return pil_img.fromarray(
+                #  cv2.cvtColor(cv2.imread(img_fn), cv2.COLOR_BGR2RGB))
+    #  with open(img_fn, 'rb') as f:
+        #  img = pil_img.open(f).convert('RGB')
+    #  return img
+    if img_fn.endswith('jpeg') or img_fn.endswith('jpg'):
+        try:
+            with open(img_fn, 'rb') as f:
+                img = np.array(jpeg.JPEG(f).decode())
+        except jpeg.JPEGRuntimeError:
+            # logger.warning('{} produced a JPEGRuntimeError', img_fn)
+            img = cv2.cvtColor(cv2.imread(img_fn), cv2.COLOR_BGR2RGB)
+    else:
+    #  elif img_fn.endswith('png') or img_fn.endswith('JPG') or img_fn.endswith(''):
+        img = cv2.cvtColor(cv2.imread(img_fn), cv2.COLOR_BGR2RGB)
+    return img.astype(np.float32)
+
+
+def generate_heatmaps_2d(joints, joints_vis, num_joints=24, heatmap_size=56, image_size=224, sigma=1.75):
+    '''
+    :param joints:  [num_joints, 3]
+    :param joints_vis: [num_joints, 3]
+    :return: target, target_weight(1: visible, 0: invisible)
+    '''
+    target_weight = np.ones((num_joints, 1), dtype=np.float32)
+    target_weight[:, 0] = joints_vis[:, 0]
+
+    target = np.zeros((num_joints, heatmap_size, heatmap_size), dtype=np.float32)
+
+    tmp_size = sigma * 3
+
+    # denormalize joint into heatmap coordinates
+    joints = (joints + 1.) * (image_size / 2.)
+
+    for joint_id in range(num_joints):
+        feat_stride = image_size / heatmap_size
+        mu_x = int(joints[joint_id][0] / feat_stride + 0.5)
+        mu_y = int(joints[joint_id][1] / feat_stride + 0.5)
+        # Check that any part of the gaussian is in-bounds
+        ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
+        br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]
+        if ul[0] >= heatmap_size or ul[1] >= heatmap_size \
+                or br[0] < 0 or br[1] < 0:
+            # If not, just return the image as is
+            target_weight[joint_id] = 0
+            continue
+
+        # # Generate gaussian
+        size = 2 * tmp_size + 1
+        x = np.arange(0, size, 1, np.float32)
+        y = x[:, np.newaxis]
+        x0 = y0 = size // 2
+        # The gaussian is not normalized, we want the center value to equal 1
+        g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
+
+        # Usable gaussian range
+        g_x = max(0, -ul[0]), min(br[0], heatmap_size) - ul[0]
+        g_y = max(0, -ul[1]), min(br[1], heatmap_size) - ul[1]
+        # Image range
+        img_x = max(0, ul[0]), min(br[0], heatmap_size)
+        img_y = max(0, ul[1]), min(br[1], heatmap_size)
+
+        v = target_weight[joint_id]
+        if v > 0.5:
+            target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \
+                g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
+
+    return target, target_weight
+
+
+def generate_part_labels(vertices, faces, cam_t, neural_renderer, body_part_texture, K, R, part_bins):
+    batch_size = vertices.shape[0]
+
+    body_parts, depth, mask = neural_renderer(
+        vertices,
+        faces.expand(batch_size, -1, -1),
+        textures=body_part_texture.expand(batch_size, -1, -1, -1, -1, -1),
+        K=K.expand(batch_size, -1, -1),
+        R=R.expand(batch_size, -1, -1),
+        t=cam_t.unsqueeze(1),
+    )
+
+    render_rgb = body_parts.clone()
+
+    body_parts = body_parts.permute(0, 2, 3, 1)
+    body_parts *= 255. # multiply it with 255 to make labels distant
+    body_parts, _ = body_parts.max(-1) # reduce to single channel
+
+    body_parts = torch.bucketize(body_parts.detach(), part_bins, right=True) # np.digitize(body_parts, bins, right=True)
+
+    # add 1 to make background label 0
+    body_parts = body_parts.long() + 1
+    body_parts = body_parts * mask.detach()
+
+    return body_parts.long(), render_rgb
+
+
+def generate_heatmaps_2d_batch(joints, num_joints=24, heatmap_size=56, image_size=224, sigma=1.75):
+    batch_size = joints.shape[0]
+
+    joints = joints.detach().cpu().numpy()
+    joints_vis = np.ones_like(joints)
+
+    heatmaps = []
+    heatmaps_vis = []
+    for i in range(batch_size):
+        hm, hm_vis = generate_heatmaps_2d(joints[i], joints_vis[i], num_joints, heatmap_size, image_size, sigma)
+        heatmaps.append(hm)
+        heatmaps_vis.append(hm_vis)
+
+    return torch.from_numpy(np.stack(heatmaps)).float().to('cuda'), \
+           torch.from_numpy(np.stack(heatmaps_vis)).float().to('cuda')      
+
+
+def get_body_part_texture(faces, model_type='smpl', non_parametric=False):
+    if model_type == 'smpl':
+        n_vertices = 6890
+        segmentation_path = 'data/smpl_vert_segmentation.json'
+    if model_type == 'smplx':
+        n_vertices = 10475
+        segmentation_path = 'data/smplx_vert_segmentation.json'
+
+    with open(segmentation_path, 'rb') as f:
+        part_segmentation = json.load(f)
+
+    # map all vertex ids to the joint ids
+    joint_names = get_smpl_joint_names()
+    smplx_extra_joint_names = ['leftEye', 'eyeballs', 'rightEye']
+    body_vert_idx = np.zeros((n_vertices), dtype=np.int32) - 1  # -1 for missing label
+    for i, (k, v) in enumerate(part_segmentation.items()):
+        if k in smplx_extra_joint_names and model_type == 'smplx':
+            k = 'head'  # map all extra smplx face joints to head
+        body_joint_idx = joint_names.index(k)
+        body_vert_idx[v] = body_joint_idx
+
+    # pare implementation
+    # import joblib
+    # part_segmentation = joblib.load('data/smpl_partSegmentation_mapping.pkl')
+    # body_vert_idx = part_segmentation['smpl_index']
+
+    n_parts = 24.
+
+    if non_parametric:
+        # reduce the number of body_parts to 14
+        # by mapping some joints to others
+        n_parts = 14.
+        joint_mapping = map_smpl_to_common()
+
+        for jm in joint_mapping:
+            for j in jm[0]:
+                body_vert_idx[body_vert_idx==j] = jm[1]
+
+    vertex_colors = np.ones((n_vertices, 4))
+    vertex_colors[:, :3] = body_vert_idx[..., None]
+
+    vertex_colors = color.to_rgba(vertex_colors)
+    vertex_colors = vertex_colors[:, :3]/255.
+
+    face_colors = vertex_colors[faces].min(axis=1)
+    texture = np.zeros((1, faces.shape[0], 1, 1, 3), dtype=np.float32)
+    # texture[0, :, 0, 0, :] = face_colors[:, :3] / n_parts
+    texture[0, :, 0, 0, :] = face_colors[:, :3]
+
+    vertex_colors = torch.from_numpy(vertex_colors).float()
+    texture = torch.from_numpy(texture).float()
+    return vertex_colors, texture
+
+
+def get_default_camera(focal_length, img_h, img_w, is_cam_batch=False):
+    if not is_cam_batch:
+        K = torch.eye(3)
+        K[0, 0] = focal_length
+        K[1, 1] = focal_length
+        K[2, 2] = 1
+        K[0, 2] = img_w / 2.
+        K[1, 2] = img_h / 2.
+        K = K[None, :, :]
+        R = torch.eye(3)[None, :, :]
+    else:
+        bs = focal_length.shape[0]
+        K = torch.eye(3)[None, :, :].repeat(bs, 1, 1)
+        K[:, 0, 0] = focal_length[:, 0]
+        K[:, 1, 1] = focal_length[:, 1]
+        K[:, 2, 2] = 1
+        K[:, 0, 2] = img_w / 2.
+        K[:, 1, 2] = img_h / 2.
+        R = torch.eye(3)[None, :, :].repeat(bs, 1, 1)
+    return K, R
+
+
+def read_exif_data(img_fname):
+    import PIL.Image
+    import PIL.ExifTags
+
+    img = PIL.Image.open(img_fname)
+    exif_data = img._getexif()
+
+    if exif_data == None:
+        return None
+
+    exif = {
+        PIL.ExifTags.TAGS[k]: v
+        for k, v in exif_data.items()
+        if k in PIL.ExifTags.TAGS
+    }
+    return exif

utils/kp_utils.py

@@ -0,0 +1,1114 @@
+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/loss.py

@@ -0,0 +1,207 @@
+import torch
+import torch.nn as nn
+from common import constants
+from models.smpl import SMPL
+from smplx import SMPLX
+import pickle as pkl
+import numpy as np
+from utils.mesh_utils import save_results_mesh
+from utils.diff_renderer import Pytorch3D
+import os
+import cv2
+
+
+class sem_loss_function(nn.Module):
+    def __init__(self):
+        super(sem_loss_function, self).__init__()
+        self.ce = nn.BCELoss()
+
+    def forward(self, y_true, y_pred):
+        loss = self.ce(y_pred, y_true)
+        return loss
+
+
+class class_loss_function(nn.Module):
+    def __init__(self):
+        super(class_loss_function, self).__init__()
+        self.ce_loss = nn.BCELoss()
+        # self.ce_loss = nn.MultiLabelSoftMarginLoss()
+        # self.ce_loss = nn.MultiLabelMarginLoss()
+
+    def forward(self, y_true, y_pred, valid_mask):
+        # y_true = torch.squeeze(y_true, 1).long()
+        # y_true = torch.squeeze(y_true, 1)
+        # y_pred = torch.squeeze(y_pred, 1)
+        bs = y_true.shape[0]
+        if bs != 1:
+            y_pred = y_pred[valid_mask == 1]
+            y_true = y_true[valid_mask == 1]
+        if len(y_pred) > 0:
+            return self.ce_loss(y_pred, y_true)
+        else:
+            return torch.tensor(0.0).to(y_pred.device)
+
+
+class pixel_anchoring_function(nn.Module):
+    def __init__(self, model_type, device='cuda'):
+        super(pixel_anchoring_function, self).__init__()
+
+        self.device = device
+
+        self.model_type = model_type
+
+        if self.model_type == 'smplx':
+            # load mapping from smpl vertices to smplx vertices
+            mapping_pkl = os.path.join(constants.CONTACT_MAPPING_PATH, "smpl_to_smplx.pkl")
+            with open(mapping_pkl, 'rb') as f:
+                smpl_to_smplx_mapping = pkl.load(f)
+                smpl_to_smplx_mapping = smpl_to_smplx_mapping["matrix"]
+            self.smpl_to_smplx_mapping = torch.from_numpy(smpl_to_smplx_mapping).float().to(self.device)
+
+
+        # Setup the SMPL model
+        if self.model_type == 'smpl':
+            self.n_vertices = 6890
+            self.body_model = SMPL(constants.SMPL_MODEL_DIR).to(self.device)
+        if self.model_type == 'smplx':
+            self.n_vertices = 10475
+            self.body_model = SMPLX(constants.SMPLX_MODEL_DIR,
+                                    num_betas=10,
+                                    use_pca=False).to(self.device)
+        self.body_faces = torch.LongTensor(self.body_model.faces.astype(np.int32)).to(self.device)
+
+        self.ce_loss = nn.BCELoss()
+
+    def get_posed_mesh(self, body_params, debug=False):
+        betas = body_params['betas']
+        pose = body_params['pose']
+        transl = body_params['transl']
+
+        # extra smplx params
+        extra_args = {'jaw_pose': torch.zeros((betas.shape[0], 3)).float().to(self.device),
+                      'leye_pose': torch.zeros((betas.shape[0], 3)).float().to(self.device),
+                      'reye_pose': torch.zeros((betas.shape[0], 3)).float().to(self.device),
+                      'expression': torch.zeros((betas.shape[0], 10)).float().to(self.device),
+                      'left_hand_pose': torch.zeros((betas.shape[0], 45)).float().to(self.device),
+                      'right_hand_pose': torch.zeros((betas.shape[0], 45)).float().to(self.device)}
+
+        smpl_output = self.body_model(betas=betas,
+                                      body_pose=pose[:, 3:],
+                                      global_orient=pose[:, :3],
+                                      pose2rot=True,
+                                      transl=transl,
+                                      **extra_args)
+        smpl_verts = smpl_output.vertices
+        smpl_joints = smpl_output.joints
+
+        if debug:
+            for mesh_i in range(smpl_verts.shape[0]):
+                out_dir = 'temp_meshes'
+                os.makedirs(out_dir, exist_ok=True)
+                out_file = os.path.join(out_dir, f'temp_mesh_{mesh_i:04d}.obj')
+                save_results_mesh(smpl_verts[mesh_i], self.body_model.faces, out_file)
+        return smpl_verts, smpl_joints
+
+
+    def render_batch(self, smpl_verts, cam_k, img_scale_factor, vertex_colors=None, face_textures=None, debug=False):
+
+        bs = smpl_verts.shape[0]
+
+        # Incorporate resizing factor into the camera
+        img_w = 256 # TODO: Remove hardcoding
+        img_h = 256 # TODO: Remove hardcoding
+        focal_length_x = cam_k[:, 0, 0] * img_scale_factor[:, 0]
+        focal_length_y = cam_k[:, 1, 1] * img_scale_factor[:, 1]
+        # convert to float for pytorch3d
+        focal_length_x, focal_length_y = focal_length_x.float(), focal_length_y.float()
+
+        # concatenate focal length
+        focal_length = torch.stack([focal_length_x, focal_length_y], dim=1)
+
+        # Setup renderer
+        renderer = Pytorch3D(img_h=img_h,
+                                  img_w=img_w,
+                                  focal_length=focal_length,
+                                  smpl_faces=self.body_faces,
+                                  texture_mode='deco',
+                                  vertex_colors=vertex_colors,
+                                  face_textures=face_textures,
+                                  is_train=True,
+                                  is_cam_batch=True)
+        front_view = renderer(smpl_verts)
+        if debug:
+            # visualize the front view as images in a temp_image folder
+            for i in range(bs):
+                front_view_rgb = front_view[i, :3, :, :].permute(1, 2, 0).detach().cpu()
+                front_view_mask = front_view[i, 3, :, :].detach().cpu()
+                out_dir = 'temp_images'
+                os.makedirs(out_dir, exist_ok=True)
+                out_file_rgb = os.path.join(out_dir, f'{i:04d}_rgb.png')
+                out_file_mask = os.path.join(out_dir, f'{i:04d}_mask.png')
+                cv2.imwrite(out_file_rgb, front_view_rgb.numpy()*255)
+                cv2.imwrite(out_file_mask, front_view_mask.numpy()*255)
+
+        return front_view
+
+    def paint_contact(self, pred_contact):
+        """
+        Paints the contact vertices on the SMPL mesh
+
+        Args:
+            pred_contact: prbabilities of contact vertices
+
+        Returns:
+            pred_rgb: RGB colors for the contact vertices
+        """
+        bs = pred_contact.shape[0]
+
+        # initialize black and while colors
+        colors = torch.tensor([[0, 0, 0], [1, 1, 1]]).float().to(self.device)
+        colors = torch.unsqueeze(colors, 0).expand(bs, -1, -1)
+
+        # add another dimension to the contact probabilities for inverse probabilities
+        pred_contact = torch.unsqueeze(pred_contact, 2)
+        pred_contact = torch.cat((1 - pred_contact, pred_contact), 2)
+
+        # get pred_rgb colors
+        pred_vert_rgb = torch.bmm(pred_contact, colors)
+        pred_face_rgb = pred_vert_rgb[:, self.body_faces, :][:, :, 0, :] # take the first vertex color
+        pred_face_texture = torch.zeros((bs, self.body_faces.shape[0], 1, 1, 3), dtype=torch.float32).to(self.device)
+        pred_face_texture[:, :, 0, 0, :] = pred_face_rgb
+        return pred_vert_rgb, pred_face_texture
+
+    def forward(self, pred_contact, body_params, cam_k, img_scale_factor, gt_contact_polygon, valid_mask):
+        """
+        Takes predicted contact labels (probabilities), transfers them to the posed mesh and
+        renders to the image. Loss is computed between the rendered contact and the ground truth
+        polygons from HOT.
+
+        Args:
+            pred_contact: predicted contact labels (probabilities)
+            body_params: SMPL parameters in camera coords
+            cam_k: camera intrinsics
+            gt_contact_polygon: ground truth polygons from HOT
+        """
+        # convert pred_contact to smplx
+        bs = pred_contact.shape[0]
+        if self.model_type == 'smplx':
+            smpl_to_smplx_mapping = self.smpl_to_smplx_mapping[None].expand(bs, -1, -1)
+            pred_contact = torch.bmm(smpl_to_smplx_mapping, pred_contact[..., None])
+            pred_contact = pred_contact.squeeze()
+
+        # get the posed mesh
+        smpl_verts, smpl_joints = self.get_posed_mesh(body_params)
+
+        # paint the contact vertices on the mesh
+        vertex_colors, face_textures = self.paint_contact(pred_contact)
+
+        # render the mesh
+        front_view = self.render_batch(smpl_verts, cam_k, img_scale_factor, vertex_colors, face_textures)
+        front_view_rgb = front_view[:, :3, :, :].permute(0, 2, 3, 1)
+        front_view_mask = front_view[:, 3, :, :]
+
+        # compute segmentation loss between rendered contact mask and ground truth contact mask
+        front_view_rgb = front_view_rgb[valid_mask == 1]
+        gt_contact_polygon = gt_contact_polygon[valid_mask == 1]
+        loss = self.ce_loss(front_view_rgb, gt_contact_polygon)
+        return loss, front_view_rgb, front_view_mask

utils/mesh_utils.py

@@ -0,0 +1,6 @@
+import trimesh
+
+def save_results_mesh(vertices, faces, filename):
+  mesh = trimesh.Trimesh(vertices, faces, process=False)
+  mesh.export(filename)
+  print(f'save results to {filename}')

utils/metrics.py

@@ -0,0 +1,106 @@
+import numpy as np
+import torch
+import monai.metrics as metrics
+from common.constants import DIST_MATRIX_PATH
+
+DIST_MATRIX = np.load(DIST_MATRIX_PATH)
+
+def metric(mask, pred, back=True):
+  iou = metrics.compute_meaniou(pred, mask, back, False)
+  iou = iou.mean()
+
+  return iou
+
+def precision_recall_f1score(gt, pred):
+    """
+    Compute precision, recall, and f1
+    """
+
+    # gt = gt.numpy()
+    # pred = pred.numpy()
+
+    precision = torch.zeros(gt.shape[0])
+    recall = torch.zeros(gt.shape[0])
+    f1 = torch.zeros(gt.shape[0])
+    
+    for b in range(gt.shape[0]):
+        tp_num = gt[b, pred[b, :] >= 0.5].sum()
+        precision_denominator = (pred[b, :] >= 0.5).sum()
+        recall_denominator = (gt[b, :]).sum()
+
+        precision_ = tp_num / precision_denominator
+        recall_ = tp_num / recall_denominator
+        if precision_denominator == 0: # if no pred
+            precision_ = 1.
+            recall_ = 0.
+            f1_ = 0.
+        elif recall_denominator == 0: # if no GT
+            precision_ = 0.
+            recall_ = 1.
+            f1_ = 0.
+        elif (precision_ + recall_) <= 1e-10:  # to avoid precision issues
+            precision_= 0.
+            recall_= 0.
+            f1_ = 0.
+        else:
+            f1_ = 2 * precision_ * recall_ / (precision_ + recall_)
+
+        precision[b] = precision_
+        recall[b] = recall_
+        f1[b] = f1_
+
+    # return precision, recall, f1
+    return precision, recall, f1
+
+def acc_precision_recall_f1score(gt, pred):
+    """
+    Compute acc, precision, recall, and f1
+    """
+
+    # gt = gt.numpy()
+    # pred = pred.numpy()
+
+    acc = torch.zeros(gt.shape[0])
+    precision = torch.zeros(gt.shape[0])
+    recall = torch.zeros(gt.shape[0])
+    f1 = torch.zeros(gt.shape[0])
+
+    for b in range(gt.shape[0]):
+        tp_num = gt[b, pred[b, :] >= 0.5].sum()
+        precision_denominator = (pred[b, :] >= 0.5).sum()
+        recall_denominator = (gt[b, :]).sum()
+        tn_num = gt.shape[-1] - precision_denominator - recall_denominator + tp_num
+
+        acc_ = (tp_num + tn_num) / gt.shape[-1]
+        precision_ = tp_num / (precision_denominator + 1e-10)
+        recall_ = tp_num / (recall_denominator + 1e-10)
+        f1_ = 2 * precision_ * recall_ / (precision_ + recall_ + 1e-10)
+
+        acc[b] = acc_
+        precision[b] = precision_
+        recall[b] = recall_
+
+    # return precision, recall, f1
+    return acc, precision, recall, f1
+
+def det_error_metric(pred, gt):
+    
+    gt = gt.detach().cpu()
+    pred = pred.detach().cpu()
+
+    dist_matrix = torch.tensor(DIST_MATRIX)
+
+    false_positive_dist = torch.zeros(gt.shape[0])
+    false_negative_dist = torch.zeros(gt.shape[0])
+    
+    for b in range(gt.shape[0]):
+        gt_columns = dist_matrix[:, gt[b, :]==1] if any(gt[b, :]==1) else dist_matrix
+        error_matrix = gt_columns[pred[b, :] >= 0.5, :] if any(pred[b, :] >= 0.5) else gt_columns
+
+        false_positive_dist_ = error_matrix.min(dim=1)[0].mean()
+        false_negative_dist_ = error_matrix.min(dim=0)[0].mean()
+
+        false_positive_dist[b] = false_positive_dist_
+        false_negative_dist[b] = false_negative_dist_
+
+    return false_positive_dist, false_negative_dist

utils/smpl_uv.py

@@ -0,0 +1,167 @@
+import torch
+import trimesh
+import numpy as np
+import skimage.io as io
+from PIL import Image
+from smplx import SMPL
+from matplotlib import cm as mpl_cm, colors as mpl_colors
+from trimesh.visual.color import face_to_vertex_color, vertex_to_face_color, to_rgba
+
+from common import constants
+from .colorwheel import make_color_wheel_image
+
+
+def get_smpl_uv():
+    uv_obj = 'data/body_models/smpl_uv_20200910/smpl_uv.obj'
+
+    uv_map = []
+    with open(uv_obj) as f:
+        for line in f.readlines():
+            if line.startswith('vt'):
+                coords = [float(x) for x in line.split(' ')[1:]]
+                uv_map.append(coords)
+
+    uv_map = np.array(uv_map)
+
+    return uv_map
+
+
+def show_uv_texture():
+    # image = io.imread('data/body_models/smpl_uv_20200910/smpl_uv_20200910.png')
+    image = make_color_wheel_image(1024, 1024)
+    image = Image.fromarray(image)
+
+    uv = np.load('data/body_models/smpl_uv_20200910/uv_table.npy') # get_smpl_uv()
+    material = trimesh.visual.texture.SimpleMaterial(image=image)
+    tex_visuals = trimesh.visual.TextureVisuals(uv=uv, image=image, material=material)
+
+    smpl = SMPL(constants.SMPL_MODEL_DIR)
+
+    faces = smpl.faces
+    verts = smpl().vertices[0].detach().numpy()
+
+    # assert(len(uv) == len(verts))
+    print(uv.shape)
+    vc = tex_visuals.to_color().vertex_colors
+    fc = trimesh.visual.color.vertex_to_face_color(vc, faces)
+    face_colors = fc.copy()
+    fc = fc.astype(float)
+    vc = vc.astype(float)
+    fc[:,:3] = fc[:,:3] / 255.
+    vc[:,:3] = vc[:,:3] / 255.
+    print(fc[:,:3].max(), fc[:,:3].min(), fc[:,:3].mean())
+    print(vc[:, :3].max(), vc[:, :3].min(), vc[:, :3].mean())
+    np.save('data/body_models/smpl/color_wheel_face_colors.npy', fc)
+    np.save('data/body_models/smpl/color_wheel_vertex_colors.npy', vc)
+    print(fc.shape)
+    mesh = trimesh.Trimesh(verts, faces, validate=True, process=False, face_colors=face_colors)
+    # mesh = trimesh.load('data/body_models/smpl_uv_20200910/smpl_uv.obj', process=False)
+    # mesh.visual = tex_visuals
+
+    # import ipdb; ipdb.set_trace()
+    # print(vc.shape)
+    mesh.show()
+
+
+def show_colored_mesh():
+    cm = mpl_cm.get_cmap('jet')
+    norm_gt = mpl_colors.Normalize()
+
+    smpl = SMPL(constants.SMPL_MODEL_DIR)
+
+    faces = smpl.faces
+    verts = smpl().vertices[0].detach().numpy()
+
+    m = trimesh.Trimesh(verts, faces, process=False)
+
+    mode = 1
+    if mode == 0:
+        # mano_segm_labels = m.triangles_center
+        face_labels = m.triangles_center
+        face_colors = (face_labels - face_labels.min()) / np.ptp(face_labels)
+
+    elif mode == 1:
+        # print(face_labels.shape)
+        face_labels = m.triangles_center
+        face_labels = np.argsort(np.linalg.norm(face_labels, axis=-1))
+        face_colors = np.ones((13776, 4))
+        face_colors[:, 3] = 1.0
+        face_colors[:, :3] = cm(norm_gt(face_labels))[:, :3]
+    elif mode == 2:
+        # breakpoint()
+        fc = np.load('data/body_models/smpl_uv_20200910/data/vertex_texture.npy')[0, :, 0, 0, 0, :]
+        face_colors = np.ones((13776, 4))
+        face_colors[:, :3] = fc
+    mesh = trimesh.Trimesh(verts, faces, process=False, face_colors=face_colors)
+    mesh.show()
+
+
+def get_tenet_texture(mode='smplpix'):
+    # mode = 'smplpix', 'decomr'
+
+    smpl = SMPL(constants.SMPL_MODEL_DIR)
+
+    faces = smpl.faces
+    verts = smpl().vertices[0].detach().numpy()
+
+    m = trimesh.Trimesh(verts, faces, process=False)
+    if mode == 'smplpix':
+        # mano_segm_labels = m.triangles_center
+        face_labels = m.triangles_center
+        face_colors = (face_labels - face_labels.min()) / np.ptp(face_labels)
+        texture = np.zeros((1, faces.shape[0], 1, 1, 1, 3), dtype=np.float32)
+        texture[0, :, 0, 0, 0, :] = face_colors[:, :3]
+        texture = torch.from_numpy(texture).float()
+    elif mode == 'decomr':
+        texture = np.load('data/body_models/smpl_uv_20200910/data/vertex_texture.npy')
+        texture = torch.from_numpy(texture).float()
+    elif mode == 'colorwheel':
+        face_colors = np.load('data/body_models/smpl/color_wheel_face_colors.npy')
+        texture = np.zeros((1, faces.shape[0], 1, 1, 1, 3), dtype=np.float32)
+        texture[0, :, 0, 0, 0, :] = face_colors[:, :3]
+        texture = torch.from_numpy(texture).float()
+    else:
+        raise ValueError(f'{mode} is not defined!')
+
+    return texture
+
+
+def save_tenet_textures(mode='smplpix'):
+    # mode = 'smplpix', 'decomr'
+
+    smpl = SMPL(constants.SMPL_MODEL_DIR)
+
+    faces = smpl.faces
+    verts = smpl().vertices[0].detach().numpy()
+
+    m = trimesh.Trimesh(verts, faces, process=False)
+
+    if mode == 'smplpix':
+        # mano_segm_labels = m.triangles_center
+        face_labels = m.triangles_center
+        face_colors = (face_labels - face_labels.min()) / np.ptp(face_labels)
+        texture = np.zeros((1, faces.shape[0], 1, 1, 1, 3), dtype=np.float32)
+        texture[0, :, 0, 0, 0, :] = face_colors[:, :3]
+        texture = torch.from_numpy(texture).float()
+
+        vert_colors = face_to_vertex_color(m, face_colors).astype(float) / 255.0
+
+    elif mode == 'decomr':
+        texture = np.load('data/body_models/smpl_uv_20200910/data/vertex_texture.npy')
+        texture = torch.from_numpy(texture).float()
+        face_colors = texture[0, :, 0, 0, 0, :]
+        vert_colors = face_to_vertex_color(m, face_colors).astype(float) / 255.0
+
+    elif mode == 'colorwheel':
+        face_colors = np.load('data/body_models/smpl/color_wheel_face_colors.npy')
+        texture = np.zeros((1, faces.shape[0], 1, 1, 1, 3), dtype=np.float32)
+        texture[0, :, 0, 0, 0, :] = face_colors[:, :3]
+        texture = torch.from_numpy(texture).float()
+        face_colors[:, :3] *= 255
+        vert_colors = face_to_vertex_color(m, face_colors).astype(float) / 255.0
+    else:
+        raise ValueError(f'{mode} is not defined!')
+
+    print(vert_colors.shape, vert_colors.max())
+    np.save(f'data/body_models/smpl/{mode}_vertex_colors.npy', vert_colors)
+    return texture