cleared files

inference.py

@@ -1,192 +0,0 @@
-import torch
-import os
-import glob
-import argparse
-import numpy as np
-import cv2
-import PIL.Image as pil_img
-from loguru import logger
-import shutil
-
-import trimesh
-import pyrender
-
-from models.deco import DECO
-from common import constants
-
-os.environ['PYOPENGL_PLATFORM'] = 'egl'
-
-if torch.cuda.is_available():
-    device = torch.device('cuda')
-else:
-    device = torch.device('cpu')
-
-def initiate_model(args):
-    deco_model = DECO('hrnet', True, device)
-
-    logger.info(f'Loading weights from {args.model_path}')
-    checkpoint = torch.load(args.model_path)
-    deco_model.load_state_dict(checkpoint['deco'], strict=True)
-
-    deco_model.eval()
-
-    return deco_model
-
-def render_image(scene, img_res, img=None, viewer=False):
-    '''
-    Render the given pyrender scene and return the image. Can also overlay the mesh on an image.
-    '''
-    if viewer:
-        pyrender.Viewer(scene, use_raymond_lighting=True)
-        return 0
-    else:
-        r = pyrender.OffscreenRenderer(viewport_width=img_res,
-                                       viewport_height=img_res,
-                                       point_size=1.0)
-        color, _ = r.render(scene, flags=pyrender.RenderFlags.RGBA)
-        color = color.astype(np.float32) / 255.0
-
-        if img is not None:
-            valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
-            input_img = img.detach().cpu().numpy()
-            output_img = (color[:, :, :-1] * valid_mask +
-                          (1 - valid_mask) * input_img)
-        else:
-            output_img = color
-        return output_img
-
-def create_scene(mesh, img, focal_length=500, camera_center=250, img_res=500):
-    # Setup the scene
-    scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0],
-                           ambient_light=(0.3, 0.3, 0.3))
-    # add mesh for camera
-    camera_pose = np.eye(4)
-    camera_rotation = np.eye(3, 3)
-    camera_translation = np.array([0., 0, 2.5])
-    camera_pose[:3, :3] = camera_rotation
-    camera_pose[:3, 3] = camera_rotation @ camera_translation
-    pyrencamera = pyrender.camera.IntrinsicsCamera(
-        fx=focal_length, fy=focal_length,
-        cx=camera_center, cy=camera_center)
-    scene.add(pyrencamera, pose=camera_pose)
-    # create and add light
-    light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=1)
-    light_pose = np.eye(4)
-    for lp in [[1, 1, 1], [-1, 1, 1], [1, -1, 1], [-1, -1, 1]]:
-        light_pose[:3, 3] = mesh.vertices.mean(0) + np.array(lp)
-        # out_mesh.vertices.mean(0) + np.array(lp)
-        scene.add(light, pose=light_pose)
-    # add body mesh
-    material = pyrender.MetallicRoughnessMaterial(
-        metallicFactor=0.0,
-        alphaMode='OPAQUE',
-        baseColorFactor=(1.0, 1.0, 0.9, 1.0))
-    mesh_images = []
-
-    # resize input image to fit the mesh image height
-    img_height = img_res
-    img_width = int(img_height * img.shape[1] / img.shape[0])
-    img = cv2.resize(img, (img_width, img_height))
-    mesh_images.append(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
-
-    for sideview_angle in [0, 90, 180, 270]:
-        out_mesh = mesh.copy()
-        rot = trimesh.transformations.rotation_matrix(
-            np.radians(sideview_angle), [0, 1, 0])
-        out_mesh.apply_transform(rot)
-        out_mesh = pyrender.Mesh.from_trimesh(
-            out_mesh,
-            material=material)
-        mesh_pose = np.eye(4)
-        scene.add(out_mesh, pose=mesh_pose, name='mesh')
-        output_img = render_image(scene, img_res)
-        output_img = pil_img.fromarray((output_img * 255).astype(np.uint8))
-        output_img = np.asarray(output_img)[:, :, :3]
-        mesh_images.append(output_img)
-        # delete the previous mesh
-        prev_mesh = scene.get_nodes(name='mesh').pop()
-        scene.remove_node(prev_mesh)
-
-    # show upside down view
-    for topview_angle in [90, 270]:
-        out_mesh = mesh.copy()
-        rot = trimesh.transformations.rotation_matrix(
-            np.radians(topview_angle), [1, 0, 0])
-        out_mesh.apply_transform(rot)
-        out_mesh = pyrender.Mesh.from_trimesh(
-            out_mesh,
-            material=material)
-        mesh_pose = np.eye(4)
-        scene.add(out_mesh, pose=mesh_pose, name='mesh')
-        output_img = render_image(scene, img_res)
-        output_img = pil_img.fromarray((output_img * 255).astype(np.uint8))
-        output_img = np.asarray(output_img)[:, :, :3]
-        mesh_images.append(output_img)
-        # delete the previous mesh
-        prev_mesh = scene.get_nodes(name='mesh').pop()
-        scene.remove_node(prev_mesh)
-
-    # stack images
-    IMG = np.hstack(mesh_images)
-    IMG = pil_img.fromarray(IMG)
-    IMG.thumbnail((3000, 3000))
-    return IMG    
-
-def main(args):
-    if os.path.isdir(args.img_src):
-        images = glob.iglob(args.img_src + '/*', recursive=True)
-    else:
-        images = [args.img_src]
-
-    deco_model = initiate_model(args)
-    
-    smpl_path = os.path.join(constants.SMPL_MODEL_DIR, 'smpl_neutral_tpose.ply')
-    
-    for img_name in images:
-        img = cv2.imread(img_name)
-        img = cv2.resize(img, (256, 256), cv2.INTER_CUBIC)
-        img = img.transpose(2,0,1)/255.0
-        img = img[np.newaxis,:,:,:]
-        img = torch.tensor(img, dtype = torch.float32).to(device)
-
-        cont, _, _ = deco_model(img)
-        cont = cont.detach().cpu().numpy().squeeze()
-        cont_smpl = []
-        for indx, i in enumerate(cont):
-            if i >= 0.5:
-                cont_smpl.append(indx)
-        
-        img = img.detach().cpu().numpy()		
-        img = np.transpose(img[0], (1, 2, 0))		
-        img = img * 255		
-        img = img.astype(np.uint8)
-        
-        contact_smpl = np.zeros((1, 1, 6890))
-        contact_smpl[0][0][cont_smpl] = 1
-
-        body_model_smpl = trimesh.load(smpl_path, process=False)
-        for vert in range(body_model_smpl.visual.vertex_colors.shape[0]):
-            body_model_smpl.visual.vertex_colors[vert] = args.mesh_colour
-        body_model_smpl.visual.vertex_colors[cont_smpl] = args.annot_colour
-
-        rend = create_scene(body_model_smpl, img)
-        os.makedirs(os.path.join(args.out_dir, 'Renders'), exist_ok=True) 
-        rend.save(os.path.join(args.out_dir, 'Renders', os.path.basename(img_name).split('.')[0] + '.png'))
-                  
-        out_dir = os.path.join(args.out_dir, 'Preds', os.path.basename(img_name).split('.')[0])
-        os.makedirs(out_dir, exist_ok=True)          
-
-        logger.info(f'Saving mesh to {out_dir}')
-        shutil.copyfile(img_name, os.path.join(out_dir, os.path.basename(img_name)))
-        body_model_smpl.export(os.path.join(out_dir, 'pred.obj'))
-
-if __name__=='__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--img_src', help='Source of image(s). Can be file or directory', default='./demo_out', type=str)
-    parser.add_argument('--out_dir', help='Where to store images', default='./demo_out', type=str)
-    parser.add_argument('--model_path', help='Path to best model weights', default='./checkpoints/Release_Checkpoint/deco_best.pth', type=str)
-    parser.add_argument('--mesh_colour', help='Colour of the mesh', nargs='+', type=int, default=[130, 130, 130, 255])
-    parser.add_argument('--annot_colour', help='Colour of the mesh', nargs='+', type=int, default=[0, 255, 0, 255])
-    args = parser.parse_args()
-
-    main(args)

tester.py

@@ -1,61 +0,0 @@
-import torch
-from torch.utils.data import DataLoader
-from loguru import logger
-
-from train.trainer_step import TrainStepper
-from train.base_trainer import evaluator
-from data.base_dataset import BaseDataset
-from models.deco import DECO
-from utils.config import parse_args, run_grid_search_experiments
-
-def test(hparams):
-    deco_model = DECO(hparams.TRAINING.ENCODER, hparams.TRAINING.CONTEXT, device)
-    pytorch_total_params = sum(p.numel() for p in deco_model.parameters() if p.requires_grad)
-    print('Total number of trainable parameters: ', pytorch_total_params)
-
-    solver = TrainStepper(deco_model, hparams.TRAINING.CONTEXT, hparams.OPTIMIZER.LR, hparams.TRAINING.LOSS_WEIGHTS, hparams.TRAINING.PAL_LOSS_WEIGHTS, device)
-
-    logger.info(f'Loading weights from {hparams.TRAINING.BEST_MODEL_PATH}')
-    _, _ = solver.load(hparams.TRAINING.BEST_MODEL_PATH)
-    
-    # Run testing
-    for test_loader in val_loaders:
-        dataset_name = test_loader.dataset.dataset
-        test_dict, total_time = evaluator(test_loader, solver, hparams, 0, dataset_name, return_dict=True)
-
-        print('Test Contact Precision: ', test_dict['cont_precision'])
-        print('Test Contact Recall: ', test_dict['cont_recall'])
-        print('Test Contact F1 Score: ', test_dict['cont_f1'])
-        print('Test Contact FP Geo. Error: ', test_dict['fp_geo_err'])
-        print('Test Contact FN Geo. Error: ', test_dict['fn_geo_err'])
-        if hparams.TRAINING.CONTEXT:
-            print('Test Contact Semantic Segmentation IoU: ', test_dict['sem_iou'])
-            print('Test Contact Part Segmentation IoU: ', test_dict['part_iou'])
-        print('\nTime taken per image for evaluation: ', total_time)
-        print('-'*50)
-
-if __name__ == '__main__':
-    args = parse_args()
-    hparams = run_grid_search_experiments(
-        args,
-        script='tester.py',
-        change_wt_name=False
-    )
-
-    if torch.cuda.is_available():
-        device = torch.device('cuda')
-    else:
-        device = torch.device('cpu')
-
-    val_datasets = []
-    for ds in hparams.VALIDATION.DATASETS:
-        if ds in ['rich', 'prox']:
-            val_datasets.append(BaseDataset(ds, 'val', model_type='smplx', normalize=hparams.DATASET.NORMALIZE_IMAGES))
-        elif ds in ['damon']:
-            val_datasets.append(BaseDataset(ds, 'val', model_type='smpl', normalize=hparams.DATASET.NORMALIZE_IMAGES))
-        else:
-            raise ValueError('Dataset not supported')
-
-    val_loaders = [DataLoader(val_dataset, batch_size=hparams.DATASET.BATCH_SIZE, shuffle=False, num_workers=hparams.DATASET.NUM_WORKERS) for val_dataset in val_datasets]
-
-    test(hparams)

train.py

@@ -1,94 +0,0 @@
-import torch
-from torch.utils.data import DataLoader
-import os
-
-from train.trainer_step import TrainStepper
-from train.base_trainer import trainer, evaluator
-from data.base_dataset import BaseDataset
-from data.mixed_dataset import MixedDataset
-from models.deco import DECO
-from utils.config import parse_args, run_grid_search_experiments
-
-def train(hparams):
-    deco_model = DECO(hparams.TRAINING.ENCODER, hparams.TRAINING.CONTEXT, device)
-
-    solver = TrainStepper(deco_model, hparams.TRAINING.CONTEXT, hparams.OPTIMIZER.LR, hparams.TRAINING.LOSS_WEIGHTS, hparams.TRAINING.PAL_LOSS_WEIGHTS, device)
-
-    vb_f1 = 0
-    start_ep = 0
-    num = 0
-    k = True
-    latest_model_path = hparams.TRAINING.BEST_MODEL_PATH.replace('best', 'latest')
-    if os.path.exists(latest_model_path):
-      _, vb_f1 = solver.load(hparams.TRAINING.BEST_MODEL_PATH)
-      start_ep, _ = solver.load(latest_model_path)
-    
-    for epoch in range(start_ep+1, hparams.TRAINING.NUM_EPOCHS + 1):
-        # Train one epoch
-        trainer(epoch, train_loader, solver, hparams)
-        # Run evaluation
-        vc_f1 = None
-        for val_loader in val_loaders:
-            dataset_name = val_loader.dataset.dataset
-            vc_f1_ds = evaluator(val_loader, solver, hparams, epoch, dataset_name, normalize=hparams.DATASET.NORMALIZE_IMAGES)
-            if dataset_name == hparams.VALIDATION.MAIN_DATASET:
-                vc_f1 = vc_f1_ds
-        if vc_f1 is None:
-            raise ValueError('Main dataset not found in validation datasets')
-
-        print('Learning rate: ', solver.lr)
-
-        print('---------------------------------------------')
-        print('---------------------------------------------')
-
-        solver.save(epoch, vc_f1, latest_model_path)
-
-        if epoch % hparams.TRAINING.CHECKPOINT_EPOCHS == 0:
-          inter_model_path = latest_model_path.replace('latest', 'epoch_'+str(epoch).zfill(3))
-          solver.save(epoch, vc_f1, inter_model_path)
-
-        if vc_f1 < vb_f1:
-          num += 1
-          print('Not Saving model: Best Val F1 = ', vb_f1, ' Current Val F1 = ', vc_f1)
-        else:
-          num = 0
-          vb_f1 = vc_f1
-          print('Saving model...')
-          solver.save(epoch, vb_f1, hparams.TRAINING.BEST_MODEL_PATH)
-
-        if num >= hparams.OPTIMIZER.NUM_UPDATE_LR: solver.update_lr()
-        if num >= hparams.TRAINING.NUM_EARLY_STOP:
-          print('Early Stop')
-          k = False
-
-        if k: continue
-        else: break
-
-
-if __name__ == '__main__':
-    args = parse_args()
-    hparams = run_grid_search_experiments(
-        args,
-        script='train.py',
-    )
-
-    if torch.cuda.is_available():
-        device = torch.device('cuda')
-    else:
-        device = torch.device('cpu')
-
-    train_dataset = MixedDataset(hparams.TRAINING.DATASETS, 'train', dataset_mix_pdf=hparams.TRAINING.DATASET_MIX_PDF,  normalize=hparams.DATASET.NORMALIZE_IMAGES)
-
-    val_datasets = []
-    for ds in hparams.VALIDATION.DATASETS:
-        if ds in ['rich', 'prox']:
-            val_datasets.append(BaseDataset(ds, 'val', model_type='smplx', normalize=hparams.DATASET.NORMALIZE_IMAGES))
-        elif ds in ['damon']:
-            val_datasets.append(BaseDataset(ds, 'val', model_type='smpl', normalize=hparams.DATASET.NORMALIZE_IMAGES))
-        else:
-            raise ValueError('Dataset not supported')
-
-    train_loader = DataLoader(train_dataset, hparams.DATASET.BATCH_SIZE, shuffle=True, num_workers=hparams.DATASET.NUM_WORKERS)
-    val_loaders = [DataLoader(val_dataset, batch_size=hparams.DATASET.BATCH_SIZE, shuffle=False, num_workers=hparams.DATASET.NUM_WORKERS) for val_dataset in val_datasets]
-
-    train(hparams)

train/base_trainer.py

@@ -1,103 +0,0 @@
-from tqdm import tqdm
-from utils.metrics import metric, precision_recall_f1score, det_error_metric
-import torch
-import numpy as np
-from vis.visualize import gen_render
-
-
-def trainer(epoch, train_loader, solver, hparams, compute_metrics=False):
-
-    total_epochs = hparams.TRAINING.NUM_EPOCHS
-    print('Training Epoch {}/{}'.format(epoch, total_epochs))
-
-    length = len(train_loader)
-    iterator = tqdm(enumerate(train_loader), total=length, leave=False, desc=f'Training Epoch: {epoch}/{total_epochs}')
-    for step, batch in iterator:
-        losses, output = solver.optimize(batch)
-    return losses, output
-
-@torch.no_grad()
-def evaluator(val_loader, solver, hparams, epoch=0, dataset_name='Unknown', normalize=True, return_dict=False):
-    total_epochs = hparams.TRAINING.NUM_EPOCHS
-
-    batch_size = val_loader.batch_size
-    dataset_size = len(val_loader.dataset)
-    print(f'Dataset size: {dataset_size}')
-
-    val_epoch_cont_pre = np.zeros(dataset_size)
-    val_epoch_cont_rec = np.zeros(dataset_size)
-    val_epoch_cont_f1 = np.zeros(dataset_size)
-    val_epoch_fp_geo_err = np.zeros(dataset_size)
-    val_epoch_fn_geo_err = np.zeros(dataset_size)
-    if hparams.TRAINING.CONTEXT:
-        val_epoch_sem_iou = np.zeros(dataset_size)
-        val_epoch_part_iou = np.zeros(dataset_size)
-
-    val_epoch_cont_loss = np.zeros(dataset_size)
-    
-    total_time = 0
-
-    rend_images = []
-
-    eval_dict = {}
-
-    length = len(val_loader)
-    iterator = tqdm(enumerate(val_loader), total=length, leave=False, desc=f'Evaluating {dataset_name.capitalize()} Epoch: {epoch}/{total_epochs}')
-    for step, batch in iterator:
-        curr_batch_size = batch['img'].shape[0]
-        losses, output, time_taken = solver.evaluate(batch)
-
-        val_epoch_cont_loss[step * batch_size:step * batch_size + curr_batch_size] = losses['cont_loss'].cpu().numpy()
-
-        # compute metrics
-        contact_labels_3d = output['contact_labels_3d_gt']
-        has_contact_3d = output['has_contact_3d']
-        # check if any value in has_contact_3d tensor is 0
-        assert torch.any(has_contact_3d == 0) == False, 'has_contact_3d tensor has 0 values'
-
-        contact_labels_3d_pred = output['contact_labels_3d_pred']
-        if hparams.TRAINING.CONTEXT:
-            sem_mask_gt = output['sem_mask_gt']
-            sem_seg_pred = output['sem_mask_pred']
-            part_mask_gt = output['part_mask_gt']
-            part_seg_pred = output['part_mask_pred']
-
-        cont_pre, cont_rec, cont_f1 = precision_recall_f1score(contact_labels_3d, contact_labels_3d_pred)
-        fp_geo_err, fn_geo_err = det_error_metric(contact_labels_3d_pred, contact_labels_3d)
-        if hparams.TRAINING.CONTEXT:
-            sem_iou = metric(sem_mask_gt, sem_seg_pred)
-            part_iou = metric(part_mask_gt, part_seg_pred)
-
-        val_epoch_cont_pre[step * batch_size:step * batch_size + curr_batch_size] = cont_pre.cpu().numpy()
-        val_epoch_cont_rec[step * batch_size:step * batch_size + curr_batch_size] = cont_rec.cpu().numpy()
-        val_epoch_cont_f1[step * batch_size:step * batch_size + curr_batch_size] = cont_f1.cpu().numpy()
-        val_epoch_fp_geo_err[step * batch_size:step * batch_size + curr_batch_size] = fp_geo_err.cpu().numpy()
-        val_epoch_fn_geo_err[step * batch_size:step * batch_size + curr_batch_size] = fn_geo_err.cpu().numpy()
-        if hparams.TRAINING.CONTEXT:
-            val_epoch_sem_iou[step * batch_size:step * batch_size + curr_batch_size] = sem_iou.cpu().numpy()
-            val_epoch_part_iou[step * batch_size:step * batch_size + curr_batch_size] = part_iou.cpu().numpy()
-        
-        total_time += time_taken
-
-        # logging every summary_steps steps
-        if step % hparams.VALIDATION.SUMMARY_STEPS == 0:
-            if hparams.TRAINING.CONTEXT:
-                rend = gen_render(output, normalize)
-                rend_images.append(rend)
-
-    eval_dict['cont_precision'] = np.sum(val_epoch_cont_pre) / dataset_size
-    eval_dict['cont_recall'] = np.sum(val_epoch_cont_rec) / dataset_size
-    eval_dict['cont_f1'] = np.sum(val_epoch_cont_f1) / dataset_size
-    eval_dict['fp_geo_err'] = np.sum(val_epoch_fp_geo_err) / dataset_size
-    eval_dict['fn_geo_err'] = np.sum(val_epoch_fn_geo_err) / dataset_size
-    if hparams.TRAINING.CONTEXT:
-        eval_dict['sem_iou'] = np.sum(val_epoch_sem_iou) / dataset_size
-        eval_dict['part_iou'] = np.sum(val_epoch_part_iou) / dataset_size
-        eval_dict['images'] = rend_images
-    
-    total_time /= dataset_size
-
-    val_epoch_cont_loss = np.sum(val_epoch_cont_loss) / dataset_size
-    if return_dict:
-        return eval_dict, total_time
-    return eval_dict['cont_f1']

train/trainer_step.py

@@ -1,291 +0,0 @@
-from utils.loss import sem_loss_function, class_loss_function, pixel_anchoring_function
-import torch
-import os
-import time
-
-
-class TrainStepper():
-    def __init__(self, deco_model, context, learning_rate, loss_weight, pal_loss_weight, device):
-        self.device = device
-
-        self.model = deco_model
-        self.context = context
-
-        if self.context:
-            self.optimizer_sem = torch.optim.Adam(params=list(self.model.encoder_sem.parameters()) + list(self.model.decoder_sem.parameters()),
-                                                lr=learning_rate, weight_decay=0.0001)
-            self.optimizer_part = torch.optim.Adam(
-                params=list(self.model.encoder_part.parameters()) + list(self.model.decoder_part.parameters()), lr=learning_rate,
-                weight_decay=0.0001)
-        self.optimizer_contact = torch.optim.Adam(
-            params=list(self.model.encoder_sem.parameters()) + list(self.model.encoder_part.parameters()) + list(
-                self.model.cross_att.parameters()) + list(self.model.classif.parameters()), lr=learning_rate, weight_decay=0.0001)
-
-        if self.context: self.sem_loss = sem_loss_function().to(device)
-        self.class_loss = class_loss_function().to(device)
-        self.pixel_anchoring_loss_smplx = pixel_anchoring_function(model_type='smplx').to(device)
-        self.pixel_anchoring_loss_smpl = pixel_anchoring_function(model_type='smpl').to(device)
-        self.lr = learning_rate
-        self.loss_weight = loss_weight
-        self.pal_loss_weight = pal_loss_weight
-
-    def optimize(self, batch):
-        self.model.train()
-
-        img_paths = batch['img_path']
-        img = batch['img'].to(self.device)
-
-        img_scale_factor = batch['img_scale_factor'].to(self.device)
-
-        pose = batch['pose'].to(self.device)
-        betas = batch['betas'].to(self.device)
-        transl = batch['transl'].to(self.device)
-        has_smpl = batch['has_smpl'].to(self.device)
-        is_smplx = batch['is_smplx'].to(self.device)
-
-        cam_k = batch['cam_k'].to(self.device)
-
-        gt_contact_labels_3d = batch['contact_label_3d'].to(self.device)
-        has_contact_3d = batch['has_contact_3d'].to(self.device)
-
-        if self.context:
-            sem_mask_gt = batch['sem_mask'].to(self.device)
-            part_mask_gt = batch['part_mask'].to(self.device)
-
-        polygon_contact_2d = batch['polygon_contact_2d'].to(self.device)
-        has_polygon_contact_2d = batch['has_polygon_contact_2d'].to(self.device)
-
-        # Forward pass
-        if self.context:
-            cont, sem_mask_pred, part_mask_pred = self.model(img)
-        else:
-            cont = self.model(img)    
-
-        if self.context:
-            loss_sem = self.sem_loss(sem_mask_gt, sem_mask_pred)
-            loss_part = self.sem_loss(part_mask_gt, part_mask_pred)
-        valid_contact_3d = has_contact_3d
-        loss_cont = self.class_loss(gt_contact_labels_3d, cont, valid_contact_3d)
-        valid_polygon_contact_2d = has_polygon_contact_2d
-
-        if self.pal_loss_weight > 0 and (is_smplx == 0).sum() > 0:
-            smpl_body_params = {'pose': pose[is_smplx == 0], 'betas': betas[is_smplx == 0],
-                                'transl': transl[is_smplx == 0],
-                                'has_smpl': has_smpl[is_smplx == 0]}
-            loss_pix_anchoring_smpl, contact_2d_pred_rgb_smpl, _ = self.pixel_anchoring_loss_smpl(cont[is_smplx == 0],
-                                                                                                  smpl_body_params,
-                                                                                                  cam_k[is_smplx == 0],
-                                                                                                  img_scale_factor[
-                                                                                                      is_smplx == 0],
-                                                                                                  polygon_contact_2d[
-                                                                                                      is_smplx == 0],
-                                                                                                  valid_polygon_contact_2d[
-                                                                                                      is_smplx == 0])
-            # weigh the smpl loss based on the number of smpl sample
-            loss_pix_anchoring = loss_pix_anchoring_smpl * (is_smplx == 0).sum() / len(is_smplx)
-            contact_2d_pred_rgb = contact_2d_pred_rgb_smpl
-        else:
-            loss_pix_anchoring = 0
-            contact_2d_pred_rgb = torch.zeros_like(polygon_contact_2d)
-
-        if self.context: loss = loss_sem + loss_part + self.loss_weight * loss_cont + self.pal_loss_weight * loss_pix_anchoring
-        else: loss = self.loss_weight * loss_cont + self.pal_loss_weight * loss_pix_anchoring
-
-        if self.context:
-            self.optimizer_sem.zero_grad()
-            self.optimizer_part.zero_grad()
-        self.optimizer_contact.zero_grad()
-
-        loss.backward()
-
-        if self.context:
-            self.optimizer_sem.step()
-            self.optimizer_part.step()
-        self.optimizer_contact.step()
-
-        if self.context:
-            losses = {'sem_loss': loss_sem,
-                    'part_loss': loss_part,
-                    'cont_loss': loss_cont,
-                    'pal_loss': loss_pix_anchoring,
-                    'total_loss': loss}
-        else:
-            losses = {'cont_loss': loss_cont,
-                    'pal_loss': loss_pix_anchoring,
-                    'total_loss': loss}         
-
-        if self.context:
-            output = {
-                'img': img,
-                'sem_mask_gt': sem_mask_gt,
-                'sem_mask_pred': sem_mask_pred,
-                'part_mask_gt': part_mask_gt,
-                'part_mask_pred': part_mask_pred,
-                'has_contact_2d': has_polygon_contact_2d,
-                'contact_2d_gt': polygon_contact_2d,
-                'contact_2d_pred_rgb': contact_2d_pred_rgb,
-                'has_contact_3d': has_contact_3d,
-                'contact_labels_3d_gt': gt_contact_labels_3d,
-                'contact_labels_3d_pred': cont}
-        else:
-            output = {
-                'img': img,
-                'has_contact_2d': has_polygon_contact_2d,
-                'contact_2d_gt': polygon_contact_2d,
-                'contact_2d_pred_rgb': contact_2d_pred_rgb,
-                'has_contact_3d': has_contact_3d,
-                'contact_labels_3d_gt': gt_contact_labels_3d,
-                'contact_labels_3d_pred': cont}   
-
-        return losses, output
-
-    @torch.no_grad()
-    def evaluate(self, batch):
-        self.model.eval()
-
-        img_paths = batch['img_path']
-        img = batch['img'].to(self.device)
-
-        img_scale_factor = batch['img_scale_factor'].to(self.device)
-
-        pose = batch['pose'].to(self.device)
-        betas = batch['betas'].to(self.device)
-        transl = batch['transl'].to(self.device)
-        has_smpl = batch['has_smpl'].to(self.device)
-        is_smplx = batch['is_smplx'].to(self.device)
-
-        cam_k = batch['cam_k'].to(self.device)
-
-        gt_contact_labels_3d = batch['contact_label_3d'].to(self.device)
-        has_contact_3d = batch['has_contact_3d'].to(self.device)
-
-        if self.context:
-            sem_mask_gt = batch['sem_mask'].to(self.device)
-            part_mask_gt = batch['part_mask'].to(self.device)
-
-        polygon_contact_2d = batch['polygon_contact_2d'].to(self.device)
-        has_polygon_contact_2d = batch['has_polygon_contact_2d'].to(self.device)
-
-        # Forward pass
-        initial_time = time.time()
-        if self.context: cont, sem_mask_pred, part_mask_pred = self.model(img)
-        else: cont = self.model(img)
-        time_taken = time.time() - initial_time
-
-        if self.context:
-            loss_sem = self.sem_loss(sem_mask_gt, sem_mask_pred)
-            loss_part = self.sem_loss(part_mask_gt, part_mask_pred)
-        valid_contact_3d = has_contact_3d
-        loss_cont = self.class_loss(gt_contact_labels_3d, cont, valid_contact_3d)
-        valid_polygon_contact_2d = has_polygon_contact_2d
-
-        if self.pal_loss_weight > 0 and (is_smplx == 0).sum() > 0: # PAL loss only on 2D contacts in HOT which only has SMPL
-            smpl_body_params = {'pose': pose[is_smplx == 0], 'betas': betas[is_smplx == 0], 'transl': transl[is_smplx == 0],
-                                'has_smpl': has_smpl[is_smplx == 0]}
-            loss_pix_anchoring_smpl, contact_2d_pred_rgb_smpl, _ = self.pixel_anchoring_loss_smpl(cont[is_smplx == 0],
-                                                                                                 smpl_body_params,
-                                                                                                 cam_k[is_smplx == 0],
-                                                                                                 img_scale_factor[
-                                                                                                     is_smplx == 0],
-                                                                                                 polygon_contact_2d[
-                                                                                                     is_smplx == 0],
-                                                                                                 valid_polygon_contact_2d[
-                                                                                                     is_smplx == 0])
-            # weight the smpl loss based on the number of smpl samples
-            contact_2d_pred_rgb = contact_2d_pred_rgb_smpl
-            loss_pix_anchoring = loss_pix_anchoring_smpl * (is_smplx == 0).sum() / len(is_smplx)
-        else:
-            loss_pix_anchoring = 0
-            contact_2d_pred_rgb = torch.zeros_like(polygon_contact_2d)
-
-        if self.context: loss = loss_sem + loss_part + self.loss_weight * loss_cont + self.pal_loss_weight * loss_pix_anchoring
-        else: loss = self.loss_weight * loss_cont + self.pal_loss_weight * loss_pix_anchoring
-
-        if self.context:
-            losses = {'sem_loss': loss_sem,
-                    'part_loss': loss_part,
-                    'cont_loss': loss_cont,
-                    'pal_loss': loss_pix_anchoring,
-                    'total_loss': loss}
-        else:
-            losses = {'cont_loss': loss_cont,
-                  'pal_loss': loss_pix_anchoring,
-                  'total_loss': loss}            
-
-        if self.context:
-            output = {
-                'img': img,
-                'sem_mask_gt': sem_mask_gt,
-                'sem_mask_pred': sem_mask_pred,
-                'part_mask_gt': part_mask_gt,
-                'part_mask_pred': part_mask_pred,
-                'has_contact_2d': has_polygon_contact_2d,
-                'contact_2d_gt': polygon_contact_2d,
-                'contact_2d_pred_rgb': contact_2d_pred_rgb,
-                'has_contact_3d': has_contact_3d,
-                'contact_labels_3d_gt': gt_contact_labels_3d,
-                'contact_labels_3d_pred': cont}
-        else:
-            output = {
-                'img': img,
-                'has_contact_2d': has_polygon_contact_2d,
-                'contact_2d_gt': polygon_contact_2d,
-                'contact_2d_pred_rgb': contact_2d_pred_rgb,
-                'has_contact_3d': has_contact_3d,
-                'contact_labels_3d_gt': gt_contact_labels_3d,
-                'contact_labels_3d_pred': cont}        
-
-        return losses, output, time_taken
-
-    def save(self, ep, f1, model_path):
-        # create model directory if it does not exist
-        os.makedirs(os.path.dirname(model_path), exist_ok=True)
-        if self.context:
-            torch.save({
-                'epoch': ep,
-                'deco': self.model.state_dict(),
-                'f1': f1,
-                'sem_optim': self.optimizer_sem.state_dict(),
-                'part_optim': self.optimizer_part.state_dict(),
-                'contact_optim': self.optimizer_contact.state_dict()
-            },
-                model_path)
-        else:
-            torch.save({
-                'epoch': ep,
-                'deco': self.model.state_dict(),
-                'f1': f1,
-                'sem_optim': self.optimizer_sem.state_dict(),
-                'part_optim': self.optimizer_part.state_dict(),
-                'contact_optim': self.optimizer_contact.state_dict()
-            },
-                model_path)    
-
-    def load(self, model_path):
-        print(f'~~~ Loading existing checkpoint from {model_path} ~~~')
-        checkpoint = torch.load(model_path)
-        self.model.load_state_dict(checkpoint['deco'], strict=True)
-
-        if self.context:
-            self.optimizer_sem.load_state_dict(checkpoint['sem_optim'])
-            self.optimizer_part.load_state_dict(checkpoint['part_optim'])
-        self.optimizer_contact.load_state_dict(checkpoint['contact_optim'])
-        epoch = checkpoint['epoch']
-        f1 = checkpoint['f1']
-        return epoch, f1
-
-    def update_lr(self, factor=2):
-        if factor:
-            new_lr = self.lr / factor
-
-        if self.context:
-            self.optimizer_sem = torch.optim.Adam(params=list(self.model.encoder_sem.parameters()) + list(self.model.decoder_sem.parameters()),
-                                                lr=new_lr, weight_decay=0.0001)
-            self.optimizer_part = torch.optim.Adam(
-                params=list(self.model.encoder_part.parameters()) + list(self.model.decoder_part.parameters()), lr=new_lr, weight_decay=0.0001)
-        self.optimizer_contact = torch.optim.Adam(
-            params=list(self.model.encoder_sem.parameters()) + list(self.model.encoder_part.parameters()) + list(
-                self.model.cross_att.parameters()) + list(self.model.classif.parameters()), lr=new_lr, weight_decay=0.0001)
-
-        print('update learning rate: %f -> %f' % (self.lr, new_lr))
-        self.lr = new_lr

utils/cluster.py

@@ -1,99 +0,0 @@
-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

@@ -1,22 +0,0 @@
-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

@@ -1,196 +0,0 @@
-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

@@ -1,45 +0,0 @@
-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

@@ -1,287 +0,0 @@
-# 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

@@ -1,17 +0,0 @@
-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

@@ -1,625 +0,0 @@
-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

@@ -1,444 +0,0 @@
-"""
-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

@@ -1,1114 +0,0 @@
-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

@@ -1,207 +0,0 @@
-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

@@ -1,6 +0,0 @@
-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

@@ -1,106 +0,0 @@
-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

@@ -1,167 +0,0 @@
-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

vis/visualize.py

@@ -1,209 +0,0 @@
-import cv2
-import os
-import trimesh
-import PIL.Image as pil_img
-import numpy as np
-import pyrender
-from common import constants
-
-os.environ['PYOPENGL_PLATFORM'] = 'egl'
-
-def render_image(scene, img_res, img=None, viewer=False):
-    '''
-    Render the given pyrender scene and return the image. Can also overlay the mesh on an image.
-    '''
-    if viewer:
-        pyrender.Viewer(scene, use_raymond_lighting=True)
-        return 0
-    else:
-        r = pyrender.OffscreenRenderer(viewport_width=img_res,
-                                       viewport_height=img_res,
-                                       point_size=1.0)
-        color, _ = r.render(scene, flags=pyrender.RenderFlags.RGBA)
-        color = color.astype(np.float32) / 255.0
-
-        if img is not None:
-            valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
-            input_img = img.detach().cpu().numpy()
-            output_img = (color[:, :, :-1] * valid_mask +
-                          (1 - valid_mask) * input_img)
-        else:
-            output_img = color
-        return output_img
-
-def create_scene(mesh, img, focal_length=500, camera_center=250, img_res=500):
-    # Setup the scene
-    scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0],
-                           ambient_light=(0.3, 0.3, 0.3))
-    # add mesh for camera
-    camera_pose = np.eye(4)
-    camera_rotation = np.eye(3, 3)
-    camera_translation = np.array([0., 0, 2.5])
-    camera_pose[:3, :3] = camera_rotation
-    camera_pose[:3, 3] = camera_rotation @ camera_translation
-    pyrencamera = pyrender.camera.IntrinsicsCamera(
-        fx=focal_length, fy=focal_length,
-        cx=camera_center, cy=camera_center)
-    scene.add(pyrencamera, pose=camera_pose)
-    # create and add light
-    light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=1)
-    light_pose = np.eye(4)
-    for lp in [[1, 1, 1], [-1, 1, 1], [1, -1, 1], [-1, -1, 1]]:
-        light_pose[:3, 3] = mesh.vertices.mean(0) + np.array(lp)
-        # out_mesh.vertices.mean(0) + np.array(lp)
-        scene.add(light, pose=light_pose)
-    # add body mesh
-    material = pyrender.MetallicRoughnessMaterial(
-        metallicFactor=0.0,
-        alphaMode='OPAQUE',
-        baseColorFactor=(1.0, 1.0, 0.9, 1.0))
-    mesh_images = []
-
-    # resize input image to fit the mesh image height
-    # print(img.shape)
-    img_height = img_res
-    img_width = int(img_height * img.shape[1] / img.shape[0])
-    img = cv2.resize(img, (img_width, img_height))
-    mesh_images.append(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
-
-    for sideview_angle in [0, 90, 180, 270]:
-        out_mesh = mesh.copy()
-        rot = trimesh.transformations.rotation_matrix(
-            np.radians(sideview_angle), [0, 1, 0])
-        out_mesh.apply_transform(rot)
-        out_mesh = pyrender.Mesh.from_trimesh(
-            out_mesh,
-            material=material)
-        mesh_pose = np.eye(4)
-        scene.add(out_mesh, pose=mesh_pose, name='mesh')
-        output_img = render_image(scene, img_res)
-        output_img = pil_img.fromarray((output_img * 255).astype(np.uint8))
-        output_img = np.asarray(output_img)[:, :, :3]
-        mesh_images.append(output_img)
-        # delete the previous mesh
-        prev_mesh = scene.get_nodes(name='mesh').pop()
-        scene.remove_node(prev_mesh)
-
-    # show upside down view
-    for topview_angle in [90, 270]:
-        out_mesh = mesh.copy()
-        rot = trimesh.transformations.rotation_matrix(
-            np.radians(topview_angle), [1, 0, 0])
-        out_mesh.apply_transform(rot)
-        out_mesh = pyrender.Mesh.from_trimesh(
-            out_mesh,
-            material=material)
-        mesh_pose = np.eye(4)
-        scene.add(out_mesh, pose=mesh_pose, name='mesh')
-        output_img = render_image(scene, img_res)
-        output_img = pil_img.fromarray((output_img * 255).astype(np.uint8))
-        output_img = np.asarray(output_img)[:, :, :3]
-        mesh_images.append(output_img)
-        # delete the previous mesh
-        prev_mesh = scene.get_nodes(name='mesh').pop()
-        scene.remove_node(prev_mesh)
-
-    # stack images
-    IMG = np.hstack(mesh_images)
-    IMG = pil_img.fromarray(IMG)
-    IMG.thumbnail((3000, 3000))
-    return IMG
-
-# img = cv2.imread('../samples/prox_N3OpenArea_03301_01_s001_frame_00694.jpg')
-# mesh = trimesh.load('../samples/mesh.ply', process=False)
-# comb_img = create_scene(mesh, img)
-# comb_img.save('../samples/combined_image.png')
-
-def unsplit(img, palette):
-    rgb_img = np.zeros((img.shape[0], img.shape[1], 3))
-    for i in range(img.shape[0]):
-        for j in range(img.shape[1]):
-            id = np.argmax(img[i, j, :])
-            rgb_img[i, j, :] = palette[id]
-
-    return rgb_img
-
-def gen_render(output, normalize=True):
-    img = output['img'].cpu().numpy()
-    contact_labels_3d = output['contact_labels_3d_gt'].cpu().numpy()
-    contact_labels_3d_pred = output['contact_labels_3d_pred'].cpu().numpy()
-    sem_mask_gt = output['sem_mask_gt'].cpu().numpy()
-    sem_mask_pred = output['sem_mask_pred'].cpu().numpy()
-    part_mask_gt = output['part_mask_gt'].cpu().numpy()
-    part_mask_pred = output['part_mask_pred'].cpu().numpy()
-    contact_2d_gt_rgb = output['contact_2d_gt'].cpu().numpy()
-    contact_2d_pred_rgb = output['contact_2d_pred_rgb'].cpu().numpy()
-
-    mesh_path = './data/smpl/smpl_neutral_tpose.ply'
-    gt_mesh = trimesh.load(mesh_path, process=False)
-    pred_mesh = trimesh.load(mesh_path, process=False)
-
-    img = np.transpose(img[0], (1, 2, 0))
-    if normalize:
-        # unnormalize the image before displaying
-        mean = np.array(constants.IMG_NORM_MEAN, dtype=np.float32)
-        std = np.array(constants.IMG_NORM_STD, dtype=np.float32)
-        img = img * std + mean
-    img = img * 255
-    img = img.astype(np.uint8)
-    color = np.array([0, 0, 0, 255])
-    th = 0.5
-
-    contact_labels_3d = contact_labels_3d[0, :]
-    for vid, val in enumerate(contact_labels_3d):
-        if val >= th:
-            gt_mesh.visual.vertex_colors[vid] = color
-
-    contact_labels_3d_pred = contact_labels_3d_pred[0, :]
-    for vid, val in enumerate(contact_labels_3d_pred):
-        if val >= th:
-            pred_mesh.visual.vertex_colors[vid] = color
-
-    gt_rend = create_scene(gt_mesh, img)
-    pred_rend = create_scene(pred_mesh, img)
-
-    sem_palette = [[220, 20, 60], [119, 11, 32], [0, 0, 142], [0, 0, 230], [106, 0, 228], [0, 60, 100], [0, 80, 100], [0, 0, 70], [0, 0, 192], [250, 170, 30], [100, 170, 30], [220, 220, 0], [175, 116, 175], [250, 0, 30], [165, 42, 42], [255, 77, 255], [0, 226, 252], [182, 182, 255], [0, 82, 0], [120, 166, 157], [110, 76, 0], [174, 57, 255], [199, 100, 0], [72, 0, 118], [255, 179, 240], [0, 125, 92], [209, 0, 151], [188, 208, 182], [0, 220, 176], [255, 99, 164], [92, 0, 73], [133, 129, 255], [78, 180, 255], [0, 228, 0], [174, 255, 243], [45, 89, 255], [134, 134, 103], [145, 148, 174], [255, 208, 186], [197, 226, 255], [171, 134, 1], [109, 63, 54], [207, 138, 255], [151, 0, 95], [9, 80, 61], [84, 105, 51], [74, 65, 105], [166, 196, 102], [208, 195, 210], [255, 109, 65], [0, 143, 149], [179, 0, 194], [209, 99, 106], [5, 121, 0], [227, 255, 205], [147, 186, 208], [153, 69, 1], [3, 95, 161], [163, 255, 0], [119, 0, 170], [0, 182, 199], [0, 165, 120], [183, 130, 88], [95, 32, 0], [130, 114, 135], [110, 129, 133], [166, 74, 118], [219, 142, 185], [79, 210, 114], [178, 90, 62], [65, 70, 15], [127, 167, 115], [59, 105, 106], [142, 108, 45], [196, 172, 0], [95, 54, 80], [128, 76, 255], [201, 57, 1], [246, 0, 122], [191, 162, 208], [255, 255, 128], [147, 211, 203], [150, 100, 100], [168, 171, 172], [146, 112, 198], [210, 170, 100], [92, 136, 89], [218, 88, 184], [241, 129, 0], [217, 17, 255], [124, 74, 181], [70, 70, 70], [255, 228, 255], [154, 208, 0], [193, 0, 92], [76, 91, 113], [255, 180, 195], [106, 154, 176], [230, 150, 140], [60, 143, 255], [128, 64, 128], [92, 82, 55], [254, 212, 124], [73, 77, 174], [255, 160, 98], [255, 255, 255], [104, 84, 109], [169, 164, 131], [225, 199, 255], [137, 54, 74], [135, 158, 223], [7, 246, 231], [107, 255, 200], [58, 41, 149], [183, 121, 142], [255, 73, 97], [107, 142, 35], [190, 153, 153], [146, 139, 141], [70, 130, 180], [134, 199, 156], [209, 226, 140], [96, 36, 108], [96, 96, 96], [64, 170, 64], [152, 251, 152], [208, 229, 228], [206, 186, 171], [152, 161, 64], [116, 112, 0], [0, 114, 143], [102, 102, 156], [250, 141, 255]]
-    # part_palette = [(0,0,0), (128,0,0), (255,0,0), (0,85,0), (170,0,51), (255,85,0), (0,0,85), (0,119,221), (85,85,0), (0,85,85), (85,51,0), (52,86,128), (0,128,0), (0,0,255), (51,170,221), (0,255,255), (85,255,170), (170,255,85), (255,255,0), (255,170,0)]
-    part_palette = [[0, 0, 0], [220, 20, 60], [119, 11, 32], [0, 0, 142], [0, 0, 230], [106, 0, 228], [0, 60, 100], [0, 80, 100], [0, 0, 70], [0, 0, 192], [250, 170, 30], [100, 170, 30], [220, 220, 0], [175, 116, 175], [250, 0, 30], [165, 42, 42], [255, 77, 255], [0, 226, 252], [182, 182, 255], [0, 82, 0], [120, 166, 157], [110, 76, 0], [174, 57, 255], [199, 100, 0], [72, 0, 118], [255, 179, 240]]
-    hot_palette = [[0, 0, 0], [220, 20, 60], [119, 11, 32], [0, 0, 142], [0, 0, 230], [106, 0, 228], [0, 60, 100], [0, 80, 100], [0, 0, 70], [0, 0, 192], [250, 170, 30], [100, 170, 30], [220, 220, 0], [175, 116, 175], [250, 0, 30], [165, 42, 42], [255, 77, 255], [0, 226, 252]]
-
-    sem_mask_gt = np.transpose(sem_mask_gt[0], (1, 2, 0))*255
-    sem_mask_gt = sem_mask_gt.astype(np.uint8)
-    sem_mask_pred = np.transpose(sem_mask_pred[0], (1, 2, 0))*255
-    sem_mask_pred = sem_mask_pred.astype(np.uint8)
-    part_mask_gt = np.transpose(part_mask_gt[0], (1, 2, 0))*255
-    part_mask_gt = part_mask_gt.astype(np.uint8)
-    part_mask_pred = np.transpose(part_mask_pred[0], (1, 2, 0))*255
-    part_mask_pred = part_mask_pred.astype(np.uint8)
-    contact_2d_gt_rgb = contact_2d_gt_rgb[0]*255
-    contact_2d_gt_rgb = contact_2d_gt_rgb.astype(np.uint8)
-    contact_2d_pred_rgb = contact_2d_pred_rgb[0]*255
-    contact_2d_pred_rgb = contact_2d_pred_rgb.astype(np.uint8)
-
-    sem_mask_rgb = unsplit(sem_mask_gt, sem_palette)
-    sem_pred_rgb = unsplit(sem_mask_pred, sem_palette)
-    part_mask_rgb = unsplit(part_mask_gt, part_palette)
-    part_pred_rgb = unsplit(part_mask_pred, part_palette)
-
-    sem_mask_rgb = sem_mask_rgb.astype(np.uint8)
-    sem_pred_rgb = sem_pred_rgb.astype(np.uint8)
-    part_mask_rgb = part_mask_rgb.astype(np.uint8)
-    part_pred_rgb = part_pred_rgb.astype(np.uint8)
-
-    sem_mask_rgb = pil_img.fromarray(sem_mask_rgb)
-    sem_pred_rgb = pil_img.fromarray(sem_pred_rgb)
-    part_mask_rgb = pil_img.fromarray(part_mask_rgb)
-    part_pred_rgb = pil_img.fromarray(part_pred_rgb)
-    contact_2d_gt_rgb = pil_img.fromarray(contact_2d_gt_rgb)
-    contact_2d_pred_rgb = pil_img.fromarray(contact_2d_pred_rgb)
-
-    tot_rend = pil_img.new('RGB', (3000, 2000))
-    tot_rend.paste(gt_rend, (0, 0))
-    tot_rend.paste(pred_rend, (0, 450))
-    tot_rend.paste(sem_mask_rgb, (0, 900))
-    tot_rend.paste(sem_pred_rgb, (400, 900))
-    tot_rend.paste(part_mask_rgb, (0, 1300))
-    tot_rend.paste(part_pred_rgb, (400, 1300))
-    tot_rend.paste(contact_2d_gt_rgb, (0, 1700))
-    tot_rend.paste(contact_2d_pred_rgb, (400, 1700))
-    return tot_rend