diff --git a/lib / pymaf / configs / pymaf_config.yaml b/lib / pymaf / configs / pymaf_config.yaml deleted file mode 100644 index 5e4bb7fff6145ebd520dc3db51b0acede783129d..0000000000000000000000000000000000000000 --- a/lib / pymaf / configs / pymaf_config.yaml +++ /dev/null @@ -1,47 +0,0 @@ -SOLVER: - MAX_ITER: 500000 - TYPE: Adam - BASE_LR: 0.00005 - GAMMA: 0.1 - STEPS: [0] - EPOCHS: [0] -DEBUG: False -LOGDIR: '' -DEVICE: cuda -NUM_WORKERS: 8 -SEED_VALUE: -1 -LOSS: - KP_2D_W: 300.0 - KP_3D_W: 300.0 - SHAPE_W: 0.06 - POSE_W: 60.0 - VERT_W: 0.0 - INDEX_WEIGHTS: 2.0 - # Loss weights for surface parts. (24 Parts) - PART_WEIGHTS: 0.3 - # Loss weights for UV regression. - POINT_REGRESSION_WEIGHTS: 0.5 -TRAIN: - NUM_WORKERS: 8 - BATCH_SIZE: 64 - PIN_MEMORY: True -TEST: - BATCH_SIZE: 32 -MODEL: - PyMAF: - BACKBONE: 'res50' - MLP_DIM: [256, 128, 64, 5] - N_ITER: 3 - AUX_SUPV_ON: True - DP_HEATMAP_SIZE: 56 -RES_MODEL: - DECONV_WITH_BIAS: False - NUM_DECONV_LAYERS: 3 - NUM_DECONV_FILTERS: - - 256 - - 256 - - 256 - NUM_DECONV_KERNELS: - - 4 - - 4 - - 4 \ No newline at end of file diff --git a/lib / pymaf /core / __init__.py b/lib / pymaf /core / __init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/lib / pymaf /core / train_options.py b/lib / pymaf /core / train_options.py deleted file mode 100644 index 15657f6d55ae39f93896c149e290f54d1a594538..0000000000000000000000000000000000000000 --- a/lib / pymaf /core / train_options.py +++ /dev/null @@ -1,135 +0,0 @@ -import argparse - - -class TrainOptions(): - def __init__(self): - self.parser = argparse.ArgumentParser() - - gen = self.parser.add_argument_group('General') - gen.add_argument( - '--resume', - dest='resume', - default=False, - action='store_true', - help='Resume from checkpoint (Use latest checkpoint by default') - - io = self.parser.add_argument_group('io') - io.add_argument('--log_dir', - default='logs', - help='Directory to store logs') - io.add_argument( - '--pretrained_checkpoint', - default=None, - help='Load a pretrained checkpoint at the beginning training') - - train = self.parser.add_argument_group('Training Options') - train.add_argument('--num_epochs', - type=int, - default=200, - help='Total number of training epochs') - train.add_argument('--regressor', - type=str, - choices=['hmr', 'pymaf_net'], - default='pymaf_net', - help='Name of the SMPL regressor.') - train.add_argument('--cfg_file', - type=str, - default='./configs/pymaf_config.yaml', - help='config file path for PyMAF.') - train.add_argument( - '--img_res', - type=int, - default=224, - help='Rescale bounding boxes to size [img_res, img_res] before feeding them in the network' - ) - train.add_argument( - '--rot_factor', - type=float, - default=30, - help='Random rotation in the range [-rot_factor, rot_factor]') - train.add_argument( - '--noise_factor', - type=float, - default=0.4, - help='Randomly multiply pixel values with factor in the range [1-noise_factor, 1+noise_factor]' - ) - train.add_argument( - '--scale_factor', - type=float, - default=0.25, - help='Rescale bounding boxes by a factor of [1-scale_factor,1+scale_factor]' - ) - train.add_argument( - '--openpose_train_weight', - default=0., - help='Weight for OpenPose keypoints during training') - train.add_argument('--gt_train_weight', - default=1., - help='Weight for GT keypoints during training') - train.add_argument('--eval_dataset', - type=str, - default='h36m-p2-mosh', - help='Name of the evaluation dataset.') - train.add_argument('--single_dataset', - default=False, - action='store_true', - help='Use a single dataset') - train.add_argument('--single_dataname', - type=str, - default='h36m', - help='Name of the single dataset.') - train.add_argument('--eval_pve', - default=False, - action='store_true', - help='evaluate PVE') - train.add_argument('--overwrite', - default=False, - action='store_true', - help='overwrite the latest checkpoint') - - train.add_argument('--distributed', - action='store_true', - help='Use distributed training') - train.add_argument('--dist_backend', - default='nccl', - type=str, - help='distributed backend') - train.add_argument('--dist_url', - default='tcp://127.0.0.1:10356', - type=str, - help='url used to set up distributed training') - train.add_argument('--world_size', - default=1, - type=int, - help='number of nodes for distributed training') - train.add_argument("--local_rank", default=0, type=int) - train.add_argument('--rank', - default=0, - type=int, - help='node rank for distributed training') - train.add_argument( - '--multiprocessing_distributed', - action='store_true', - help='Use multi-processing distributed training to launch ' - 'N processes per node, which has N GPUs. This is the ' - 'fastest way to use PyTorch for either single node or ' - 'multi node data parallel training') - - misc = self.parser.add_argument_group('Misc Options') - misc.add_argument('--misc', - help="Modify config options using the command-line", - default=None, - nargs=argparse.REMAINDER) - return - - def parse_args(self): - """Parse input arguments.""" - self.args = self.parser.parse_args() - self.save_dump() - return self.args - - def save_dump(self): - """Store all argument values to a json file. - The default location is logs/expname/args.json. - """ - pass \ No newline at end of file diff --git a/lib / pymaf /core /base_trainer.py b/lib / pymaf /core /base_trainer.py deleted file mode 100644 index f06856c7038fbd652a3dc077a872da66d3585dec..0000000000000000000000000000000000000000 --- a/lib / pymaf /core /base_trainer.py +++ /dev/null @@ -1,107 +0,0 @@ -# This script is borrowed and extended from https://github.com/nkolot/SPIN/blob/master/utils/base_trainer.py -from __future__ import division -import logging -from utils import CheckpointSaver -from tensorboardX import SummaryWriter - -import torch -from tqdm import tqdm - -tqdm.monitor_interval = 0 - - -logger = logging.getLogger(__name__) - - -class BaseTrainer(object): - """Base class for Trainer objects. - Takes care of checkpointing/logging/resuming training. - """ - - def __init__(self, options): - self.options = options - if options.multiprocessing_distributed: - self.device = torch.device('cuda', options.gpu) - else: - self.device = torch.device( - 'cuda' if torch.cuda.is_available() else 'cpu') - # override this function to define your model, optimizers etc. - self.saver = CheckpointSaver(save_dir=options.checkpoint_dir, - overwrite=options.overwrite) - if options.rank == 0: - self.summary_writer = SummaryWriter(self.options.summary_dir) - self.init_fn() - - self.checkpoint = None - if options.resume and self.saver.exists_checkpoint(): - self.checkpoint = self.saver.load_checkpoint( - self.models_dict, self.optimizers_dict) - - if self.checkpoint is None: - self.epoch_count = 0 - self.step_count = 0 - else: - self.epoch_count = self.checkpoint['epoch'] - self.step_count = self.checkpoint['total_step_count'] - - if self.checkpoint is not None: - self.checkpoint_batch_idx = self.checkpoint['batch_idx'] - else: - self.checkpoint_batch_idx = 0 - - self.best_performance = float('inf') - - def load_pretrained(self, checkpoint_file=None): - """Load a pretrained checkpoint. - This is different from resuming training using --resume. - """ - if checkpoint_file is not None: - checkpoint = torch.load(checkpoint_file) - for model in self.models_dict: - if model in checkpoint: - self.models_dict[model].load_state_dict(checkpoint[model], - strict=True) - print(f'Checkpoint {model} loaded') - - def move_dict_to_device(self, dict, device, tensor2float=False): - for k, v in dict.items(): - if isinstance(v, torch.Tensor): - if tensor2float: - dict[k] = v.float().to(device) - else: - dict[k] = v.to(device) - - # The following methods (with the possible exception of test) have to be implemented in the derived classes - def train(self, epoch): - raise NotImplementedError('You need to provide an train method') - - def init_fn(self): - raise NotImplementedError('You need to provide an _init_fn method') - - def train_step(self, input_batch): - raise NotImplementedError('You need to provide a _train_step method') - - def train_summaries(self, input_batch): - raise NotImplementedError( - 'You need to provide a _train_summaries method') - - def visualize(self, input_batch): - raise NotImplementedError('You need to provide a visualize method') - - def validate(self): - pass - - def test(self): - pass - - def evaluate(self): - pass - - def fit(self): - # Run training for num_epochs epochs - for epoch in tqdm(range(self.epoch_count, self.options.num_epochs), - total=self.options.num_epochs, - initial=self.epoch_count): - self.epoch_count = epoch - self.train(epoch) - return \ No newline at end of file diff --git a/lib / pymaf /core /cfgs.py b/lib / pymaf /core /cfgs.py deleted file mode 100644 index 09ac4fa48483aa9e595b7e4b27dfa7426cb11d33..0000000000000000000000000000000000000000 --- a/lib / pymaf /core /cfgs.py +++ /dev/null @@ -1,100 +0,0 @@ -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import os -import json -from yacs.config import CfgNode as CN - -# Configuration variables -cfg = CN(new_allowed=True) - -cfg.OUTPUT_DIR = 'results' -cfg.DEVICE = 'cuda' -cfg.DEBUG = False -cfg.LOGDIR = '' -cfg.VAL_VIS_BATCH_FREQ = 200 -cfg.TRAIN_VIS_ITER_FERQ = 1000 -cfg.SEED_VALUE = -1 - -cfg.TRAIN = CN(new_allowed=True) - -cfg.LOSS = CN(new_allowed=True) -cfg.LOSS.KP_2D_W = 300.0 -cfg.LOSS.KP_3D_W = 300.0 -cfg.LOSS.SHAPE_W = 0.06 -cfg.LOSS.POSE_W = 60.0 -cfg.LOSS.VERT_W = 0.0 - -# Loss weights for dense correspondences -cfg.LOSS.INDEX_WEIGHTS = 2.0 -# Loss weights for surface parts. (24 Parts) -cfg.LOSS.PART_WEIGHTS = 0.3 -# Loss weights for UV regression. -cfg.LOSS.POINT_REGRESSION_WEIGHTS = 0.5 - -cfg.MODEL = CN(new_allowed=True) - -cfg.MODEL.PyMAF = CN(new_allowed=True) - -# switch -cfg.TRAIN.VAL_LOOP = True - -cfg.TEST = CN(new_allowed=True) - - -def get_cfg_defaults(): - """Get a yacs CfgNode 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 cfg.clone() - return cfg - - -def update_cfg(cfg_file): - # cfg = get_cfg_defaults() - cfg.merge_from_file(cfg_file) - # return cfg.clone() - return cfg - - -def parse_args(args): - cfg_file = args.cfg_file - if args.cfg_file is not None: - cfg = update_cfg(args.cfg_file) - else: - cfg = get_cfg_defaults() - - # if args.misc is not None: - # cfg.merge_from_list(args.misc) - - return cfg - - -def parse_args_extend(args): - if args.resume: - if not os.path.exists(args.log_dir): - raise ValueError( - 'Experiment are set to resume mode, but log directory does not exist.' - ) - - # load log's cfg - cfg_file = os.path.join(args.log_dir, 'cfg.yaml') - cfg = update_cfg(cfg_file) - - if args.misc is not None: - cfg.merge_from_list(args.misc) - else: - parse_args(args) diff --git a/lib / pymaf /core /constants.py b/lib / pymaf /core /constants.py deleted file mode 100644 index 69f507432f3adfd73164a15d03f86f02c3db5287..0000000000000000000000000000000000000000 --- a/lib / pymaf /core /constants.py +++ /dev/null @@ -1,153 +0,0 @@ -# This script is borrowed and extended from https://github.com/nkolot/SPIN/blob/master/constants.py -FOCAL_LENGTH = 5000. -IMG_RES = 224 - -# Mean and standard deviation for normalizing input image -IMG_NORM_MEAN = [0.485, 0.456, 0.406] -IMG_NORM_STD = [0.229, 0.224, 0.225] -""" -We create a superset of joints containing the OpenPose joints together with the ones that each dataset provides. -We keep a superset of 24 joints such that we include all joints from every dataset. -If a dataset doesn't provide annotations for a specific joint, we simply ignore it. -The joints used here are the following: -""" -JOINT_NAMES = [ - # 25 OpenPose joints (in the order provided by OpenPose) - 'OP Nose', - 'OP Neck', - 'OP RShoulder', - 'OP RElbow', - 'OP RWrist', - 'OP LShoulder', - 'OP LElbow', - 'OP LWrist', - 'OP MidHip', - 'OP RHip', - 'OP RKnee', - 'OP RAnkle', - 'OP LHip', - 'OP LKnee', - 'OP LAnkle', - 'OP REye', - 'OP LEye', - 'OP REar', - 'OP LEar', - 'OP LBigToe', - 'OP LSmallToe', - 'OP LHeel', - 'OP RBigToe', - 'OP RSmallToe', - 'OP RHeel', - # 24 Ground Truth joints (superset of joints from different datasets) - 'Right Ankle', - 'Right Knee', - 'Right Hip', # 2 - 'Left Hip', - 'Left Knee', # 4 - 'Left Ankle', - 'Right Wrist', # 6 - 'Right Elbow', - 'Right Shoulder', # 8 - 'Left Shoulder', - 'Left Elbow', # 10 - 'Left Wrist', - 'Neck (LSP)', # 12 - 'Top of Head (LSP)', - 'Pelvis (MPII)', # 14 - 'Thorax (MPII)', - 'Spine (H36M)', # 16 - 'Jaw (H36M)', - 'Head (H36M)', # 18 - 'Nose', - 'Left Eye', - 'Right Eye', - 'Left Ear', - 'Right Ear' -] - -# Dict containing the joints in numerical order -JOINT_IDS = {JOINT_NAMES[i]: i for i in range(len(JOINT_NAMES))} - -# Map joints to SMPL joints -JOINT_MAP = { - 'OP Nose': 24, - 'OP Neck': 12, - 'OP RShoulder': 17, - 'OP RElbow': 19, - 'OP RWrist': 21, - 'OP LShoulder': 16, - 'OP LElbow': 18, - 'OP LWrist': 20, - 'OP MidHip': 0, - 'OP RHip': 2, - 'OP RKnee': 5, - 'OP RAnkle': 8, - 'OP LHip': 1, - 'OP LKnee': 4, - 'OP LAnkle': 7, - 'OP REye': 25, - 'OP LEye': 26, - 'OP REar': 27, - 'OP LEar': 28, - 'OP LBigToe': 29, - 'OP LSmallToe': 30, - 'OP LHeel': 31, - 'OP RBigToe': 32, - 'OP RSmallToe': 33, - 'OP RHeel': 34, - 'Right Ankle': 8, - 'Right Knee': 5, - 'Right Hip': 45, - 'Left Hip': 46, - 'Left Knee': 4, - 'Left Ankle': 7, - 'Right Wrist': 21, - 'Right Elbow': 19, - 'Right Shoulder': 17, - 'Left Shoulder': 16, - 'Left Elbow': 18, - 'Left Wrist': 20, - 'Neck (LSP)': 47, - 'Top of Head (LSP)': 48, - 'Pelvis (MPII)': 49, - 'Thorax (MPII)': 50, - 'Spine (H36M)': 51, - 'Jaw (H36M)': 52, - 'Head (H36M)': 53, - 'Nose': 24, - 'Left Eye': 26, - 'Right Eye': 25, - 'Left Ear': 28, - 'Right Ear': 27 -} - -# Joint selectors -# Indices to get the 14 LSP joints from the 17 H36M joints -H36M_TO_J17 = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9] -H36M_TO_J14 = H36M_TO_J17[:14] -# Indices to get the 14 LSP joints from the ground truth joints -J24_TO_J17 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 18, 14, 16, 17] -J24_TO_J14 = J24_TO_J17[:14] -J24_TO_J19 = J24_TO_J17[:14] + [19, 20, 21, 22, 23] -J24_TO_JCOCO = [19, 20, 21, 22, 23, 9, 8, 10, 7, 11, 6, 3, 2, 4, 1, 5, 0] - -# Permutation of SMPL pose parameters when flipping the shape -SMPL_JOINTS_FLIP_PERM = [ - 0, 2, 1, 3, 5, 4, 6, 8, 7, 9, 11, 10, 12, 14, 13, 15, 17, 16, 19, 18, 21, - 20, 23, 22 -] -SMPL_POSE_FLIP_PERM = [] -for i in SMPL_JOINTS_FLIP_PERM: - SMPL_POSE_FLIP_PERM.append(3 * i) - SMPL_POSE_FLIP_PERM.append(3 * i + 1) - SMPL_POSE_FLIP_PERM.append(3 * i + 2) -# Permutation indices for the 24 ground truth joints -J24_FLIP_PERM = [ - 5, 4, 3, 2, 1, 0, 11, 10, 9, 8, 7, 6, 12, 13, 14, 15, 16, 17, 18, 19, 21, - 20, 23, 22 -] -# Permutation indices for the full set of 49 joints -J49_FLIP_PERM = [0, 1, 5, 6, 7, 2, 3, 4, 8, 12, 13, 14, 9, 10, 11, 16, 15, 18, 17, 22, 23, 24, 19, 20, 21]\ - + [25+i for i in J24_FLIP_PERM] -SMPL_J49_FLIP_PERM = [0, 1, 5, 6, 7, 2, 3, 4, 8, 12, 13, 14, 9, 10, 11, 16, 15, 18, 17, 22, 23, 24, 19, 20, 21]\ - + [25+i for i in SMPL_JOINTS_FLIP_PERM] \ No newline at end of file diff --git a/lib / pymaf /core /fits_dict.py b/lib / pymaf /core /fits_dict.py deleted file mode 100644 index 43a3415dd0395ac3d481a3497b0397a60bc5ef90..0000000000000000000000000000000000000000 --- a/lib / pymaf /core /fits_dict.py +++ /dev/null @@ -1,133 +0,0 @@ -''' -This script is borrowed and extended from https://github.com/nkolot/SPIN/blob/master/train/fits_dict.py -''' -import os -import cv2 -import torch -import numpy as np -from torchgeometry import angle_axis_to_rotation_matrix, rotation_matrix_to_angle_axis - -from core import path_config, constants - -import logging - -logger = logging.getLogger(__name__) - - -class FitsDict(): - """ Dictionary keeping track of the best fit per image in the training set """ - - def __init__(self, options, train_dataset): - self.options = options - self.train_dataset = train_dataset - self.fits_dict = {} - self.valid_fit_state = {} - # array used to flip SMPL pose parameters - self.flipped_parts = torch.tensor(constants.SMPL_POSE_FLIP_PERM, - dtype=torch.int64) - # Load dictionary state - for ds_name, ds in train_dataset.dataset_dict.items(): - if ds_name in ['h36m']: - dict_file = os.path.join(path_config.FINAL_FITS_DIR, - ds_name + '.npy') - self.fits_dict[ds_name] = torch.from_numpy(np.load(dict_file)) - self.valid_fit_state[ds_name] = torch.ones(len( - self.fits_dict[ds_name]), - dtype=torch.uint8) - else: - dict_file = os.path.join(path_config.FINAL_FITS_DIR, - ds_name + '.npz') - fits_dict = np.load(dict_file) - opt_pose = torch.from_numpy(fits_dict['pose']) - opt_betas = torch.from_numpy(fits_dict['betas']) - opt_valid_fit = torch.from_numpy(fits_dict['valid_fit']).to( - torch.uint8) - self.fits_dict[ds_name] = torch.cat([opt_pose, opt_betas], - dim=1) - self.valid_fit_state[ds_name] = opt_valid_fit - - if not options.single_dataset: - for ds in train_dataset.datasets: - if ds.dataset not in ['h36m']: - ds.pose = self.fits_dict[ds.dataset][:, :72].numpy() - ds.betas = self.fits_dict[ds.dataset][:, 72:].numpy() - ds.has_smpl = self.valid_fit_state[ds.dataset].numpy() - - def save(self): - """ Save dictionary state to disk """ - for ds_name in self.train_dataset.dataset_dict.keys(): - dict_file = os.path.join(self.options.checkpoint_dir, - ds_name + '_fits.npy') - np.save(dict_file, self.fits_dict[ds_name].cpu().numpy()) - - def __getitem__(self, x): - """ Retrieve dictionary entries """ - dataset_name, ind, rot, is_flipped = x - batch_size = len(dataset_name) - pose = torch.zeros((batch_size, 72)) - betas = torch.zeros((batch_size, 10)) - for ds, i, n in zip(dataset_name, ind, range(batch_size)): - params = self.fits_dict[ds][i] - pose[n, :] = params[:72] - betas[n, :] = params[72:] - pose = pose.clone() - # Apply flipping and rotation - pose = self.flip_pose(self.rotate_pose(pose, rot), is_flipped) - betas = betas.clone() - return pose, betas - - def get_vaild_state(self, dataset_name, ind): - batch_size = len(dataset_name) - valid_fit = torch.zeros(batch_size, dtype=torch.uint8) - for ds, i, n in zip(dataset_name, ind, range(batch_size)): - valid_fit[n] = self.valid_fit_state[ds][i] - valid_fit = valid_fit.clone() - return valid_fit - - def __setitem__(self, x, val): - """ Update dictionary entries """ - dataset_name, ind, rot, is_flipped, update = x - pose, betas = val - batch_size = len(dataset_name) - # Undo flipping and rotation - pose = self.rotate_pose(self.flip_pose(pose, is_flipped), -rot) - params = torch.cat((pose, betas), dim=-1).cpu() - for ds, i, n in zip(dataset_name, ind, range(batch_size)): - if update[n]: - self.fits_dict[ds][i] = params[n] - - def flip_pose(self, pose, is_flipped): - """flip SMPL pose parameters""" - is_flipped = is_flipped.byte() - pose_f = pose.clone() - pose_f[is_flipped, :] = pose[is_flipped][:, self.flipped_parts] - # we also negate the second and the third dimension of the axis-angle representation - pose_f[is_flipped, 1::3] *= -1 - pose_f[is_flipped, 2::3] *= -1 - return pose_f - - def rotate_pose(self, pose, rot): - """Rotate SMPL pose parameters by rot degrees""" - pose = pose.clone() - cos = torch.cos(-np.pi * rot / 180.) - sin = torch.sin(-np.pi * rot / 180.) - zeros = torch.zeros_like(cos) - r3 = torch.zeros(cos.shape[0], 1, 3, device=cos.device) - r3[:, 0, -1] = 1 - R = torch.cat([ - torch.stack([cos, -sin, zeros], dim=-1).unsqueeze(1), - torch.stack([sin, cos, zeros], dim=-1).unsqueeze(1), r3 - ], - dim=1) - global_pose = pose[:, :3] - global_pose_rotmat = angle_axis_to_rotation_matrix(global_pose) - global_pose_rotmat_3b3 = global_pose_rotmat[:, :3, :3] - global_pose_rotmat_3b3 = torch.matmul(R, global_pose_rotmat_3b3) - global_pose_rotmat[:, :3, :3] = global_pose_rotmat_3b3 - global_pose_rotmat = global_pose_rotmat[:, :-1, :-1].cpu().numpy() - global_pose_np = np.zeros((global_pose.shape[0], 3)) - for i in range(global_pose.shape[0]): - aa, _ = cv2.Rodrigues(global_pose_rotmat[i]) - global_pose_np[i, :] = aa.squeeze() - pose[:, :3] = torch.from_numpy(global_pose_np).to(pose.device) - return pose \ No newline at end of file diff --git a/lib / pymaf /core /path_config.py b/lib / pymaf /core /path_config.py deleted file mode 100644 index a053b972c515f83133c62b753eb8d31f1c16588a..0000000000000000000000000000000000000000 --- a/lib / pymaf /core /path_config.py +++ /dev/null @@ -1,24 +0,0 @@ -""" -This script is borrowed and extended from https://github.com/nkolot/SPIN/blob/master/path_config.py -path configuration -This file contains definitions of useful data stuctures and the paths -for the datasets and data files necessary to run the code. -Things you need to change: *_ROOT that indicate the path to each dataset -""" -import os -from huggingface_hub import hf_hub_url, cached_download - -# pymaf -pymaf_data_dir = hf_hub_url('Yuliang/PyMAF', '') -smpl_data_dir = hf_hub_url('Yuliang/SMPL', '') -SMPL_MODEL_DIR = os.path.join(smpl_data_dir, 'models/smpl') - -SMPL_MEAN_PARAMS = cached_download(os.path.join(pymaf_data_dir, 'smpl_mean_params.npz'), use_auth_token=os.environ['ICON']) -MESH_DOWNSAMPLEING = cached_download(os.path.join(pymaf_data_dir, 'mesh_downsampling.npz'), use_auth_token=os.environ['ICON']) -CUBE_PARTS_FILE = cached_download(os.path.join(pymaf_data_dir, 'cube_parts.npy'), use_auth_token=os.environ['ICON']) -JOINT_REGRESSOR_TRAIN_EXTRA = cached_download(os.path.join(pymaf_data_dir, 'J_regressor_extra.npy'), use_auth_token=os.environ['ICON']) -JOINT_REGRESSOR_H36M = cached_download(os.path.join(pymaf_data_dir, 'J_regressor_h36m.npy'), use_auth_token=os.environ['ICON']) -VERTEX_TEXTURE_FILE = cached_download(os.path.join(pymaf_data_dir, 'vertex_texture.npy'), use_auth_token=os.environ['ICON']) -SMPL_MEAN_PARAMS = cached_download(os.path.join(pymaf_data_dir, 'smpl_mean_params.npz'), use_auth_token=os.environ['ICON']) -CHECKPOINT_FILE = cached_download(os.path.join(pymaf_data_dir, 'pretrained_model/PyMAF_model_checkpoint.pt'), use_auth_token=os.environ['ICON']) - diff --git a/lib / pymaf /models / __init__.py b/lib / pymaf /models / __init__.py deleted file mode 100644 index af3f32ff51b1640179b30dfe22c4e3cf50c7c98c..0000000000000000000000000000000000000000 --- a/lib / pymaf /models / __init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .hmr import hmr -from .pymaf_net import pymaf_net -from .smpl import SMPL \ No newline at end of file diff --git a/lib / pymaf /models / pymaf_net.py b/lib / pymaf /models / pymaf_net.py deleted file mode 100644 index 2807abaa3c7da0be6913d2fd68cb0ad1721e2bf1..0000000000000000000000000000000000000000 --- a/lib / pymaf /models / pymaf_net.py +++ /dev/null @@ -1,362 +0,0 @@ -import torch -import torch.nn as nn -import numpy as np - -from lib.pymaf.utils.geometry import rot6d_to_rotmat, projection, rotation_matrix_to_angle_axis -from .maf_extractor import MAF_Extractor -from .smpl import SMPL, SMPL_MODEL_DIR, SMPL_MEAN_PARAMS, H36M_TO_J14 -from .hmr import ResNet_Backbone -from .res_module import IUV_predict_layer -from lib.common.config import cfg -import logging - -logger = logging.getLogger(__name__) - -BN_MOMENTUM = 0.1 - - -class Regressor(nn.Module): - def __init__(self, feat_dim, smpl_mean_params): - super().__init__() - - npose = 24 * 6 - - self.fc1 = nn.Linear(feat_dim + npose + 13, 1024) - self.drop1 = nn.Dropout() - self.fc2 = nn.Linear(1024, 1024) - self.drop2 = nn.Dropout() - self.decpose = nn.Linear(1024, npose) - self.decshape = nn.Linear(1024, 10) - self.deccam = nn.Linear(1024, 3) - nn.init.xavier_uniform_(self.decpose.weight, gain=0.01) - nn.init.xavier_uniform_(self.decshape.weight, gain=0.01) - nn.init.xavier_uniform_(self.deccam.weight, gain=0.01) - - self.smpl = SMPL(SMPL_MODEL_DIR, batch_size=64, create_transl=False) - - mean_params = np.load(smpl_mean_params) - init_pose = torch.from_numpy(mean_params['pose'][:]).unsqueeze(0) - init_shape = torch.from_numpy( - mean_params['shape'][:].astype('float32')).unsqueeze(0) - init_cam = torch.from_numpy(mean_params['cam']).unsqueeze(0) - self.register_buffer('init_pose', init_pose) - self.register_buffer('init_shape', init_shape) - self.register_buffer('init_cam', init_cam) - - def forward(self, - x, - init_pose=None, - init_shape=None, - init_cam=None, - n_iter=1, - J_regressor=None): - batch_size = x.shape[0] - - if init_pose is None: - init_pose = self.init_pose.expand(batch_size, -1) - if init_shape is None: - init_shape = self.init_shape.expand(batch_size, -1) - if init_cam is None: - init_cam = self.init_cam.expand(batch_size, -1) - - pred_pose = init_pose - pred_shape = init_shape - pred_cam = init_cam - for i in range(n_iter): - xc = torch.cat([x, pred_pose, pred_shape, pred_cam], 1) - xc = self.fc1(xc) - xc = self.drop1(xc) - xc = self.fc2(xc) - xc = self.drop2(xc) - pred_pose = self.decpose(xc) + pred_pose - pred_shape = self.decshape(xc) + pred_shape - pred_cam = self.deccam(xc) + pred_cam - - pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3) - - pred_output = self.smpl(betas=pred_shape, - body_pose=pred_rotmat[:, 1:], - global_orient=pred_rotmat[:, 0].unsqueeze(1), - pose2rot=False) - - pred_vertices = pred_output.vertices - pred_joints = pred_output.joints - pred_smpl_joints = pred_output.smpl_joints - pred_keypoints_2d = projection(pred_joints, pred_cam) - pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, - 3)).reshape( - -1, 72) - - if J_regressor is not None: - pred_joints = torch.matmul(J_regressor, pred_vertices) - pred_pelvis = pred_joints[:, [0], :].clone() - pred_joints = pred_joints[:, H36M_TO_J14, :] - pred_joints = pred_joints - pred_pelvis - - output = { - 'theta': torch.cat([pred_cam, pred_shape, pose], dim=1), - 'verts': pred_vertices, - 'kp_2d': pred_keypoints_2d, - 'kp_3d': pred_joints, - 'smpl_kp_3d': pred_smpl_joints, - 'rotmat': pred_rotmat, - 'pred_cam': pred_cam, - 'pred_shape': pred_shape, - 'pred_pose': pred_pose, - } - return output - - def forward_init(self, - x, - init_pose=None, - init_shape=None, - init_cam=None, - n_iter=1, - J_regressor=None): - batch_size = x.shape[0] - - if init_pose is None: - init_pose = self.init_pose.expand(batch_size, -1) - if init_shape is None: - init_shape = self.init_shape.expand(batch_size, -1) - if init_cam is None: - init_cam = self.init_cam.expand(batch_size, -1) - - pred_pose = init_pose - pred_shape = init_shape - pred_cam = init_cam - - pred_rotmat = rot6d_to_rotmat(pred_pose.contiguous()).view( - batch_size, 24, 3, 3) - - pred_output = self.smpl(betas=pred_shape, - body_pose=pred_rotmat[:, 1:], - global_orient=pred_rotmat[:, 0].unsqueeze(1), - pose2rot=False) - - pred_vertices = pred_output.vertices - pred_joints = pred_output.joints - pred_smpl_joints = pred_output.smpl_joints - pred_keypoints_2d = projection(pred_joints, pred_cam) - pose = rotation_matrix_to_angle_axis(pred_rotmat.reshape(-1, 3, - 3)).reshape( - -1, 72) - - if J_regressor is not None: - pred_joints = torch.matmul(J_regressor, pred_vertices) - pred_pelvis = pred_joints[:, [0], :].clone() - pred_joints = pred_joints[:, H36M_TO_J14, :] - pred_joints = pred_joints - pred_pelvis - - output = { - 'theta': torch.cat([pred_cam, pred_shape, pose], dim=1), - 'verts': pred_vertices, - 'kp_2d': pred_keypoints_2d, - 'kp_3d': pred_joints, - 'smpl_kp_3d': pred_smpl_joints, - 'rotmat': pred_rotmat, - 'pred_cam': pred_cam, - 'pred_shape': pred_shape, - 'pred_pose': pred_pose, - } - return output - - -class PyMAF(nn.Module): - """ PyMAF based Deep Regressor for Human Mesh Recovery - PyMAF: 3D Human Pose and Shape Regression with Pyramidal Mesh Alignment Feedback Loop, in ICCV, 2021 - """ - - def __init__(self, smpl_mean_params=SMPL_MEAN_PARAMS, pretrained=True): - super().__init__() - self.feature_extractor = ResNet_Backbone( - model=cfg.MODEL.PyMAF.BACKBONE, pretrained=pretrained) - - # deconv layers - self.inplanes = self.feature_extractor.inplanes - self.deconv_with_bias = cfg.RES_MODEL.DECONV_WITH_BIAS - self.deconv_layers = self._make_deconv_layer( - cfg.RES_MODEL.NUM_DECONV_LAYERS, - cfg.RES_MODEL.NUM_DECONV_FILTERS, - cfg.RES_MODEL.NUM_DECONV_KERNELS, - ) - - self.maf_extractor = nn.ModuleList() - for _ in range(cfg.MODEL.PyMAF.N_ITER): - self.maf_extractor.append(MAF_Extractor()) - ma_feat_len = self.maf_extractor[-1].Dmap.shape[ - 0] * cfg.MODEL.PyMAF.MLP_DIM[-1] - - grid_size = 21 - xv, yv = torch.meshgrid([ - torch.linspace(-1, 1, grid_size), - torch.linspace(-1, 1, grid_size) - ]) - points_grid = torch.stack([xv.reshape(-1), - yv.reshape(-1)]).unsqueeze(0) - self.register_buffer('points_grid', points_grid) - grid_feat_len = grid_size * grid_size * cfg.MODEL.PyMAF.MLP_DIM[-1] - - self.regressor = nn.ModuleList() - for i in range(cfg.MODEL.PyMAF.N_ITER): - if i == 0: - ref_infeat_dim = grid_feat_len - else: - ref_infeat_dim = ma_feat_len - self.regressor.append( - Regressor(feat_dim=ref_infeat_dim, - smpl_mean_params=smpl_mean_params)) - - dp_feat_dim = 256 - self.with_uv = cfg.LOSS.POINT_REGRESSION_WEIGHTS > 0 - if cfg.MODEL.PyMAF.AUX_SUPV_ON: - self.dp_head = IUV_predict_layer(feat_dim=dp_feat_dim) - - 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), - ) - - 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_deconv_layer(self, num_layers, num_filters, num_kernels): - """ - Deconv_layer used in Simple Baselines: - Xiao et al. Simple Baselines for Human Pose Estimation and Tracking - https://github.com/microsoft/human-pose-estimation.pytorch - """ - assert num_layers == len(num_filters), \ - 'ERROR: num_deconv_layers is different len(num_deconv_filters)' - assert num_layers == len(num_kernels), \ - 'ERROR: num_deconv_layers is different len(num_deconv_filters)' - - def _get_deconv_cfg(deconv_kernel, index): - if deconv_kernel == 4: - padding = 1 - output_padding = 0 - elif deconv_kernel == 3: - padding = 1 - output_padding = 1 - elif deconv_kernel == 2: - padding = 0 - output_padding = 0 - - return deconv_kernel, padding, output_padding - - layers = [] - for i in range(num_layers): - kernel, padding, output_padding = _get_deconv_cfg( - num_kernels[i], i) - - planes = num_filters[i] - layers.append( - nn.ConvTranspose2d(in_channels=self.inplanes, - out_channels=planes, - kernel_size=kernel, - stride=2, - padding=padding, - output_padding=output_padding, - bias=self.deconv_with_bias)) - layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)) - layers.append(nn.ReLU(inplace=True)) - self.inplanes = planes - - return nn.Sequential(*layers) - - def forward(self, x, J_regressor=None): - - batch_size = x.shape[0] - - # spatial features and global features - s_feat, g_feat = self.feature_extractor(x) - - assert cfg.MODEL.PyMAF.N_ITER >= 0 and cfg.MODEL.PyMAF.N_ITER <= 3 - if cfg.MODEL.PyMAF.N_ITER == 1: - deconv_blocks = [self.deconv_layers] - elif cfg.MODEL.PyMAF.N_ITER == 2: - deconv_blocks = [self.deconv_layers[0:6], self.deconv_layers[6:9]] - elif cfg.MODEL.PyMAF.N_ITER == 3: - deconv_blocks = [ - self.deconv_layers[0:3], self.deconv_layers[3:6], - self.deconv_layers[6:9] - ] - - out_list = {} - - # initial parameters - # TODO: remove the initial mesh generation during forward to reduce runtime - # by generating initial mesh the beforehand: smpl_output = self.init_smpl - smpl_output = self.regressor[0].forward_init(g_feat, - J_regressor=J_regressor) - - out_list['smpl_out'] = [smpl_output] - out_list['dp_out'] = [] - - # for visulization - vis_feat_list = [s_feat.detach()] - - # parameter predictions - for rf_i in range(cfg.MODEL.PyMAF.N_ITER): - pred_cam = smpl_output['pred_cam'] - pred_shape = smpl_output['pred_shape'] - pred_pose = smpl_output['pred_pose'] - - pred_cam = pred_cam.detach() - pred_shape = pred_shape.detach() - pred_pose = pred_pose.detach() - - s_feat_i = deconv_blocks[rf_i](s_feat) - s_feat = s_feat_i - vis_feat_list.append(s_feat_i.detach()) - - self.maf_extractor[rf_i].im_feat = s_feat_i - self.maf_extractor[rf_i].cam = pred_cam - - if rf_i == 0: - sample_points = torch.transpose( - self.points_grid.expand(batch_size, -1, -1), 1, 2) - ref_feature = self.maf_extractor[rf_i].sampling(sample_points) - else: - pred_smpl_verts = smpl_output['verts'].detach() - # TODO: use a more sparse SMPL implementation (with 431 vertices) for acceleration - pred_smpl_verts_ds = torch.matmul( - self.maf_extractor[rf_i].Dmap.unsqueeze(0), - pred_smpl_verts) # [B, 431, 3] - ref_feature = self.maf_extractor[rf_i]( - pred_smpl_verts_ds) # [B, 431 * n_feat] - - smpl_output = self.regressor[rf_i](ref_feature, - pred_pose, - pred_shape, - pred_cam, - n_iter=1, - J_regressor=J_regressor) - out_list['smpl_out'].append(smpl_output) - - if self.training and cfg.MODEL.PyMAF.AUX_SUPV_ON: - iuv_out_dict = self.dp_head(s_feat) - out_list['dp_out'].append(iuv_out_dict) - - return out_list - - -def pymaf_net(smpl_mean_params, pretrained=True): - """ Constructs an PyMAF model with ResNet50 backbone. - Args: - pretrained (bool): If True, returns a model pre-trained on ImageNet - """ - model = PyMAF(smpl_mean_params, pretrained) - return model diff --git a/lib / pymaf /models / smpl.py b/lib / pymaf /models / smpl.py deleted file mode 100644 index ad0059acc3d88d7bf13d6bca25ed9da1b82bb5fe..0000000000000000000000000000000000000000 --- a/lib / pymaf /models / smpl.py +++ /dev/null @@ -1,92 +0,0 @@ -# This script is borrowed from https://github.com/nkolot/SPIN/blob/master/models/smpl.py - -import torch -import numpy as np -from lib.smplx import SMPL as _SMPL -from lib.smplx.body_models import ModelOutput -from lib.smplx.lbs import vertices2joints -from collections import namedtuple - -from lib.pymaf.core import path_config, constants - -SMPL_MEAN_PARAMS = path_config.SMPL_MEAN_PARAMS -SMPL_MODEL_DIR = path_config.SMPL_MODEL_DIR - -# Indices to get the 14 LSP joints from the 17 H36M joints -H36M_TO_J17 = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9] -H36M_TO_J14 = H36M_TO_J17[:14] - - -class SMPL(_SMPL): - """ Extension of the official SMPL implementation to support more joints """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - joints = [constants.JOINT_MAP[i] for i in constants.JOINT_NAMES] - J_regressor_extra = np.load(path_config.JOINT_REGRESSOR_TRAIN_EXTRA) - self.register_buffer( - 'J_regressor_extra', - torch.tensor(J_regressor_extra, dtype=torch.float32)) - self.joint_map = torch.tensor(joints, dtype=torch.long) - self.ModelOutput = namedtuple( - 'ModelOutput_', ModelOutput._fields + ( - 'smpl_joints', - 'joints_J19', - )) - self.ModelOutput.__new__.__defaults__ = (None, ) * len( - self.ModelOutput._fields) - - def forward(self, *args, **kwargs): - kwargs['get_skin'] = True - smpl_output = super().forward(*args, **kwargs) - extra_joints = vertices2joints(self.J_regressor_extra, - smpl_output.vertices) - # smpl_output.joints: [B, 45, 3] extra_joints: [B, 9, 3] - vertices = smpl_output.vertices - joints = torch.cat([smpl_output.joints, extra_joints], dim=1) - smpl_joints = smpl_output.joints[:, :24] - joints = joints[:, self.joint_map, :] # [B, 49, 3] - joints_J24 = joints[:, -24:, :] - joints_J19 = joints_J24[:, constants.J24_TO_J19, :] - output = self.ModelOutput(vertices=vertices, - global_orient=smpl_output.global_orient, - body_pose=smpl_output.body_pose, - joints=joints, - joints_J19=joints_J19, - smpl_joints=smpl_joints, - betas=smpl_output.betas, - full_pose=smpl_output.full_pose) - return output - - -def get_smpl_faces(): - smpl = SMPL(SMPL_MODEL_DIR, batch_size=1, create_transl=False) - return smpl.faces - - -def get_part_joints(smpl_joints): - batch_size = smpl_joints.shape[0] - - # part_joints = torch.zeros().to(smpl_joints.device) - - one_seg_pairs = [(0, 1), (0, 2), (0, 3), (3, 6), (9, 12), (9, 13), (9, 14), - (12, 15), (13, 16), (14, 17)] - two_seg_pairs = [(1, 4), (2, 5), (4, 7), (5, 8), (16, 18), (17, 19), - (18, 20), (19, 21)] - - one_seg_pairs.extend(two_seg_pairs) - - single_joints = [(10), (11), (15), (22), (23)] - - part_joints = [] - - for j_p in one_seg_pairs: - new_joint = torch.mean(smpl_joints[:, j_p], dim=1, keepdim=True) - part_joints.append(new_joint) - - for j_p in single_joints: - part_joints.append(smpl_joints[:, j_p:j_p + 1]) - - part_joints = torch.cat(part_joints, dim=1) - - return part_joints diff --git a/lib / pymaf /models /hmr.py b/lib / pymaf /models /hmr.py deleted file mode 100644 index 9fb1cd6b7e2e4581f2c5d9cb5e952049b5a075e1..0000000000000000000000000000000000000000 --- a/lib / pymaf /models /hmr.py +++ /dev/null @@ -1,303 +0,0 @@ -# This script is borrowed from https://github.com/nkolot/SPIN/blob/master/models/hmr.py - -import torch -import torch.nn as nn -import torchvision.models.resnet as resnet -import numpy as np -import math -from lib.pymaf.utils.geometry import rot6d_to_rotmat - -import logging - -logger = logging.getLogger(__name__) - -BN_MOMENTUM = 0.1 - - -class Bottleneck(nn.Module): - """ Redefinition of Bottleneck residual block - Adapted from the official PyTorch implementation - """ - expansion = 4 - - def __init__(self, inplanes, planes, stride=1, downsample=None): - super().__init__() - self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) - self.bn1 = nn.BatchNorm2d(planes) - self.conv2 = nn.Conv2d(planes, - planes, - kernel_size=3, - stride=stride, - padding=1, - bias=False) - self.bn2 = nn.BatchNorm2d(planes) - self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) - self.bn3 = nn.BatchNorm2d(planes * 4) - 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 ResNet_Backbone(nn.Module): - """ Feature Extrator with ResNet backbone - """ - - def __init__(self, model='res50', pretrained=True): - if model == 'res50': - block, layers = Bottleneck, [3, 4, 6, 3] - else: - pass # TODO - - self.inplanes = 64 - super().__init__() - npose = 24 * 6 - self.conv1 = nn.Conv2d(3, - 64, - kernel_size=7, - stride=2, - padding=3, - bias=False) - self.bn1 = nn.BatchNorm2d(64) - self.relu = nn.ReLU(inplace=True) - self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) - self.layer1 = self._make_layer(block, 64, layers[0]) - self.layer2 = self._make_layer(block, 128, layers[1], stride=2) - self.layer3 = self._make_layer(block, 256, layers[2], stride=2) - self.layer4 = self._make_layer(block, 512, layers[3], stride=2) - self.avgpool = nn.AvgPool2d(7, stride=1) - - if pretrained: - resnet_imagenet = resnet.resnet50(pretrained=True) - self.load_state_dict(resnet_imagenet.state_dict(), strict=False) - logger.info('loaded resnet50 imagenet pretrained model') - - 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), - ) - - 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_deconv_layer(self, num_layers, num_filters, num_kernels): - assert num_layers == len(num_filters), \ - 'ERROR: num_deconv_layers is different len(num_deconv_filters)' - assert num_layers == len(num_kernels), \ - 'ERROR: num_deconv_layers is different len(num_deconv_filters)' - - def _get_deconv_cfg(deconv_kernel, index): - if deconv_kernel == 4: - padding = 1 - output_padding = 0 - elif deconv_kernel == 3: - padding = 1 - output_padding = 1 - elif deconv_kernel == 2: - padding = 0 - output_padding = 0 - - return deconv_kernel, padding, output_padding - - layers = [] - for i in range(num_layers): - kernel, padding, output_padding = _get_deconv_cfg( - num_kernels[i], i) - - planes = num_filters[i] - layers.append( - nn.ConvTranspose2d(in_channels=self.inplanes, - out_channels=planes, - kernel_size=kernel, - stride=2, - padding=padding, - output_padding=output_padding, - bias=self.deconv_with_bias)) - layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)) - layers.append(nn.ReLU(inplace=True)) - self.inplanes = planes - - return nn.Sequential(*layers) - - def forward(self, x): - - batch_size = x.shape[0] - - x = self.conv1(x) - x = self.bn1(x) - x = self.relu(x) - x = self.maxpool(x) - - x1 = self.layer1(x) - x2 = self.layer2(x1) - x3 = self.layer3(x2) - x4 = self.layer4(x3) - - xf = self.avgpool(x4) - xf = xf.view(xf.size(0), -1) - - x_featmap = x4 - - return x_featmap, xf - - -class HMR(nn.Module): - """ SMPL Iterative Regressor with ResNet50 backbone - """ - - def __init__(self, block, layers, smpl_mean_params): - self.inplanes = 64 - super().__init__() - npose = 24 * 6 - self.conv1 = nn.Conv2d(3, - 64, - kernel_size=7, - stride=2, - padding=3, - bias=False) - self.bn1 = nn.BatchNorm2d(64) - self.relu = nn.ReLU(inplace=True) - self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) - self.layer1 = self._make_layer(block, 64, layers[0]) - self.layer2 = self._make_layer(block, 128, layers[1], stride=2) - self.layer3 = self._make_layer(block, 256, layers[2], stride=2) - self.layer4 = self._make_layer(block, 512, layers[3], stride=2) - self.avgpool = nn.AvgPool2d(7, stride=1) - self.fc1 = nn.Linear(512 * block.expansion + npose + 13, 1024) - self.drop1 = nn.Dropout() - self.fc2 = nn.Linear(1024, 1024) - self.drop2 = nn.Dropout() - self.decpose = nn.Linear(1024, npose) - self.decshape = nn.Linear(1024, 10) - self.deccam = nn.Linear(1024, 3) - nn.init.xavier_uniform_(self.decpose.weight, gain=0.01) - nn.init.xavier_uniform_(self.decshape.weight, gain=0.01) - nn.init.xavier_uniform_(self.deccam.weight, gain=0.01) - - for m in self.modules(): - if isinstance(m, nn.Conv2d): - n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels - m.weight.data.normal_(0, math.sqrt(2. / n)) - elif isinstance(m, nn.BatchNorm2d): - m.weight.data.fill_(1) - m.bias.data.zero_() - - mean_params = np.load(smpl_mean_params) - init_pose = torch.from_numpy(mean_params['pose'][:]).unsqueeze(0) - init_shape = torch.from_numpy( - mean_params['shape'][:].astype('float32')).unsqueeze(0) - init_cam = torch.from_numpy(mean_params['cam']).unsqueeze(0) - self.register_buffer('init_pose', init_pose) - self.register_buffer('init_shape', init_shape) - self.register_buffer('init_cam', init_cam) - - 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), - ) - - 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 forward(self, - x, - init_pose=None, - init_shape=None, - init_cam=None, - n_iter=3): - - batch_size = x.shape[0] - - if init_pose is None: - init_pose = self.init_pose.expand(batch_size, -1) - if init_shape is None: - init_shape = self.init_shape.expand(batch_size, -1) - if init_cam is None: - init_cam = self.init_cam.expand(batch_size, -1) - - x = self.conv1(x) - x = self.bn1(x) - x = self.relu(x) - x = self.maxpool(x) - - x1 = self.layer1(x) - x2 = self.layer2(x1) - x3 = self.layer3(x2) - x4 = self.layer4(x3) - - xf = self.avgpool(x4) - xf = xf.view(xf.size(0), -1) - - pred_pose = init_pose - pred_shape = init_shape - pred_cam = init_cam - for i in range(n_iter): - xc = torch.cat([xf, pred_pose, pred_shape, pred_cam], 1) - xc = self.fc1(xc) - xc = self.drop1(xc) - xc = self.fc2(xc) - xc = self.drop2(xc) - pred_pose = self.decpose(xc) + pred_pose - pred_shape = self.decshape(xc) + pred_shape - pred_cam = self.deccam(xc) + pred_cam - - pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3) - - return pred_rotmat, pred_shape, pred_cam - - -def hmr(smpl_mean_params, pretrained=True, **kwargs): - """ Constructs an HMR model with ResNet50 backbone. - Args: - pretrained (bool): If True, returns a model pre-trained on ImageNet - """ - model = HMR(Bottleneck, [3, 4, 6, 3], smpl_mean_params, **kwargs) - if pretrained: - resnet_imagenet = resnet.resnet50(pretrained=True) - model.load_state_dict(resnet_imagenet.state_dict(), strict=False) - return model diff --git a/lib / pymaf /models /maf_extractor.py b/lib / pymaf /models /maf_extractor.py deleted file mode 100644 index 66a228d7992ea284f2374b6cfb7343c5a47c902a..0000000000000000000000000000000000000000 --- a/lib / pymaf /models /maf_extractor.py +++ /dev/null @@ -1,135 +0,0 @@ -# This script is borrowed and extended from https://github.com/shunsukesaito/PIFu/blob/master/lib/model/SurfaceClassifier.py - -from packaging import version -import torch -import scipy -import numpy as np -import torch.nn as nn -import torch.nn.functional as F - -from lib.common.config import cfg -from lib.pymaf.utils.geometry import projection -from lib.pymaf.core.path_config import MESH_DOWNSAMPLEING - -import logging - -logger = logging.getLogger(__name__) - - -class MAF_Extractor(nn.Module): - ''' Mesh-aligned Feature Extrator - As discussed in the paper, we extract mesh-aligned features based on 2D projection of the mesh vertices. - The features extrated from spatial feature maps will go through a MLP for dimension reduction. - ''' - - def __init__(self, device=torch.device('cuda')): - super().__init__() - - self.device = device - self.filters = [] - self.num_views = 1 - filter_channels = cfg.MODEL.PyMAF.MLP_DIM - self.last_op = nn.ReLU(True) - - for l in range(0, len(filter_channels) - 1): - if 0 != l: - self.filters.append( - nn.Conv1d(filter_channels[l] + filter_channels[0], - filter_channels[l + 1], 1)) - else: - self.filters.append( - nn.Conv1d(filter_channels[l], filter_channels[l + 1], 1)) - - self.add_module("conv%d" % l, self.filters[l]) - - self.im_feat = None - self.cam = None - - # downsample SMPL mesh and assign part labels - # from https://github.com/nkolot/GraphCMR/blob/master/data/mesh_downsampling.npz - smpl_mesh_graph = np.load(MESH_DOWNSAMPLEING, - allow_pickle=True, - encoding='latin1') - - A = smpl_mesh_graph['A'] - U = smpl_mesh_graph['U'] - D = smpl_mesh_graph['D'] # shape: (2,) - - # downsampling - ptD = [] - for i in range(len(D)): - d = scipy.sparse.coo_matrix(D[i]) - i = torch.LongTensor(np.array([d.row, d.col])) - v = torch.FloatTensor(d.data) - ptD.append(torch.sparse.FloatTensor(i, v, d.shape)) - - # downsampling mapping from 6890 points to 431 points - # ptD[0].to_dense() - Size: [1723, 6890] - # ptD[1].to_dense() - Size: [431. 1723] - Dmap = torch.matmul(ptD[1].to_dense(), - ptD[0].to_dense()) # 6890 -> 431 - self.register_buffer('Dmap', Dmap) - - def reduce_dim(self, feature): - ''' - Dimension reduction by multi-layer perceptrons - :param feature: list of [B, C_s, N] point-wise features before dimension reduction - :return: [B, C_p x N] concatantion of point-wise features after dimension reduction - ''' - y = feature - tmpy = feature - for i, f in enumerate(self.filters): - y = self._modules['conv' + - str(i)](y if i == 0 else torch.cat([y, tmpy], 1)) - if i != len(self.filters) - 1: - y = F.leaky_relu(y) - if self.num_views > 1 and i == len(self.filters) // 2: - y = y.view(-1, self.num_views, y.shape[1], - y.shape[2]).mean(dim=1) - tmpy = feature.view(-1, self.num_views, feature.shape[1], - feature.shape[2]).mean(dim=1) - - y = self.last_op(y) - - y = y.view(y.shape[0], -1) - return y - - def sampling(self, points, im_feat=None, z_feat=None): - ''' - Given 2D points, sample the point-wise features for each point, - the dimension of point-wise features will be reduced from C_s to C_p by MLP. - Image features should be pre-computed before this call. - :param points: [B, N, 2] image coordinates of points - :im_feat: [B, C_s, H_s, W_s] spatial feature maps - :return: [B, C_p x N] concatantion of point-wise features after dimension reduction - ''' - if im_feat is None: - im_feat = self.im_feat - - batch_size = im_feat.shape[0] - - if version.parse(torch.__version__) >= version.parse('1.3.0'): - # Default grid_sample behavior has changed to align_corners=False since 1.3.0. - point_feat = torch.nn.functional.grid_sample( - im_feat, points.unsqueeze(2), align_corners=True)[..., 0] - else: - point_feat = torch.nn.functional.grid_sample( - im_feat, points.unsqueeze(2))[..., 0] - - mesh_align_feat = self.reduce_dim(point_feat) - return mesh_align_feat - - def forward(self, p, s_feat=None, cam=None, **kwargs): - ''' Returns mesh-aligned features for the 3D mesh points. - Args: - p (tensor): [B, N_m, 3] mesh vertices - s_feat (tensor): [B, C_s, H_s, W_s] spatial feature maps - cam (tensor): [B, 3] camera - Return: - mesh_align_feat (tensor): [B, C_p x N_m] mesh-aligned features - ''' - if cam is None: - cam = self.cam - p_proj_2d = projection(p, cam, retain_z=False) - mesh_align_feat = self.sampling(p_proj_2d, s_feat) - return mesh_align_feat diff --git a/lib / pymaf /models /res_module.py b/lib / pymaf /models /res_module.py deleted file mode 100644 index 60c6a1a87326e80c68fa0c058206af4062c39d00..0000000000000000000000000000000000000000 --- a/lib / pymaf /models /res_module.py +++ /dev/null @@ -1,385 +0,0 @@ -# code brought in part from https://github.com/microsoft/human-pose-estimation.pytorch/blob/master/lib/models/pose_resnet.py - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import torch -import torch.nn as nn -import torch.nn.functional as F -from collections import OrderedDict -import os -from lib.pymaf.core.cfgs import cfg - -import logging - -logger = logging.getLogger(__name__) - -BN_MOMENTUM = 0.1 - - -def conv3x3(in_planes, out_planes, stride=1, bias=False, groups=1): - """3x3 convolution with padding""" - return nn.Conv2d(in_planes * groups, - out_planes * groups, - kernel_size=3, - stride=stride, - padding=1, - bias=bias, - groups=groups) - - -class BasicBlock(nn.Module): - expansion = 1 - - def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1): - super().__init__() - self.conv1 = conv3x3(inplanes, planes, stride, groups=groups) - self.bn1 = nn.BatchNorm2d(planes * groups, momentum=BN_MOMENTUM) - self.relu = nn.ReLU(inplace=True) - self.conv2 = conv3x3(planes, planes, groups=groups) - self.bn2 = nn.BatchNorm2d(planes * groups, 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, groups=1): - super().__init__() - self.conv1 = nn.Conv2d(inplanes * groups, - planes * groups, - kernel_size=1, - bias=False, - groups=groups) - self.bn1 = nn.BatchNorm2d(planes * groups, momentum=BN_MOMENTUM) - self.conv2 = nn.Conv2d(planes * groups, - planes * groups, - kernel_size=3, - stride=stride, - padding=1, - bias=False, - groups=groups) - self.bn2 = nn.BatchNorm2d(planes * groups, momentum=BN_MOMENTUM) - self.conv3 = nn.Conv2d(planes * groups, - planes * self.expansion * groups, - kernel_size=1, - bias=False, - groups=groups) - self.bn3 = nn.BatchNorm2d(planes * self.expansion * groups, - 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 - - -resnet_spec = { - 18: (BasicBlock, [2, 2, 2, 2]), - 34: (BasicBlock, [3, 4, 6, 3]), - 50: (Bottleneck, [3, 4, 6, 3]), - 101: (Bottleneck, [3, 4, 23, 3]), - 152: (Bottleneck, [3, 8, 36, 3]) -} - - -class IUV_predict_layer(nn.Module): - def __init__(self, - feat_dim=256, - final_cov_k=3, - part_out_dim=25, - with_uv=True): - super().__init__() - - self.with_uv = with_uv - if self.with_uv: - self.predict_u = nn.Conv2d(in_channels=feat_dim, - out_channels=25, - kernel_size=final_cov_k, - stride=1, - padding=1 if final_cov_k == 3 else 0) - - self.predict_v = nn.Conv2d(in_channels=feat_dim, - out_channels=25, - kernel_size=final_cov_k, - stride=1, - padding=1 if final_cov_k == 3 else 0) - - self.predict_ann_index = nn.Conv2d( - in_channels=feat_dim, - out_channels=15, - kernel_size=final_cov_k, - stride=1, - padding=1 if final_cov_k == 3 else 0) - - self.predict_uv_index = nn.Conv2d(in_channels=feat_dim, - out_channels=25, - kernel_size=final_cov_k, - stride=1, - padding=1 if final_cov_k == 3 else 0) - - self.inplanes = feat_dim - - 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), - ) - - 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 forward(self, x): - return_dict = {} - - predict_uv_index = self.predict_uv_index(x) - predict_ann_index = self.predict_ann_index(x) - - return_dict['predict_uv_index'] = predict_uv_index - return_dict['predict_ann_index'] = predict_ann_index - - if self.with_uv: - predict_u = self.predict_u(x) - predict_v = self.predict_v(x) - return_dict['predict_u'] = predict_u - return_dict['predict_v'] = predict_v - else: - return_dict['predict_u'] = None - return_dict['predict_v'] = None - # return_dict['predict_u'] = torch.zeros(predict_uv_index.shape).to(predict_uv_index.device) - # return_dict['predict_v'] = torch.zeros(predict_uv_index.shape).to(predict_uv_index.device) - - return return_dict - - -class SmplResNet(nn.Module): - def __init__(self, - resnet_nums, - in_channels=3, - num_classes=229, - last_stride=2, - n_extra_feat=0, - truncate=0, - **kwargs): - super().__init__() - - self.inplanes = 64 - self.truncate = truncate - # extra = cfg.MODEL.EXTRA - # self.deconv_with_bias = extra.DECONV_WITH_BIAS - block, layers = resnet_spec[resnet_nums] - - self.conv1 = nn.Conv2d(in_channels, - 64, - kernel_size=7, - stride=2, - padding=3, - bias=False) - self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM) - self.relu = nn.ReLU(inplace=True) - self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) - self.layer1 = self._make_layer(block, 64, layers[0]) - self.layer2 = self._make_layer(block, 128, layers[1], stride=2) - self.layer3 = self._make_layer(block, 256, layers[2], - stride=2) if truncate < 2 else None - self.layer4 = self._make_layer( - block, 512, layers[3], - stride=last_stride) if truncate < 1 else None - - self.avg_pooling = nn.AdaptiveAvgPool2d(1) - - self.num_classes = num_classes - if num_classes > 0: - self.final_layer = nn.Linear(512 * block.expansion, num_classes) - nn.init.xavier_uniform_(self.final_layer.weight, gain=0.01) - - self.n_extra_feat = n_extra_feat - if n_extra_feat > 0: - self.trans_conv = nn.Sequential( - nn.Conv2d(n_extra_feat + 512 * block.expansion, - 512 * block.expansion, - kernel_size=1, - bias=False), - nn.BatchNorm2d(512 * block.expansion, momentum=BN_MOMENTUM), - nn.ReLU(True)) - - 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 forward(self, x, infeat=None): - x = self.conv1(x) - x = self.bn1(x) - x = self.relu(x) - x = self.maxpool(x) - - x1 = self.layer1(x) - x2 = self.layer2(x1) - x3 = self.layer3(x2) if self.truncate < 2 else x2 - x4 = self.layer4(x3) if self.truncate < 1 else x3 - - if infeat is not None: - x4 = self.trans_conv(torch.cat([infeat, x4], 1)) - - if self.num_classes > 0: - xp = self.avg_pooling(x4) - cls = self.final_layer(xp.view(xp.size(0), -1)) - if not cfg.DANET.USE_MEAN_PARA: - # for non-negative scale - scale = F.relu(cls[:, 0]).unsqueeze(1) - cls = torch.cat((scale, cls[:, 1:]), dim=1) - else: - cls = None - - return cls, {'x4': x4} - - def init_weights(self, pretrained=''): - if os.path.isfile(pretrained): - logger.info('=> loading pretrained model {}'.format(pretrained)) - # self.load_state_dict(pretrained_state_dict, strict=False) - checkpoint = torch.load(pretrained) - if isinstance(checkpoint, OrderedDict): - # state_dict = checkpoint - state_dict_old = self.state_dict() - for key in state_dict_old.keys(): - if key in checkpoint.keys(): - if state_dict_old[key].shape != checkpoint[key].shape: - del checkpoint[key] - state_dict = checkpoint - elif isinstance(checkpoint, dict) and 'state_dict' in checkpoint: - state_dict_old = checkpoint['state_dict'] - state_dict = OrderedDict() - # delete 'module.' because it is saved from DataParallel module - for key in state_dict_old.keys(): - if key.startswith('module.'): - # state_dict[key[7:]] = state_dict[key] - # state_dict.pop(key) - state_dict[key[7:]] = state_dict_old[key] - else: - state_dict[key] = state_dict_old[key] - else: - raise RuntimeError( - 'No state_dict found in checkpoint file {}'.format( - pretrained)) - self.load_state_dict(state_dict, strict=False) - else: - logger.error('=> imagenet pretrained model dose not exist') - logger.error('=> please download it first') - raise ValueError('imagenet pretrained model does not exist') - - -class LimbResLayers(nn.Module): - def __init__(self, - resnet_nums, - inplanes, - outplanes=None, - groups=1, - **kwargs): - super().__init__() - - self.inplanes = inplanes - block, layers = resnet_spec[resnet_nums] - self.outplanes = 512 if outplanes == None else outplanes - self.layer4 = self._make_layer(block, - self.outplanes, - layers[3], - stride=2, - groups=groups) - - self.avg_pooling = nn.AdaptiveAvgPool2d(1) - - def _make_layer(self, block, planes, blocks, stride=1, groups=1): - downsample = None - if stride != 1 or self.inplanes != planes * block.expansion: - downsample = nn.Sequential( - nn.Conv2d(self.inplanes * groups, - planes * block.expansion * groups, - kernel_size=1, - stride=stride, - bias=False, - groups=groups), - nn.BatchNorm2d(planes * block.expansion * groups, - momentum=BN_MOMENTUM), - ) - - layers = [] - layers.append( - block(self.inplanes, planes, stride, downsample, groups=groups)) - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append(block(self.inplanes, planes, groups=groups)) - - return nn.Sequential(*layers) - - def forward(self, x): - x = self.layer4(x) - x = self.avg_pooling(x) - - return x \ No newline at end of file diff --git a/lib / pymaf /utils / __init__.py b/lib / pymaf /utils / __init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/lib / pymaf /utils / geometry.py b/lib / pymaf /utils / geometry.py deleted file mode 100644 index d65db5895f5546d9a119955703243c4d4917892c..0000000000000000000000000000000000000000 --- a/lib / pymaf /utils / geometry.py +++ /dev/null @@ -1,435 +0,0 @@ -import torch -import numpy as np -from torch.nn import functional as F -""" -Useful geometric operations, e.g. Perspective projection and a differentiable Rodrigues formula -Parts of the code are taken from https://github.com/MandyMo/pytorch_HMR -""" - - -def batch_rodrigues(theta): - """Convert axis-angle representation to rotation matrix. - Args: - theta: size = [B, 3] - Returns: - Rotation matrix corresponding to the quaternion -- size = [B, 3, 3] - """ - l1norm = torch.norm(theta + 1e-8, p=2, dim=1) - angle = torch.unsqueeze(l1norm, -1) - normalized = torch.div(theta, angle) - angle = angle * 0.5 - v_cos = torch.cos(angle) - v_sin = torch.sin(angle) - quat = torch.cat([v_cos, v_sin * normalized], dim=1) - return quat_to_rotmat(quat) - - -def quat_to_rotmat(quat): - """Convert quaternion coefficients to rotation matrix. - Args: - quat: size = [B, 4] 4 <===>(w, x, y, z) - Returns: - Rotation matrix corresponding to the quaternion -- size = [B, 3, 3] - """ - norm_quat = quat - norm_quat = norm_quat / norm_quat.norm(p=2, dim=1, keepdim=True) - w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:, - 2], norm_quat[:, - 3] - - B = quat.size(0) - - w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2) - wx, wy, wz = w * x, w * y, w * z - xy, xz, yz = x * y, x * z, y * z - - rotMat = torch.stack([ - w2 + x2 - y2 - z2, 2 * xy - 2 * wz, 2 * wy + 2 * xz, 2 * wz + 2 * xy, - w2 - x2 + y2 - z2, 2 * yz - 2 * wx, 2 * xz - 2 * wy, 2 * wx + 2 * yz, - w2 - x2 - y2 + z2 - ], - dim=1).view(B, 3, 3) - return rotMat - - -def rotation_matrix_to_angle_axis(rotation_matrix): - """ - This function is borrowed from https://github.com/kornia/kornia - Convert 3x4 rotation matrix to Rodrigues vector - Args: - rotation_matrix (Tensor): rotation matrix. - Returns: - Tensor: Rodrigues vector transformation. - Shape: - - Input: :math:`(N, 3, 4)` - - Output: :math:`(N, 3)` - Example: - >>> input = torch.rand(2, 3, 4) # Nx4x4 - >>> output = tgm.rotation_matrix_to_angle_axis(input) # Nx3 - """ - if rotation_matrix.shape[1:] == (3, 3): - rot_mat = rotation_matrix.reshape(-1, 3, 3) - hom = torch.tensor([0, 0, 1], - dtype=torch.float32, - device=rotation_matrix.device).reshape( - 1, 3, 1).expand(rot_mat.shape[0], -1, -1) - rotation_matrix = torch.cat([rot_mat, hom], dim=-1) - - quaternion = rotation_matrix_to_quaternion(rotation_matrix) - aa = quaternion_to_angle_axis(quaternion) - aa[torch.isnan(aa)] = 0.0 - return aa - - -def quaternion_to_angle_axis(quaternion: torch.Tensor) -> torch.Tensor: - """ - This function is borrowed from https://github.com/kornia/kornia - Convert quaternion vector to angle axis of rotation. - Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h - Args: - quaternion (torch.Tensor): tensor with quaternions. - Return: - torch.Tensor: tensor with angle axis of rotation. - Shape: - - Input: :math:`(*, 4)` where `*` means, any number of dimensions - - Output: :math:`(*, 3)` - Example: - >>> quaternion = torch.rand(2, 4) # Nx4 - >>> angle_axis = tgm.quaternion_to_angle_axis(quaternion) # Nx3 - """ - if not torch.is_tensor(quaternion): - raise TypeError("Input type is not a torch.Tensor. Got {}".format( - type(quaternion))) - - if not quaternion.shape[-1] == 4: - raise ValueError( - "Input must be a tensor of shape Nx4 or 4. Got {}".format( - quaternion.shape)) - # unpack input and compute conversion - q1: torch.Tensor = quaternion[..., 1] - q2: torch.Tensor = quaternion[..., 2] - q3: torch.Tensor = quaternion[..., 3] - sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3 - - sin_theta: torch.Tensor = torch.sqrt(sin_squared_theta) - cos_theta: torch.Tensor = quaternion[..., 0] - two_theta: torch.Tensor = 2.0 * torch.where( - cos_theta < 0.0, torch.atan2(-sin_theta, -cos_theta), - torch.atan2(sin_theta, cos_theta)) - - k_pos: torch.Tensor = two_theta / sin_theta - k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta) - k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg) - - angle_axis: torch.Tensor = torch.zeros_like(quaternion)[..., :3] - angle_axis[..., 0] += q1 * k - angle_axis[..., 1] += q2 * k - angle_axis[..., 2] += q3 * k - return angle_axis - - -def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6): - """ - This function is borrowed from https://github.com/kornia/kornia - Convert 3x4 rotation matrix to 4d quaternion vector - This algorithm is based on algorithm described in - https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201 - Args: - rotation_matrix (Tensor): the rotation matrix to convert. - Return: - Tensor: the rotation in quaternion - Shape: - - Input: :math:`(N, 3, 4)` - - Output: :math:`(N, 4)` - Example: - >>> input = torch.rand(4, 3, 4) # Nx3x4 - >>> output = tgm.rotation_matrix_to_quaternion(input) # Nx4 - """ - if not torch.is_tensor(rotation_matrix): - raise TypeError("Input type is not a torch.Tensor. Got {}".format( - type(rotation_matrix))) - - if len(rotation_matrix.shape) > 3: - raise ValueError( - "Input size must be a three dimensional tensor. Got {}".format( - rotation_matrix.shape)) - if not rotation_matrix.shape[-2:] == (3, 4): - raise ValueError( - "Input size must be a N x 3 x 4 tensor. Got {}".format( - rotation_matrix.shape)) - - rmat_t = torch.transpose(rotation_matrix, 1, 2) - - mask_d2 = rmat_t[:, 2, 2] < eps - - mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1] - mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1] - - t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2] - q0 = torch.stack([ - rmat_t[:, 1, 2] - rmat_t[:, 2, 1], t0, - rmat_t[:, 0, 1] + rmat_t[:, 1, 0], rmat_t[:, 2, 0] + rmat_t[:, 0, 2] - ], -1) - t0_rep = t0.repeat(4, 1).t() - - t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2] - q1 = torch.stack([ - rmat_t[:, 2, 0] - rmat_t[:, 0, 2], rmat_t[:, 0, 1] + rmat_t[:, 1, 0], - t1, rmat_t[:, 1, 2] + rmat_t[:, 2, 1] - ], -1) - t1_rep = t1.repeat(4, 1).t() - - t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2] - q2 = torch.stack([ - rmat_t[:, 0, 1] - rmat_t[:, 1, 0], rmat_t[:, 2, 0] + rmat_t[:, 0, 2], - rmat_t[:, 1, 2] + rmat_t[:, 2, 1], t2 - ], -1) - t2_rep = t2.repeat(4, 1).t() - - t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2] - q3 = torch.stack([ - t3, rmat_t[:, 1, 2] - rmat_t[:, 2, 1], - rmat_t[:, 2, 0] - rmat_t[:, 0, 2], rmat_t[:, 0, 1] - rmat_t[:, 1, 0] - ], -1) - t3_rep = t3.repeat(4, 1).t() - - mask_c0 = mask_d2 * mask_d0_d1 - mask_c1 = mask_d2 * ~mask_d0_d1 - mask_c2 = ~mask_d2 * mask_d0_nd1 - mask_c3 = ~mask_d2 * ~mask_d0_nd1 - mask_c0 = mask_c0.view(-1, 1).type_as(q0) - mask_c1 = mask_c1.view(-1, 1).type_as(q1) - mask_c2 = mask_c2.view(-1, 1).type_as(q2) - mask_c3 = mask_c3.view(-1, 1).type_as(q3) - - q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3 - q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 + # noqa - t2_rep * mask_c2 + t3_rep * mask_c3) # noqa - q *= 0.5 - return q - - -def rot6d_to_rotmat(x): - """Convert 6D rotation representation to 3x3 rotation matrix. - Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019 - Input: - (B,6) Batch of 6-D rotation representations - Output: - (B,3,3) Batch of corresponding rotation matrices - """ - x = x.view(-1, 3, 2) - a1 = x[:, :, 0] - a2 = x[:, :, 1] - b1 = F.normalize(a1) - b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1) - b3 = torch.cross(b1, b2) - return torch.stack((b1, b2, b3), dim=-1) - - -def projection(pred_joints, pred_camera, retain_z=False): - pred_cam_t = torch.stack([ - pred_camera[:, 1], pred_camera[:, 2], 2 * 5000. / - (224. * pred_camera[:, 0] + 1e-9) - ], - dim=-1) - batch_size = pred_joints.shape[0] - camera_center = torch.zeros(batch_size, 2) - pred_keypoints_2d = perspective_projection( - pred_joints, - rotation=torch.eye(3).unsqueeze(0).expand(batch_size, -1, - -1).to(pred_joints.device), - translation=pred_cam_t, - focal_length=5000., - camera_center=camera_center, - retain_z=retain_z) - # Normalize keypoints to [-1,1] - pred_keypoints_2d = pred_keypoints_2d / (224. / 2.) - return pred_keypoints_2d - - -def perspective_projection(points, - rotation, - translation, - focal_length, - camera_center, - retain_z=False): - """ - This function computes the perspective projection of a set of points. - Input: - points (bs, N, 3): 3D points - rotation (bs, 3, 3): Camera rotation - translation (bs, 3): Camera translation - focal_length (bs,) or scalar: Focal length - camera_center (bs, 2): Camera center - """ - batch_size = points.shape[0] - K = torch.zeros([batch_size, 3, 3], device=points.device) - K[:, 0, 0] = focal_length - K[:, 1, 1] = focal_length - K[:, 2, 2] = 1. - K[:, :-1, -1] = camera_center - - # Transform points - points = torch.einsum('bij,bkj->bki', rotation, points) - points = points + translation.unsqueeze(1) - - # Apply perspective distortion - projected_points = points / points[:, :, -1].unsqueeze(-1) - - # Apply camera intrinsics - projected_points = torch.einsum('bij,bkj->bki', K, projected_points) - - if retain_z: - return projected_points - else: - return projected_points[:, :, :-1] - - -def estimate_translation_np(S, - joints_2d, - joints_conf, - focal_length=5000, - img_size=224): - """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d. - Input: - S: (25, 3) 3D joint locations - joints: (25, 3) 2D joint locations and confidence - Returns: - (3,) camera translation vector - """ - - num_joints = S.shape[0] - # focal length - f = np.array([focal_length, focal_length]) - # optical center - center = np.array([img_size / 2., img_size / 2.]) - - # transformations - Z = np.reshape(np.tile(S[:, 2], (2, 1)).T, -1) - XY = np.reshape(S[:, 0:2], -1) - O = np.tile(center, num_joints) - F = np.tile(f, num_joints) - weight2 = np.reshape(np.tile(np.sqrt(joints_conf), (2, 1)).T, -1) - - # least squares - Q = np.array([ - F * np.tile(np.array([1, 0]), num_joints), - F * np.tile(np.array([0, 1]), num_joints), - O - np.reshape(joints_2d, -1) - ]).T - c = (np.reshape(joints_2d, -1) - O) * Z - F * XY - - # weighted least squares - W = np.diagflat(weight2) - Q = np.dot(W, Q) - c = np.dot(W, c) - - # square matrix - A = np.dot(Q.T, Q) - b = np.dot(Q.T, c) - - # solution - trans = np.linalg.solve(A, b) - - return trans - - -def estimate_translation(S, joints_2d, focal_length=5000., img_size=224.): - """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d. - Input: - S: (B, 49, 3) 3D joint locations - joints: (B, 49, 3) 2D joint locations and confidence - Returns: - (B, 3) camera translation vectors - """ - - device = S.device - # Use only joints 25:49 (GT joints) - S = S[:, 25:, :].cpu().numpy() - joints_2d = joints_2d[:, 25:, :].cpu().numpy() - joints_conf = joints_2d[:, :, -1] - joints_2d = joints_2d[:, :, :-1] - trans = np.zeros((S.shape[0], 3), dtype=np.float32) - # Find the translation for each example in the batch - for i in range(S.shape[0]): - S_i = S[i] - joints_i = joints_2d[i] - conf_i = joints_conf[i] - trans[i] = estimate_translation_np(S_i, - joints_i, - conf_i, - focal_length=focal_length, - img_size=img_size) - return torch.from_numpy(trans).to(device) - - -def Rot_y(angle, category='torch', prepend_dim=True, device=None): - '''Rotate around y-axis by angle - Args: - category: 'torch' or 'numpy' - prepend_dim: prepend an extra dimension - Return: Rotation matrix with shape [1, 3, 3] (prepend_dim=True) - ''' - m = np.array([[np.cos(angle), 0., np.sin(angle)], [0., 1., 0.], - [-np.sin(angle), 0., np.cos(angle)]]) - if category == 'torch': - if prepend_dim: - return torch.tensor(m, dtype=torch.float, - device=device).unsqueeze(0) - else: - return torch.tensor(m, dtype=torch.float, device=device) - elif category == 'numpy': - if prepend_dim: - return np.expand_dims(m, 0) - else: - return m - else: - raise ValueError("category must be 'torch' or 'numpy'") - - -def Rot_x(angle, category='torch', prepend_dim=True, device=None): - '''Rotate around x-axis by angle - Args: - category: 'torch' or 'numpy' - prepend_dim: prepend an extra dimension - Return: Rotation matrix with shape [1, 3, 3] (prepend_dim=True) - ''' - m = np.array([[1., 0., 0.], [0., np.cos(angle), -np.sin(angle)], - [0., np.sin(angle), np.cos(angle)]]) - if category == 'torch': - if prepend_dim: - return torch.tensor(m, dtype=torch.float, - device=device).unsqueeze(0) - else: - return torch.tensor(m, dtype=torch.float, device=device) - elif category == 'numpy': - if prepend_dim: - return np.expand_dims(m, 0) - else: - return m - else: - raise ValueError("category must be 'torch' or 'numpy'") - - -def Rot_z(angle, category='torch', prepend_dim=True, device=None): - '''Rotate around z-axis by angle - Args: - category: 'torch' or 'numpy' - prepend_dim: prepend an extra dimension - Return: Rotation matrix with shape [1, 3, 3] (prepend_dim=True) - ''' - m = np.array([[np.cos(angle), -np.sin(angle), 0.], - [np.sin(angle), np.cos(angle), 0.], [0., 0., 1.]]) - if category == 'torch': - if prepend_dim: - return torch.tensor(m, dtype=torch.float, - device=device).unsqueeze(0) - else: - return torch.tensor(m, dtype=torch.float, device=device) - elif category == 'numpy': - if prepend_dim: - return np.expand_dims(m, 0) - else: - return m - else: - raise ValueError("category must be 'torch' or 'numpy'") \ No newline at end of file diff --git a/lib / pymaf /utils / imutils.py b/lib / pymaf /utils / imutils.py deleted file mode 100644 index 000c9357842036abe496bdafe8223458ca12bb75..0000000000000000000000000000000000000000 --- a/lib / pymaf /utils / imutils.py +++ /dev/null @@ -1,491 +0,0 @@ -""" -This file contains functions that are used to perform data augmentation. -""" -import cv2 -import io -import torch -import numpy as np -from PIL import Image -from rembg import remove -from rembg.session_factory import new_session -from torchvision.models import detection - -from lib.pymaf.core import constants -from lib.pymaf.utils.streamer import aug_matrix -from lib.common.cloth_extraction import load_segmentation -from torchvision import transforms - - -def load_img(img_file): - - img = cv2.imread(img_file, cv2.IMREAD_UNCHANGED) - if len(img.shape) == 2: - img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) - - if not img_file.endswith("png"): - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - else: - img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR) - - return img - - -def get_bbox(img, det): - - input = np.float32(img) - input = (input / 255.0 - - (0.5, 0.5, 0.5)) / (0.5, 0.5, 0.5) # TO [-1.0, 1.0] - input = input.transpose(2, 0, 1) # TO [3 x H x W] - bboxes, probs = det(torch.from_numpy(input).float().unsqueeze(0)) - - probs = probs.unsqueeze(3) - bboxes = (bboxes * probs).sum(dim=1, keepdim=True) / probs.sum( - dim=1, keepdim=True) - bbox = bboxes[0, 0, 0].cpu().numpy() - - return bbox -# Michael Black is - - -def get_transformer(input_res): - - image_to_tensor = transforms.Compose([ - transforms.Resize(input_res), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ]) - - mask_to_tensor = transforms.Compose([ - transforms.Resize(input_res), - transforms.ToTensor(), - transforms.Normalize((0.0, ), (1.0, )) - ]) - - image_to_pymaf_tensor = transforms.Compose([ - transforms.Resize(size=224), - transforms.Normalize(mean=constants.IMG_NORM_MEAN, - std=constants.IMG_NORM_STD) - ]) - - image_to_pixie_tensor = transforms.Compose([ - transforms.Resize(224) - ]) - - def image_to_hybrik_tensor(img): - # mean - img[0].add_(-0.406) - img[1].add_(-0.457) - img[2].add_(-0.480) - - # std - img[0].div_(0.225) - img[1].div_(0.224) - img[2].div_(0.229) - return img - - return [image_to_tensor, mask_to_tensor, image_to_pymaf_tensor, image_to_pixie_tensor, image_to_hybrik_tensor] - - -def process_image(img_file, hps_type, input_res=512, device=None, seg_path=None): - """Read image, do preprocessing and possibly crop it according to the bounding box. - If there are bounding box annotations, use them to crop the image. - If no bounding box is specified but openpose detections are available, use them to get the bounding box. - """ - - [image_to_tensor, mask_to_tensor, image_to_pymaf_tensor, - image_to_pixie_tensor, image_to_hybrik_tensor] = get_transformer(input_res) - - img_ori = load_img(img_file) - - in_height, in_width, _ = img_ori.shape - M = aug_matrix(in_width, in_height, input_res*2, input_res*2) - - # from rectangle to square - img_for_crop = cv2.warpAffine(img_ori, M[0:2, :], - (input_res*2, input_res*2), flags=cv2.INTER_CUBIC) - - # detection for bbox - detector = detection.maskrcnn_resnet50_fpn(pretrained=True) - detector.eval() - predictions = detector( - [torch.from_numpy(img_for_crop).permute(2, 0, 1) / 255.])[0] - human_ids = torch.where( - predictions["scores"] == predictions["scores"][predictions['labels'] == 1].max()) - bbox = predictions["boxes"][human_ids, :].flatten().detach().cpu().numpy() - - width = bbox[2] - bbox[0] - height = bbox[3] - bbox[1] - center = np.array([(bbox[0] + bbox[2]) / 2.0, - (bbox[1] + bbox[3]) / 2.0]) - - scale = max(height, width) / 180 - - if hps_type == 'hybrik': - img_np = crop_for_hybrik(img_for_crop, center, - np.array([scale * 180, scale * 180])) - else: - img_np, cropping_parameters = crop( - img_for_crop, center, scale, (input_res, input_res)) - - img_pil = Image.fromarray(remove(img_np, post_process_mask=True, session=new_session("u2net"))) - - # for icon - img_rgb = image_to_tensor(img_pil.convert("RGB")) - img_mask = torch.tensor(1.0) - (mask_to_tensor(img_pil.split()[-1]) < - torch.tensor(0.5)).float() - img_tensor = img_rgb * img_mask - - # for hps - img_hps = img_np.astype(np.float32) / 255. - img_hps = torch.from_numpy(img_hps).permute(2, 0, 1) - - if hps_type == 'bev': - img_hps = img_np[:, :, [2, 1, 0]] - elif hps_type == 'hybrik': - img_hps = image_to_hybrik_tensor(img_hps).unsqueeze(0).to(device) - elif hps_type != 'pixie': - img_hps = image_to_pymaf_tensor(img_hps).unsqueeze(0).to(device) - else: - img_hps = image_to_pixie_tensor(img_hps).unsqueeze(0).to(device) - - # uncrop params - uncrop_param = {'center': center, - 'scale': scale, - 'ori_shape': img_ori.shape, - 'box_shape': img_np.shape, - 'crop_shape': img_for_crop.shape, - 'M': M} - - if not (seg_path is None): - segmentations = load_segmentation(seg_path, (in_height, in_width)) - seg_coord_normalized = [] - for seg in segmentations: - coord_normalized = [] - for xy in seg['coordinates']: - xy_h = np.vstack((xy[:, 0], xy[:, 1], np.ones(len(xy)))).T - warped_indeces = M[0:2, :] @ xy_h[:, :, None] - warped_indeces = np.array(warped_indeces).astype(int) - warped_indeces.resize((warped_indeces.shape[:2])) - - # cropped_indeces = crop_segmentation(warped_indeces, center, scale, (input_res, input_res), img_np.shape) - cropped_indeces = crop_segmentation( - warped_indeces, (input_res, input_res), cropping_parameters) - - indices = np.vstack( - (cropped_indeces[:, 0], cropped_indeces[:, 1])).T - - # Convert to NDC coordinates - seg_cropped_normalized = 2*(indices / input_res) - 1 - # Don't know why we need to divide by 50 but it works ¯\_(ツ)_/¯ (probably some scaling factor somewhere) - # Divide only by 45 on the horizontal axis to take the curve of the human body into account - seg_cropped_normalized[:, 0] = ( - 1/40) * seg_cropped_normalized[:, 0] - seg_cropped_normalized[:, 1] = ( - 1/50) * seg_cropped_normalized[:, 1] - coord_normalized.append(seg_cropped_normalized) - - seg['coord_normalized'] = coord_normalized - seg_coord_normalized.append(seg) - - return img_tensor, img_hps, img_ori, img_mask, uncrop_param, seg_coord_normalized - - return img_tensor, img_hps, img_ori, img_mask, uncrop_param - - -def get_transform(center, scale, res): - """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 - - return t - - -def transform(pt, center, scale, res, invert=0): - """Transform pixel location to different reference.""" - t = get_transform(center, scale, res) - 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 np.around(new_pt[:2]).astype(np.int16) - - -def crop(img, center, scale, res): - """Crop image according to the supplied bounding box.""" - - # Upper left point - ul = np.array(transform([0, 0], center, scale, res, invert=1)) - - # Bottom right point - br = np.array(transform(res, center, scale, res, invert=1)) - - 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 len(img.shape) == 2: - new_img = np.array(Image.fromarray(new_img).resize(res)) - else: - new_img = np.array(Image.fromarray( - new_img.astype(np.uint8)).resize(res)) - - return new_img, (old_x, new_x, old_y, new_y, new_shape) - - -def crop_segmentation(org_coord, res, cropping_parameters): - old_x, new_x, old_y, new_y, new_shape = cropping_parameters - - new_coord = np.zeros((org_coord.shape)) - new_coord[:, 0] = new_x[0] + (org_coord[:, 0] - old_x[0]) - new_coord[:, 1] = new_y[0] + (org_coord[:, 1] - old_y[0]) - - new_coord[:, 0] = res[0] * (new_coord[:, 0] / new_shape[1]) - new_coord[:, 1] = res[1] * (new_coord[:, 1] / new_shape[0]) - - return new_coord - - -def crop_for_hybrik(img, center, scale): - inp_h, inp_w = (256, 256) - trans = get_affine_transform(center, scale, 0, [inp_w, inp_h]) - new_img = cv2.warpAffine( - img, trans, (int(inp_w), int(inp_h)), flags=cv2.INTER_LINEAR) - return new_img - - -def get_affine_transform(center, - scale, - rot, - output_size, - shift=np.array([0, 0], dtype=np.float32), - inv=0): - - def get_dir(src_point, rot_rad): - """Rotate the point by `rot_rad` degree.""" - sn, cs = np.sin(rot_rad), np.cos(rot_rad) - - src_result = [0, 0] - src_result[0] = src_point[0] * cs - src_point[1] * sn - src_result[1] = src_point[0] * sn + src_point[1] * cs - - return src_result - - def get_3rd_point(a, b): - """Return vector c that perpendicular to (a - b).""" - direct = a - b - return b + np.array([-direct[1], direct[0]], dtype=np.float32) - - if not isinstance(scale, np.ndarray) and not isinstance(scale, list): - scale = np.array([scale, scale]) - - scale_tmp = scale - src_w = scale_tmp[0] - dst_w = output_size[0] - dst_h = output_size[1] - - rot_rad = np.pi * rot / 180 - src_dir = get_dir([0, src_w * -0.5], rot_rad) - dst_dir = np.array([0, dst_w * -0.5], np.float32) - - src = np.zeros((3, 2), dtype=np.float32) - dst = np.zeros((3, 2), dtype=np.float32) - src[0, :] = center + scale_tmp * shift - src[1, :] = center + src_dir + scale_tmp * shift - dst[0, :] = [dst_w * 0.5, dst_h * 0.5] - dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir - - src[2:, :] = get_3rd_point(src[0, :], src[1, :]) - dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) - - if inv: - trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) - else: - trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) - - return trans - - -def corner_align(ul, br): - - if ul[1]-ul[0] != br[1]-br[0]: - ul[1] = ul[0]+br[1]-br[0] - - return ul, br - - -def uncrop(img, center, scale, orig_shape): - """'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([0, 0], center, scale, res, invert=1)) - # Bottom right point - br = np.array(transform(res, center, scale, res, invert=1)) - - # quick fix - ul, br = corner_align(ul, br) - - # size of cropped image - crop_shape = [br[1] - ul[1], br[0] - ul[0]] - 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 = np.array(Image.fromarray(img.astype(np.uint8)).resize(crop_shape)) - - 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, is_smpl=False): - """Flip keypoints.""" - if len(kp) == 24: - if is_smpl: - flipped_parts = constants.SMPL_JOINTS_FLIP_PERM - else: - flipped_parts = constants.J24_FLIP_PERM - elif len(kp) == 49: - if is_smpl: - flipped_parts = constants.SMPL_J49_FLIP_PERM - else: - 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 normalize_2d_kp(kp_2d, crop_size=224, inv=False): - # Normalize keypoints between -1, 1 - if not inv: - ratio = 1.0 / crop_size - kp_2d = 2.0 * kp_2d * ratio - 1.0 - else: - ratio = 1.0 / crop_size - kp_2d = (kp_2d + 1.0) / (2 * ratio) - - return kp_2d - - -def generate_heatmap(joints, heatmap_size, sigma=1, joints_vis=None): - ''' - param joints: [num_joints, 3] - param joints_vis: [num_joints, 3] - return: target, target_weight(1: visible, 0: invisible) - ''' - num_joints = joints.shape[0] - device = joints.device - cur_device = torch.device(device.type, device.index) - if not hasattr(heatmap_size, '__len__'): - # width height - heatmap_size = [heatmap_size, heatmap_size] - assert len(heatmap_size) == 2 - target_weight = np.ones((num_joints, 1), dtype=np.float32) - if joints_vis is not None: - target_weight[:, 0] = joints_vis[:, 0] - target = torch.zeros((num_joints, heatmap_size[1], heatmap_size[0]), - dtype=torch.float32, - device=cur_device) - - tmp_size = sigma * 3 - - for joint_id in range(num_joints): - mu_x = int(joints[joint_id][0] * heatmap_size[0] + 0.5) - mu_y = int(joints[joint_id][1] * heatmap_size[1] + 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[0] or ul[1] >= heatmap_size[1] \ - 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)) - # g = torch.from_numpy(g.astype(np.float32)) - - x = torch.arange(0, size, dtype=torch.float32, device=cur_device) - y = x.unsqueeze(-1) - x0 = y0 = size // 2 - # The gaussian is not normalized, we want the center value to equal 1 - g = torch.exp(-((x - x0)**2 + (y - y0)**2) / (2 * sigma**2)) - - # Usable gaussian range - g_x = max(0, -ul[0]), min(br[0], heatmap_size[0]) - ul[0] - g_y = max(0, -ul[1]), min(br[1], heatmap_size[1]) - ul[1] - # Image range - img_x = max(0, ul[0]), min(br[0], heatmap_size[0]) - img_y = max(0, ul[1]), min(br[1], heatmap_size[1]) - - 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 \ No newline at end of file diff --git a/lib / pymaf /utils / streamer.py b/lib / pymaf /utils / streamer.py deleted file mode 100644 index 4ab3074b2cf5b87f9332afa6239d0911efeff1ac..0000000000000000000000000000000000000000 --- a/lib / pymaf /utils / streamer.py +++ /dev/null @@ -1,142 +0,0 @@ -import cv2 -import torch -import numpy as np -import imageio - - -def aug_matrix(w1, h1, w2, h2): - dx = (w2 - w1) / 2.0 - dy = (h2 - h1) / 2.0 - - matrix_trans = np.array([[1.0, 0, dx], - [0, 1.0, dy], - [0, 0, 1.0]]) - - scale = np.min([float(w2)/w1, float(h2)/h1]) - - M = get_affine_matrix( - center=(w2 / 2.0, h2 / 2.0), - translate=(0, 0), - scale=scale) - - M = np.array(M + [0., 0., 1.]).reshape(3, 3) - M = M.dot(matrix_trans) - - return M - - -def get_affine_matrix(center, translate, scale): - cx, cy = center - tx, ty = translate - - M = [1, 0, 0, - 0, 1, 0] - M = [x * scale for x in M] - - # Apply translation and of center translation: RSS * C^-1 - M[2] += M[0] * (-cx) + M[1] * (-cy) - M[5] += M[3] * (-cx) + M[4] * (-cy) - - # Apply center translation: T * C * RSS * C^-1 - M[2] += cx + tx - M[5] += cy + ty - return M - - -class BaseStreamer(): - """This streamer will return images at 512x512 size. - """ - - def __init__(self, - width=512, height=512, pad=True, - mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), - **kwargs): - self.width = width - self.height = height - self.pad = pad - self.mean = np.array(mean) - self.std = np.array(std) - - self.loader = self.create_loader() - - def create_loader(self): - raise NotImplementedError - yield np.zeros((600, 400, 3)) # in RGB (0, 255) - - def __getitem__(self, index): - image = next(self.loader) - in_height, in_width, _ = image.shape - M = aug_matrix(in_width, in_height, self.width, self.height, self.pad) - image = cv2.warpAffine( - image, M[0:2, :], (self.width, self.height), flags=cv2.INTER_CUBIC) - - input = np.float32(image) - input = (input / 255.0 - self.mean) / self.std # TO [-1.0, 1.0] - input = input.transpose(2, 0, 1) # TO [3 x H x W] - return torch.from_numpy(input).float() - - def __len__(self): - raise NotImplementedError - - -class CaptureStreamer(BaseStreamer): - """This streamer takes webcam as input. - """ - - def __init__(self, id=0, width=512, height=512, pad=True, **kwargs): - super().__init__(width, height, pad, **kwargs) - self.capture = cv2.VideoCapture(id) - - def create_loader(self): - while True: - _, image = self.capture.read() - image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # RGB - yield image - - def __len__(self): - return 100_000_000 - - def __del__(self): - self.capture.release() - - -class VideoListStreamer(BaseStreamer): - """This streamer takes a list of video files as input. - """ - - def __init__(self, files, width=512, height=512, pad=True, **kwargs): - super().__init__(width, height, pad, **kwargs) - self.files = files - self.captures = [imageio.get_reader(f) for f in files] - self.nframes = sum([int(cap._meta["fps"] * cap._meta["duration"]) - for cap in self.captures]) - - def create_loader(self): - for capture in self.captures: - for image in capture: # RGB - yield image - - def __len__(self): - return self.nframes - - def __del__(self): - for capture in self.captures: - capture.close() - - -class ImageListStreamer(BaseStreamer): - """This streamer takes a list of image files as input. - """ - - def __init__(self, files, width=512, height=512, pad=True, **kwargs): - super().__init__(width, height, pad, **kwargs) - self.files = files - - def create_loader(self): - for f in self.files: - image = cv2.imread(f, cv2.IMREAD_UNCHANGED)[:, :, 0:3] - image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # RGB - yield image - - def __len__(self): - return len(self.files) \ No newline at end of file diff --git a/lib / pymaf /utils /transforms.py b/lib / pymaf /utils /transforms.py deleted file mode 100644 index db0bf42829bf07a1a22d962c1378098ba6e26e9d..0000000000000000000000000000000000000000 --- a/lib / pymaf /utils /transforms.py +++ /dev/null @@ -1,78 +0,0 @@ -# ------------------------------------------------------------------------------ -# Copyright (c) Microsoft -# Licensed under the MIT License. -# Written by Bin Xiao (Bin.Xiao@microsoft.com) -# ------------------------------------------------------------------------------ - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import cv2 -import numpy as np - - -def transform_preds(coords, center, scale, output_size): - target_coords = np.zeros(coords.shape) - trans = get_affine_transform(center, scale, 0, output_size, inv=1) - for p in range(coords.shape[0]): - target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans) - return target_coords - - -def get_affine_transform(center, - scale, - rot, - output_size, - shift=np.array([0, 0], dtype=np.float32), - inv=0): - if not isinstance(scale, np.ndarray) and not isinstance(scale, list): - # print(scale) - scale = np.array([scale, scale]) - - scale_tmp = scale * 200.0 - src_w = scale_tmp[0] - dst_w = output_size[0] - dst_h = output_size[1] - - rot_rad = np.pi * rot / 180 - src_dir = get_dir([0, src_w * -0.5], rot_rad) - dst_dir = np.array([0, dst_w * -0.5], np.float32) - - src = np.zeros((3, 2), dtype=np.float32) - dst = np.zeros((3, 2), dtype=np.float32) - src[0, :] = center + scale_tmp * shift - src[1, :] = center + src_dir + scale_tmp * shift - dst[0, :] = [dst_w * 0.5, dst_h * 0.5] - dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir - - src[2:, :] = get_3rd_point(src[0, :], src[1, :]) - dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) - - if inv: - trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) - else: - trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) - - return trans - - -def affine_transform(pt, t): - new_pt = np.array([pt[0], pt[1], 1.]).T - new_pt = np.dot(t, new_pt) - return new_pt[:2] - - -def get_3rd_point(a, b): - direct = a - b - return b + np.array([-direct[1], direct[0]], dtype=np.float32) - - -def get_dir(src_point, rot_rad): - sn, cs = np.sin(rot_rad), np.cos(rot_rad) - - src_result = [0, 0] - src_result[0] = src_point[0] * cs - src_point[1] * sn - src_result[1] = src_point[0] * sn + src_point[1] * cs - - return src_result \ No newline at end of file diff --git a/lib / renderer / __init__.py b/lib / renderer / __init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/lib / renderer / camera.py b/lib / renderer / camera.py deleted file mode 100644 index fde488826c8eba933a00485d39cdc21d4d2dae8f..0000000000000000000000000000000000000000 --- a/lib / renderer / camera.py +++ /dev/null @@ -1,226 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import cv2 -import numpy as np - -from .glm import ortho - - -class Camera: - def __init__(self, width=1600, height=1200): - # Focal Length - # equivalent 50mm - focal = np.sqrt(width * width + height * height) - self.focal_x = focal - self.focal_y = focal - # Principal Point Offset - self.principal_x = width / 2 - self.principal_y = height / 2 - # Axis Skew - self.skew = 0 - # Image Size - self.width = width - self.height = height - - self.near = 1 - self.far = 10 - - # Camera Center - self.center = np.array([0, 0, 1.6]) - self.direction = np.array([0, 0, -1]) - self.right = np.array([1, 0, 0]) - self.up = np.array([0, 1, 0]) - - self.ortho_ratio = None - - def sanity_check(self): - self.center = self.center.reshape([-1]) - self.direction = self.direction.reshape([-1]) - self.right = self.right.reshape([-1]) - self.up = self.up.reshape([-1]) - - assert len(self.center) == 3 - assert len(self.direction) == 3 - assert len(self.right) == 3 - assert len(self.up) == 3 - - @staticmethod - def normalize_vector(v): - v_norm = np.linalg.norm(v) - return v if v_norm == 0 else v / v_norm - - def get_real_z_value(self, z): - z_near = self.near - z_far = self.far - z_n = 2.0 * z - 1.0 - z_e = 2.0 * z_near * z_far / (z_far + z_near - z_n * (z_far - z_near)) - return z_e - - def get_rotation_matrix(self): - rot_mat = np.eye(3) - s = self.right - s = self.normalize_vector(s) - rot_mat[0, :] = s - u = self.up - u = self.normalize_vector(u) - rot_mat[1, :] = -u - rot_mat[2, :] = self.normalize_vector(self.direction) - - return rot_mat - - def get_translation_vector(self): - rot_mat = self.get_rotation_matrix() - trans = -np.dot(rot_mat, self.center) - return trans - - def get_intrinsic_matrix(self): - int_mat = np.eye(3) - - int_mat[0, 0] = self.focal_x - int_mat[1, 1] = self.focal_y - int_mat[0, 1] = self.skew - int_mat[0, 2] = self.principal_x - int_mat[1, 2] = self.principal_y - - return int_mat - - def get_projection_matrix(self): - ext_mat = self.get_extrinsic_matrix() - int_mat = self.get_intrinsic_matrix() - - return np.matmul(int_mat, ext_mat) - - def get_extrinsic_matrix(self): - rot_mat = self.get_rotation_matrix() - int_mat = self.get_intrinsic_matrix() - trans = self.get_translation_vector() - - extrinsic = np.eye(4) - extrinsic[:3, :3] = rot_mat - extrinsic[:3, 3] = trans - - return extrinsic[:3, :] - - def set_rotation_matrix(self, rot_mat): - self.direction = rot_mat[2, :] - self.up = -rot_mat[1, :] - self.right = rot_mat[0, :] - - def set_intrinsic_matrix(self, int_mat): - self.focal_x = int_mat[0, 0] - self.focal_y = int_mat[1, 1] - self.skew = int_mat[0, 1] - self.principal_x = int_mat[0, 2] - self.principal_y = int_mat[1, 2] - - def set_projection_matrix(self, proj_mat): - res = cv2.decomposeProjectionMatrix(proj_mat) - int_mat, rot_mat, camera_center_homo = res[0], res[1], res[2] - camera_center = camera_center_homo[0:3] / camera_center_homo[3] - camera_center = camera_center.reshape(-1) - int_mat = int_mat / int_mat[2][2] - - self.set_intrinsic_matrix(int_mat) - self.set_rotation_matrix(rot_mat) - self.center = camera_center - - self.sanity_check() - - def get_gl_matrix(self): - z_near = self.near - z_far = self.far - rot_mat = self.get_rotation_matrix() - int_mat = self.get_intrinsic_matrix() - trans = self.get_translation_vector() - - extrinsic = np.eye(4) - extrinsic[:3, :3] = rot_mat - extrinsic[:3, 3] = trans - axis_adj = np.eye(4) - axis_adj[2, 2] = -1 - axis_adj[1, 1] = -1 - model_view = np.matmul(axis_adj, extrinsic) - - projective = np.zeros([4, 4]) - projective[:2, :2] = int_mat[:2, :2] - projective[:2, 2:3] = -int_mat[:2, 2:3] - projective[3, 2] = -1 - projective[2, 2] = (z_near + z_far) - projective[2, 3] = (z_near * z_far) - - if self.ortho_ratio is None: - ndc = ortho(0, self.width, 0, self.height, z_near, z_far) - perspective = np.matmul(ndc, projective) - else: - perspective = ortho(-self.width * self.ortho_ratio / 2, - self.width * self.ortho_ratio / 2, - -self.height * self.ortho_ratio / 2, - self.height * self.ortho_ratio / 2, z_near, - z_far) - - return perspective, model_view - - -def KRT_from_P(proj_mat, normalize_K=True): - res = cv2.decomposeProjectionMatrix(proj_mat) - K, Rot, camera_center_homog = res[0], res[1], res[2] - camera_center = camera_center_homog[0:3] / camera_center_homog[3] - trans = -Rot.dot(camera_center) - if normalize_K: - K = K / K[2][2] - return K, Rot, trans - - -def MVP_from_P(proj_mat, width, height, near=0.1, far=10000): - ''' - Convert OpenCV camera calibration matrix to OpenGL projection and model view matrix - :param proj_mat: OpenCV camera projeciton matrix - :param width: Image width - :param height: Image height - :param near: Z near value - :param far: Z far value - :return: OpenGL projection matrix and model view matrix - ''' - res = cv2.decomposeProjectionMatrix(proj_mat) - K, Rot, camera_center_homog = res[0], res[1], res[2] - camera_center = camera_center_homog[0:3] / camera_center_homog[3] - trans = -Rot.dot(camera_center) - K = K / K[2][2] - - extrinsic = np.eye(4) - extrinsic[:3, :3] = Rot - extrinsic[:3, 3:4] = trans - axis_adj = np.eye(4) - axis_adj[2, 2] = -1 - axis_adj[1, 1] = -1 - model_view = np.matmul(axis_adj, extrinsic) - - zFar = far - zNear = near - projective = np.zeros([4, 4]) - projective[:2, :2] = K[:2, :2] - projective[:2, 2:3] = -K[:2, 2:3] - projective[3, 2] = -1 - projective[2, 2] = (zNear + zFar) - projective[2, 3] = (zNear * zFar) - - ndc = ortho(0, width, 0, height, zNear, zFar) - - perspective = np.matmul(ndc, projective) - - return perspective, model_view diff --git a/lib / renderer / gl / __init__.py b/lib / renderer / gl / __init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/lib / renderer / gl / data / color.fs b/lib / renderer / gl / data / color.fs deleted file mode 100644 index e28b6f65605bb380428c84c326d00d9e1445a899..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / data / color.fs +++ /dev/null @@ -1,20 +0,0 @@ -#version 330 core - -layout (location = 0) out vec4 FragColor; -layout (location = 1) out vec4 FragNormal; -layout (location = 2) out vec4 FragDepth; - -in vec3 Color; -in vec3 CamNormal; -in vec3 depth; - - -void main() -{ - FragColor = vec4(Color,1.0); - - vec3 cam_norm_normalized = normalize(CamNormal); - vec3 rgb = (cam_norm_normalized + 1.0) / 2.0; - FragNormal = vec4(rgb, 1.0); - FragDepth = vec4(depth.xyz, 1.0); -} \ No newline at end of file diff --git a/lib / renderer / gl / data /color.vs b/lib / renderer / gl / data /color.vs deleted file mode 100644 index 1256f7eb3f3605b8f848d452fcf601d5a18b95e2..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / data /color.vs +++ /dev/null @@ -1,29 +0,0 @@ -#version 330 core - -layout (location = 0) in vec3 a_Position; -layout (location = 1) in vec3 a_Color; -layout (location = 2) in vec3 a_Normal; - -out vec3 CamNormal; -out vec3 CamPos; -out vec3 Color; -out vec3 depth; - - -uniform mat3 RotMat; -uniform mat4 NormMat; -uniform mat4 ModelMat; -uniform mat4 PerspMat; - -void main() -{ - vec3 a_Position = (NormMat * vec4(a_Position,1.0)).xyz; - gl_Position = PerspMat * ModelMat * vec4(RotMat * a_Position, 1.0); - Color = a_Color; - - mat3 R = mat3(ModelMat) * RotMat; - CamNormal = (R * a_Normal); - - depth = vec3(gl_Position.z / gl_Position.w); - -} \ No newline at end of file diff --git a/lib / renderer / gl / data /normal.fs b/lib / renderer / gl / data /normal.fs deleted file mode 100644 index 9e2770952e27d9265ccb100833245beed3ebebe5..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / data /normal.fs +++ /dev/null @@ -1,12 +0,0 @@ -#version 330 - -out vec4 FragColor; - -in vec3 CamNormal; - -void main() -{ - vec3 cam_norm_normalized = normalize(CamNormal); - vec3 rgb = (cam_norm_normalized + 1.0) / 2.0; - FragColor = vec4(rgb, 1.0); -} \ No newline at end of file diff --git a/lib / renderer / gl / data /normal.vs b/lib / renderer / gl / data /normal.vs deleted file mode 100644 index a0f7f50b1cedfd677843b2a60cf9051b2134c347..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / data /normal.vs +++ /dev/null @@ -1,15 +0,0 @@ -#version 330 - -layout (location = 0) in vec3 Position; -layout (location = 1) in vec3 Normal; - -out vec3 CamNormal; - -uniform mat4 ModelMat; -uniform mat4 PerspMat; - -void main() -{ - gl_Position = PerspMat * ModelMat * vec4(Position, 1.0); - CamNormal = (ModelMat * vec4(Normal, 0.0)).xyz; -} \ No newline at end of file diff --git a/lib / renderer / gl / data /prt.fs b/lib / renderer / gl / data /prt.fs deleted file mode 100644 index 3737e2a5b51d5b001cdb2f3d8157793caa98fd54..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / data /prt.fs +++ /dev/null @@ -1,157 +0,0 @@ -#version 330 - -uniform vec3 SHCoeffs[9]; -uniform uint analytic; - -uniform uint hasNormalMap; -uniform uint hasAlbedoMap; - -uniform sampler2D AlbedoMap; -uniform sampler2D NormalMap; - -in VertexData { - vec3 Position; - vec3 Depth; - vec3 ModelNormal; - vec2 Texcoord; - vec3 Tangent; - vec3 Bitangent; - vec3 PRT1; - vec3 PRT2; - vec3 PRT3; - vec3 Label; -} VertexIn; - -layout (location = 0) out vec4 FragColor; -layout (location = 1) out vec4 FragNormal; -layout (location = 2) out vec4 FragPosition; -layout (location = 3) out vec4 FragAlbedo; -layout (location = 4) out vec4 FragShading; -layout (location = 5) out vec4 FragPRT1; -layout (location = 6) out vec4 FragPRT2; -// layout (location = 7) out vec4 FragPRT3; -layout (location = 7) out vec4 FragLabel; - - -vec4 gammaCorrection(vec4 vec, float g) -{ - return vec4(pow(vec.x, 1.0/g), pow(vec.y, 1.0/g), pow(vec.z, 1.0/g), vec.w); -} - -vec3 gammaCorrection(vec3 vec, float g) -{ - return vec3(pow(vec.x, 1.0/g), pow(vec.y, 1.0/g), pow(vec.z, 1.0/g)); -} - -void evaluateH(vec3 n, out float H[9]) -{ - float c1 = 0.429043, c2 = 0.511664, - c3 = 0.743125, c4 = 0.886227, c5 = 0.247708; - - H[0] = c4; - H[1] = 2.0 * c2 * n[1]; - H[2] = 2.0 * c2 * n[2]; - H[3] = 2.0 * c2 * n[0]; - H[4] = 2.0 * c1 * n[0] * n[1]; - H[5] = 2.0 * c1 * n[1] * n[2]; - H[6] = c3 * n[2] * n[2] - c5; - H[7] = 2.0 * c1 * n[2] * n[0]; - H[8] = c1 * (n[0] * n[0] - n[1] * n[1]); -} - -vec3 evaluateLightingModel(vec3 normal) -{ - float H[9]; - evaluateH(normal, H); - vec3 res = vec3(0.0); - for (int i = 0; i < 9; i++) { - res += H[i] * SHCoeffs[i]; - } - return res; -} - -// nC: coarse geometry normal, nH: fine normal from normal map -vec3 evaluateLightingModelHybrid(vec3 nC, vec3 nH, mat3 prt) -{ - float HC[9], HH[9]; - evaluateH(nC, HC); - evaluateH(nH, HH); - - vec3 res = vec3(0.0); - vec3 shadow = vec3(0.0); - vec3 unshadow = vec3(0.0); - for(int i = 0; i < 3; ++i){ - for(int j = 0; j < 3; ++j){ - int id = i*3+j; - res += HH[id]* SHCoeffs[id]; - shadow += prt[i][j] * SHCoeffs[id]; - unshadow += HC[id] * SHCoeffs[id]; - } - } - vec3 ratio = clamp(shadow/unshadow,0.0,1.0); - res = ratio * res; - - return res; -} - -vec3 evaluateLightingModelPRT(mat3 prt) -{ - vec3 res = vec3(0.0); - for(int i = 0; i < 3; ++i){ - for(int j = 0; j < 3; ++j){ - res += prt[i][j] * SHCoeffs[i*3+j]; - } - } - - return res; -} - -void main() -{ - vec2 uv = VertexIn.Texcoord; - vec3 nC = normalize(VertexIn.ModelNormal); - vec3 nml = nC; - mat3 prt = mat3(VertexIn.PRT1, VertexIn.PRT2, VertexIn.PRT3); - - if(hasAlbedoMap == uint(0)) - FragAlbedo = vec4(1.0); - else - FragAlbedo = texture(AlbedoMap, uv);//gammaCorrection(texture(AlbedoMap, uv), 1.0/2.2); - - if(hasNormalMap == uint(0)) - { - if(analytic == uint(0)) - FragShading = vec4(evaluateLightingModelPRT(prt), 1.0f); - else - FragShading = vec4(evaluateLightingModel(nC), 1.0f); - } - else - { - vec3 n_tan = normalize(texture(NormalMap, uv).rgb*2.0-vec3(1.0)); - - mat3 TBN = mat3(normalize(VertexIn.Tangent),normalize(VertexIn.Bitangent),nC); - vec3 nH = normalize(TBN * n_tan); - - if(analytic == uint(0)) - FragShading = vec4(evaluateLightingModelHybrid(nC,nH,prt),1.0f); - else - FragShading = vec4(evaluateLightingModel(nH), 1.0f); - - nml = nH; - } - - FragShading = gammaCorrection(FragShading, 2.2); - FragColor = clamp(FragAlbedo * FragShading, 0.0, 1.0); - FragNormal = vec4(0.5*(nml+vec3(1.0)), 1.0); - FragPosition = vec4(VertexIn.Depth.xyz, 1.0); - FragShading = vec4(clamp(0.5*FragShading.xyz, 0.0, 1.0),1.0); - // FragColor = gammaCorrection(clamp(FragAlbedo * FragShading, 0.0, 1.0),2.2); - // FragNormal = vec4(0.5*(nml+vec3(1.0)), 1.0); - // FragPosition = vec4(VertexIn.Position,VertexIn.Depth.x); - // FragShading = vec4(gammaCorrection(clamp(0.5*FragShading.xyz, 0.0, 1.0),2.2),1.0); - // FragAlbedo = gammaCorrection(FragAlbedo,2.2); - FragPRT1 = vec4(VertexIn.PRT1,1.0); - FragPRT2 = vec4(VertexIn.PRT2,1.0); - // FragPRT3 = vec4(VertexIn.PRT3,1.0); - FragLabel = vec4(VertexIn.Label,1.0); -} \ No newline at end of file diff --git a/lib / renderer / gl / data /prt.vs b/lib / renderer / gl / data /prt.vs deleted file mode 100644 index 60cc800f20a5cc765a433e08e634684f014f7513..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / data /prt.vs +++ /dev/null @@ -1,156 +0,0 @@ -#version 330 - -layout (location = 0) in vec3 a_Position; -layout (location = 1) in vec3 a_Normal; -layout (location = 2) in vec2 a_TextureCoord; -layout (location = 3) in vec3 a_Tangent; -layout (location = 4) in vec3 a_Bitangent; -layout (location = 5) in vec3 a_PRT1; -layout (location = 6) in vec3 a_PRT2; -layout (location = 7) in vec3 a_PRT3; -layout (location = 8) in vec3 a_Label; -out VertexData { - vec3 Position; - vec3 Depth; - vec3 ModelNormal; - vec2 Texcoord; - vec3 Tangent; - vec3 Bitangent; - vec3 PRT1; - vec3 PRT2; - vec3 PRT3; - vec3 Label; -} VertexOut; -uniform mat3 RotMat; -uniform mat4 NormMat; -uniform mat4 ModelMat; -uniform mat4 PerspMat; -float s_c3 = 0.94617469575; // (3*sqrt(5))/(4*sqrt(pi)) -float s_c4 = -0.31539156525;// (-sqrt(5))/(4*sqrt(pi)) -float s_c5 = 0.54627421529; // (sqrt(15))/(4*sqrt(pi)) -float s_c_scale = 1.0/0.91529123286551084; -float s_c_scale_inv = 0.91529123286551084; -float s_rc2 = 1.5853309190550713*s_c_scale; -float s_c4_div_c3 = s_c4/s_c3; -float s_c4_div_c3_x2 = (s_c4/s_c3)*2.0; -float s_scale_dst2 = s_c3 * s_c_scale_inv; -float s_scale_dst4 = s_c5 * s_c_scale_inv; -void OptRotateBand0(float x[1], mat3 R, out float dst[1]) -{ - dst[0] = x[0]; -} -// 9 multiplies -void OptRotateBand1(float x[3], mat3 R, out float dst[3]) -{ - // derived from SlowRotateBand1 - dst[0] = ( R[1][1])*x[0] + (-R[1][2])*x[1] + ( R[1][0])*x[2]; - dst[1] = (-R[2][1])*x[0] + ( R[2][2])*x[1] + (-R[2][0])*x[2]; - dst[2] = ( R[0][1])*x[0] + (-R[0][2])*x[1] + ( R[0][0])*x[2]; -} -// 48 multiplies -void OptRotateBand2(float x[5], mat3 R, out float dst[5]) -{ - // Sparse matrix multiply - float sh0 = x[3] + x[4] + x[4] - x[1]; - float sh1 = x[0] + s_rc2*x[2] + x[3] + x[4]; - float sh2 = x[0]; - float sh3 = -x[3]; - float sh4 = -x[1]; - - // Rotations. R0 and R1 just use the raw matrix columns - float r2x = R[0][0] + R[0][1]; - float r2y = R[1][0] + R[1][1]; - float r2z = R[2][0] + R[2][1]; - - float r3x = R[0][0] + R[0][2]; - float r3y = R[1][0] + R[1][2]; - float r3z = R[2][0] + R[2][2]; - - float r4x = R[0][1] + R[0][2]; - float r4y = R[1][1] + R[1][2]; - float r4z = R[2][1] + R[2][2]; - - // dense matrix multiplication one column at a time - - // column 0 - float sh0_x = sh0 * R[0][0]; - float sh0_y = sh0 * R[1][0]; - float d0 = sh0_x * R[1][0]; - float d1 = sh0_y * R[2][0]; - float d2 = sh0 * (R[2][0] * R[2][0] + s_c4_div_c3); - float d3 = sh0_x * R[2][0]; - float d4 = sh0_x * R[0][0] - sh0_y * R[1][0]; - - // column 1 - float sh1_x = sh1 * R[0][2]; - float sh1_y = sh1 * R[1][2]; - d0 += sh1_x * R[1][2]; - d1 += sh1_y * R[2][2]; - d2 += sh1 * (R[2][2] * R[2][2] + s_c4_div_c3); - d3 += sh1_x * R[2][2]; - d4 += sh1_x * R[0][2] - sh1_y * R[1][2]; - - // column 2 - float sh2_x = sh2 * r2x; - float sh2_y = sh2 * r2y; - d0 += sh2_x * r2y; - d1 += sh2_y * r2z; - d2 += sh2 * (r2z * r2z + s_c4_div_c3_x2); - d3 += sh2_x * r2z; - d4 += sh2_x * r2x - sh2_y * r2y; - - // column 3 - float sh3_x = sh3 * r3x; - float sh3_y = sh3 * r3y; - d0 += sh3_x * r3y; - d1 += sh3_y * r3z; - d2 += sh3 * (r3z * r3z + s_c4_div_c3_x2); - d3 += sh3_x * r3z; - d4 += sh3_x * r3x - sh3_y * r3y; - - // column 4 - float sh4_x = sh4 * r4x; - float sh4_y = sh4 * r4y; - d0 += sh4_x * r4y; - d1 += sh4_y * r4z; - d2 += sh4 * (r4z * r4z + s_c4_div_c3_x2); - d3 += sh4_x * r4z; - d4 += sh4_x * r4x - sh4_y * r4y; - - // extra multipliers - dst[0] = d0; - dst[1] = -d1; - dst[2] = d2 * s_scale_dst2; - dst[3] = -d3; - dst[4] = d4 * s_scale_dst4; -} -void main() -{ - // normalization - vec3 pos = (NormMat * vec4(a_Position,1.0)).xyz; - mat3 R = mat3(ModelMat) * RotMat; - VertexOut.ModelNormal = (R * a_Normal); - VertexOut.Position = R * pos; - VertexOut.Texcoord = a_TextureCoord; - VertexOut.Tangent = (R * a_Tangent); - VertexOut.Bitangent = (R * a_Bitangent); - VertexOut.Label = a_Label; - float PRT0, PRT1[3], PRT2[5]; - PRT0 = a_PRT1[0]; - PRT1[0] = a_PRT1[1]; - PRT1[1] = a_PRT1[2]; - PRT1[2] = a_PRT2[0]; - PRT2[0] = a_PRT2[1]; - PRT2[1] = a_PRT2[2]; - PRT2[2] = a_PRT3[0]; - PRT2[3] = a_PRT3[1]; - PRT2[4] = a_PRT3[2]; - OptRotateBand1(PRT1, R, PRT1); - OptRotateBand2(PRT2, R, PRT2); - VertexOut.PRT1 = vec3(PRT0,PRT1[0],PRT1[1]); - VertexOut.PRT2 = vec3(PRT1[2],PRT2[0],PRT2[1]); - VertexOut.PRT3 = vec3(PRT2[2],PRT2[3],PRT2[4]); - gl_Position = PerspMat * ModelMat * vec4(RotMat * pos, 1.0); - - VertexOut.Depth = vec3(gl_Position.z / gl_Position.w); -} \ No newline at end of file diff --git a/lib / renderer / gl / data /prt_uv.fs b/lib / renderer / gl / data /prt_uv.fs deleted file mode 100644 index 6e90b25c62b41c8cf61afd29333372193047d5f1..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / data /prt_uv.fs +++ /dev/null @@ -1,141 +0,0 @@ -#version 330 - -uniform vec3 SHCoeffs[9]; -uniform uint analytic; - -uniform uint hasNormalMap; -uniform uint hasAlbedoMap; - -uniform sampler2D AlbedoMap; -uniform sampler2D NormalMap; - -in VertexData { - vec3 Position; - vec3 ModelNormal; - vec3 CameraNormal; - vec2 Texcoord; - vec3 Tangent; - vec3 Bitangent; - vec3 PRT1; - vec3 PRT2; - vec3 PRT3; -} VertexIn; - -layout (location = 0) out vec4 FragColor; -layout (location = 1) out vec4 FragPosition; -layout (location = 2) out vec4 FragNormal; - -vec4 gammaCorrection(vec4 vec, float g) -{ - return vec4(pow(vec.x, 1.0/g), pow(vec.y, 1.0/g), pow(vec.z, 1.0/g), vec.w); -} - -vec3 gammaCorrection(vec3 vec, float g) -{ - return vec3(pow(vec.x, 1.0/g), pow(vec.y, 1.0/g), pow(vec.z, 1.0/g)); -} - -void evaluateH(vec3 n, out float H[9]) -{ - float c1 = 0.429043, c2 = 0.511664, - c3 = 0.743125, c4 = 0.886227, c5 = 0.247708; - - H[0] = c4; - H[1] = 2.0 * c2 * n[1]; - H[2] = 2.0 * c2 * n[2]; - H[3] = 2.0 * c2 * n[0]; - H[4] = 2.0 * c1 * n[0] * n[1]; - H[5] = 2.0 * c1 * n[1] * n[2]; - H[6] = c3 * n[2] * n[2] - c5; - H[7] = 2.0 * c1 * n[2] * n[0]; - H[8] = c1 * (n[0] * n[0] - n[1] * n[1]); -} - -vec3 evaluateLightingModel(vec3 normal) -{ - float H[9]; - evaluateH(normal, H); - vec3 res = vec3(0.0); - for (int i = 0; i < 9; i++) { - res += H[i] * SHCoeffs[i]; - } - return res; -} - -// nC: coarse geometry normal, nH: fine normal from normal map -vec3 evaluateLightingModelHybrid(vec3 nC, vec3 nH, mat3 prt) -{ - float HC[9], HH[9]; - evaluateH(nC, HC); - evaluateH(nH, HH); - - vec3 res = vec3(0.0); - vec3 shadow = vec3(0.0); - vec3 unshadow = vec3(0.0); - for(int i = 0; i < 3; ++i){ - for(int j = 0; j < 3; ++j){ - int id = i*3+j; - res += HH[id]* SHCoeffs[id]; - shadow += prt[i][j] * SHCoeffs[id]; - unshadow += HC[id] * SHCoeffs[id]; - } - } - vec3 ratio = clamp(shadow/unshadow,0.0,1.0); - res = ratio * res; - - return res; -} - -vec3 evaluateLightingModelPRT(mat3 prt) -{ - vec3 res = vec3(0.0); - for(int i = 0; i < 3; ++i){ - for(int j = 0; j < 3; ++j){ - res += prt[i][j] * SHCoeffs[i*3+j]; - } - } - - return res; -} - -void main() -{ - vec2 uv = VertexIn.Texcoord; - vec3 nM = normalize(VertexIn.ModelNormal); - vec3 nC = normalize(VertexIn.CameraNormal); - vec3 nml = nC; - mat3 prt = mat3(VertexIn.PRT1, VertexIn.PRT2, VertexIn.PRT3); - - vec4 albedo, shading; - if(hasAlbedoMap == uint(0)) - albedo = vec4(1.0); - else - albedo = texture(AlbedoMap, uv);//gammaCorrection(texture(AlbedoMap, uv), 1.0/2.2); - - if(hasNormalMap == uint(0)) - { - if(analytic == uint(0)) - shading = vec4(evaluateLightingModelPRT(prt), 1.0f); - else - shading = vec4(evaluateLightingModel(nC), 1.0f); - } - else - { - vec3 n_tan = normalize(texture(NormalMap, uv).rgb*2.0-vec3(1.0)); - - mat3 TBN = mat3(normalize(VertexIn.Tangent),normalize(VertexIn.Bitangent),nC); - vec3 nH = normalize(TBN * n_tan); - - if(analytic == uint(0)) - shading = vec4(evaluateLightingModelHybrid(nC,nH,prt),1.0f); - else - shading = vec4(evaluateLightingModel(nH), 1.0f); - - nml = nH; - } - - shading = gammaCorrection(shading, 2.2); - FragColor = clamp(albedo * shading, 0.0, 1.0); - FragPosition = vec4(VertexIn.Position,1.0); - FragNormal = vec4(0.5*(nM+vec3(1.0)),1.0); -} \ No newline at end of file diff --git a/lib / renderer / gl / data /prt_uv.vs b/lib / renderer / gl / data /prt_uv.vs deleted file mode 100644 index c3cc0a2f2b381032780e450020a96d353168cded..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / data /prt_uv.vs +++ /dev/null @@ -1,168 +0,0 @@ -#version 330 - -layout (location = 0) in vec3 a_Position; -layout (location = 1) in vec3 a_Normal; -layout (location = 2) in vec2 a_TextureCoord; -layout (location = 3) in vec3 a_Tangent; -layout (location = 4) in vec3 a_Bitangent; -layout (location = 5) in vec3 a_PRT1; -layout (location = 6) in vec3 a_PRT2; -layout (location = 7) in vec3 a_PRT3; - -out VertexData { - vec3 Position; - vec3 ModelNormal; - vec3 CameraNormal; - vec2 Texcoord; - vec3 Tangent; - vec3 Bitangent; - vec3 PRT1; - vec3 PRT2; - vec3 PRT3; -} VertexOut; - -uniform mat3 RotMat; -uniform mat4 NormMat; -uniform mat4 ModelMat; -uniform mat4 PerspMat; - -#define pi 3.1415926535897932384626433832795 - -float s_c3 = 0.94617469575; // (3*sqrt(5))/(4*sqrt(pi)) -float s_c4 = -0.31539156525;// (-sqrt(5))/(4*sqrt(pi)) -float s_c5 = 0.54627421529; // (sqrt(15))/(4*sqrt(pi)) - -float s_c_scale = 1.0/0.91529123286551084; -float s_c_scale_inv = 0.91529123286551084; - -float s_rc2 = 1.5853309190550713*s_c_scale; -float s_c4_div_c3 = s_c4/s_c3; -float s_c4_div_c3_x2 = (s_c4/s_c3)*2.0; - -float s_scale_dst2 = s_c3 * s_c_scale_inv; -float s_scale_dst4 = s_c5 * s_c_scale_inv; - -void OptRotateBand0(float x[1], mat3 R, out float dst[1]) -{ - dst[0] = x[0]; -} - -// 9 multiplies -void OptRotateBand1(float x[3], mat3 R, out float dst[3]) -{ - // derived from SlowRotateBand1 - dst[0] = ( R[1][1])*x[0] + (-R[1][2])*x[1] + ( R[1][0])*x[2]; - dst[1] = (-R[2][1])*x[0] + ( R[2][2])*x[1] + (-R[2][0])*x[2]; - dst[2] = ( R[0][1])*x[0] + (-R[0][2])*x[1] + ( R[0][0])*x[2]; -} - -// 48 multiplies -void OptRotateBand2(float x[5], mat3 R, out float dst[5]) -{ - // Sparse matrix multiply - float sh0 = x[3] + x[4] + x[4] - x[1]; - float sh1 = x[0] + s_rc2*x[2] + x[3] + x[4]; - float sh2 = x[0]; - float sh3 = -x[3]; - float sh4 = -x[1]; - - // Rotations. R0 and R1 just use the raw matrix columns - float r2x = R[0][0] + R[0][1]; - float r2y = R[1][0] + R[1][1]; - float r2z = R[2][0] + R[2][1]; - - float r3x = R[0][0] + R[0][2]; - float r3y = R[1][0] + R[1][2]; - float r3z = R[2][0] + R[2][2]; - - float r4x = R[0][1] + R[0][2]; - float r4y = R[1][1] + R[1][2]; - float r4z = R[2][1] + R[2][2]; - - // dense matrix multiplication one column at a time - - // column 0 - float sh0_x = sh0 * R[0][0]; - float sh0_y = sh0 * R[1][0]; - float d0 = sh0_x * R[1][0]; - float d1 = sh0_y * R[2][0]; - float d2 = sh0 * (R[2][0] * R[2][0] + s_c4_div_c3); - float d3 = sh0_x * R[2][0]; - float d4 = sh0_x * R[0][0] - sh0_y * R[1][0]; - - // column 1 - float sh1_x = sh1 * R[0][2]; - float sh1_y = sh1 * R[1][2]; - d0 += sh1_x * R[1][2]; - d1 += sh1_y * R[2][2]; - d2 += sh1 * (R[2][2] * R[2][2] + s_c4_div_c3); - d3 += sh1_x * R[2][2]; - d4 += sh1_x * R[0][2] - sh1_y * R[1][2]; - - // column 2 - float sh2_x = sh2 * r2x; - float sh2_y = sh2 * r2y; - d0 += sh2_x * r2y; - d1 += sh2_y * r2z; - d2 += sh2 * (r2z * r2z + s_c4_div_c3_x2); - d3 += sh2_x * r2z; - d4 += sh2_x * r2x - sh2_y * r2y; - - // column 3 - float sh3_x = sh3 * r3x; - float sh3_y = sh3 * r3y; - d0 += sh3_x * r3y; - d1 += sh3_y * r3z; - d2 += sh3 * (r3z * r3z + s_c4_div_c3_x2); - d3 += sh3_x * r3z; - d4 += sh3_x * r3x - sh3_y * r3y; - - // column 4 - float sh4_x = sh4 * r4x; - float sh4_y = sh4 * r4y; - d0 += sh4_x * r4y; - d1 += sh4_y * r4z; - d2 += sh4 * (r4z * r4z + s_c4_div_c3_x2); - d3 += sh4_x * r4z; - d4 += sh4_x * r4x - sh4_y * r4y; - - // extra multipliers - dst[0] = d0; - dst[1] = -d1; - dst[2] = d2 * s_scale_dst2; - dst[3] = -d3; - dst[4] = d4 * s_scale_dst4; -} - -void main() -{ - // normalization - mat3 R = mat3(ModelMat) * RotMat; - VertexOut.ModelNormal = a_Normal; - VertexOut.CameraNormal = (R * a_Normal); - VertexOut.Position = a_Position; - VertexOut.Texcoord = a_TextureCoord; - VertexOut.Tangent = (R * a_Tangent); - VertexOut.Bitangent = (R * a_Bitangent); - float PRT0, PRT1[3], PRT2[5]; - PRT0 = a_PRT1[0]; - PRT1[0] = a_PRT1[1]; - PRT1[1] = a_PRT1[2]; - PRT1[2] = a_PRT2[0]; - PRT2[0] = a_PRT2[1]; - PRT2[1] = a_PRT2[2]; - PRT2[2] = a_PRT3[0]; - PRT2[3] = a_PRT3[1]; - PRT2[4] = a_PRT3[2]; - - OptRotateBand1(PRT1, R, PRT1); - OptRotateBand2(PRT2, R, PRT2); - - VertexOut.PRT1 = vec3(PRT0,PRT1[0],PRT1[1]); - VertexOut.PRT2 = vec3(PRT1[2],PRT2[0],PRT2[1]); - VertexOut.PRT3 = vec3(PRT2[2],PRT2[3],PRT2[4]); - - gl_Position = vec4(a_TextureCoord, 0.0, 1.0) - vec4(0.5, 0.5, 0, 0); - gl_Position[0] *= 2.0; - gl_Position[1] *= 2.0; -} \ No newline at end of file diff --git a/lib / renderer / gl / data /quad.fs b/lib / renderer / gl / data /quad.fs deleted file mode 100644 index f43502f2352ca2adf19d11e809946b51498df5a5..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / data /quad.fs +++ /dev/null @@ -1,11 +0,0 @@ -#version 330 core -out vec4 FragColor; - -in vec2 TexCoord; - -uniform sampler2D screenTexture; - -void main() -{ - FragColor = texture(screenTexture, TexCoord); -} \ No newline at end of file diff --git a/lib / renderer / gl / data /quad.vs b/lib / renderer / gl / data /quad.vs deleted file mode 100644 index 811044631a1f29f5b45c490b2d40297f3127b6ea..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / data /quad.vs +++ /dev/null @@ -1,11 +0,0 @@ -#version 330 core -layout (location = 0) in vec2 aPos; -layout (location = 1) in vec2 aTexCoord; - -out vec2 TexCoord; - -void main() -{ - gl_Position = vec4(aPos.x, aPos.y, 0.0, 1.0); - TexCoord = aTexCoord; -} \ No newline at end of file diff --git a/lib / renderer / gl / framework.py b/lib / renderer / gl / framework.py deleted file mode 100644 index 90bc4c64a9f855d455da70f085358118a2f49b15..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / framework.py +++ /dev/null @@ -1,95 +0,0 @@ -# Mario Rosasco, 2016 -# adapted from framework.cpp, Copyright (C) 2010-2012 by Jason L. McKesson -# This file is licensed under the MIT License. -# -# NB: Unlike in the framework.cpp organization, the main loop is contained -# in the tutorial files, not in this framework file. Additionally, a copy of -# this module file must exist in the same directory as the tutorial files -# to be imported properly. - -import os -from OpenGL.GL import * - - -# Function that creates and compiles shaders according to the given type (a GL enum value) and -# shader program (a file containing a GLSL program). -def loadShader(shaderType, shaderFile): - # check if file exists, get full path name - strFilename = findFileOrThrow(shaderFile) - shaderData = None - with open(strFilename, 'r') as f: - shaderData = f.read() - - shader = glCreateShader(shaderType) - glShaderSource( - shader, - shaderData) # note that this is a simpler function call than in C - - # This shader compilation is more explicit than the one used in - # framework.cpp, which relies on a glutil wrapper function. - # This is made explicit here mainly to decrease dependence on pyOpenGL - # utilities and wrappers, which docs caution may change in future versions. - glCompileShader(shader) - - status = glGetShaderiv(shader, GL_COMPILE_STATUS) - if status == GL_FALSE: - # Note that getting the error log is much simpler in Python than in C/C++ - # and does not require explicit handling of the string buffer - strInfoLog = glGetShaderInfoLog(shader) - strShaderType = "" - if shaderType is GL_VERTEX_SHADER: - strShaderType = "vertex" - elif shaderType is GL_GEOMETRY_SHADER: - strShaderType = "geometry" - elif shaderType is GL_FRAGMENT_SHADER: - strShaderType = "fragment" - - print("Compilation failure for " + strShaderType + " shader:\n" + - str(strInfoLog)) - - return shader - - -# Function that accepts a list of shaders, compiles them, and returns a handle to the compiled program -def createProgram(shaderList): - program = glCreateProgram() - - for shader in shaderList: - glAttachShader(program, shader) - - glLinkProgram(program) - - status = glGetProgramiv(program, GL_LINK_STATUS) - if status == GL_FALSE: - # Note that getting the error log is much simpler in Python than in C/C++ - # and does not require explicit handling of the string buffer - strInfoLog = glGetProgramInfoLog(program) - print("Linker failure: \n" + str(strInfoLog)) - - for shader in shaderList: - glDetachShader(program, shader) - - return program - - -# Helper function to locate and open the target file (passed in as a string). -# Returns the full path to the file as a string. -def findFileOrThrow(strBasename): - # Keep constant names in C-style convention, for readability - # when comparing to C(/C++) code. - if os.path.isfile(strBasename): - return strBasename - - LOCAL_FILE_DIR = "data" + os.sep - GLOBAL_FILE_DIR = os.path.dirname( - os.path.abspath(__file__)) + os.sep + "data" + os.sep - - strFilename = LOCAL_FILE_DIR + strBasename - if os.path.isfile(strFilename): - return strFilename - - strFilename = GLOBAL_FILE_DIR + strBasename - if os.path.isfile(strFilename): - return strFilename - - raise IOError('Could not find target file ' + strBasename) \ No newline at end of file diff --git a/lib / renderer / gl / glcontext.py b/lib / renderer / gl / glcontext.py deleted file mode 100644 index 0601b79d856e83675d219dde3feade7a9f2c836c..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / glcontext.py +++ /dev/null @@ -1,136 +0,0 @@ -"""Headless GPU-accelerated OpenGL context creation on Google Colaboratory. - -Typical usage: - - # Optional PyOpenGL configuratiopn can be done here. - # import OpenGL - # OpenGL.ERROR_CHECKING = True - - # 'glcontext' must be imported before any OpenGL.* API. - from lucid.misc.gl.glcontext import create_opengl_context - - # Now it's safe to import OpenGL and EGL functions - import OpenGL.GL as gl - - # create_opengl_context() creates a GL context that is attached to an - # offscreen surface of the specified size. Note that rendering to buffers - # of other sizes and formats is still possible with OpenGL Framebuffers. - # - # Users are expected to directly use the EGL API in case more advanced - # context management is required. - width, height = 640, 480 - create_opengl_context((width, height)) - - # OpenGL context is available here. - -""" - -from __future__ import print_function - -# pylint: disable=unused-import,g-import-not-at-top,g-statement-before-imports - -try: - import OpenGL -except: - print('This module depends on PyOpenGL.') - print('Please run "\033[1m!pip install -q pyopengl\033[0m" ' - 'prior importing this module.') - raise - -import ctypes -from ctypes import pointer, util -import os - -os.environ['PYOPENGL_PLATFORM'] = 'egl' - -# OpenGL loading workaround. -# -# * PyOpenGL tries to load libGL, but we need libOpenGL, see [1,2]. -# This could have been solved by a symlink libGL->libOpenGL, but: -# -# * Python 2.7 can't find libGL and linEGL due to a bug (see [3]) -# in ctypes.util, that was only wixed in Python 3.6. -# -# So, the only solution I've found is to monkeypatch ctypes.util -# [1] https://devblogs.nvidia.com/egl-eye-opengl-visualization-without-x-server/ -# [2] https://devblogs.nvidia.com/linking-opengl-server-side-rendering/ -# [3] https://bugs.python.org/issue9998 -_find_library_old = ctypes.util.find_library -try: - - def _find_library_new(name): - return { - 'GL': 'libOpenGL.so', - 'EGL': 'libEGL.so', - }.get(name, _find_library_old(name)) - - util.find_library = _find_library_new - import OpenGL.GL as gl - import OpenGL.EGL as egl -except: - print('Unable to load OpenGL libraries. ' - 'Make sure you use GPU-enabled backend.') - print('Press "Runtime->Change runtime type" and set ' - '"Hardware accelerator" to GPU.') - raise -finally: - util.find_library = _find_library_old - - -def create_opengl_context(surface_size=(640, 480)): - """Create offscreen OpenGL context and make it current. - - Users are expected to directly use EGL API in case more advanced - context management is required. - - Args: - surface_size: (width, height), size of the offscreen rendering surface. - """ - egl_display = egl.eglGetDisplay(egl.EGL_DEFAULT_DISPLAY) - - major, minor = egl.EGLint(), egl.EGLint() - egl.eglInitialize(egl_display, pointer(major), pointer(minor)) - - config_attribs = [ - egl.EGL_SURFACE_TYPE, egl.EGL_PBUFFER_BIT, egl.EGL_BLUE_SIZE, 8, - egl.EGL_GREEN_SIZE, 8, egl.EGL_RED_SIZE, 8, egl.EGL_DEPTH_SIZE, 24, - egl.EGL_RENDERABLE_TYPE, egl.EGL_OPENGL_BIT, egl.EGL_NONE - ] - config_attribs = (egl.EGLint * len(config_attribs))(*config_attribs) - - num_configs = egl.EGLint() - egl_cfg = egl.EGLConfig() - egl.eglChooseConfig(egl_display, config_attribs, pointer(egl_cfg), 1, - pointer(num_configs)) - - width, height = surface_size - pbuffer_attribs = [ - egl.EGL_WIDTH, - width, - egl.EGL_HEIGHT, - height, - egl.EGL_NONE, - ] - pbuffer_attribs = (egl.EGLint * len(pbuffer_attribs))(*pbuffer_attribs) - egl_surf = egl.eglCreatePbufferSurface(egl_display, egl_cfg, - pbuffer_attribs) - - egl.eglBindAPI(egl.EGL_OPENGL_API) - - context_attribs = None - # context_attribs = [ - # egl.EGL_CONTEXT_MAJOR_VERSION, - # 4, - # egl.EGL_CONTEXT_MINOR_VERSION, - # 1, - # egl.EGL_NONE, - # ] - - egl_context = egl.eglCreateContext(egl_display, egl_cfg, - egl.EGL_NO_CONTEXT, context_attribs) - egl.eglMakeCurrent(egl_display, egl_surf, egl_surf, egl_context) - - buffer_type = egl.EGLint() - out = egl.eglQueryContext(egl_display, egl_context, - egl.EGL_CONTEXT_CLIENT_VERSION, buffer_type) - # print(buffer_type) \ No newline at end of file diff --git a/lib / renderer / gl / init_gl.py b/lib / renderer / gl / init_gl.py deleted file mode 100644 index 92d7e1a90b54d7a568e5849d6400bd960e14958a..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / init_gl.py +++ /dev/null @@ -1,24 +0,0 @@ -_glut_window = None -_context_inited = None - - -def initialize_GL_context(width=512, height=512, egl=False): - ''' - default context uses GLUT - ''' - if not egl: - import OpenGL.GLUT as GLUT - display_mode = GLUT.GLUT_DOUBLE | GLUT.GLUT_RGB | GLUT.GLUT_DEPTH - global _glut_window - if _glut_window is None: - GLUT.glutInit() - GLUT.glutInitDisplayMode(display_mode) - GLUT.glutInitWindowSize(width, height) - GLUT.glutInitWindowPosition(0, 0) - _glut_window = GLUT.glutCreateWindow("My Render.") - else: - from .glcontext import create_opengl_context - global _context_inited - if _context_inited is None: - create_opengl_context((width, height)) - _context_inited = True \ No newline at end of file diff --git a/lib / renderer / gl / norm_render.py b/lib / renderer / gl / norm_render.py deleted file mode 100644 index fa1ca7ea15139f3eae8d375a9a52223fee507dda..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / norm_render.py +++ /dev/null @@ -1,75 +0,0 @@ -''' -MIT License -Copyright (c) 2019 Shunsuke Saito, Zeng Huang, and Ryota Natsume -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. -''' -from OpenGL.GLUT import * - -from .render2 import Render - - -class NormRender(Render): - def __init__(self, - width=1600, - height=1200, - name='Cam Renderer', - program_files=['simple.fs', 'simple.vs'], - color_size=1, - ms_rate=1): - Render.__init__(self, width, height, name, program_files, color_size, - ms_rate) - self.camera = None - - glutDisplayFunc(self.display) - glutKeyboardFunc(self.keyboard) - - def set_camera(self, camera): - self.camera = camera - self.projection_matrix, self.model_view_matrix = camera.get_gl_matrix() - - def set_matrices(self, projection, modelview): - self.projection_matrix = projection - self.model_view_matrix = modelview - - def keyboard(self, key, x, y): - # up - eps = 1 - # print(key) - if key == b'w': - self.camera.center += eps * self.camera.direction - elif key == b's': - self.camera.center -= eps * self.camera.direction - if key == b'a': - self.camera.center -= eps * self.camera.right - elif key == b'd': - self.camera.center += eps * self.camera.right - if key == b' ': - self.camera.center += eps * self.camera.up - elif key == b'x': - self.camera.center -= eps * self.camera.up - elif key == b'i': - self.camera.near += 0.1 * eps - self.camera.far += 0.1 * eps - elif key == b'o': - self.camera.near -= 0.1 * eps - self.camera.far -= 0.1 * eps - - self.projection_matrix, self.model_view_matrix = self.camera.get_gl_matrix( - ) - - def show(self): - glutMainLoop() \ No newline at end of file diff --git a/lib / renderer / gl / render.py b/lib / renderer / gl / render.py deleted file mode 100644 index 3a7e8da2bb1d041ecb74596f4cd9b73a4b30db27..0000000000000000000000000000000000000000 --- a/lib / renderer / gl / render.py +++ /dev/null @@ -1,380 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from ctypes import * - -import numpy as np -from .framework import * - -GLUT = None - - -# NOTE: Render class assumes GL context is created already. -class Render: - def __init__(self, - width=1600, - height=1200, - name='GL Renderer', - program_files=['simple.fs', 'simple.vs'], - color_size=1, - ms_rate=1, - egl=False): - self.width = width - self.height = height - self.name = name - self.use_inverse_depth = False - self.egl = egl - - glEnable(GL_DEPTH_TEST) - - glClampColor(GL_CLAMP_READ_COLOR, GL_FALSE) - glClampColor(GL_CLAMP_FRAGMENT_COLOR, GL_FALSE) - glClampColor(GL_CLAMP_VERTEX_COLOR, GL_FALSE) - - # init program - shader_list = [] - - for program_file in program_files: - _, ext = os.path.splitext(program_file) - if ext == '.vs': - shader_list.append(loadShader(GL_VERTEX_SHADER, program_file)) - elif ext == '.fs': - shader_list.append(loadShader(GL_FRAGMENT_SHADER, - program_file)) - elif ext == '.gs': - shader_list.append(loadShader(GL_GEOMETRY_SHADER, - program_file)) - - self.program = createProgram(shader_list) - - for shader in shader_list: - glDeleteShader(shader) - - # Init uniform variables - self.model_mat_unif = glGetUniformLocation(self.program, 'ModelMat') - self.persp_mat_unif = glGetUniformLocation(self.program, 'PerspMat') - - self.vertex_buffer = glGenBuffers(1) - - # Init screen quad program and buffer - self.quad_program, self.quad_buffer = self.init_quad_program() - - # Configure frame buffer - self.frame_buffer = glGenFramebuffers(1) - glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) - - self.intermediate_fbo = None - if ms_rate > 1: - # Configure texture buffer to render to - self.color_buffer = [] - for i in range(color_size): - color_buffer = glGenTextures(1) - multi_sample_rate = ms_rate - glBindTexture(GL_TEXTURE_2D_MULTISAMPLE, color_buffer) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, - GL_CLAMP_TO_EDGE) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, - GL_CLAMP_TO_EDGE) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, - GL_LINEAR) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, - GL_LINEAR) - glTexImage2DMultisample(GL_TEXTURE_2D_MULTISAMPLE, - multi_sample_rate, GL_RGBA32F, - self.width, self.height, GL_TRUE) - glBindTexture(GL_TEXTURE_2D_MULTISAMPLE, 0) - glFramebufferTexture2D(GL_FRAMEBUFFER, - GL_COLOR_ATTACHMENT0 + i, - GL_TEXTURE_2D_MULTISAMPLE, color_buffer, - 0) - self.color_buffer.append(color_buffer) - - self.render_buffer = glGenRenderbuffers(1) - glBindRenderbuffer(GL_RENDERBUFFER, self.render_buffer) - glRenderbufferStorageMultisample(GL_RENDERBUFFER, - multi_sample_rate, - GL_DEPTH24_STENCIL8, self.width, - self.height) - glBindRenderbuffer(GL_RENDERBUFFER, 0) - glFramebufferRenderbuffer(GL_FRAMEBUFFER, - GL_DEPTH_STENCIL_ATTACHMENT, - GL_RENDERBUFFER, self.render_buffer) - - attachments = [] - for i in range(color_size): - attachments.append(GL_COLOR_ATTACHMENT0 + i) - glDrawBuffers(color_size, attachments) - glBindFramebuffer(GL_FRAMEBUFFER, 0) - - self.intermediate_fbo = glGenFramebuffers(1) - glBindFramebuffer(GL_FRAMEBUFFER, self.intermediate_fbo) - - self.screen_texture = [] - for i in range(color_size): - screen_texture = glGenTextures(1) - glBindTexture(GL_TEXTURE_2D, screen_texture) - glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, self.width, - self.height, 0, GL_RGBA, GL_FLOAT, None) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, - GL_LINEAR) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, - GL_LINEAR) - glFramebufferTexture2D(GL_FRAMEBUFFER, - GL_COLOR_ATTACHMENT0 + i, GL_TEXTURE_2D, - screen_texture, 0) - self.screen_texture.append(screen_texture) - - glDrawBuffers(color_size, attachments) - glBindFramebuffer(GL_FRAMEBUFFER, 0) - else: - self.color_buffer = [] - for i in range(color_size): - color_buffer = glGenTextures(1) - glBindTexture(GL_TEXTURE_2D, color_buffer) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, - GL_CLAMP_TO_EDGE) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, - GL_CLAMP_TO_EDGE) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, - GL_NEAREST) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, - GL_NEAREST) - glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, self.width, - self.height, 0, GL_RGBA, GL_FLOAT, None) - glFramebufferTexture2D(GL_FRAMEBUFFER, - GL_COLOR_ATTACHMENT0 + i, GL_TEXTURE_2D, - color_buffer, 0) - self.color_buffer.append(color_buffer) - - # Configure depth texture map to render to - self.depth_buffer = glGenTextures(1) - glBindTexture(GL_TEXTURE_2D, self.depth_buffer) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST) - glTexParameteri(GL_TEXTURE_2D, GL_DEPTH_TEXTURE_MODE, GL_INTENSITY) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_MODE, - GL_COMPARE_R_TO_TEXTURE) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_FUNC, GL_LEQUAL) - glTexImage2D(GL_TEXTURE_2D, 0, GL_DEPTH_COMPONENT, self.width, - self.height, 0, GL_DEPTH_COMPONENT, GL_FLOAT, None) - glFramebufferTexture2D(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, - GL_TEXTURE_2D, self.depth_buffer, 0) - - attachments = [] - for i in range(color_size): - attachments.append(GL_COLOR_ATTACHMENT0 + i) - glDrawBuffers(color_size, attachments) - self.screen_texture = self.color_buffer - - glBindFramebuffer(GL_FRAMEBUFFER, 0) - - # Configure texture buffer if needed - self.render_texture = None - - # NOTE: original render_texture only support one input - # this is tentative member of this issue - self.render_texture_v2 = {} - - # Inner storage for buffer data - self.vertex_data = None - self.vertex_dim = None - self.n_vertices = None - - self.model_view_matrix = None - self.projection_matrix = None - - if not egl: - global GLUT - import OpenGL.GLUT as GLUT - GLUT.glutDisplayFunc(self.display) - - def init_quad_program(self): - shader_list = [] - - shader_list.append(loadShader(GL_VERTEX_SHADER, "quad.vs")) - shader_list.append(loadShader(GL_FRAGMENT_SHADER, "quad.fs")) - - the_program = createProgram(shader_list) - - for shader in shader_list: - glDeleteShader(shader) - - # vertex attributes for a quad that fills the entire screen in Normalized Device Coordinates. - # positions # texCoords - quad_vertices = np.array([ - -1.0, 1.0, 0.0, 1.0, -1.0, -1.0, 0.0, 0.0, 1.0, -1.0, 1.0, 0.0, - -1.0, 1.0, 0.0, 1.0, 1.0, -1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0 - ]) - - quad_buffer = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, quad_buffer) - glBufferData(GL_ARRAY_BUFFER, quad_vertices, GL_STATIC_DRAW) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - return the_program, quad_buffer - - def set_mesh(self, vertices, faces): - self.vertex_data = vertices[faces.reshape([-1])] - self.vertex_dim = self.vertex_data.shape[1] - self.n_vertices = self.vertex_data.shape[0] - - glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) - glBufferData(GL_ARRAY_BUFFER, self.vertex_data, GL_STATIC_DRAW) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - def set_viewpoint(self, projection, model_view): - self.projection_matrix = projection - self.model_view_matrix = model_view - - def draw_init(self): - glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) - glEnable(GL_DEPTH_TEST) - - glClearColor(0.0, 0.0, 0.0, 0.0) - if self.use_inverse_depth: - glDepthFunc(GL_GREATER) - glClearDepth(0.0) - else: - glDepthFunc(GL_LESS) - glClearDepth(1.0) - glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT) - - def draw_end(self): - if self.intermediate_fbo is not None: - for i in range(len(self.color_buffer)): - glBindFramebuffer(GL_READ_FRAMEBUFFER, self.frame_buffer) - glReadBuffer(GL_COLOR_ATTACHMENT0 + i) - glBindFramebuffer(GL_DRAW_FRAMEBUFFER, self.intermediate_fbo) - glDrawBuffer(GL_COLOR_ATTACHMENT0 + i) - glBlitFramebuffer(0, 0, self.width, self.height, 0, 0, - self.width, self.height, GL_COLOR_BUFFER_BIT, - GL_NEAREST) - - glBindFramebuffer(GL_FRAMEBUFFER, 0) - glDepthFunc(GL_LESS) - glClearDepth(1.0) - - def draw(self): - self.draw_init() - - glUseProgram(self.program) - glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, - self.model_view_matrix.transpose()) - glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, - self.projection_matrix.transpose()) - - glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) - - glEnableVertexAttribArray(0) - glVertexAttribPointer(0, self.vertex_dim, GL_DOUBLE, GL_FALSE, 0, None) - - glDrawArrays(GL_TRIANGLES, 0, self.n_vertices) - - glDisableVertexAttribArray(0) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - glUseProgram(0) - - self.draw_end() - - def get_color(self, color_id=0): - glBindFramebuffer( - GL_FRAMEBUFFER, self.intermediate_fbo - if self.intermediate_fbo is not None else self.frame_buffer) - glReadBuffer(GL_COLOR_ATTACHMENT0 + color_id) - data = glReadPixels(0, - 0, - self.width, - self.height, - GL_RGBA, - GL_FLOAT, - outputType=None) - glBindFramebuffer(GL_FRAMEBUFFER, 0) - rgb = data.reshape(self.height, self.width, -1) - rgb = np.flip(rgb, 0) - return rgb - - def get_z_value(self): - glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) - data = glReadPixels(0, - 0, - self.width, - self.height, - GL_DEPTH_COMPONENT, - GL_FLOAT, - outputType=None) - glBindFramebuffer(GL_FRAMEBUFFER, 0) - z = data.reshape(self.height, self.width) - z = np.flip(z, 0) - return z - - def display(self): - self.draw() - - if not self.egl: - # First we draw a scene. - # Notice the result is stored in the texture buffer. - - # Then we return to the default frame buffer since we will display on the screen. - glBindFramebuffer(GL_FRAMEBUFFER, 0) - - # Do the clean-up. - glClearColor(0.0, 0.0, 0.0, 0.0) - glClear(GL_COLOR_BUFFER_BIT) - - # We draw a rectangle which covers the whole screen. - glUseProgram(self.quad_program) - glBindBuffer(GL_ARRAY_BUFFER, self.quad_buffer) - - size_of_double = 8 - glEnableVertexAttribArray(0) - glVertexAttribPointer(0, 2, GL_DOUBLE, GL_FALSE, - 4 * size_of_double, None) - glEnableVertexAttribArray(1) - glVertexAttribPointer(1, 2, GL_DOUBLE, GL_FALSE, - 4 * size_of_double, - c_void_p(2 * size_of_double)) - - glDisable(GL_DEPTH_TEST) - - # The stored texture is then mapped to this rectangle. - # properly assing color buffer texture - glActiveTexture(GL_TEXTURE0) - glBindTexture(GL_TEXTURE_2D, self.screen_texture[0]) - glUniform1i( - glGetUniformLocation(self.quad_program, 'screenTexture'), 0) - - glDrawArrays(GL_TRIANGLES, 0, 6) - - glDisableVertexAttribArray(1) - glDisableVertexAttribArray(0) - - glEnable(GL_DEPTH_TEST) - glBindBuffer(GL_ARRAY_BUFFER, 0) - glUseProgram(0) - - GLUT.glutSwapBuffers() - GLUT.glutPostRedisplay() - - def show(self): - if not self.egl: - GLUT.glutMainLoop() \ No newline at end of file diff --git a/lib / renderer / gl /cam_render.py b/lib / renderer / gl /cam_render.py deleted file mode 100644 index 6379b5d8a8b443fb1281eb7e025d4fc4d63c523d..0000000000000000000000000000000000000000 --- a/lib / renderer / gl /cam_render.py +++ /dev/null @@ -1,80 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from .render import Render - -GLUT = None - - -class CamRender(Render): - def __init__(self, - width=1600, - height=1200, - name='Cam Renderer', - program_files=['simple.fs', 'simple.vs'], - color_size=1, - ms_rate=1, - egl=False): - Render.__init__(self, - width, - height, - name, - program_files, - color_size, - ms_rate=ms_rate, - egl=egl) - self.camera = None - - if not egl: - global GLUT - import OpenGL.GLUT as GLUT - GLUT.glutDisplayFunc(self.display) - GLUT.glutKeyboardFunc(self.keyboard) - - def set_camera(self, camera): - self.camera = camera - self.projection_matrix, self.model_view_matrix = camera.get_gl_matrix() - - def keyboard(self, key, x, y): - # up - eps = 1 - # print(key) - if key == b'w': - self.camera.center += eps * self.camera.direction - elif key == b's': - self.camera.center -= eps * self.camera.direction - if key == b'a': - self.camera.center -= eps * self.camera.right - elif key == b'd': - self.camera.center += eps * self.camera.right - if key == b' ': - self.camera.center += eps * self.camera.up - elif key == b'x': - self.camera.center -= eps * self.camera.up - elif key == b'i': - self.camera.near += 0.1 * eps - self.camera.far += 0.1 * eps - elif key == b'o': - self.camera.near -= 0.1 * eps - self.camera.far -= 0.1 * eps - - self.projection_matrix, self.model_view_matrix = self.camera.get_gl_matrix( - ) - - def show(self): - if GLUT is not None: - GLUT.glutMainLoop() \ No newline at end of file diff --git a/lib / renderer / gl /color_render.py b/lib / renderer / gl /color_render.py deleted file mode 100644 index 6fd7b51db4816e8feafde843335acaa60856c174..0000000000000000000000000000000000000000 --- a/lib / renderer / gl /color_render.py +++ /dev/null @@ -1,158 +0,0 @@ -''' -MIT License -Copyright (c) 2019 Shunsuke Saito, Zeng Huang, and Ryota Natsume -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. -''' -import numpy as np -import random - -from .framework import * -from .cam_render import CamRender - - -class ColorRender(CamRender): - def __init__(self, width=1600, height=1200, name='Color Renderer', egl=False): - program_files = ['color.vs', 'color.fs'] - CamRender.__init__(self, - width, - height, - name, - program_files=program_files, - color_size=3, egl=egl) - - # WARNING: this differs from vertex_buffer and vertex_data in Render - self.vert_buffer = {} - self.vert_data = {} - - # normal - self.norm_buffer = {} - self.norm_data = {} - - self.color_buffer = {} - self.color_data = {} - - self.vertex_dim = {} - self.n_vertices = {} - - self.rot_mat_unif = glGetUniformLocation(self.program, 'RotMat') - self.rot_matrix = np.eye(3) - - self.norm_mat_unif = glGetUniformLocation(self.program, 'NormMat') - self.normalize_matrix = np.eye(4) - - def set_norm_mat(self, scale, center): - N = np.eye(4) - N[:3, :3] = scale * np.eye(3) - N[:3, 3] = -scale * center - - self.normalize_matrix = N - - def set_mesh(self, vertices, faces, color, normals, mat_name='all'): - - self.vert_data[mat_name] = vertices[faces.reshape([-1])] - self.n_vertices[mat_name] = self.vert_data[mat_name].shape[0] - self.vertex_dim[mat_name] = self.vert_data[mat_name].shape[1] - self.color_data[mat_name] = color[faces.reshape([-1])] - self.norm_data[mat_name] = normals[faces.reshape([-1])] - - if mat_name not in self.vert_buffer.keys(): - self.vert_buffer[mat_name] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.vert_buffer[mat_name]) - glBufferData(GL_ARRAY_BUFFER, self.vert_data[mat_name], GL_STATIC_DRAW) - - if mat_name not in self.color_buffer.keys(): - self.color_buffer[mat_name] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.color_buffer[mat_name]) - glBufferData(GL_ARRAY_BUFFER, self.color_data[mat_name], - GL_STATIC_DRAW) - - if mat_name not in self.norm_buffer.keys(): - self.norm_buffer[mat_name] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer[mat_name]) - glBufferData(GL_ARRAY_BUFFER, self.norm_data[mat_name], GL_STATIC_DRAW) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - def cleanup(self): - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - for key in self.vert_data: - glDeleteBuffers(1, [self.vert_buffer[key]]) - glDeleteBuffers(1, [self.color_buffer[key]]) - glDeleteBuffers(1, [self.norm_buffer[key]]) - - self.norm_buffer = {} - self.norm_data = {} - - self.vert_buffer = {} - self.vert_data = {} - - self.color_buffer = {} - self.color_data = {} - - self.render_texture_mat = {} - - self.vertex_dim = {} - self.n_vertices = {} - - def draw(self): - self.draw_init() - - glEnable(GL_MULTISAMPLE) - - glUseProgram(self.program) - glUniformMatrix4fv(self.norm_mat_unif, 1, GL_FALSE, - self.normalize_matrix.transpose()) - glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, - self.model_view_matrix.transpose()) - glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, - self.projection_matrix.transpose()) - glUniformMatrix3fv(self.rot_mat_unif, 1, GL_FALSE, - self.rot_matrix.transpose()) - - for mat in self.vert_buffer: - - # Handle vertex buffer - glBindBuffer(GL_ARRAY_BUFFER, self.vert_buffer[mat]) - glEnableVertexAttribArray(0) - glVertexAttribPointer(0, self.vertex_dim[mat], GL_DOUBLE, GL_FALSE, - 0, None) - - # Handle color buffer - glBindBuffer(GL_ARRAY_BUFFER, self.color_buffer[mat]) - glEnableVertexAttribArray(1) - glVertexAttribPointer(1, 3, GL_DOUBLE, GL_FALSE, 0, None) - - # Handle normal buffer - glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer[mat]) - glEnableVertexAttribArray(2) - glVertexAttribPointer(2, 3, GL_DOUBLE, GL_FALSE, 0, None) - - glDrawArrays(GL_TRIANGLES, 0, self.n_vertices[mat]) - - glDisableVertexAttribArray(2) - glDisableVertexAttribArray(1) - glDisableVertexAttribArray(0) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - glUseProgram(0) - - glDisable(GL_MULTISAMPLE) - - self.draw_end() \ No newline at end of file diff --git a/lib / renderer / gl /normal_render.py b/lib / renderer / gl /normal_render.py deleted file mode 100644 index 02cc72311564d615793c2e53863b7230d0a17a5b..0000000000000000000000000000000000000000 --- a/lib / renderer / gl /normal_render.py +++ /dev/null @@ -1,93 +0,0 @@ -''' -MIT License -Copyright (c) 2019 Shunsuke Saito, Zeng Huang, and Ryota Natsume -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. -''' -import numpy as np -import math - -from .framework import * -from .norm_render import NormRender - - -class NormalRender(NormRender): - def __init__(self, width=1600, height=1200, name='Normal Renderer'): - NormRender.__init__(self, - width, - height, - name, - program_files=['normal.vs', 'normal.fs']) - - self.norm_buffer = glGenBuffers(1) - - self.norm_data = None - - def set_normal_mesh(self, vertices, faces, norms, face_normals): - NormRender.set_mesh(self, vertices, faces) - - self.norm_data = norms[face_normals.reshape([-1])] - - glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer) - glBufferData(GL_ARRAY_BUFFER, self.norm_data, GL_STATIC_DRAW) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - def euler_to_rot_mat(self, r_x, r_y, r_z): - R_x = np.array([[1, 0, 0], [0, math.cos(r_x), -math.sin(r_x)], - [0, math.sin(r_x), math.cos(r_x)]]) - - R_y = np.array([[math.cos(r_y), 0, math.sin(r_y)], [0, 1, 0], - [-math.sin(r_y), 0, math.cos(r_y)]]) - - R_z = np.array([[math.cos(r_z), -math.sin(r_z), 0], - [math.sin(r_z), math.cos(r_z), 0], [0, 0, 1]]) - - R = np.dot(R_z, np.dot(R_y, R_x)) - - return R - - def draw(self): - self.draw_init() - - glUseProgram(self.program) - glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, - self.model_view_matrix.transpose()) - glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, - self.projection_matrix.transpose()) - - # Handle vertex buffer - glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) - - glEnableVertexAttribArray(0) - glVertexAttribPointer(0, self.vertex_dim, GL_DOUBLE, GL_FALSE, 0, None) - - # Handle normal buffer - glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer) - - glEnableVertexAttribArray(1) - glVertexAttribPointer(1, 3, GL_DOUBLE, GL_FALSE, 0, None) - - glDrawArrays(GL_TRIANGLES, 0, self.n_vertices) - - glDisableVertexAttribArray(1) - glDisableVertexAttribArray(0) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - glUseProgram(0) - - self.draw_end() \ No newline at end of file diff --git a/lib / renderer / gl /prt_render.py b/lib / renderer / gl /prt_render.py deleted file mode 100644 index a8b3ce290c6d8b7e0a3eb894230a99b0329a7818..0000000000000000000000000000000000000000 --- a/lib / renderer / gl /prt_render.py +++ /dev/null @@ -1,450 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import numpy as np -import random - -from .framework import * -from .cam_render import CamRender - - -class PRTRender(CamRender): - def __init__(self, - width=1600, - height=1200, - name='PRT Renderer', - uv_mode=False, - ms_rate=1, - egl=False): - program_files = ['prt.vs', 'prt.fs' - ] if not uv_mode else ['prt_uv.vs', 'prt_uv.fs'] - CamRender.__init__(self, - width, - height, - name, - program_files=program_files, - color_size=8, - ms_rate=ms_rate, - egl=egl) - - # WARNING: this differs from vertex_buffer and vertex_data in Render - self.vert_buffer = {} - self.vert_data = {} - - self.vert_label_buffer = {} - self.vert_label_data = {} - - self.norm_buffer = {} - self.norm_data = {} - - self.tan_buffer = {} - self.tan_data = {} - - self.btan_buffer = {} - self.btan_data = {} - - self.prt1_buffer = {} - self.prt1_data = {} - - self.prt2_buffer = {} - self.prt2_data = {} - - self.prt3_buffer = {} - self.prt3_data = {} - - self.uv_buffer = {} - self.uv_data = {} - - self.render_texture_mat = {} - - self.vertex_dim = {} - self.n_vertices = {} - self.label_dim = {} - - self.norm_mat_unif = glGetUniformLocation(self.program, 'NormMat') - self.normalize_matrix = np.eye(4) - - self.shcoeff_unif = glGetUniformLocation(self.program, 'SHCoeffs') - self.shcoeffs = np.zeros((9, 3)) - self.shcoeffs[0, :] = 1.0 - #self.shcoeffs[1:,:] = np.random.rand(8,3) - - self.hasAlbedoUnif = glGetUniformLocation(self.program, 'hasAlbedoMap') - self.hasNormalUnif = glGetUniformLocation(self.program, 'hasNormalMap') - - self.analyticUnif = glGetUniformLocation(self.program, 'analytic') - self.analytic = False - - self.rot_mat_unif = glGetUniformLocation(self.program, 'RotMat') - self.rot_matrix = np.eye(3) - - def set_texture(self, mat_name, smplr_name, texture): - # texture_image: H x W x 3 - width = texture.shape[1] - height = texture.shape[0] - texture = np.flip(texture, 0) - img_data = np.fromstring(texture.tostring(), np.uint8) - - if mat_name not in self.render_texture_mat: - self.render_texture_mat[mat_name] = {} - if smplr_name in self.render_texture_mat[mat_name].keys(): - glDeleteTextures([self.render_texture_mat[mat_name][smplr_name]]) - del self.render_texture_mat[mat_name][smplr_name] - - self.render_texture_mat[mat_name][smplr_name] = glGenTextures(1) - glActiveTexture(GL_TEXTURE0) - - glPixelStorei(GL_UNPACK_ALIGNMENT, 1) - glBindTexture(GL_TEXTURE_2D, - self.render_texture_mat[mat_name][smplr_name]) - - glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, width, height, 0, GL_RGB, - GL_UNSIGNED_BYTE, img_data) - - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAX_LEVEL, 3) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, - GL_LINEAR_MIPMAP_LINEAR) - - glGenerateMipmap(GL_TEXTURE_2D) - - def set_albedo(self, texture_image, mat_name='all'): - self.set_texture(mat_name, 'AlbedoMap', texture_image) - - def set_normal_map(self, texture_image, mat_name='all'): - self.set_texture(mat_name, 'NormalMap', texture_image) - - def set_mesh(self, - vertices, - faces, - norms, - faces_nml, - uvs, - faces_uvs, - prt, - faces_prt, - tans, - bitans, - verts_label=None, - mat_name='all'): - - self.vert_data[mat_name] = vertices[faces.reshape([-1])] - self.vert_label_data[mat_name] = verts_label[faces.reshape([-1])] - self.n_vertices[mat_name] = self.vert_data[mat_name].shape[0] - self.vertex_dim[mat_name] = self.vert_data[mat_name].shape[1] - self.label_dim[mat_name] = self.vert_label_data[mat_name].shape[1] - - if mat_name not in self.vert_buffer.keys(): - self.vert_buffer[mat_name] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.vert_buffer[mat_name]) - glBufferData(GL_ARRAY_BUFFER, self.vert_data[mat_name], GL_STATIC_DRAW) - - if mat_name not in self.vert_label_buffer.keys(): - self.vert_label_buffer[mat_name] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.vert_label_buffer[mat_name]) - glBufferData(GL_ARRAY_BUFFER, self.vert_label_data[mat_name], - GL_STATIC_DRAW) - - self.uv_data[mat_name] = uvs[faces_uvs.reshape([-1])] - if mat_name not in self.uv_buffer.keys(): - self.uv_buffer[mat_name] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.uv_buffer[mat_name]) - glBufferData(GL_ARRAY_BUFFER, self.uv_data[mat_name], GL_STATIC_DRAW) - - self.norm_data[mat_name] = norms[faces_nml.reshape([-1])] - if mat_name not in self.norm_buffer.keys(): - self.norm_buffer[mat_name] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer[mat_name]) - glBufferData(GL_ARRAY_BUFFER, self.norm_data[mat_name], GL_STATIC_DRAW) - - self.tan_data[mat_name] = tans[faces_nml.reshape([-1])] - if mat_name not in self.tan_buffer.keys(): - self.tan_buffer[mat_name] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.tan_buffer[mat_name]) - glBufferData(GL_ARRAY_BUFFER, self.tan_data[mat_name], GL_STATIC_DRAW) - - self.btan_data[mat_name] = bitans[faces_nml.reshape([-1])] - if mat_name not in self.btan_buffer.keys(): - self.btan_buffer[mat_name] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.btan_buffer[mat_name]) - glBufferData(GL_ARRAY_BUFFER, self.btan_data[mat_name], GL_STATIC_DRAW) - - self.prt1_data[mat_name] = prt[faces_prt.reshape([-1])][:, :3] - self.prt2_data[mat_name] = prt[faces_prt.reshape([-1])][:, 3:6] - self.prt3_data[mat_name] = prt[faces_prt.reshape([-1])][:, 6:] - - if mat_name not in self.prt1_buffer.keys(): - self.prt1_buffer[mat_name] = glGenBuffers(1) - if mat_name not in self.prt2_buffer.keys(): - self.prt2_buffer[mat_name] = glGenBuffers(1) - if mat_name not in self.prt3_buffer.keys(): - self.prt3_buffer[mat_name] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.prt1_buffer[mat_name]) - glBufferData(GL_ARRAY_BUFFER, self.prt1_data[mat_name], GL_STATIC_DRAW) - glBindBuffer(GL_ARRAY_BUFFER, self.prt2_buffer[mat_name]) - glBufferData(GL_ARRAY_BUFFER, self.prt2_data[mat_name], GL_STATIC_DRAW) - glBindBuffer(GL_ARRAY_BUFFER, self.prt3_buffer[mat_name]) - glBufferData(GL_ARRAY_BUFFER, self.prt3_data[mat_name], GL_STATIC_DRAW) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - def set_mesh_mtl(self, - vertices, - faces, - norms, - faces_nml, - uvs, - faces_uvs, - tans, - bitans, - prt, - verts_label=None): - for key in faces: - self.vert_data[key] = vertices[faces[key].reshape([-1])] - self.vert_label_data[key] = verts_label[faces[key].reshape([-1])] - self.n_vertices[key] = self.vert_data[key].shape[0] - self.vertex_dim[key] = self.vert_data[key].shape[1] - self.label_dim[key] = self.vert_label_data[key].shape[1] - - if key not in self.vert_buffer.keys(): - self.vert_buffer[key] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.vert_buffer[key]) - glBufferData(GL_ARRAY_BUFFER, self.vert_data[key], GL_STATIC_DRAW) - - if key not in self.vert_label_buffer.keys(): - self.vert_label_buffer[key] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.vert_label_buffer[key]) - glBufferData(GL_ARRAY_BUFFER, self.vert_label_data[key], - GL_STATIC_DRAW) - - self.uv_data[key] = uvs[faces_uvs[key].reshape([-1])] - if key not in self.uv_buffer.keys(): - self.uv_buffer[key] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.uv_buffer[key]) - glBufferData(GL_ARRAY_BUFFER, self.uv_data[key], GL_STATIC_DRAW) - - self.norm_data[key] = norms[faces_nml[key].reshape([-1])] - if key not in self.norm_buffer.keys(): - self.norm_buffer[key] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer[key]) - glBufferData(GL_ARRAY_BUFFER, self.norm_data[key], GL_STATIC_DRAW) - - self.tan_data[key] = tans[faces_nml[key].reshape([-1])] - if key not in self.tan_buffer.keys(): - self.tan_buffer[key] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.tan_buffer[key]) - glBufferData(GL_ARRAY_BUFFER, self.tan_data[key], GL_STATIC_DRAW) - - self.btan_data[key] = bitans[faces_nml[key].reshape([-1])] - if key not in self.btan_buffer.keys(): - self.btan_buffer[key] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.btan_buffer[key]) - glBufferData(GL_ARRAY_BUFFER, self.btan_data[key], GL_STATIC_DRAW) - - self.prt1_data[key] = prt[faces[key].reshape([-1])][:, :3] - self.prt2_data[key] = prt[faces[key].reshape([-1])][:, 3:6] - self.prt3_data[key] = prt[faces[key].reshape([-1])][:, 6:] - - if key not in self.prt1_buffer.keys(): - self.prt1_buffer[key] = glGenBuffers(1) - if key not in self.prt2_buffer.keys(): - self.prt2_buffer[key] = glGenBuffers(1) - if key not in self.prt3_buffer.keys(): - self.prt3_buffer[key] = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, self.prt1_buffer[key]) - glBufferData(GL_ARRAY_BUFFER, self.prt1_data[key], GL_STATIC_DRAW) - glBindBuffer(GL_ARRAY_BUFFER, self.prt2_buffer[key]) - glBufferData(GL_ARRAY_BUFFER, self.prt2_data[key], GL_STATIC_DRAW) - glBindBuffer(GL_ARRAY_BUFFER, self.prt3_buffer[key]) - glBufferData(GL_ARRAY_BUFFER, self.prt3_data[key], GL_STATIC_DRAW) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - def cleanup(self): - - glBindBuffer(GL_ARRAY_BUFFER, 0) - for key in self.vert_data: - glDeleteBuffers(1, [self.vert_buffer[key]]) - glDeleteBuffers(1, [self.norm_buffer[key]]) - glDeleteBuffers(1, [self.uv_buffer[key]]) - glDeleteBuffers(1, [self.vert_label_buffer[key]]) - - glDeleteBuffers(1, [self.tan_buffer[key]]) - glDeleteBuffers(1, [self.btan_buffer[key]]) - glDeleteBuffers(1, [self.prt1_buffer[key]]) - glDeleteBuffers(1, [self.prt2_buffer[key]]) - glDeleteBuffers(1, [self.prt3_buffer[key]]) - - glDeleteBuffers(1, []) - - for smplr in self.render_texture_mat[key]: - glDeleteTextures([self.render_texture_mat[key][smplr]]) - - self.vert_buffer = {} - self.vert_data = {} - - self.vert_label_buffer = {} - self.vert_label_data = {} - - self.norm_buffer = {} - self.norm_data = {} - - self.tan_buffer = {} - self.tan_data = {} - - self.btan_buffer = {} - self.btan_data = {} - - self.prt1_buffer = {} - self.prt1_data = {} - - self.prt2_buffer = {} - self.prt2_data = {} - - self.prt3_buffer = {} - self.prt3_data = {} - - self.uv_buffer = {} - self.uv_data = {} - - self.render_texture_mat = {} - - self.vertex_dim = {} - self.n_vertices = {} - self.label_dim = {} - - def randomize_sh(self): - self.shcoeffs[0, :] = 0.8 - self.shcoeffs[1:, :] = 1.0 * np.random.rand(8, 3) - - def set_sh(self, sh): - self.shcoeffs = sh - - def set_norm_mat(self, scale, center): - N = np.eye(4) - N[:3, :3] = scale * np.eye(3) - N[:3, 3] = -scale * center - - self.normalize_matrix = N - - def draw(self): - self.draw_init() - - glDisable(GL_BLEND) - #glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA) - glEnable(GL_MULTISAMPLE) - - glUseProgram(self.program) - glUniformMatrix4fv(self.norm_mat_unif, 1, GL_FALSE, - self.normalize_matrix.transpose()) - glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, - self.model_view_matrix.transpose()) - glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, - self.projection_matrix.transpose()) - - if 'AlbedoMap' in self.render_texture_mat['all']: - glUniform1ui(self.hasAlbedoUnif, GLuint(1)) - else: - glUniform1ui(self.hasAlbedoUnif, GLuint(0)) - - if 'NormalMap' in self.render_texture_mat['all']: - glUniform1ui(self.hasNormalUnif, GLuint(1)) - else: - glUniform1ui(self.hasNormalUnif, GLuint(0)) - - glUniform1ui(self.analyticUnif, - GLuint(1) if self.analytic else GLuint(0)) - - glUniform3fv(self.shcoeff_unif, 9, self.shcoeffs) - - glUniformMatrix3fv(self.rot_mat_unif, 1, GL_FALSE, - self.rot_matrix.transpose()) - - for mat in self.vert_buffer: - # Handle vertex buffer - glBindBuffer(GL_ARRAY_BUFFER, self.vert_buffer[mat]) - glEnableVertexAttribArray(0) - glVertexAttribPointer(0, self.vertex_dim[mat], GL_DOUBLE, GL_FALSE, - 0, None) - - # Handle normal buffer - glBindBuffer(GL_ARRAY_BUFFER, self.norm_buffer[mat]) - glEnableVertexAttribArray(1) - glVertexAttribPointer(1, 3, GL_DOUBLE, GL_FALSE, 0, None) - - # Handle uv buffer - glBindBuffer(GL_ARRAY_BUFFER, self.uv_buffer[mat]) - glEnableVertexAttribArray(2) - glVertexAttribPointer(2, 2, GL_DOUBLE, GL_FALSE, 0, None) - - # Handle tan buffer - glBindBuffer(GL_ARRAY_BUFFER, self.tan_buffer[mat]) - glEnableVertexAttribArray(3) - glVertexAttribPointer(3, 3, GL_DOUBLE, GL_FALSE, 0, None) - - # Handle btan buffer - glBindBuffer(GL_ARRAY_BUFFER, self.btan_buffer[mat]) - glEnableVertexAttribArray(4) - glVertexAttribPointer(4, 3, GL_DOUBLE, GL_FALSE, 0, None) - - # Handle PTR buffer - glBindBuffer(GL_ARRAY_BUFFER, self.prt1_buffer[mat]) - glEnableVertexAttribArray(5) - glVertexAttribPointer(5, 3, GL_DOUBLE, GL_FALSE, 0, None) - - glBindBuffer(GL_ARRAY_BUFFER, self.prt2_buffer[mat]) - glEnableVertexAttribArray(6) - glVertexAttribPointer(6, 3, GL_DOUBLE, GL_FALSE, 0, None) - - glBindBuffer(GL_ARRAY_BUFFER, self.prt3_buffer[mat]) - glEnableVertexAttribArray(7) - glVertexAttribPointer(7, 3, GL_DOUBLE, GL_FALSE, 0, None) - - # Handle vertex label buffer - glBindBuffer(GL_ARRAY_BUFFER, self.vert_label_buffer[mat]) - glEnableVertexAttribArray(8) - glVertexAttribPointer(8, self.label_dim[mat], GL_DOUBLE, GL_FALSE, - 0, None) - - for i, smplr in enumerate(self.render_texture_mat[mat]): - glActiveTexture(GL_TEXTURE0 + i) - glBindTexture(GL_TEXTURE_2D, - self.render_texture_mat[mat][smplr]) - glUniform1i(glGetUniformLocation(self.program, smplr), i) - - glDrawArrays(GL_TRIANGLES, 0, self.n_vertices[mat]) - - glDisableVertexAttribArray(8) - glDisableVertexAttribArray(7) - glDisableVertexAttribArray(6) - glDisableVertexAttribArray(5) - glDisableVertexAttribArray(4) - glDisableVertexAttribArray(3) - glDisableVertexAttribArray(2) - glDisableVertexAttribArray(1) - glDisableVertexAttribArray(0) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - glUseProgram(0) - - glDisable(GL_BLEND) - glDisable(GL_MULTISAMPLE) - - self.draw_end() diff --git a/lib / renderer / gl /render2.py b/lib / renderer / gl /render2.py deleted file mode 100644 index 8f9f66beda68adaa141af27591c131c5a7999c1e..0000000000000000000000000000000000000000 --- a/lib / renderer / gl /render2.py +++ /dev/null @@ -1,384 +0,0 @@ -''' -MIT License -Copyright (c) 2019 Shunsuke Saito, Zeng Huang, and Ryota Natsume -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. -''' -import numpy as np -from OpenGL.GLUT import * -from .framework import * - -_glut_window = None - - -class Render: - def __init__(self, - width=1600, - height=1200, - name='GL Renderer', - program_files=['simple.fs', 'simple.vs'], - color_size=1, - ms_rate=1): - self.width = width - self.height = height - self.name = name - self.display_mode = GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH - self.use_inverse_depth = False - - global _glut_window - if _glut_window is None: - glutInit() - glutInitDisplayMode(self.display_mode) - glutInitWindowSize(self.width, self.height) - glutInitWindowPosition(0, 0) - _glut_window = glutCreateWindow("My Render.") - - # glEnable(GL_DEPTH_CLAMP) - glEnable(GL_DEPTH_TEST) - - glClampColor(GL_CLAMP_READ_COLOR, GL_FALSE) - glClampColor(GL_CLAMP_FRAGMENT_COLOR, GL_FALSE) - glClampColor(GL_CLAMP_VERTEX_COLOR, GL_FALSE) - - # init program - shader_list = [] - - for program_file in program_files: - _, ext = os.path.splitext(program_file) - if ext == '.vs': - shader_list.append(loadShader(GL_VERTEX_SHADER, program_file)) - elif ext == '.fs': - shader_list.append(loadShader(GL_FRAGMENT_SHADER, - program_file)) - elif ext == '.gs': - shader_list.append(loadShader(GL_GEOMETRY_SHADER, - program_file)) - - self.program = createProgram(shader_list) - - for shader in shader_list: - glDeleteShader(shader) - - # Init uniform variables - self.model_mat_unif = glGetUniformLocation(self.program, 'ModelMat') - self.persp_mat_unif = glGetUniformLocation(self.program, 'PerspMat') - - self.vertex_buffer = glGenBuffers(1) - - # Init screen quad program and buffer - self.quad_program, self.quad_buffer = self.init_quad_program() - - # Configure frame buffer - self.frame_buffer = glGenFramebuffers(1) - glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) - - self.intermediate_fbo = None - if ms_rate > 1: - # Configure texture buffer to render to - self.color_buffer = [] - for i in range(color_size): - color_buffer = glGenTextures(1) - multi_sample_rate = ms_rate - glBindTexture(GL_TEXTURE_2D_MULTISAMPLE, color_buffer) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, - GL_CLAMP_TO_EDGE) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, - GL_CLAMP_TO_EDGE) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, - GL_LINEAR) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, - GL_LINEAR) - glTexImage2DMultisample(GL_TEXTURE_2D_MULTISAMPLE, - multi_sample_rate, GL_RGBA32F, - self.width, self.height, GL_TRUE) - glBindTexture(GL_TEXTURE_2D_MULTISAMPLE, 0) - glFramebufferTexture2D(GL_FRAMEBUFFER, - GL_COLOR_ATTACHMENT0 + i, - GL_TEXTURE_2D_MULTISAMPLE, color_buffer, - 0) - self.color_buffer.append(color_buffer) - - self.render_buffer = glGenRenderbuffers(1) - glBindRenderbuffer(GL_RENDERBUFFER, self.render_buffer) - glRenderbufferStorageMultisample(GL_RENDERBUFFER, - multi_sample_rate, - GL_DEPTH24_STENCIL8, self.width, - self.height) - glBindRenderbuffer(GL_RENDERBUFFER, 0) - glFramebufferRenderbuffer(GL_FRAMEBUFFER, - GL_DEPTH_STENCIL_ATTACHMENT, - GL_RENDERBUFFER, self.render_buffer) - - attachments = [] - for i in range(color_size): - attachments.append(GL_COLOR_ATTACHMENT0 + i) - glDrawBuffers(color_size, attachments) - glBindFramebuffer(GL_FRAMEBUFFER, 0) - - self.intermediate_fbo = glGenFramebuffers(1) - glBindFramebuffer(GL_FRAMEBUFFER, self.intermediate_fbo) - - self.screen_texture = [] - for i in range(color_size): - screen_texture = glGenTextures(1) - glBindTexture(GL_TEXTURE_2D, screen_texture) - glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, self.width, - self.height, 0, GL_RGBA, GL_FLOAT, None) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, - GL_LINEAR) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, - GL_LINEAR) - glFramebufferTexture2D(GL_FRAMEBUFFER, - GL_COLOR_ATTACHMENT0 + i, GL_TEXTURE_2D, - screen_texture, 0) - self.screen_texture.append(screen_texture) - - glDrawBuffers(color_size, attachments) - glBindFramebuffer(GL_FRAMEBUFFER, 0) - else: - self.color_buffer = [] - for i in range(color_size): - color_buffer = glGenTextures(1) - glBindTexture(GL_TEXTURE_2D, color_buffer) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, - GL_CLAMP_TO_EDGE) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, - GL_CLAMP_TO_EDGE) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, - GL_NEAREST) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, - GL_NEAREST) - glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, self.width, - self.height, 0, GL_RGBA, GL_FLOAT, None) - glFramebufferTexture2D(GL_FRAMEBUFFER, - GL_COLOR_ATTACHMENT0 + i, GL_TEXTURE_2D, - color_buffer, 0) - self.color_buffer.append(color_buffer) - - # Configure depth texture map to render to - self.depth_buffer = glGenTextures(1) - glBindTexture(GL_TEXTURE_2D, self.depth_buffer) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST) - glTexParameteri(GL_TEXTURE_2D, GL_DEPTH_TEXTURE_MODE, GL_INTENSITY) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_MODE, - GL_COMPARE_R_TO_TEXTURE) - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_FUNC, GL_LEQUAL) - glTexImage2D(GL_TEXTURE_2D, 0, GL_DEPTH_COMPONENT, self.width, - self.height, 0, GL_DEPTH_COMPONENT, GL_FLOAT, None) - glFramebufferTexture2D(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, - GL_TEXTURE_2D, self.depth_buffer, 0) - - attachments = [] - for i in range(color_size): - attachments.append(GL_COLOR_ATTACHMENT0 + i) - glDrawBuffers(color_size, attachments) - self.screen_texture = self.color_buffer - - glBindFramebuffer(GL_FRAMEBUFFER, 0) - - # Configure texture buffer if needed - self.render_texture = None - - # NOTE: original render_texture only support one input - # this is tentative member of this issue - self.render_texture_v2 = {} - - # Inner storage for buffer data - self.vertex_data = None - self.vertex_dim = None - self.n_vertices = None - - self.model_view_matrix = None - self.projection_matrix = None - - glutDisplayFunc(self.display) - - def init_quad_program(self): - shader_list = [] - - shader_list.append(loadShader(GL_VERTEX_SHADER, "quad.vs")) - shader_list.append(loadShader(GL_FRAGMENT_SHADER, "quad.fs")) - - the_program = createProgram(shader_list) - - for shader in shader_list: - glDeleteShader(shader) - - # vertex attributes for a quad that fills the entire screen in Normalized Device Coordinates. - # positions # texCoords - quad_vertices = np.array([ - -1.0, 1.0, 0.0, 1.0, -1.0, -1.0, 0.0, 0.0, 1.0, -1.0, 1.0, 0.0, - -1.0, 1.0, 0.0, 1.0, 1.0, -1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0 - ]) - - quad_buffer = glGenBuffers(1) - glBindBuffer(GL_ARRAY_BUFFER, quad_buffer) - glBufferData(GL_ARRAY_BUFFER, quad_vertices, GL_STATIC_DRAW) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - return the_program, quad_buffer - - def set_mesh(self, vertices, faces): - self.vertex_data = vertices[faces.reshape([-1])] - self.vertex_dim = self.vertex_data.shape[1] - self.n_vertices = self.vertex_data.shape[0] - - glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) - glBufferData(GL_ARRAY_BUFFER, self.vertex_data, GL_STATIC_DRAW) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - def set_viewpoint(self, projection, model_view): - self.projection_matrix = projection - self.model_view_matrix = model_view - - def draw_init(self): - glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) - glEnable(GL_DEPTH_TEST) - - # glClearColor(0.0, 0.0, 0.0, 0.0) - glClearColor(1.0, 1.0, 1.0, 0.0) # Black background - - if self.use_inverse_depth: - glDepthFunc(GL_GREATER) - glClearDepth(0.0) - else: - glDepthFunc(GL_LESS) - glClearDepth(1.0) - glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT) - - def draw_end(self): - if self.intermediate_fbo is not None: - for i in range(len(self.color_buffer)): - glBindFramebuffer(GL_READ_FRAMEBUFFER, self.frame_buffer) - glReadBuffer(GL_COLOR_ATTACHMENT0 + i) - glBindFramebuffer(GL_DRAW_FRAMEBUFFER, self.intermediate_fbo) - glDrawBuffer(GL_COLOR_ATTACHMENT0 + i) - glBlitFramebuffer(0, 0, self.width, self.height, 0, 0, - self.width, self.height, GL_COLOR_BUFFER_BIT, - GL_NEAREST) - - glBindFramebuffer(GL_FRAMEBUFFER, 0) - glDepthFunc(GL_LESS) - glClearDepth(1.0) - - def draw(self): - self.draw_init() - - glUseProgram(self.program) - glUniformMatrix4fv(self.model_mat_unif, 1, GL_FALSE, - self.model_view_matrix.transpose()) - glUniformMatrix4fv(self.persp_mat_unif, 1, GL_FALSE, - self.projection_matrix.transpose()) - - glBindBuffer(GL_ARRAY_BUFFER, self.vertex_buffer) - - glEnableVertexAttribArray(0) - glVertexAttribPointer(0, self.vertex_dim, GL_DOUBLE, GL_FALSE, 0, None) - - glDrawArrays(GL_TRIANGLES, 0, self.n_vertices) - - glDisableVertexAttribArray(0) - - glBindBuffer(GL_ARRAY_BUFFER, 0) - - glUseProgram(0) - - self.draw_end() - - def get_color(self, color_id=0): - glBindFramebuffer( - GL_FRAMEBUFFER, self.intermediate_fbo - if self.intermediate_fbo is not None else self.frame_buffer) - glReadBuffer(GL_COLOR_ATTACHMENT0 + color_id) - data = glReadPixels(0, - 0, - self.width, - self.height, - GL_RGBA, - GL_FLOAT, - outputType=None) - glBindFramebuffer(GL_FRAMEBUFFER, 0) - rgb = data.reshape(self.height, self.width, -1) - rgb = np.flip(rgb, 0) - return rgb - - def get_z_value(self): - glBindFramebuffer(GL_FRAMEBUFFER, self.frame_buffer) - data = glReadPixels(0, - 0, - self.width, - self.height, - GL_DEPTH_COMPONENT, - GL_FLOAT, - outputType=None) - glBindFramebuffer(GL_FRAMEBUFFER, 0) - z = data.reshape(self.height, self.width) - z = np.flip(z, 0) - return z - - def display(self): - # First we draw a scene. - # Notice the result is stored in the texture buffer. - self.draw() - - # Then we return to the default frame buffer since we will display on the screen. - glBindFramebuffer(GL_FRAMEBUFFER, 0) - - # Do the clean-up. - # glClearColor(0.0, 0.0, 0.0, 0.0) #Black background - glClearColor(1.0, 1.0, 1.0, 0.0) # Black background - glClear(GL_COLOR_BUFFER_BIT) - - # We draw a rectangle which covers the whole screen. - glUseProgram(self.quad_program) - glBindBuffer(GL_ARRAY_BUFFER, self.quad_buffer) - - size_of_double = 8 - glEnableVertexAttribArray(0) - glVertexAttribPointer(0, 2, GL_DOUBLE, GL_FALSE, 4 * size_of_double, - None) - glEnableVertexAttribArray(1) - glVertexAttribPointer(1, 2, GL_DOUBLE, GL_FALSE, 4 * size_of_double, - c_void_p(2 * size_of_double)) - - glDisable(GL_DEPTH_TEST) - - # The stored texture is then mapped to this rectangle. - # properly assing color buffer texture - glActiveTexture(GL_TEXTURE0) - glBindTexture(GL_TEXTURE_2D, self.screen_texture[0]) - glUniform1i(glGetUniformLocation(self.quad_program, 'screenTexture'), - 0) - - glDrawArrays(GL_TRIANGLES, 0, 6) - - glDisableVertexAttribArray(1) - glDisableVertexAttribArray(0) - - glEnable(GL_DEPTH_TEST) - glBindBuffer(GL_ARRAY_BUFFER, 0) - glUseProgram(0) - - glutSwapBuffers() - glutPostRedisplay() - - def show(self): - glutMainLoop() \ No newline at end of file diff --git a/lib / renderer /glm.py b/lib / renderer /glm.py deleted file mode 100644 index 5068d464c398710e3480078aad3ad13925d3a73b..0000000000000000000000000000000000000000 --- a/lib / renderer /glm.py +++ /dev/null @@ -1,143 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import numpy as np - - -def vec3(x, y, z): - return np.array([x, y, z], dtype=np.float32) - - -def radians(v): - return np.radians(v) - - -def identity(): - return np.identity(4, dtype=np.float32) - - -def empty(): - return np.zeros([4, 4], dtype=np.float32) - - -def magnitude(v): - return np.linalg.norm(v) - - -def normalize(v): - m = magnitude(v) - return v if m == 0 else v / m - - -def dot(u, v): - return np.sum(u * v) - - -def cross(u, v): - res = vec3(0, 0, 0) - res[0] = u[1] * v[2] - u[2] * v[1] - res[1] = u[2] * v[0] - u[0] * v[2] - res[2] = u[0] * v[1] - u[1] * v[0] - return res - - -# below functions can be optimized - - -def translate(m, v): - res = np.copy(m) - res[:, 3] = m[:, 0] * v[0] + m[:, 1] * v[1] + m[:, 2] * v[2] + m[:, 3] - return res - - -def rotate(m, angle, v): - a = angle - c = np.cos(a) - s = np.sin(a) - - axis = normalize(v) - temp = (1 - c) * axis - - rot = empty() - rot[0][0] = c + temp[0] * axis[0] - rot[0][1] = temp[0] * axis[1] + s * axis[2] - rot[0][2] = temp[0] * axis[2] - s * axis[1] - - rot[1][0] = temp[1] * axis[0] - s * axis[2] - rot[1][1] = c + temp[1] * axis[1] - rot[1][2] = temp[1] * axis[2] + s * axis[0] - - rot[2][0] = temp[2] * axis[0] + s * axis[1] - rot[2][1] = temp[2] * axis[1] - s * axis[0] - rot[2][2] = c + temp[2] * axis[2] - - res = empty() - res[:, 0] = m[:, 0] * rot[0][0] + m[:, 1] * rot[0][1] + m[:, 2] * rot[0][2] - res[:, 1] = m[:, 0] * rot[1][0] + m[:, 1] * rot[1][1] + m[:, 2] * rot[1][2] - res[:, 2] = m[:, 0] * rot[2][0] + m[:, 1] * rot[2][1] + m[:, 2] * rot[2][2] - res[:, 3] = m[:, 3] - return res - - -def perspective(fovy, aspect, zNear, zFar): - tanHalfFovy = np.tan(fovy / 2) - - res = empty() - res[0][0] = 1 / (aspect * tanHalfFovy) - res[1][1] = 1 / (tanHalfFovy) - res[2][3] = -1 - res[2][2] = -(zFar + zNear) / (zFar - zNear) - res[3][2] = -(2 * zFar * zNear) / (zFar - zNear) - - return res.T - - -def ortho(left, right, bottom, top, zNear, zFar): - # res = np.ones([4, 4], dtype=np.float32) - res = identity() - res[0][0] = 2 / (right - left) - res[1][1] = 2 / (top - bottom) - res[2][2] = -2 / (zFar - zNear) - res[3][0] = -(right + left) / (right - left) - res[3][1] = -(top + bottom) / (top - bottom) - res[3][2] = -(zFar + zNear) / (zFar - zNear) - return res.T - - -def lookat(eye, center, up): - f = normalize(center - eye) - s = normalize(cross(f, up)) - u = cross(s, f) - - res = identity() - res[0][0] = s[0] - res[1][0] = s[1] - res[2][0] = s[2] - res[0][1] = u[0] - res[1][1] = u[1] - res[2][1] = u[2] - res[0][2] = -f[0] - res[1][2] = -f[1] - res[2][2] = -f[2] - res[3][0] = -dot(s, eye) - res[3][1] = -dot(u, eye) - res[3][2] = -dot(f, eye) - return res.T - - -def transform(d, m): - return np.dot(m, d.T).T \ No newline at end of file diff --git a/lib / renderer /mesh.py b/lib / renderer /mesh.py deleted file mode 100644 index 1bba90625694abd908c86089914956b63afe0ed6..0000000000000000000000000000000000000000 --- a/lib / renderer /mesh.py +++ /dev/null @@ -1,526 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from lib.dataset.mesh_util import SMPLX -from lib.common.render_utils import face_vertices -import numpy as np -import lib.smplx as smplx -import trimesh -import torch -import torch.nn.functional as F - -model_init_params = dict( - gender='male', - model_type='smplx', - model_path=SMPLX().model_dir, - create_global_orient=False, - create_body_pose=False, - create_betas=False, - create_left_hand_pose=False, - create_right_hand_pose=False, - create_expression=False, - create_jaw_pose=False, - create_leye_pose=False, - create_reye_pose=False, - create_transl=False, - num_pca_comps=12) - - -def get_smpl_model(model_type, gender): return smplx.create( - **model_init_params) - - -def normalization(data): - _range = np.max(data) - np.min(data) - return ((data - np.min(data)) / _range) - - -def sigmoid(x): - z = 1 / (1 + np.exp(-x)) - return z - - -def load_fit_body(fitted_path, scale, smpl_type='smplx', smpl_gender='neutral', noise_dict=None): - - param = np.load(fitted_path, allow_pickle=True) - for key in param.keys(): - param[key] = torch.as_tensor(param[key]) - - smpl_model = get_smpl_model(smpl_type, smpl_gender) - model_forward_params = dict(betas=param['betas'], - global_orient=param['global_orient'], - body_pose=param['body_pose'], - left_hand_pose=param['left_hand_pose'], - right_hand_pose=param['right_hand_pose'], - jaw_pose=param['jaw_pose'], - leye_pose=param['leye_pose'], - reye_pose=param['reye_pose'], - expression=param['expression'], - return_verts=True) - - if noise_dict is not None: - model_forward_params.update(noise_dict) - - smpl_out = smpl_model(**model_forward_params) - - smpl_verts = ( - (smpl_out.vertices[0] * param['scale'] + param['translation']) * scale).detach() - smpl_joints = ( - (smpl_out.joints[0] * param['scale'] + param['translation']) * scale).detach() - smpl_mesh = trimesh.Trimesh(smpl_verts, - smpl_model.faces, - process=False, maintain_order=True) - - return smpl_mesh, smpl_joints - - -def load_ori_fit_body(fitted_path, smpl_type='smplx', smpl_gender='neutral'): - - param = np.load(fitted_path, allow_pickle=True) - for key in param.keys(): - param[key] = torch.as_tensor(param[key]) - - smpl_model = get_smpl_model(smpl_type, smpl_gender) - model_forward_params = dict(betas=param['betas'], - global_orient=param['global_orient'], - body_pose=param['body_pose'], - left_hand_pose=param['left_hand_pose'], - right_hand_pose=param['right_hand_pose'], - jaw_pose=param['jaw_pose'], - leye_pose=param['leye_pose'], - reye_pose=param['reye_pose'], - expression=param['expression'], - return_verts=True) - - smpl_out = smpl_model(**model_forward_params) - - smpl_verts = smpl_out.vertices[0].detach() - smpl_mesh = trimesh.Trimesh(smpl_verts, - smpl_model.faces, - process=False, maintain_order=True) - - return smpl_mesh - - -def save_obj_mesh(mesh_path, verts, faces): - file = open(mesh_path, 'w') - for v in verts: - file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2])) - for f in faces: - f_plus = f + 1 - file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2])) - file.close() - - -# https://github.com/ratcave/wavefront_reader -def read_mtlfile(fname): - materials = {} - with open(fname) as f: - lines = f.read().splitlines() - - for line in lines: - if line: - split_line = line.strip().split(' ', 1) - if len(split_line) < 2: - continue - - prefix, data = split_line[0], split_line[1] - if 'newmtl' in prefix: - material = {} - materials[data] = material - elif materials: - if data: - split_data = data.strip().split(' ') - - # assume texture maps are in the same level - # WARNING: do not include space in your filename!! - if 'map' in prefix: - material[prefix] = split_data[-1].split('\\')[-1] - elif len(split_data) > 1: - material[prefix] = tuple(float(d) for d in split_data) - else: - try: - material[prefix] = int(data) - except ValueError: - material[prefix] = float(data) - - return materials - - -def load_obj_mesh_mtl(mesh_file): - vertex_data = [] - norm_data = [] - uv_data = [] - - face_data = [] - face_norm_data = [] - face_uv_data = [] - - # face per material - face_data_mat = {} - face_norm_data_mat = {} - face_uv_data_mat = {} - - # current material name - mtl_data = None - cur_mat = None - - if isinstance(mesh_file, str): - f = open(mesh_file, "r") - else: - f = mesh_file - for line in f: - if isinstance(line, bytes): - line = line.decode("utf-8") - if line.startswith('#'): - continue - values = line.split() - if not values: - continue - - if values[0] == 'v': - v = list(map(float, values[1:4])) - vertex_data.append(v) - elif values[0] == 'vn': - vn = list(map(float, values[1:4])) - norm_data.append(vn) - elif values[0] == 'vt': - vt = list(map(float, values[1:3])) - uv_data.append(vt) - elif values[0] == 'mtllib': - mtl_data = read_mtlfile( - mesh_file.replace(mesh_file.split('/')[-1], values[1])) - elif values[0] == 'usemtl': - cur_mat = values[1] - elif values[0] == 'f': - # local triangle data - l_face_data = [] - l_face_uv_data = [] - l_face_norm_data = [] - - # quad mesh - if len(values) > 4: - f = list( - map( - lambda x: int(x.split('/')[0]) if int(x.split('/')[0]) - < 0 else int(x.split('/')[0]) - 1, values[1:4])) - l_face_data.append(f) - f = list( - map( - lambda x: int(x.split('/')[0]) - if int(x.split('/')[0]) < 0 else int(x.split('/')[0]) - - 1, [values[3], values[4], values[1]])) - l_face_data.append(f) - # tri mesh - else: - f = list( - map( - lambda x: int(x.split('/')[0]) if int(x.split('/')[0]) - < 0 else int(x.split('/')[0]) - 1, values[1:4])) - l_face_data.append(f) - # deal with texture - if len(values[1].split('/')) >= 2: - # quad mesh - if len(values) > 4: - f = list( - map( - lambda x: int(x.split('/')[1]) - if int(x.split('/')[1]) < 0 else int( - x.split('/')[1]) - 1, values[1:4])) - l_face_uv_data.append(f) - f = list( - map( - lambda x: int(x.split('/')[1]) - if int(x.split('/')[1]) < 0 else int( - x.split('/')[1]) - 1, - [values[3], values[4], values[1]])) - l_face_uv_data.append(f) - # tri mesh - elif len(values[1].split('/')[1]) != 0: - f = list( - map( - lambda x: int(x.split('/')[1]) - if int(x.split('/')[1]) < 0 else int( - x.split('/')[1]) - 1, values[1:4])) - l_face_uv_data.append(f) - # deal with normal - if len(values[1].split('/')) == 3: - # quad mesh - if len(values) > 4: - f = list( - map( - lambda x: int(x.split('/')[2]) - if int(x.split('/')[2]) < 0 else int( - x.split('/')[2]) - 1, values[1:4])) - l_face_norm_data.append(f) - f = list( - map( - lambda x: int(x.split('/')[2]) - if int(x.split('/')[2]) < 0 else int( - x.split('/')[2]) - 1, - [values[3], values[4], values[1]])) - l_face_norm_data.append(f) - # tri mesh - elif len(values[1].split('/')[2]) != 0: - f = list( - map( - lambda x: int(x.split('/')[2]) - if int(x.split('/')[2]) < 0 else int( - x.split('/')[2]) - 1, values[1:4])) - l_face_norm_data.append(f) - - face_data += l_face_data - face_uv_data += l_face_uv_data - face_norm_data += l_face_norm_data - - if cur_mat is not None: - if cur_mat not in face_data_mat.keys(): - face_data_mat[cur_mat] = [] - if cur_mat not in face_uv_data_mat.keys(): - face_uv_data_mat[cur_mat] = [] - if cur_mat not in face_norm_data_mat.keys(): - face_norm_data_mat[cur_mat] = [] - face_data_mat[cur_mat] += l_face_data - face_uv_data_mat[cur_mat] += l_face_uv_data - face_norm_data_mat[cur_mat] += l_face_norm_data - - vertices = np.array(vertex_data) - faces = np.array(face_data) - - norms = np.array(norm_data) - norms = normalize_v3(norms) - face_normals = np.array(face_norm_data) - - uvs = np.array(uv_data) - face_uvs = np.array(face_uv_data) - - out_tuple = (vertices, faces, norms, face_normals, uvs, face_uvs) - - if cur_mat is not None and mtl_data is not None: - for key in face_data_mat: - face_data_mat[key] = np.array(face_data_mat[key]) - face_uv_data_mat[key] = np.array(face_uv_data_mat[key]) - face_norm_data_mat[key] = np.array(face_norm_data_mat[key]) - - out_tuple += (face_data_mat, face_norm_data_mat, face_uv_data_mat, - mtl_data) - - return out_tuple - - -def load_scan(mesh_file, with_normal=False, with_texture=False): - vertex_data = [] - norm_data = [] - uv_data = [] - - face_data = [] - face_norm_data = [] - face_uv_data = [] - - if isinstance(mesh_file, str): - f = open(mesh_file, "r") - else: - f = mesh_file - for line in f: - if isinstance(line, bytes): - line = line.decode("utf-8") - if line.startswith('#'): - continue - values = line.split() - if not values: - continue - - if values[0] == 'v': - v = list(map(float, values[1:4])) - vertex_data.append(v) - elif values[0] == 'vn': - vn = list(map(float, values[1:4])) - norm_data.append(vn) - elif values[0] == 'vt': - vt = list(map(float, values[1:3])) - uv_data.append(vt) - - elif values[0] == 'f': - # quad mesh - if len(values) > 4: - f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) - face_data.append(f) - f = list( - map(lambda x: int(x.split('/')[0]), - [values[3], values[4], values[1]])) - face_data.append(f) - # tri mesh - else: - f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) - face_data.append(f) - - # deal with texture - if len(values[1].split('/')) >= 2: - # quad mesh - if len(values) > 4: - f = list(map(lambda x: int(x.split('/')[1]), values[1:4])) - face_uv_data.append(f) - f = list( - map(lambda x: int(x.split('/')[1]), - [values[3], values[4], values[1]])) - face_uv_data.append(f) - # tri mesh - elif len(values[1].split('/')[1]) != 0: - f = list(map(lambda x: int(x.split('/')[1]), values[1:4])) - face_uv_data.append(f) - # deal with normal - if len(values[1].split('/')) == 3: - # quad mesh - if len(values) > 4: - f = list(map(lambda x: int(x.split('/')[2]), values[1:4])) - face_norm_data.append(f) - f = list( - map(lambda x: int(x.split('/')[2]), - [values[3], values[4], values[1]])) - face_norm_data.append(f) - # tri mesh - elif len(values[1].split('/')[2]) != 0: - f = list(map(lambda x: int(x.split('/')[2]), values[1:4])) - face_norm_data.append(f) - - vertices = np.array(vertex_data) - faces = np.array(face_data) - 1 - - if with_texture and with_normal: - uvs = np.array(uv_data) - face_uvs = np.array(face_uv_data) - 1 - norms = np.array(norm_data) - if norms.shape[0] == 0: - norms = compute_normal(vertices, faces) - face_normals = faces - else: - norms = normalize_v3(norms) - face_normals = np.array(face_norm_data) - 1 - return vertices, faces, norms, face_normals, uvs, face_uvs - - if with_texture: - uvs = np.array(uv_data) - face_uvs = np.array(face_uv_data) - 1 - return vertices, faces, uvs, face_uvs - - if with_normal: - norms = np.array(norm_data) - norms = normalize_v3(norms) - face_normals = np.array(face_norm_data) - 1 - return vertices, faces, norms, face_normals - - return vertices, faces - - -def normalize_v3(arr): - ''' Normalize a numpy array of 3 component vectors shape=(n,3) ''' - lens = np.sqrt(arr[:, 0]**2 + arr[:, 1]**2 + arr[:, 2]**2) - eps = 0.00000001 - lens[lens < eps] = eps - arr[:, 0] /= lens - arr[:, 1] /= lens - arr[:, 2] /= lens - return arr - - -def compute_normal(vertices, faces): - # Create a zeroed array with the same type and shape as our vertices i.e., per vertex normal - norm = np.zeros(vertices.shape, dtype=vertices.dtype) - # Create an indexed view into the vertex array using the array of three indices for triangles - tris = vertices[faces] - # Calculate the normal for all the triangles, by taking the cross product of the vectors v1-v0, and v2-v0 in each triangle - n = np.cross(tris[::, 1] - tris[::, 0], tris[::, 2] - tris[::, 0]) - # n is now an array of normals per triangle. The length of each normal is dependent the vertices, - # we need to normalize these, so that our next step weights each normal equally. - normalize_v3(n) - # now we have a normalized array of normals, one per triangle, i.e., per triangle normals. - # But instead of one per triangle (i.e., flat shading), we add to each vertex in that triangle, - # the triangles' normal. Multiple triangles would then contribute to every vertex, so we need to normalize again afterwards. - # The cool part, we can actually add the normals through an indexed view of our (zeroed) per vertex normal array - norm[faces[:, 0]] += n - norm[faces[:, 1]] += n - norm[faces[:, 2]] += n - normalize_v3(norm) - - return norm - - -def compute_normal_batch(vertices, faces): - - bs, nv = vertices.shape[:2] - bs, nf = faces.shape[:2] - - vert_norm = torch.zeros(bs * nv, 3).type_as(vertices) - tris = face_vertices(vertices, faces) - face_norm = F.normalize(torch.cross(tris[:, :, 1] - tris[:, :, 0], - tris[:, :, 2] - tris[:, :, 0]), - dim=-1) - - faces = (faces + - (torch.arange(bs).type_as(faces) * nv)[:, None, None]).view( - -1, 3) - - vert_norm[faces[:, 0]] += face_norm.view(-1, 3) - vert_norm[faces[:, 1]] += face_norm.view(-1, 3) - vert_norm[faces[:, 2]] += face_norm.view(-1, 3) - - vert_norm = F.normalize(vert_norm, dim=-1).view(bs, nv, 3) - - return vert_norm - - -# compute tangent and bitangent -def compute_tangent(vertices, faces, normals, uvs, faceuvs): - # NOTE: this could be numerically unstable around [0,0,1] - # but other current solutions are pretty freaky somehow - c1 = np.cross(normals, np.array([0, 1, 0.0])) - tan = c1 - normalize_v3(tan) - btan = np.cross(normals, tan) - - # NOTE: traditional version is below - - # pts_tris = vertices[faces] - # uv_tris = uvs[faceuvs] - - # W = np.stack([pts_tris[::, 1] - pts_tris[::, 0], pts_tris[::, 2] - pts_tris[::, 0]],2) - # UV = np.stack([uv_tris[::, 1] - uv_tris[::, 0], uv_tris[::, 2] - uv_tris[::, 0]], 1) - - # for i in range(W.shape[0]): - # W[i,::] = W[i,::].dot(np.linalg.inv(UV[i,::])) - - # tan = np.zeros(vertices.shape, dtype=vertices.dtype) - # tan[faces[:,0]] += W[:,:,0] - # tan[faces[:,1]] += W[:,:,0] - # tan[faces[:,2]] += W[:,:,0] - - # btan = np.zeros(vertices.shape, dtype=vertices.dtype) - # btan[faces[:,0]] += W[:,:,1] - # btan[faces[:,1]] += W[:,:,1] - # btan[faces[:,2]] += W[:,:,1] - - # normalize_v3(tan) - - # ndott = np.sum(normals*tan, 1, keepdims=True) - # tan = tan - ndott * normals - - # normalize_v3(btan) - # normalize_v3(tan) - - # tan[np.sum(np.cross(normals, tan) * btan, 1) < 0,:] *= -1.0 - - return tan, btan diff --git a/lib / renderer /opengl_util.py b/lib / renderer /opengl_util.py deleted file mode 100644 index e68bdc2b3c069f6035bcb9c89da5d92ef1a73c54..0000000000000000000000000000000000000000 --- a/lib / renderer /opengl_util.py +++ /dev/null @@ -1,369 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import os - -from lib.renderer.mesh import load_scan, compute_tangent -from lib.renderer.camera import Camera -import cv2 -import math -import random -import numpy as np - - -def render_result(rndr, shader_id, path, mask=False): - - cam_render = rndr.get_color(shader_id) - cam_render = cv2.cvtColor(cam_render, cv2.COLOR_RGBA2BGRA) - - os.makedirs(os.path.dirname(path), exist_ok=True) - if shader_id != 2: - cv2.imwrite(path, np.uint8(255.0 * cam_render)) - else: - cam_render[:, :, -1] -= 0.5 - cam_render[:, :, -1] *= 2.0 - if not mask: - cv2.imwrite(path, np.uint8(255.0 / 2.0 * (cam_render + 1.0))) - else: - cv2.imwrite(path, np.uint8(-1.0 * cam_render[:, :, [3]])) - - -def make_rotate(rx, ry, rz): - sinX = np.sin(rx) - sinY = np.sin(ry) - sinZ = np.sin(rz) - - cosX = np.cos(rx) - cosY = np.cos(ry) - cosZ = np.cos(rz) - - Rx = np.zeros((3, 3)) - Rx[0, 0] = 1.0 - Rx[1, 1] = cosX - Rx[1, 2] = -sinX - Rx[2, 1] = sinX - Rx[2, 2] = cosX - - Ry = np.zeros((3, 3)) - Ry[0, 0] = cosY - Ry[0, 2] = sinY - Ry[1, 1] = 1.0 - Ry[2, 0] = -sinY - Ry[2, 2] = cosY - - Rz = np.zeros((3, 3)) - Rz[0, 0] = cosZ - Rz[0, 1] = -sinZ - Rz[1, 0] = sinZ - Rz[1, 1] = cosZ - Rz[2, 2] = 1.0 - - R = np.matmul(np.matmul(Rz, Ry), Rx) - return R - - -def rotateSH(SH, R): - SHn = SH - - # 1st order - SHn[1] = R[1, 1] * SH[1] - R[1, 2] * SH[2] + R[1, 0] * SH[3] - SHn[2] = -R[2, 1] * SH[1] + R[2, 2] * SH[2] - R[2, 0] * SH[3] - SHn[3] = R[0, 1] * SH[1] - R[0, 2] * SH[2] + R[0, 0] * SH[3] - - # 2nd order - SHn[4:, 0] = rotateBand2(SH[4:, 0], R) - SHn[4:, 1] = rotateBand2(SH[4:, 1], R) - SHn[4:, 2] = rotateBand2(SH[4:, 2], R) - - return SHn - - -def rotateBand2(x, R): - s_c3 = 0.94617469575 - s_c4 = -0.31539156525 - s_c5 = 0.54627421529 - - s_c_scale = 1.0 / 0.91529123286551084 - s_c_scale_inv = 0.91529123286551084 - - s_rc2 = 1.5853309190550713 * s_c_scale - s_c4_div_c3 = s_c4 / s_c3 - s_c4_div_c3_x2 = (s_c4 / s_c3) * 2.0 - - s_scale_dst2 = s_c3 * s_c_scale_inv - s_scale_dst4 = s_c5 * s_c_scale_inv - - sh0 = x[3] + x[4] + x[4] - x[1] - sh1 = x[0] + s_rc2 * x[2] + x[3] + x[4] - sh2 = x[0] - sh3 = -x[3] - sh4 = -x[1] - - r2x = R[0][0] + R[0][1] - r2y = R[1][0] + R[1][1] - r2z = R[2][0] + R[2][1] - - r3x = R[0][0] + R[0][2] - r3y = R[1][0] + R[1][2] - r3z = R[2][0] + R[2][2] - - r4x = R[0][1] + R[0][2] - r4y = R[1][1] + R[1][2] - r4z = R[2][1] + R[2][2] - - sh0_x = sh0 * R[0][0] - sh0_y = sh0 * R[1][0] - d0 = sh0_x * R[1][0] - d1 = sh0_y * R[2][0] - d2 = sh0 * (R[2][0] * R[2][0] + s_c4_div_c3) - d3 = sh0_x * R[2][0] - d4 = sh0_x * R[0][0] - sh0_y * R[1][0] - - sh1_x = sh1 * R[0][2] - sh1_y = sh1 * R[1][2] - d0 += sh1_x * R[1][2] - d1 += sh1_y * R[2][2] - d2 += sh1 * (R[2][2] * R[2][2] + s_c4_div_c3) - d3 += sh1_x * R[2][2] - d4 += sh1_x * R[0][2] - sh1_y * R[1][2] - - sh2_x = sh2 * r2x - sh2_y = sh2 * r2y - d0 += sh2_x * r2y - d1 += sh2_y * r2z - d2 += sh2 * (r2z * r2z + s_c4_div_c3_x2) - d3 += sh2_x * r2z - d4 += sh2_x * r2x - sh2_y * r2y - - sh3_x = sh3 * r3x - sh3_y = sh3 * r3y - d0 += sh3_x * r3y - d1 += sh3_y * r3z - d2 += sh3 * (r3z * r3z + s_c4_div_c3_x2) - d3 += sh3_x * r3z - d4 += sh3_x * r3x - sh3_y * r3y - - sh4_x = sh4 * r4x - sh4_y = sh4 * r4y - d0 += sh4_x * r4y - d1 += sh4_y * r4z - d2 += sh4 * (r4z * r4z + s_c4_div_c3_x2) - d3 += sh4_x * r4z - d4 += sh4_x * r4x - sh4_y * r4y - - dst = x - dst[0] = d0 - dst[1] = -d1 - dst[2] = d2 * s_scale_dst2 - dst[3] = -d3 - dst[4] = d4 * s_scale_dst4 - - return dst - - -def load_calib(param, render_size=512): - # pixel unit / world unit - ortho_ratio = param['ortho_ratio'] - # world unit / model unit - scale = param['scale'] - # camera center world coordinate - center = param['center'] - # model rotation - R = param['R'] - - translate = -np.matmul(R, center).reshape(3, 1) - extrinsic = np.concatenate([R, translate], axis=1) - extrinsic = np.concatenate( - [extrinsic, np.array([0, 0, 0, 1]).reshape(1, 4)], 0) - # Match camera space to image pixel space - scale_intrinsic = np.identity(4) - scale_intrinsic[0, 0] = scale / ortho_ratio - scale_intrinsic[1, 1] = -scale / ortho_ratio - scale_intrinsic[2, 2] = scale / ortho_ratio - # Match image pixel space to image uv space - uv_intrinsic = np.identity(4) - uv_intrinsic[0, 0] = 1.0 / float(render_size // 2) - uv_intrinsic[1, 1] = 1.0 / float(render_size // 2) - uv_intrinsic[2, 2] = 1.0 / float(render_size // 2) - - intrinsic = np.matmul(uv_intrinsic, scale_intrinsic) - calib = np.concatenate([extrinsic, intrinsic], axis=0) - return calib - - -def render_prt_ortho(out_path, - folder_name, - subject_name, - shs, - rndr, - rndr_uv, - im_size, - angl_step=4, - n_light=1, - pitch=[0]): - cam = Camera(width=im_size, height=im_size) - cam.ortho_ratio = 0.4 * (512 / im_size) - cam.near = -100 - cam.far = 100 - cam.sanity_check() - - # set path for obj, prt - mesh_file = os.path.join(folder_name, subject_name + '_100k.obj') - if not os.path.exists(mesh_file): - print('ERROR: obj file does not exist!!', mesh_file) - return - prt_file = os.path.join(folder_name, 'bounce', 'bounce0.txt') - if not os.path.exists(prt_file): - print('ERROR: prt file does not exist!!!', prt_file) - return - face_prt_file = os.path.join(folder_name, 'bounce', 'face.npy') - if not os.path.exists(face_prt_file): - print('ERROR: face prt file does not exist!!!', prt_file) - return - text_file = os.path.join(folder_name, 'tex', subject_name + '_dif_2k.jpg') - if not os.path.exists(text_file): - print('ERROR: dif file does not exist!!', text_file) - return - - texture_image = cv2.imread(text_file) - texture_image = cv2.cvtColor(texture_image, cv2.COLOR_BGR2RGB) - - vertices, faces, normals, faces_normals, textures, face_textures = load_scan( - mesh_file, with_normal=True, with_texture=True) - vmin = vertices.min(0) - vmax = vertices.max(0) - up_axis = 1 if (vmax - vmin).argmax() == 1 else 2 - - vmed = np.median(vertices, 0) - vmed[up_axis] = 0.5 * (vmax[up_axis] + vmin[up_axis]) - y_scale = 180 / (vmax[up_axis] - vmin[up_axis]) - - rndr.set_norm_mat(y_scale, vmed) - rndr_uv.set_norm_mat(y_scale, vmed) - - tan, bitan = compute_tangent(vertices, faces, normals, textures, - face_textures) - prt = np.loadtxt(prt_file) - face_prt = np.load(face_prt_file) - rndr.set_mesh(vertices, faces, normals, faces_normals, textures, - face_textures, prt, face_prt, tan, bitan) - rndr.set_albedo(texture_image) - - rndr_uv.set_mesh(vertices, faces, normals, faces_normals, textures, - face_textures, prt, face_prt, tan, bitan) - rndr_uv.set_albedo(texture_image) - - os.makedirs(os.path.join(out_path, 'GEO', 'OBJ', subject_name), - exist_ok=True) - os.makedirs(os.path.join(out_path, 'PARAM', subject_name), exist_ok=True) - os.makedirs(os.path.join(out_path, 'RENDER', subject_name), exist_ok=True) - os.makedirs(os.path.join(out_path, 'MASK', subject_name), exist_ok=True) - os.makedirs(os.path.join(out_path, 'UV_RENDER', subject_name), - exist_ok=True) - os.makedirs(os.path.join(out_path, 'UV_MASK', subject_name), exist_ok=True) - os.makedirs(os.path.join(out_path, 'UV_POS', subject_name), exist_ok=True) - os.makedirs(os.path.join(out_path, 'UV_NORMAL', subject_name), - exist_ok=True) - - if not os.path.exists(os.path.join(out_path, 'val.txt')): - f = open(os.path.join(out_path, 'val.txt'), 'w') - f.close() - - # copy obj file - cmd = 'cp %s %s' % (mesh_file, - os.path.join(out_path, 'GEO', 'OBJ', subject_name)) - print(cmd) - os.system(cmd) - - for p in pitch: - for y in tqdm(range(0, 360, angl_step)): - R = np.matmul(make_rotate(math.radians(p), 0, 0), - make_rotate(0, math.radians(y), 0)) - if up_axis == 2: - R = np.matmul(R, make_rotate(math.radians(90), 0, 0)) - - rndr.rot_matrix = R - rndr_uv.rot_matrix = R - rndr.set_camera(cam) - rndr_uv.set_camera(cam) - - for j in range(n_light): - sh_id = random.randint(0, shs.shape[0] - 1) - sh = shs[sh_id] - sh_angle = 0.2 * np.pi * (random.random() - 0.5) - sh = rotateSH(sh, make_rotate(0, sh_angle, 0).T) - - dic = { - 'sh': sh, - 'ortho_ratio': cam.ortho_ratio, - 'scale': y_scale, - 'center': vmed, - 'R': R - } - - rndr.set_sh(sh) - rndr.analytic = False - rndr.use_inverse_depth = False - rndr.display() - - out_all_f = rndr.get_color(0) - out_mask = out_all_f[:, :, 3] - out_all_f = cv2.cvtColor(out_all_f, cv2.COLOR_RGBA2BGR) - - np.save( - os.path.join(out_path, 'PARAM', subject_name, - '%d_%d_%02d.npy' % (y, p, j)), dic) - cv2.imwrite( - os.path.join(out_path, 'RENDER', subject_name, - '%d_%d_%02d.jpg' % (y, p, j)), - 255.0 * out_all_f) - cv2.imwrite( - os.path.join(out_path, 'MASK', subject_name, - '%d_%d_%02d.png' % (y, p, j)), - 255.0 * out_mask) - - rndr_uv.set_sh(sh) - rndr_uv.analytic = False - rndr_uv.use_inverse_depth = False - rndr_uv.display() - - uv_color = rndr_uv.get_color(0) - uv_color = cv2.cvtColor(uv_color, cv2.COLOR_RGBA2BGR) - cv2.imwrite( - os.path.join(out_path, 'UV_RENDER', subject_name, - '%d_%d_%02d.jpg' % (y, p, j)), - 255.0 * uv_color) - - if y == 0 and j == 0 and p == pitch[0]: - uv_pos = rndr_uv.get_color(1) - uv_mask = uv_pos[:, :, 3] - cv2.imwrite( - os.path.join(out_path, 'UV_MASK', subject_name, - '00.png'), 255.0 * uv_mask) - - data = { - 'default': uv_pos[:, :, :3] - } # default is a reserved name - pyexr.write( - os.path.join(out_path, 'UV_POS', subject_name, - '00.exr'), data) - - uv_nml = rndr_uv.get_color(2) - uv_nml = cv2.cvtColor(uv_nml, cv2.COLOR_RGBA2BGR) - cv2.imwrite( - os.path.join(out_path, 'UV_NORMAL', subject_name, - '00.png'), 255.0 * uv_nml) \ No newline at end of file diff --git a/lib / renderer /prt_util.py b/lib / renderer /prt_util.py deleted file mode 100644 index de8290ca7a326774b319f09bbb444dc82e9161d7..0000000000000000000000000000000000000000 --- a/lib / renderer /prt_util.py +++ /dev/null @@ -1,199 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import os -import trimesh -import numpy as np -import math -from scipy.special import sph_harm -import argparse -from tqdm import tqdm -from trimesh.util import bounds_tree - - -def factratio(N, D): - if N >= D: - prod = 1.0 - for i in range(D + 1, N + 1): - prod *= i - return prod - else: - prod = 1.0 - for i in range(N + 1, D + 1): - prod *= i - return 1.0 / prod - - -def KVal(M, L): - return math.sqrt(((2 * L + 1) / (4 * math.pi)) * (factratio(L - M, L + M))) - - -def AssociatedLegendre(M, L, x): - if M < 0 or M > L or np.max(np.abs(x)) > 1.0: - return np.zeros_like(x) - - pmm = np.ones_like(x) - if M > 0: - somx2 = np.sqrt((1.0 + x) * (1.0 - x)) - fact = 1.0 - for i in range(1, M + 1): - pmm = -pmm * fact * somx2 - fact = fact + 2 - - if L == M: - return pmm - else: - pmmp1 = x * (2 * M + 1) * pmm - if L == M + 1: - return pmmp1 - else: - pll = np.zeros_like(x) - for i in range(M + 2, L + 1): - pll = (x * (2 * i - 1) * pmmp1 - (i + M - 1) * pmm) / (i - M) - pmm = pmmp1 - pmmp1 = pll - return pll - - -def SphericalHarmonic(M, L, theta, phi): - if M > 0: - return math.sqrt(2.0) * KVal(M, L) * np.cos( - M * phi) * AssociatedLegendre(M, L, np.cos(theta)) - elif M < 0: - return math.sqrt(2.0) * KVal(-M, L) * np.sin( - -M * phi) * AssociatedLegendre(-M, L, np.cos(theta)) - else: - return KVal(0, L) * AssociatedLegendre(0, L, np.cos(theta)) - - -def save_obj(mesh_path, verts): - file = open(mesh_path, 'w') - for v in verts: - file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2])) - file.close() - - -def sampleSphericalDirections(n): - xv = np.random.rand(n, n) - yv = np.random.rand(n, n) - theta = np.arccos(1 - 2 * xv) - phi = 2.0 * math.pi * yv - - phi = phi.reshape(-1) - theta = theta.reshape(-1) - - vx = -np.sin(theta) * np.cos(phi) - vy = -np.sin(theta) * np.sin(phi) - vz = np.cos(theta) - return np.stack([vx, vy, vz], 1), phi, theta - - -def getSHCoeffs(order, phi, theta): - shs = [] - for n in range(0, order + 1): - for m in range(-n, n + 1): - s = SphericalHarmonic(m, n, theta, phi) - shs.append(s) - - return np.stack(shs, 1) - - -def computePRT(mesh_path, scale, n, order): - - prt_dir = os.path.join(os.path.dirname(mesh_path), "prt") - bounce_path = os.path.join(prt_dir, "bounce.npy") - face_path = os.path.join(prt_dir, "face.npy") - - os.makedirs(prt_dir, exist_ok=True) - - PRT = None - F = None - - if os.path.exists(bounce_path) and os.path.exists(face_path): - - PRT = np.load(bounce_path) - F = np.load(face_path) - - else: - - mesh = trimesh.load(mesh_path, - skip_materials=True, - process=False, - maintain_order=True) - mesh.vertices *= scale - - vectors_orig, phi, theta = sampleSphericalDirections(n) - SH_orig = getSHCoeffs(order, phi, theta) - - w = 4.0 * math.pi / (n * n) - - origins = mesh.vertices - normals = mesh.vertex_normals - n_v = origins.shape[0] - - origins = np.repeat(origins[:, None], n, axis=1).reshape(-1, 3) - normals = np.repeat(normals[:, None], n, axis=1).reshape(-1, 3) - PRT_all = None - for i in range(n): - SH = np.repeat(SH_orig[None, (i * n):((i + 1) * n)], n_v, - axis=0).reshape(-1, SH_orig.shape[1]) - vectors = np.repeat(vectors_orig[None, (i * n):((i + 1) * n)], - n_v, - axis=0).reshape(-1, 3) - - dots = (vectors * normals).sum(1) - front = (dots > 0.0) - - delta = 1e-3 * min(mesh.bounding_box.extents) - - hits = mesh.ray.intersects_any(origins + delta * normals, vectors) - nohits = np.logical_and(front, np.logical_not(hits)) - - PRT = (nohits.astype(np.float) * dots)[:, None] * SH - - if PRT_all is not None: - PRT_all += (PRT.reshape(-1, n, SH.shape[1]).sum(1)) - else: - PRT_all = (PRT.reshape(-1, n, SH.shape[1]).sum(1)) - - PRT = w * PRT_all - F = mesh.faces - - np.save(bounce_path, PRT) - np.save(face_path, F) - - # NOTE: trimesh sometimes break the original vertex order, but topology will not change. - # when loading PRT in other program, use the triangle list from trimesh. - - return PRT, F - - -def testPRT(obj_path, n=40): - - os.makedirs(os.path.join(os.path.dirname(obj_path), - f'../bounce/{os.path.basename(obj_path)[:-4]}'), - exist_ok=True) - - PRT, F = computePRT(obj_path, n, 2) - np.savetxt( - os.path.join(os.path.dirname(obj_path), - f'../bounce/{os.path.basename(obj_path)[:-4]}', - 'bounce.npy'), PRT) - np.save( - os.path.join(os.path.dirname(obj_path), - f'../bounce/{os.path.basename(obj_path)[:-4]}', - 'face.npy'), F) \ No newline at end of file diff --git a/lib /common / __init__.py b/lib /common / __init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/lib /common / seg3d_utils.py b/lib /common / seg3d_utils.py deleted file mode 100644 index a86a6d37daa3586a30066a856a6b9a8d361d1d5d..0000000000000000000000000000000000000000 --- a/lib /common / seg3d_utils.py +++ /dev/null @@ -1,390 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import torch -import torch.nn as nn -import torch.nn.functional as F -import matplotlib.pyplot as plt - - -def plot_mask2D(mask, - title="", - point_coords=None, - figsize=10, - point_marker_size=5): - ''' - Simple plotting tool to show intermediate mask predictions and points - where PointRend is applied. - Args: - mask (Tensor): mask prediction of shape HxW - title (str): title for the plot - point_coords ((Tensor, Tensor)): x and y point coordinates - figsize (int): size of the figure to plot - point_marker_size (int): marker size for points - ''' - - H, W = mask.shape - plt.figure(figsize=(figsize, figsize)) - if title: - title += ", " - plt.title("{}resolution {}x{}".format(title, H, W), fontsize=30) - plt.ylabel(H, fontsize=30) - plt.xlabel(W, fontsize=30) - plt.xticks([], []) - plt.yticks([], []) - plt.imshow(mask.detach(), - interpolation="nearest", - cmap=plt.get_cmap('gray')) - if point_coords is not None: - plt.scatter(x=point_coords[0], - y=point_coords[1], - color="red", - s=point_marker_size, - clip_on=True) - plt.xlim(-0.5, W - 0.5) - plt.ylim(H - 0.5, -0.5) - plt.show() - - -def plot_mask3D(mask=None, - title="", - point_coords=None, - figsize=1500, - point_marker_size=8, - interactive=True): - ''' - Simple plotting tool to show intermediate mask predictions and points - where PointRend is applied. - Args: - mask (Tensor): mask prediction of shape DxHxW - title (str): title for the plot - point_coords ((Tensor, Tensor, Tensor)): x and y and z point coordinates - figsize (int): size of the figure to plot - point_marker_size (int): marker size for points - ''' - import trimesh - import vtkplotter - from skimage import measure - - vp = vtkplotter.Plotter(title=title, size=(figsize, figsize)) - vis_list = [] - - if mask is not None: - mask = mask.detach().to("cpu").numpy() - mask = mask.transpose(2, 1, 0) - - # marching cube to find surface - verts, faces, normals, values = measure.marching_cubes_lewiner( - mask, 0.5, gradient_direction='ascent') - - # create a mesh - mesh = trimesh.Trimesh(verts, faces) - mesh.visual.face_colors = [200, 200, 250, 100] - vis_list.append(mesh) - - if point_coords is not None: - point_coords = torch.stack(point_coords, 1).to("cpu").numpy() - - # import numpy as np - # select_x = np.logical_and(point_coords[:, 0] >= 16, point_coords[:, 0] <= 112) - # select_y = np.logical_and(point_coords[:, 1] >= 48, point_coords[:, 1] <= 272) - # select_z = np.logical_and(point_coords[:, 2] >= 16, point_coords[:, 2] <= 112) - # select = np.logical_and(np.logical_and(select_x, select_y), select_z) - # point_coords = point_coords[select, :] - - pc = vtkplotter.Points(point_coords, r=point_marker_size, c='red') - vis_list.append(pc) - - vp.show(*vis_list, - bg="white", - axes=1, - interactive=interactive, - azimuth=30, - elevation=30) - - -def create_grid3D(min, max, steps): - if type(min) is int: - min = (min, min, min) # (x, y, z) - if type(max) is int: - max = (max, max, max) # (x, y) - if type(steps) is int: - steps = (steps, steps, steps) # (x, y, z) - arrangeX = torch.linspace(min[0], max[0], steps[0]).long() - arrangeY = torch.linspace(min[1], max[1], steps[1]).long() - arrangeZ = torch.linspace(min[2], max[2], steps[2]).long() - gridD, girdH, gridW = torch.meshgrid([arrangeZ, arrangeY, arrangeX]) - coords = torch.stack([gridW, girdH, - gridD]) # [2, steps[0], steps[1], steps[2]] - coords = coords.view(3, -1).t() # [N, 3] - return coords - - -def create_grid2D(min, max, steps): - if type(min) is int: - min = (min, min) # (x, y) - if type(max) is int: - max = (max, max) # (x, y) - if type(steps) is int: - steps = (steps, steps) # (x, y) - arrangeX = torch.linspace(min[0], max[0], steps[0]).long() - arrangeY = torch.linspace(min[1], max[1], steps[1]).long() - girdH, gridW = torch.meshgrid([arrangeY, arrangeX]) - coords = torch.stack([gridW, girdH]) # [2, steps[0], steps[1]] - coords = coords.view(2, -1).t() # [N, 2] - return coords - - -class SmoothConv2D(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size=3): - super().__init__() - assert kernel_size % 2 == 1, "kernel_size for smooth_conv must be odd: {3, 5, ...}" - self.padding = (kernel_size - 1) // 2 - - weight = torch.ones( - (in_channels, out_channels, kernel_size, kernel_size), - dtype=torch.float32) / (kernel_size**2) - self.register_buffer('weight', weight) - - def forward(self, input): - return F.conv2d(input, self.weight, padding=self.padding) - - -class SmoothConv3D(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size=3): - super().__init__() - assert kernel_size % 2 == 1, "kernel_size for smooth_conv must be odd: {3, 5, ...}" - self.padding = (kernel_size - 1) // 2 - - weight = torch.ones( - (in_channels, out_channels, kernel_size, kernel_size, kernel_size), - dtype=torch.float32) / (kernel_size**3) - self.register_buffer('weight', weight) - - def forward(self, input): - return F.conv3d(input, self.weight, padding=self.padding) - - -def build_smooth_conv3D(in_channels=1, - out_channels=1, - kernel_size=3, - padding=1): - smooth_conv = torch.nn.Conv3d(in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - padding=padding) - smooth_conv.weight.data = torch.ones( - (in_channels, out_channels, kernel_size, kernel_size, kernel_size), - dtype=torch.float32) / (kernel_size**3) - smooth_conv.bias.data = torch.zeros(out_channels) - return smooth_conv - - -def build_smooth_conv2D(in_channels=1, - out_channels=1, - kernel_size=3, - padding=1): - smooth_conv = torch.nn.Conv2d(in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - padding=padding) - smooth_conv.weight.data = torch.ones( - (in_channels, out_channels, kernel_size, kernel_size), - dtype=torch.float32) / (kernel_size**2) - smooth_conv.bias.data = torch.zeros(out_channels) - return smooth_conv - - -def get_uncertain_point_coords_on_grid3D(uncertainty_map, num_points, - **kwargs): - """ - Find `num_points` most uncertain points from `uncertainty_map` grid. - Args: - uncertainty_map (Tensor): A tensor of shape (N, 1, H, W, D) that contains uncertainty - values for a set of points on a regular H x W x D grid. - num_points (int): The number of points P to select. - Returns: - point_indices (Tensor): A tensor of shape (N, P) that contains indices from - [0, H x W x D) of the most uncertain points. - point_coords (Tensor): A tensor of shape (N, P, 3) that contains [0, 1] x [0, 1] normalized - coordinates of the most uncertain points from the H x W x D grid. - """ - R, _, D, H, W = uncertainty_map.shape - # h_step = 1.0 / float(H) - # w_step = 1.0 / float(W) - # d_step = 1.0 / float(D) - - num_points = min(D * H * W, num_points) - point_scores, point_indices = torch.topk(uncertainty_map.view( - R, D * H * W), - k=num_points, - dim=1) - point_coords = torch.zeros(R, - num_points, - 3, - dtype=torch.float, - device=uncertainty_map.device) - # point_coords[:, :, 0] = h_step / 2.0 + (point_indices // (W * D)).to(torch.float) * h_step - # point_coords[:, :, 1] = w_step / 2.0 + (point_indices % (W * D) // D).to(torch.float) * w_step - # point_coords[:, :, 2] = d_step / 2.0 + (point_indices % D).to(torch.float) * d_step - point_coords[:, :, 0] = (point_indices % W).to(torch.float) # x - point_coords[:, :, 1] = (point_indices % (H * W) // W).to(torch.float) # y - point_coords[:, :, 2] = (point_indices // (H * W)).to(torch.float) # z - print(f"resolution {D} x {H} x {W}", point_scores.min(), - point_scores.max()) - return point_indices, point_coords - - -def get_uncertain_point_coords_on_grid3D_faster(uncertainty_map, num_points, - clip_min): - """ - Find `num_points` most uncertain points from `uncertainty_map` grid. - Args: - uncertainty_map (Tensor): A tensor of shape (N, 1, H, W, D) that contains uncertainty - values for a set of points on a regular H x W x D grid. - num_points (int): The number of points P to select. - Returns: - point_indices (Tensor): A tensor of shape (N, P) that contains indices from - [0, H x W x D) of the most uncertain points. - point_coords (Tensor): A tensor of shape (N, P, 3) that contains [0, 1] x [0, 1] normalized - coordinates of the most uncertain points from the H x W x D grid. - """ - R, _, D, H, W = uncertainty_map.shape - # h_step = 1.0 / float(H) - # w_step = 1.0 / float(W) - # d_step = 1.0 / float(D) - - assert R == 1, "batchsize > 1 is not implemented!" - uncertainty_map = uncertainty_map.view(D * H * W) - indices = (uncertainty_map >= clip_min).nonzero().squeeze(1) - num_points = min(num_points, indices.size(0)) - point_scores, point_indices = torch.topk(uncertainty_map[indices], - k=num_points, - dim=0) - point_indices = indices[point_indices].unsqueeze(0) - - point_coords = torch.zeros(R, - num_points, - 3, - dtype=torch.float, - device=uncertainty_map.device) - # point_coords[:, :, 0] = h_step / 2.0 + (point_indices // (W * D)).to(torch.float) * h_step - # point_coords[:, :, 1] = w_step / 2.0 + (point_indices % (W * D) // D).to(torch.float) * w_step - # point_coords[:, :, 2] = d_step / 2.0 + (point_indices % D).to(torch.float) * d_step - point_coords[:, :, 0] = (point_indices % W).to(torch.float) # x - point_coords[:, :, 1] = (point_indices % (H * W) // W).to(torch.float) # y - point_coords[:, :, 2] = (point_indices // (H * W)).to(torch.float) # z - # print (f"resolution {D} x {H} x {W}", point_scores.min(), point_scores.max()) - return point_indices, point_coords - - -def get_uncertain_point_coords_on_grid2D(uncertainty_map, num_points, - **kwargs): - """ - Find `num_points` most uncertain points from `uncertainty_map` grid. - Args: - uncertainty_map (Tensor): A tensor of shape (N, 1, H, W) that contains uncertainty - values for a set of points on a regular H x W grid. - num_points (int): The number of points P to select. - Returns: - point_indices (Tensor): A tensor of shape (N, P) that contains indices from - [0, H x W) of the most uncertain points. - point_coords (Tensor): A tensor of shape (N, P, 2) that contains [0, 1] x [0, 1] normalized - coordinates of the most uncertain points from the H x W grid. - """ - R, _, H, W = uncertainty_map.shape - # h_step = 1.0 / float(H) - # w_step = 1.0 / float(W) - - num_points = min(H * W, num_points) - point_scores, point_indices = torch.topk(uncertainty_map.view(R, H * W), - k=num_points, - dim=1) - point_coords = torch.zeros(R, - num_points, - 2, - dtype=torch.long, - device=uncertainty_map.device) - # point_coords[:, :, 0] = w_step / 2.0 + (point_indices % W).to(torch.float) * w_step - # point_coords[:, :, 1] = h_step / 2.0 + (point_indices // W).to(torch.float) * h_step - point_coords[:, :, 0] = (point_indices % W).to(torch.long) - point_coords[:, :, 1] = (point_indices // W).to(torch.long) - # print (point_scores.min(), point_scores.max()) - return point_indices, point_coords - - -def get_uncertain_point_coords_on_grid2D_faster(uncertainty_map, num_points, - clip_min): - """ - Find `num_points` most uncertain points from `uncertainty_map` grid. - Args: - uncertainty_map (Tensor): A tensor of shape (N, 1, H, W) that contains uncertainty - values for a set of points on a regular H x W grid. - num_points (int): The number of points P to select. - Returns: - point_indices (Tensor): A tensor of shape (N, P) that contains indices from - [0, H x W) of the most uncertain points. - point_coords (Tensor): A tensor of shape (N, P, 2) that contains [0, 1] x [0, 1] normalized - coordinates of the most uncertain points from the H x W grid. - """ - R, _, H, W = uncertainty_map.shape - # h_step = 1.0 / float(H) - # w_step = 1.0 / float(W) - - assert R == 1, "batchsize > 1 is not implemented!" - uncertainty_map = uncertainty_map.view(H * W) - indices = (uncertainty_map >= clip_min).nonzero().squeeze(1) - num_points = min(num_points, indices.size(0)) - point_scores, point_indices = torch.topk(uncertainty_map[indices], - k=num_points, - dim=0) - point_indices = indices[point_indices].unsqueeze(0) - - point_coords = torch.zeros(R, - num_points, - 2, - dtype=torch.long, - device=uncertainty_map.device) - # point_coords[:, :, 0] = w_step / 2.0 + (point_indices % W).to(torch.float) * w_step - # point_coords[:, :, 1] = h_step / 2.0 + (point_indices // W).to(torch.float) * h_step - point_coords[:, :, 0] = (point_indices % W).to(torch.long) - point_coords[:, :, 1] = (point_indices // W).to(torch.long) - # print (point_scores.min(), point_scores.max()) - return point_indices, point_coords - - -def calculate_uncertainty(logits, classes=None, balance_value=0.5): - """ - We estimate uncerainty as L1 distance between 0.0 and the logit prediction in 'logits' for the - foreground class in `classes`. - Args: - logits (Tensor): A tensor of shape (R, C, ...) or (R, 1, ...) for class-specific or - class-agnostic, where R is the total number of predicted masks in all images and C is - the number of foreground classes. The values are logits. - classes (list): A list of length R that contains either predicted of ground truth class - for eash predicted mask. - Returns: - scores (Tensor): A tensor of shape (R, 1, ...) that contains uncertainty scores with - the most uncertain locations having the highest uncertainty score. - """ - if logits.shape[1] == 1: - gt_class_logits = logits - else: - gt_class_logits = logits[ - torch.arange(logits.shape[0], device=logits.device), - classes].unsqueeze(1) - return -torch.abs(gt_class_logits - balance_value) \ No newline at end of file diff --git a/lib /common /cloth_extraction.py b/lib /common /cloth_extraction.py deleted file mode 100644 index 934c2d3b572e7b5ca476ea1147c1d1bb51463c56..0000000000000000000000000000000000000000 --- a/lib /common /cloth_extraction.py +++ /dev/null @@ -1,170 +0,0 @@ -import numpy as np -import json -import os -import itertools -import trimesh -from matplotlib.path import Path -from collections import Counter -from sklearn.neighbors import KNeighborsClassifier - - -def load_segmentation(path, shape): - """ - Get a segmentation mask for a given image - Arguments: - path: path to the segmentation json file - shape: shape of the output mask - Returns: - Returns a segmentation mask - """ - with open(path) as json_file: - dict = json.load(json_file) - segmentations = [] - for key, val in dict.items(): - if not key.startswith('item'): - continue - - # Each item can have multiple polygons. Combine them to one - # segmentation_coord = list(itertools.chain.from_iterable(val['segmentation'])) - # segmentation_coord = np.round(np.array(segmentation_coord)).astype(int) - - coordinates = [] - for segmentation_coord in val['segmentation']: - # The format before is [x1,y1, x2, y2, ....] - x = segmentation_coord[::2] - y = segmentation_coord[1::2] - xy = np.vstack((x, y)).T - coordinates.append(xy) - - segmentations.append( - {'type': val['category_name'], 'type_id': val['category_id'], 'coordinates': coordinates}) - - return segmentations - - -def smpl_to_recon_labels(recon, smpl, k=1): - """ - Get the bodypart labels for the recon object by using the labels from the corresponding smpl object - Arguments: - recon: trimesh object (fully clothed model) - shape: trimesh object (smpl model) - k: number of nearest neighbours to use - Returns: - Returns a dictionary containing the bodypart and the corresponding indices - """ - smpl_vert_segmentation = json.load( - open(os.path.join(os.path.dirname(__file__), 'smpl_vert_segmentation.json'))) - n = smpl.vertices.shape[0] - y = np.array([None] * n) - for key, val in smpl_vert_segmentation.items(): - y[val] = key - - classifier = KNeighborsClassifier(n_neighbors=1) - classifier.fit(smpl.vertices, y) - - y_pred = classifier.predict(recon.vertices) - - recon_labels = {} - for key in smpl_vert_segmentation.keys(): - recon_labels[key] = list(np.argwhere( - y_pred == key).flatten().astype(int)) - - return recon_labels - - -def extract_cloth(recon, segmentation, K, R, t, smpl=None): - """ - Extract a portion of a mesh using 2d segmentation coordinates - Arguments: - recon: fully clothed mesh - seg_coord: segmentation coordinates in 2D (NDC) - K: intrinsic matrix of the projection - R: rotation matrix of the projection - t: translation vector of the projection - Returns: - Returns a submesh using the segmentation coordinates - """ - seg_coord = segmentation['coord_normalized'] - mesh = trimesh.Trimesh(recon.vertices, recon.faces) - extrinsic = np.zeros((3, 4)) - extrinsic[:3, :3] = R - extrinsic[:, 3] = t - P = K[:3, :3] @ extrinsic - - P_inv = np.linalg.pinv(P) - - # Each segmentation can contain multiple polygons - # We need to check them separately - points_so_far = [] - faces = recon.faces - for polygon in seg_coord: - n = len(polygon) - coords_h = np.hstack((polygon, np.ones((n, 1)))) - # Apply the inverse projection on homogeneus 2D coordinates to get the corresponding 3d Coordinates - XYZ = P_inv @ coords_h[:, :, None] - XYZ = XYZ.reshape((XYZ.shape[0], XYZ.shape[1])) - XYZ = XYZ[:, :3] / XYZ[:, 3, None] - - p = Path(XYZ[:, :2]) - - grid = p.contains_points(recon.vertices[:, :2]) - indeces = np.argwhere(grid == True) - points_so_far += list(indeces.flatten()) - - if smpl is not None: - num_verts = recon.vertices.shape[0] - recon_labels = smpl_to_recon_labels(recon, smpl) - body_parts_to_remove = ['rightHand', 'leftToeBase', 'leftFoot', 'rightFoot', 'head', - 'leftHandIndex1', 'rightHandIndex1', 'rightToeBase', 'leftHand', 'rightHand'] - type = segmentation['type_id'] - - # Remove additional bodyparts that are most likely not part of the segmentation but might intersect (e.g. hand in front of torso) - # https://github.com/switchablenorms/DeepFashion2 - # Short sleeve clothes - if type == 1 or type == 3 or type == 10: - body_parts_to_remove += ['leftForeArm', 'rightForeArm'] - # No sleeves at all or lower body clothes - elif type == 5 or type == 6 or type == 12 or type == 13 or type == 8 or type == 9: - body_parts_to_remove += ['leftForeArm', - 'rightForeArm', 'leftArm', 'rightArm'] - # Shorts - elif type == 7: - body_parts_to_remove += ['leftLeg', 'rightLeg', - 'leftForeArm', 'rightForeArm', 'leftArm', 'rightArm'] - - verts_to_remove = list(itertools.chain.from_iterable( - [recon_labels[part] for part in body_parts_to_remove])) - - label_mask = np.zeros(num_verts, dtype=bool) - label_mask[verts_to_remove] = True - - seg_mask = np.zeros(num_verts, dtype=bool) - seg_mask[points_so_far] = True - - # Remove points that belong to other bodyparts - # If a vertice in pointsSoFar is included in the bodyparts to remove, then these points should be removed - extra_verts_to_remove = np.array(list(seg_mask) and list(label_mask)) - - combine_mask = np.zeros(num_verts, dtype=bool) - combine_mask[points_so_far] = True - combine_mask[extra_verts_to_remove] = False - - all_indices = np.argwhere(combine_mask == True).flatten() - - i_x = np.where(np.in1d(faces[:, 0], all_indices))[0] - i_y = np.where(np.in1d(faces[:, 1], all_indices))[0] - i_z = np.where(np.in1d(faces[:, 2], all_indices))[0] - - faces_to_keep = np.array(list(set(i_x).union(i_y).union(i_z))) - mask = np.zeros(len(recon.faces), dtype=bool) - if len(faces_to_keep) > 0: - mask[faces_to_keep] = True - - mesh.update_faces(mask) - mesh.remove_unreferenced_vertices() - - # mesh.rezero() - - return mesh - - return None \ No newline at end of file diff --git a/lib /common /config.py b/lib /common /config.py deleted file mode 100644 index 2a69e8ca0ab7e42ae38055bb0b31e07ba42a4aed..0000000000000000000000000000000000000000 --- a/lib /common /config.py +++ /dev/null @@ -1,218 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from yacs.config import CfgNode as CN -import os - -_C = CN(new_allowed=True) - -# needed by trainer -_C.name = 'default' -_C.gpus = [0] -_C.test_gpus = [1] -_C.root = "./data/" -_C.ckpt_dir = './data/ckpt/' -_C.resume_path = '' -_C.normal_path = '' -_C.corr_path = '' -_C.results_path = './data/results/' -_C.projection_mode = 'orthogonal' -_C.num_views = 1 -_C.sdf = False -_C.sdf_clip = 5.0 - -_C.lr_G = 1e-3 -_C.lr_C = 1e-3 -_C.lr_N = 2e-4 -_C.weight_decay = 0.0 -_C.momentum = 0.0 -_C.optim = 'RMSprop' -_C.schedule = [5, 10, 15] -_C.gamma = 0.1 - -_C.overfit = False -_C.resume = False -_C.test_mode = False -_C.test_uv = False -_C.draw_geo_thres = 0.60 -_C.num_sanity_val_steps = 2 -_C.fast_dev = 0 -_C.get_fit = False -_C.agora = False -_C.optim_cloth = False -_C.optim_body = False -_C.mcube_res = 256 -_C.clean_mesh = True -_C.remesh = False - -_C.batch_size = 4 -_C.num_threads = 8 - -_C.num_epoch = 10 -_C.freq_plot = 0.01 -_C.freq_show_train = 0.1 -_C.freq_show_val = 0.2 -_C.freq_eval = 0.5 -_C.accu_grad_batch = 4 - -_C.test_items = ['sv', 'mv', 'mv-fusion', 'hybrid', 'dc-pred', 'gt'] - -_C.net = CN() -_C.net.gtype = 'HGPIFuNet' -_C.net.ctype = 'resnet18' -_C.net.classifierIMF = 'MultiSegClassifier' -_C.net.netIMF = 'resnet18' -_C.net.norm = 'group' -_C.net.norm_mlp = 'group' -_C.net.norm_color = 'group' -_C.net.hg_down = 'ave_pool' -_C.net.num_views = 1 - -# kernel_size, stride, dilation, padding - -_C.net.conv1 = [7, 2, 1, 3] -_C.net.conv3x3 = [3, 1, 1, 1] - -_C.net.num_stack = 4 -_C.net.num_hourglass = 2 -_C.net.hourglass_dim = 256 -_C.net.voxel_dim = 32 -_C.net.resnet_dim = 120 -_C.net.mlp_dim = [320, 1024, 512, 256, 128, 1] -_C.net.mlp_dim_knn = [320, 1024, 512, 256, 128, 3] -_C.net.mlp_dim_color = [513, 1024, 512, 256, 128, 3] -_C.net.mlp_dim_multiseg = [1088, 2048, 1024, 500] -_C.net.res_layers = [2, 3, 4] -_C.net.filter_dim = 256 -_C.net.smpl_dim = 3 - -_C.net.cly_dim = 3 -_C.net.soft_dim = 64 -_C.net.z_size = 200.0 -_C.net.N_freqs = 10 -_C.net.geo_w = 0.1 -_C.net.norm_w = 0.1 -_C.net.dc_w = 0.1 -_C.net.C_cat_to_G = False - -_C.net.skip_hourglass = True -_C.net.use_tanh = True -_C.net.soft_onehot = True -_C.net.no_residual = True -_C.net.use_attention = False - -_C.net.prior_type = "sdf" -_C.net.smpl_feats = ['sdf', 'cmap', 'norm', 'vis'] -_C.net.use_filter = True -_C.net.use_cc = False -_C.net.use_PE = False -_C.net.use_IGR = False -_C.net.in_geo = () -_C.net.in_nml = () - -_C.dataset = CN() -_C.dataset.root = '' -_C.dataset.set_splits = [0.95, 0.04] -_C.dataset.types = [ - "3dpeople", "axyz", "renderpeople", "renderpeople_p27", "humanalloy" -] -_C.dataset.scales = [1.0, 100.0, 1.0, 1.0, 100.0 / 39.37] -_C.dataset.rp_type = "pifu900" -_C.dataset.th_type = 'train' -_C.dataset.input_size = 512 -_C.dataset.rotation_num = 3 -_C.dataset.num_precomp = 10 # Number of segmentation classifiers -_C.dataset.num_multiseg = 500 # Number of categories per classifier -_C.dataset.num_knn = 10 # for loss/error -_C.dataset.num_knn_dis = 20 # for accuracy -_C.dataset.num_verts_max = 20000 -_C.dataset.zray_type = False -_C.dataset.online_smpl = False -_C.dataset.noise_type = ['z-trans', 'pose', 'beta'] -_C.dataset.noise_scale = [0.0, 0.0, 0.0] -_C.dataset.num_sample_geo = 10000 -_C.dataset.num_sample_color = 0 -_C.dataset.num_sample_seg = 0 -_C.dataset.num_sample_knn = 10000 - -_C.dataset.sigma_geo = 5.0 -_C.dataset.sigma_color = 0.10 -_C.dataset.sigma_seg = 0.10 -_C.dataset.thickness_threshold = 20.0 -_C.dataset.ray_sample_num = 2 -_C.dataset.semantic_p = False -_C.dataset.remove_outlier = False - -_C.dataset.train_bsize = 1.0 -_C.dataset.val_bsize = 1.0 -_C.dataset.test_bsize = 1.0 - - -def get_cfg_defaults(): - """Get a yacs CfgNode 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 _C.clone() - - -# Alternatively, provide a way to import the defaults as -# a global singleton: -cfg = _C # users can `from config import cfg` - -# cfg = get_cfg_defaults() -# cfg.merge_from_file('./configs/example.yaml') - -# # Now override from a list (opts could come from the command line) -# opts = ['dataset.root', './data/XXXX', 'learning_rate', '1e-2'] -# cfg.merge_from_list(opts) - - -def update_cfg(cfg_file): - # cfg = get_cfg_defaults() - _C.merge_from_file(cfg_file) - # return cfg.clone() - return _C - - -def parse_args(args): - cfg_file = args.cfg_file - if args.cfg_file is not None: - cfg = update_cfg(args.cfg_file) - else: - cfg = get_cfg_defaults() - - # if args.misc is not None: - # cfg.merge_from_list(args.misc) - - return cfg - - -def parse_args_extend(args): - if args.resume: - if not os.path.exists(args.log_dir): - raise ValueError( - 'Experiment are set to resume mode, but log directory does not exist.' - ) - - # load log's cfg - cfg_file = os.path.join(args.log_dir, 'cfg.yaml') - cfg = update_cfg(cfg_file) - - if args.misc is not None: - cfg.merge_from_list(args.misc) - else: - parse_args(args) diff --git a/lib /common /render.py b/lib /common /render.py deleted file mode 100644 index 4d8719232bbfd851102cfdd604eb5ea5ffba07a9..0000000000000000000000000000000000000000 --- a/lib /common /render.py +++ /dev/null @@ -1,389 +0,0 @@ -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from pytorch3d.renderer import ( - BlendParams, - blending, - look_at_view_transform, - FoVOrthographicCameras, - PointLights, - RasterizationSettings, - PointsRasterizationSettings, - PointsRenderer, - AlphaCompositor, - PointsRasterizer, - MeshRenderer, - MeshRasterizer, - SoftPhongShader, - SoftSilhouetteShader, - TexturesVertex, -) -from pytorch3d.renderer.mesh import TexturesVertex -from pytorch3d.structures import Meshes - -import os, subprocess - -from lib.dataset.mesh_util import SMPLX, get_visibility -import lib.common.render_utils as util -import torch -import numpy as np -from PIL import Image -from tqdm import tqdm -import cv2 -import math -from termcolor import colored - - -def image2vid(images, vid_path): - - w, h = images[0].size - videodims = (w, h) - fourcc = cv2.VideoWriter_fourcc(*'XVID') - video = cv2.VideoWriter(vid_path, fourcc, 30, videodims) - for image in images: - video.write(cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)) - video.release() - - -def query_color(verts, faces, image, device): - """query colors from points and image - Args: - verts ([B, 3]): [query verts] - faces ([M, 3]): [query faces] - image ([B, 3, H, W]): [full image] - Returns: - [np.float]: [return colors] - """ - - verts = verts.float().to(device) - faces = faces.long().to(device) - - (xy, z) = verts.split([2, 1], dim=1) - visibility = get_visibility(xy, z, faces[:, [0, 2, 1]]).flatten() - uv = xy.unsqueeze(0).unsqueeze(2) # [B, N, 2] - uv = uv * torch.tensor([1.0, -1.0]).type_as(uv) - colors = (torch.nn.functional.grid_sample(image, uv, align_corners=True)[ - 0, :, :, 0].permute(1, 0) + 1.0) * 0.5 * 255.0 - colors[visibility == 0.0] = ((Meshes(verts.unsqueeze(0), faces.unsqueeze( - 0)).verts_normals_padded().squeeze(0) + 1.0) * 0.5 * 255.0)[visibility == 0.0] - - return colors.detach().cpu() - - -class cleanShader(torch.nn.Module): - def __init__(self, device="cpu", cameras=None, blend_params=None): - super().__init__() - self.cameras = cameras - self.blend_params = blend_params if blend_params is not None else BlendParams() - - def forward(self, fragments, meshes, **kwargs): - cameras = kwargs.get("cameras", self.cameras) - if cameras is None: - msg = "Cameras must be specified either at initialization \ - or in the forward pass of TexturedSoftPhongShader" - - raise ValueError(msg) - - # get renderer output - blend_params = kwargs.get("blend_params", self.blend_params) - texels = meshes.sample_textures(fragments) - images = blending.softmax_rgb_blend( - texels, fragments, blend_params, znear=-256, zfar=256 - ) - - return images - - -class Render: - def __init__(self, size=512, device=torch.device("cuda:0")): - self.device = device - self.mesh_y_center = 100.0 - self.dis = 100.0 - self.scale = 1.0 - self.size = size - self.cam_pos = [(0, 100, 100)] - - self.mesh = None - self.deform_mesh = None - self.pcd = None - self.renderer = None - self.meshRas = None - self.type = None - self.knn = None - self.knn_inverse = None - - self.smpl_seg = None - self.smpl_cmap = None - - self.smplx = SMPLX() - - self.uv_rasterizer = util.Pytorch3dRasterizer(self.size) - - def get_camera(self, cam_id): - - R, T = look_at_view_transform( - eye=[self.cam_pos[cam_id]], - at=((0, self.mesh_y_center, 0),), - up=((0, 1, 0),), - ) - - camera = FoVOrthographicCameras( - device=self.device, - R=R, - T=T, - znear=100.0, - zfar=-100.0, - max_y=100.0, - min_y=-100.0, - max_x=100.0, - min_x=-100.0, - scale_xyz=(self.scale * np.ones(3),), - ) - - return camera - - def init_renderer(self, camera, type="clean_mesh", bg="gray"): - - if "mesh" in type: - - # rasterizer - self.raster_settings_mesh = RasterizationSettings( - image_size=self.size, - blur_radius=np.log(1.0 / 1e-4) * 1e-7, - faces_per_pixel=30, - ) - self.meshRas = MeshRasterizer( - cameras=camera, raster_settings=self.raster_settings_mesh - ) - - if bg == "black": - blendparam = BlendParams(1e-4, 1e-4, (0.0, 0.0, 0.0)) - elif bg == "white": - blendparam = BlendParams(1e-4, 1e-8, (1.0, 1.0, 1.0)) - elif bg == "gray": - blendparam = BlendParams(1e-4, 1e-8, (0.5, 0.5, 0.5)) - - if type == "ori_mesh": - - lights = PointLights( - device=self.device, - ambient_color=((0.8, 0.8, 0.8),), - diffuse_color=((0.2, 0.2, 0.2),), - specular_color=((0.0, 0.0, 0.0),), - location=[[0.0, 200.0, 0.0]], - ) - - self.renderer = MeshRenderer( - rasterizer=self.meshRas, - shader=SoftPhongShader( - device=self.device, - cameras=camera, - lights=lights, - blend_params=blendparam, - ), - ) - - if type == "silhouette": - self.raster_settings_silhouette = RasterizationSettings( - image_size=self.size, - blur_radius=np.log(1.0 / 1e-4 - 1.0) * 5e-5, - faces_per_pixel=50, - cull_backfaces=True, - ) - - self.silhouetteRas = MeshRasterizer( - cameras=camera, raster_settings=self.raster_settings_silhouette - ) - self.renderer = MeshRenderer( - rasterizer=self.silhouetteRas, shader=SoftSilhouetteShader() - ) - - if type == "pointcloud": - self.raster_settings_pcd = PointsRasterizationSettings( - image_size=self.size, radius=0.006, points_per_pixel=10 - ) - - self.pcdRas = PointsRasterizer( - cameras=camera, raster_settings=self.raster_settings_pcd - ) - self.renderer = PointsRenderer( - rasterizer=self.pcdRas, - compositor=AlphaCompositor(background_color=(0, 0, 0)), - ) - - if type == "clean_mesh": - - self.renderer = MeshRenderer( - rasterizer=self.meshRas, - shader=cleanShader( - device=self.device, cameras=camera, blend_params=blendparam - ), - ) - - def VF2Mesh(self, verts, faces): - - if not torch.is_tensor(verts): - verts = torch.tensor(verts) - if not torch.is_tensor(faces): - faces = torch.tensor(faces) - - if verts.ndimension() == 2: - verts = verts.unsqueeze(0).float() - if faces.ndimension() == 2: - faces = faces.unsqueeze(0).long() - - verts = verts.to(self.device) - faces = faces.to(self.device) - - mesh = Meshes(verts, faces).to(self.device) - - mesh.textures = TexturesVertex( - verts_features=(mesh.verts_normals_padded() + 1.0) * 0.5 - ) - - return mesh - - def load_meshes(self, verts, faces): - """load mesh into the pytorch3d renderer - Args: - verts ([N,3]): verts - faces ([N,3]): faces - offset ([N,3]): offset - """ - - # camera setting - self.scale = 100.0 - self.mesh_y_center = 0.0 - - self.cam_pos = [ - (0, self.mesh_y_center, 100.0), - (100.0, self.mesh_y_center, 0), - (0, self.mesh_y_center, -100.0), - (-100.0, self.mesh_y_center, 0), - ] - - self.type = "color" - - if isinstance(verts, list): - self.meshes = [] - for V, F in zip(verts, faces): - self.meshes.append(self.VF2Mesh(V, F)) - else: - self.meshes = [self.VF2Mesh(verts, faces)] - - def get_depth_map(self, cam_ids=[0, 2]): - - depth_maps = [] - for cam_id in cam_ids: - self.init_renderer(self.get_camera(cam_id), "clean_mesh", "gray") - fragments = self.meshRas(self.meshes[0]) - depth_map = fragments.zbuf[..., 0].squeeze(0) - if cam_id == 2: - depth_map = torch.fliplr(depth_map) - depth_maps.append(depth_map) - - return depth_maps - - def get_rgb_image(self, cam_ids=[0, 2]): - - images = [] - for cam_id in range(len(self.cam_pos)): - if cam_id in cam_ids: - self.init_renderer(self.get_camera( - cam_id), "clean_mesh", "gray") - if len(cam_ids) == 4: - rendered_img = ( - self.renderer(self.meshes[0])[ - 0:1, :, :, :3].permute(0, 3, 1, 2) - - 0.5 - ) * 2.0 - else: - rendered_img = ( - self.renderer(self.meshes[0])[ - 0:1, :, :, :3].permute(0, 3, 1, 2) - - 0.5 - ) * 2.0 - if cam_id == 2 and len(cam_ids) == 2: - rendered_img = torch.flip(rendered_img, dims=[3]) - images.append(rendered_img) - - return images - - def get_rendered_video(self, images, save_path): - - tmp_path = save_path.replace('cloth', 'tmp') - - self.cam_pos = [] - for angle in range(0, 360, 3): - self.cam_pos.append( - ( - 100.0 * math.cos(np.pi / 180 * angle), - self.mesh_y_center, - 100.0 * math.sin(np.pi / 180 * angle), - ) - ) - - old_shape = np.array(images[0].shape[:2]) - new_shape = np.around( - (self.size / old_shape[0]) * old_shape).astype(np.int) - - fourcc = cv2.VideoWriter_fourcc(*"mp4v") - video = cv2.VideoWriter( - tmp_path, fourcc, 30, (self.size * len(self.meshes) + - new_shape[1] * len(images), self.size) - ) - - pbar = tqdm(range(len(self.cam_pos))) - pbar.set_description(colored(f"exporting video {os.path.basename(save_path)}...", "blue")) - for cam_id in pbar: - self.init_renderer(self.get_camera(cam_id), "clean_mesh", "gray") - - img_lst = [ - np.array(Image.fromarray(img).resize(new_shape[::-1])).astype(np.uint8)[ - :, :, [2, 1, 0] - ] - for img in images - ] - - for mesh in self.meshes: - rendered_img = ( - (self.renderer(mesh)[0, :, :, :3] * 255.0) - .detach() - .cpu() - .numpy() - .astype(np.uint8) - ) - - img_lst.append(rendered_img) - final_img = np.concatenate(img_lst, axis=1) - video.write(final_img) - - video.release() - - os.system(f'ffmpeg -y -loglevel quiet -stats -i {tmp_path} -c:v libx264 {save_path}') - - def get_silhouette_image(self, cam_ids=[0, 2]): - - images = [] - for cam_id in range(len(self.cam_pos)): - if cam_id in cam_ids: - self.init_renderer(self.get_camera(cam_id), "silhouette") - rendered_img = self.renderer(self.meshes[0])[0:1, :, :, 3] - if cam_id == 2 and len(cam_ids) == 2: - rendered_img = torch.flip(rendered_img, dims=[2]) - images.append(rendered_img) - - return images diff --git a/lib /common /render_utils.py b/lib /common /render_utils.py deleted file mode 100644 index 00437ac2c3b243fa01d3916cd130e1437ae5b4d8..0000000000000000000000000000000000000000 --- a/lib /common /render_utils.py +++ /dev/null @@ -1,221 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import torch -from torch import nn -import trimesh -import math -from typing import NewType -from pytorch3d.structures import Meshes -from pytorch3d.renderer.mesh import rasterize_meshes - -Tensor = NewType('Tensor', torch.Tensor) - - -def solid_angles(points: Tensor, - triangles: Tensor, - thresh: float = 1e-8) -> Tensor: - ''' Compute solid angle between the input points and triangles - Follows the method described in: - The Solid Angle of a Plane Triangle - A. VAN OOSTEROM AND J. STRACKEE - IEEE TRANSACTIONS ON BIOMEDICAL ENGINEERING, - VOL. BME-30, NO. 2, FEBRUARY 1983 - Parameters - ----------- - points: BxQx3 - Tensor of input query points - triangles: BxFx3x3 - Target triangles - thresh: float - float threshold - Returns - ------- - solid_angles: BxQxF - A tensor containing the solid angle between all query points - and input triangles - ''' - # Center the triangles on the query points. Size should be BxQxFx3x3 - centered_tris = triangles[:, None] - points[:, :, None, None] - - # BxQxFx3 - norms = torch.norm(centered_tris, dim=-1) - - # Should be BxQxFx3 - cross_prod = torch.cross(centered_tris[:, :, :, 1], - centered_tris[:, :, :, 2], - dim=-1) - # Should be BxQxF - numerator = (centered_tris[:, :, :, 0] * cross_prod).sum(dim=-1) - del cross_prod - - dot01 = (centered_tris[:, :, :, 0] * centered_tris[:, :, :, 1]).sum(dim=-1) - dot12 = (centered_tris[:, :, :, 1] * centered_tris[:, :, :, 2]).sum(dim=-1) - dot02 = (centered_tris[:, :, :, 0] * centered_tris[:, :, :, 2]).sum(dim=-1) - del centered_tris - - denominator = (norms.prod(dim=-1) + dot01 * norms[:, :, :, 2] + - dot02 * norms[:, :, :, 1] + dot12 * norms[:, :, :, 0]) - del dot01, dot12, dot02, norms - - # Should be BxQ - solid_angle = torch.atan2(numerator, denominator) - del numerator, denominator - - torch.cuda.empty_cache() - - return 2 * solid_angle - - -def winding_numbers(points: Tensor, - triangles: Tensor, - thresh: float = 1e-8) -> Tensor: - ''' Uses winding_numbers to compute inside/outside - Robust inside-outside segmentation using generalized winding numbers - Alec Jacobson, - Ladislav Kavan, - Olga Sorkine-Hornung - Fast Winding Numbers for Soups and Clouds SIGGRAPH 2018 - Gavin Barill - NEIL G. Dickson - Ryan Schmidt - David I.W. Levin - and Alec Jacobson - Parameters - ----------- - points: BxQx3 - Tensor of input query points - triangles: BxFx3x3 - Target triangles - thresh: float - float threshold - Returns - ------- - winding_numbers: BxQ - A tensor containing the Generalized winding numbers - ''' - # The generalized winding number is the sum of solid angles of the point - # with respect to all triangles. - return 1 / (4 * math.pi) * solid_angles(points, triangles, - thresh=thresh).sum(dim=-1) - - -def batch_contains(verts, faces, points): - - B = verts.shape[0] - N = points.shape[1] - - verts = verts.detach().cpu() - faces = faces.detach().cpu() - points = points.detach().cpu() - contains = torch.zeros(B, N) - - for i in range(B): - contains[i] = torch.as_tensor( - trimesh.Trimesh(verts[i], faces[i]).contains(points[i])) - - return 2.0 * (contains - 0.5) - - -def dict2obj(d): - # if isinstance(d, list): - # d = [dict2obj(x) for x in d] - if not isinstance(d, dict): - return d - - class C(object): - pass - - o = C() - for k in d: - o.__dict__[k] = dict2obj(d[k]) - return o - - -def face_vertices(vertices, faces): - """ - :param vertices: [batch size, number of vertices, 3] - :param faces: [batch size, number of faces, 3] - :return: [batch size, number of faces, 3, 3] - """ - - bs, nv = vertices.shape[:2] - bs, nf = faces.shape[:2] - device = vertices.device - faces = faces + (torch.arange(bs, dtype=torch.int32).to(device) * - nv)[:, None, None] - vertices = vertices.reshape((bs * nv, vertices.shape[-1])) - - return vertices[faces.long()] - - -class Pytorch3dRasterizer(nn.Module): - """ Borrowed from https://github.com/facebookresearch/pytorch3d - Notice: - x,y,z are in image space, normalized - can only render squared image now - """ - - def __init__(self, image_size=224): - """ - use fixed raster_settings for rendering faces - """ - super().__init__() - raster_settings = { - 'image_size': image_size, - 'blur_radius': 0.0, - 'faces_per_pixel': 1, - 'bin_size': None, - 'max_faces_per_bin': None, - 'perspective_correct': True, - 'cull_backfaces': True, - } - raster_settings = dict2obj(raster_settings) - self.raster_settings = raster_settings - - def forward(self, vertices, faces, attributes=None): - fixed_vertices = vertices.clone() - fixed_vertices[..., :2] = -fixed_vertices[..., :2] - meshes_screen = Meshes(verts=fixed_vertices.float(), - faces=faces.long()) - raster_settings = self.raster_settings - pix_to_face, zbuf, bary_coords, dists = rasterize_meshes( - meshes_screen, - image_size=raster_settings.image_size, - blur_radius=raster_settings.blur_radius, - faces_per_pixel=raster_settings.faces_per_pixel, - bin_size=raster_settings.bin_size, - max_faces_per_bin=raster_settings.max_faces_per_bin, - perspective_correct=raster_settings.perspective_correct, - ) - vismask = (pix_to_face > -1).float() - D = attributes.shape[-1] - attributes = attributes.clone() - attributes = attributes.view(attributes.shape[0] * attributes.shape[1], - 3, attributes.shape[-1]) - N, H, W, K, _ = bary_coords.shape - mask = pix_to_face == -1 - pix_to_face = pix_to_face.clone() - pix_to_face[mask] = 0 - idx = pix_to_face.view(N * H * W * K, 1, 1).expand(N * H * W * K, 3, D) - pixel_face_vals = attributes.gather(0, idx).view(N, H, W, K, 3, D) - pixel_vals = (bary_coords[..., None] * pixel_face_vals).sum(dim=-2) - pixel_vals[mask] = 0 # Replace masked values in output. - pixel_vals = pixel_vals[:, :, :, 0].permute(0, 3, 1, 2) - pixel_vals = torch.cat( - [pixel_vals, vismask[:, :, :, 0][:, None, :, :]], dim=1) - return pixel_vals \ No newline at end of file diff --git a/lib /common /seg3d_lossless.py b/lib /common /seg3d_lossless.py deleted file mode 100644 index eb8dff2e1ba459f53dc5e27a4ee76c6168486203..0000000000000000000000000000000000000000 --- a/lib /common /seg3d_lossless.py +++ /dev/null @@ -1,604 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - - -from .seg3d_utils import ( - create_grid3D, - plot_mask3D, - SmoothConv3D, -) - -import torch -import torch.nn as nn -import numpy as np -import torch.nn.functional as F -import mcubes -from kaolin.ops.conversions import voxelgrids_to_trianglemeshes -import logging - -logging.getLogger("lightning").setLevel(logging.ERROR) - - -class Seg3dLossless(nn.Module): - def __init__(self, - query_func, - b_min, - b_max, - resolutions, - channels=1, - balance_value=0.5, - align_corners=False, - visualize=False, - debug=False, - use_cuda_impl=False, - faster=False, - use_shadow=False, - **kwargs): - """ - align_corners: same with how you process gt. (grid_sample / interpolate) - """ - super().__init__() - self.query_func = query_func - self.register_buffer( - 'b_min', - torch.tensor(b_min).float().unsqueeze(1)) # [bz, 1, 3] - self.register_buffer( - 'b_max', - torch.tensor(b_max).float().unsqueeze(1)) # [bz, 1, 3] - - # ti.init(arch=ti.cuda) - # self.mciso_taichi = MCISO(dim=3, N=resolutions[-1]-1) - - if type(resolutions[0]) is int: - resolutions = torch.tensor([(res, res, res) - for res in resolutions]) - else: - resolutions = torch.tensor(resolutions) - self.register_buffer('resolutions', resolutions) - self.batchsize = self.b_min.size(0) - assert self.batchsize == 1 - self.balance_value = balance_value - self.channels = channels - assert self.channels == 1 - self.align_corners = align_corners - self.visualize = visualize - self.debug = debug - self.use_cuda_impl = use_cuda_impl - self.faster = faster - self.use_shadow = use_shadow - - for resolution in resolutions: - assert resolution[0] % 2 == 1 and resolution[1] % 2 == 1, \ - f"resolution {resolution} need to be odd becuase of align_corner." - - # init first resolution - init_coords = create_grid3D(0, - resolutions[-1] - 1, - steps=resolutions[0]) # [N, 3] - init_coords = init_coords.unsqueeze(0).repeat(self.batchsize, 1, - 1) # [bz, N, 3] - self.register_buffer('init_coords', init_coords) - - # some useful tensors - calculated = torch.zeros( - (self.resolutions[-1][2], self.resolutions[-1][1], - self.resolutions[-1][0]), - dtype=torch.bool) - self.register_buffer('calculated', calculated) - - gird8_offsets = torch.stack( - torch.meshgrid([ - torch.tensor([-1, 0, 1]), - torch.tensor([-1, 0, 1]), - torch.tensor([-1, 0, 1]) - ])).int().view(3, -1).t() # [27, 3] - self.register_buffer('gird8_offsets', gird8_offsets) - - # smooth convs - self.smooth_conv3x3 = SmoothConv3D(in_channels=1, - out_channels=1, - kernel_size=3) - self.smooth_conv5x5 = SmoothConv3D(in_channels=1, - out_channels=1, - kernel_size=5) - self.smooth_conv7x7 = SmoothConv3D(in_channels=1, - out_channels=1, - kernel_size=7) - self.smooth_conv9x9 = SmoothConv3D(in_channels=1, - out_channels=1, - kernel_size=9) - - def batch_eval(self, coords, **kwargs): - """ - coords: in the coordinates of last resolution - **kwargs: for query_func - """ - coords = coords.detach() - # normalize coords to fit in [b_min, b_max] - if self.align_corners: - coords2D = coords.float() / (self.resolutions[-1] - 1) - else: - step = 1.0 / self.resolutions[-1].float() - coords2D = coords.float() / self.resolutions[-1] + step / 2 - coords2D = coords2D * (self.b_max - self.b_min) + self.b_min - # query function - occupancys = self.query_func(**kwargs, points=coords2D) - if type(occupancys) is list: - occupancys = torch.stack(occupancys) # [bz, C, N] - assert len(occupancys.size()) == 3, \ - "query_func should return a occupancy with shape of [bz, C, N]" - return occupancys - - def forward(self, **kwargs): - if self.faster: - return self._forward_faster(**kwargs) - else: - return self._forward(**kwargs) - - def _forward_faster(self, **kwargs): - """ - In faster mode, we make following changes to exchange accuracy for speed: - 1. no conflict checking: 4.88 fps -> 6.56 fps - 2. smooth_conv9x9 ~ smooth_conv3x3 for different resolution - 3. last step no examine - """ - final_W = self.resolutions[-1][0] - final_H = self.resolutions[-1][1] - final_D = self.resolutions[-1][2] - - for resolution in self.resolutions: - W, H, D = resolution - stride = (self.resolutions[-1] - 1) / (resolution - 1) - - # first step - if torch.equal(resolution, self.resolutions[0]): - coords = self.init_coords.clone() # torch.long - occupancys = self.batch_eval(coords, **kwargs) - occupancys = occupancys.view(self.batchsize, self.channels, D, - H, W) - if (occupancys > 0.5).sum() == 0: - # return F.interpolate( - # occupancys, size=(final_D, final_H, final_W), - # mode="linear", align_corners=True) - return None - - if self.visualize: - self.plot(occupancys, coords, final_D, final_H, final_W) - - with torch.no_grad(): - coords_accum = coords / stride - - # last step - elif torch.equal(resolution, self.resolutions[-1]): - - with torch.no_grad(): - # here true is correct! - valid = F.interpolate( - (occupancys > self.balance_value).float(), - size=(D, H, W), - mode="trilinear", - align_corners=True) - - # here true is correct! - occupancys = F.interpolate(occupancys.float(), - size=(D, H, W), - mode="trilinear", - align_corners=True) - - # is_boundary = (valid > 0.0) & (valid < 1.0) - is_boundary = valid == 0.5 - - # next steps - else: - coords_accum *= 2 - - with torch.no_grad(): - # here true is correct! - valid = F.interpolate( - (occupancys > self.balance_value).float(), - size=(D, H, W), - mode="trilinear", - align_corners=True) - - # here true is correct! - occupancys = F.interpolate(occupancys.float(), - size=(D, H, W), - mode="trilinear", - align_corners=True) - - is_boundary = (valid > 0.0) & (valid < 1.0) - - with torch.no_grad(): - if torch.equal(resolution, self.resolutions[1]): - is_boundary = (self.smooth_conv9x9(is_boundary.float()) - > 0)[0, 0] - elif torch.equal(resolution, self.resolutions[2]): - is_boundary = (self.smooth_conv7x7(is_boundary.float()) - > 0)[0, 0] - else: - is_boundary = (self.smooth_conv3x3(is_boundary.float()) - > 0)[0, 0] - - coords_accum = coords_accum.long() - is_boundary[coords_accum[0, :, 2], coords_accum[0, :, 1], - coords_accum[0, :, 0]] = False - point_coords = is_boundary.permute( - 2, 1, 0).nonzero(as_tuple=False).unsqueeze(0) - point_indices = (point_coords[:, :, 2] * H * W + - point_coords[:, :, 1] * W + - point_coords[:, :, 0]) - - R, C, D, H, W = occupancys.shape - - # inferred value - coords = point_coords * stride - - if coords.size(1) == 0: - continue - occupancys_topk = self.batch_eval(coords, **kwargs) - - # put mask point predictions to the right places on the upsampled grid. - R, C, D, H, W = occupancys.shape - point_indices = point_indices.unsqueeze(1).expand(-1, C, -1) - occupancys = (occupancys.reshape(R, C, D * H * W).scatter_( - 2, point_indices, occupancys_topk).view(R, C, D, H, W)) - - with torch.no_grad(): - voxels = coords / stride - coords_accum = torch.cat([voxels, coords_accum], - dim=1).unique(dim=1) - - return occupancys[0, 0] - - def _forward(self, **kwargs): - """ - output occupancy field would be: - (bz, C, res, res) - """ - final_W = self.resolutions[-1][0] - final_H = self.resolutions[-1][1] - final_D = self.resolutions[-1][2] - - calculated = self.calculated.clone() - - for resolution in self.resolutions: - W, H, D = resolution - stride = (self.resolutions[-1] - 1) / (resolution - 1) - - if self.visualize: - this_stage_coords = [] - - # first step - if torch.equal(resolution, self.resolutions[0]): - coords = self.init_coords.clone() # torch.long - occupancys = self.batch_eval(coords, **kwargs) - occupancys = occupancys.view(self.batchsize, self.channels, D, - H, W) - - if self.visualize: - self.plot(occupancys, coords, final_D, final_H, final_W) - - with torch.no_grad(): - coords_accum = coords / stride - calculated[coords[0, :, 2], coords[0, :, 1], - coords[0, :, 0]] = True - - # next steps - else: - coords_accum *= 2 - - with torch.no_grad(): - # here true is correct! - valid = F.interpolate( - (occupancys > self.balance_value).float(), - size=(D, H, W), - mode="trilinear", - align_corners=True) - - # here true is correct! - occupancys = F.interpolate(occupancys.float(), - size=(D, H, W), - mode="trilinear", - align_corners=True) - - is_boundary = (valid > 0.0) & (valid < 1.0) - - with torch.no_grad(): - # TODO - if self.use_shadow and torch.equal(resolution, - self.resolutions[-1]): - # larger z means smaller depth here - depth_res = resolution[2].item() - depth_index = torch.linspace(0, - depth_res - 1, - steps=depth_res).type_as( - occupancys.device) - depth_index_max = torch.max( - (occupancys > self.balance_value) * - (depth_index + 1), - dim=-1, - keepdim=True)[0] - 1 - shadow = depth_index < depth_index_max - is_boundary[shadow] = False - is_boundary = is_boundary[0, 0] - else: - is_boundary = (self.smooth_conv3x3(is_boundary.float()) - > 0)[0, 0] - # is_boundary = is_boundary[0, 0] - - is_boundary[coords_accum[0, :, 2], coords_accum[0, :, 1], - coords_accum[0, :, 0]] = False - point_coords = is_boundary.permute( - 2, 1, 0).nonzero(as_tuple=False).unsqueeze(0) - point_indices = (point_coords[:, :, 2] * H * W + - point_coords[:, :, 1] * W + - point_coords[:, :, 0]) - - R, C, D, H, W = occupancys.shape - # interpolated value - occupancys_interp = torch.gather( - occupancys.reshape(R, C, D * H * W), 2, - point_indices.unsqueeze(1)) - - # inferred value - coords = point_coords * stride - - if coords.size(1) == 0: - continue - occupancys_topk = self.batch_eval(coords, **kwargs) - if self.visualize: - this_stage_coords.append(coords) - - # put mask point predictions to the right places on the upsampled grid. - R, C, D, H, W = occupancys.shape - point_indices = point_indices.unsqueeze(1).expand(-1, C, -1) - occupancys = (occupancys.reshape(R, C, D * H * W).scatter_( - 2, point_indices, occupancys_topk).view(R, C, D, H, W)) - - with torch.no_grad(): - # conflicts - conflicts = ((occupancys_interp - self.balance_value) * - (occupancys_topk - self.balance_value) < 0)[0, - 0] - - if self.visualize: - self.plot(occupancys, coords, final_D, final_H, - final_W) - - voxels = coords / stride - coords_accum = torch.cat([voxels, coords_accum], - dim=1).unique(dim=1) - calculated[coords[0, :, 2], coords[0, :, 1], - coords[0, :, 0]] = True - - while conflicts.sum() > 0: - if self.use_shadow and torch.equal(resolution, - self.resolutions[-1]): - break - - with torch.no_grad(): - conflicts_coords = coords[0, conflicts, :] - - if self.debug: - self.plot(occupancys, - conflicts_coords.unsqueeze(0), - final_D, - final_H, - final_W, - title='conflicts') - - conflicts_boundary = (conflicts_coords.int() + - self.gird8_offsets.unsqueeze(1) * - stride.int()).reshape( - -1, 3).long().unique(dim=0) - conflicts_boundary[:, 0] = ( - conflicts_boundary[:, 0].clamp( - 0, - calculated.size(2) - 1)) - conflicts_boundary[:, 1] = ( - conflicts_boundary[:, 1].clamp( - 0, - calculated.size(1) - 1)) - conflicts_boundary[:, 2] = ( - conflicts_boundary[:, 2].clamp( - 0, - calculated.size(0) - 1)) - - coords = conflicts_boundary[calculated[ - conflicts_boundary[:, 2], conflicts_boundary[:, 1], - conflicts_boundary[:, 0]] == False] - - if self.debug: - self.plot(occupancys, - coords.unsqueeze(0), - final_D, - final_H, - final_W, - title='coords') - - coords = coords.unsqueeze(0) - point_coords = coords / stride - point_indices = (point_coords[:, :, 2] * H * W + - point_coords[:, :, 1] * W + - point_coords[:, :, 0]) - - R, C, D, H, W = occupancys.shape - # interpolated value - occupancys_interp = torch.gather( - occupancys.reshape(R, C, D * H * W), 2, - point_indices.unsqueeze(1)) - - # inferred value - coords = point_coords * stride - - if coords.size(1) == 0: - break - occupancys_topk = self.batch_eval(coords, **kwargs) - if self.visualize: - this_stage_coords.append(coords) - - with torch.no_grad(): - # conflicts - conflicts = ((occupancys_interp - self.balance_value) * - (occupancys_topk - self.balance_value) < - 0)[0, 0] - - # put mask point predictions to the right places on the upsampled grid. - point_indices = point_indices.unsqueeze(1).expand( - -1, C, -1) - occupancys = (occupancys.reshape(R, C, D * H * W).scatter_( - 2, point_indices, occupancys_topk).view(R, C, D, H, W)) - - with torch.no_grad(): - voxels = coords / stride - coords_accum = torch.cat([voxels, coords_accum], - dim=1).unique(dim=1) - calculated[coords[0, :, 2], coords[0, :, 1], - coords[0, :, 0]] = True - - if self.visualize: - this_stage_coords = torch.cat(this_stage_coords, dim=1) - self.plot(occupancys, this_stage_coords, final_D, final_H, - final_W) - - return occupancys[0, 0] - - def plot(self, - occupancys, - coords, - final_D, - final_H, - final_W, - title='', - **kwargs): - final = F.interpolate(occupancys.float(), - size=(final_D, final_H, final_W), - mode="trilinear", - align_corners=True) # here true is correct! - x = coords[0, :, 0].to("cpu") - y = coords[0, :, 1].to("cpu") - z = coords[0, :, 2].to("cpu") - - plot_mask3D(final[0, 0].to("cpu"), title, (x, y, z), **kwargs) - - def find_vertices(self, sdf, direction="front"): - ''' - - direction: "front" | "back" | "left" | "right" - ''' - resolution = sdf.size(2) - if direction == "front": - pass - elif direction == "left": - sdf = sdf.permute(2, 1, 0) - elif direction == "back": - inv_idx = torch.arange(sdf.size(2) - 1, -1, -1).long() - sdf = sdf[inv_idx, :, :] - elif direction == "right": - inv_idx = torch.arange(sdf.size(2) - 1, -1, -1).long() - sdf = sdf[:, :, inv_idx] - sdf = sdf.permute(2, 1, 0) - - inv_idx = torch.arange(sdf.size(2) - 1, -1, -1).long() - sdf = sdf[inv_idx, :, :] - sdf_all = sdf.permute(2, 1, 0) - - # shadow - grad_v = (sdf_all > 0.5) * torch.linspace( - resolution, 1, steps=resolution).to(sdf.device) - grad_c = torch.ones_like(sdf_all) * torch.linspace( - 0, resolution - 1, steps=resolution).to(sdf.device) - max_v, max_c = grad_v.max(dim=2) - shadow = grad_c > max_c.view(resolution, resolution, 1) - keep = (sdf_all > 0.5) & (~shadow) - - p1 = keep.nonzero(as_tuple=False).t() # [3, N] - p2 = p1.clone() # z - p2[2, :] = (p2[2, :] - 2).clamp(0, resolution) - p3 = p1.clone() # y - p3[1, :] = (p3[1, :] - 2).clamp(0, resolution) - p4 = p1.clone() # x - p4[0, :] = (p4[0, :] - 2).clamp(0, resolution) - - v1 = sdf_all[p1[0, :], p1[1, :], p1[2, :]] - v2 = sdf_all[p2[0, :], p2[1, :], p2[2, :]] - v3 = sdf_all[p3[0, :], p3[1, :], p3[2, :]] - v4 = sdf_all[p4[0, :], p4[1, :], p4[2, :]] - - X = p1[0, :].long() # [N,] - Y = p1[1, :].long() # [N,] - Z = p2[2, :].float() * (0.5 - v1) / (v2 - v1) + \ - p1[2, :].float() * (v2 - 0.5) / (v2 - v1) # [N,] - Z = Z.clamp(0, resolution) - - # normal - norm_z = v2 - v1 - norm_y = v3 - v1 - norm_x = v4 - v1 - # print (v2.min(dim=0)[0], v2.max(dim=0)[0], v3.min(dim=0)[0], v3.max(dim=0)[0]) - - norm = torch.stack([norm_x, norm_y, norm_z], dim=1) - norm = norm / torch.norm(norm, p=2, dim=1, keepdim=True) - - return X, Y, Z, norm - - def render_normal(self, resolution, X, Y, Z, norm): - image = torch.ones((1, 3, resolution, resolution), - dtype=torch.float32).to(norm.device) - color = (norm + 1) / 2.0 - color = color.clamp(0, 1) - image[0, :, Y, X] = color.t() - return image - - def display(self, sdf): - - # render - X, Y, Z, norm = self.find_vertices(sdf, direction="front") - image1 = self.render_normal(self.resolutions[-1, -1], X, Y, Z, norm) - X, Y, Z, norm = self.find_vertices(sdf, direction="left") - image2 = self.render_normal(self.resolutions[-1, -1], X, Y, Z, norm) - X, Y, Z, norm = self.find_vertices(sdf, direction="right") - image3 = self.render_normal(self.resolutions[-1, -1], X, Y, Z, norm) - X, Y, Z, norm = self.find_vertices(sdf, direction="back") - image4 = self.render_normal(self.resolutions[-1, -1], X, Y, Z, norm) - - image = torch.cat([image1, image2, image3, image4], axis=3) - image = image.detach().cpu().numpy()[0].transpose(1, 2, 0) * 255.0 - - return np.uint8(image) - - def export_mesh(self, occupancys): - - final = occupancys[1:, 1:, 1:].contiguous() - - if final.shape[0] > 256: - # for voxelgrid larger than 256^3, the required GPU memory will be > 9GB - # thus we use CPU marching_cube to avoid "CUDA out of memory" - occu_arr = final.detach().cpu().numpy() # non-smooth surface - # occu_arr = mcubes.smooth(final.detach().cpu().numpy()) # smooth surface - vertices, triangles = mcubes.marching_cubes( - occu_arr, self.balance_value) - verts = torch.as_tensor(vertices[:, [2, 1, 0]]) - faces = torch.as_tensor(triangles.astype( - np.long), dtype=torch.long)[:, [0, 2, 1]] - else: - torch.cuda.empty_cache() - vertices, triangles = voxelgrids_to_trianglemeshes( - final.unsqueeze(0)) - verts = vertices[0][:, [2, 1, 0]].cpu() - faces = triangles[0][:, [0, 2, 1]].cpu() - - return verts, faces \ No newline at end of file diff --git a/lib /common /smpl_vert_segmentation.json b/lib /common /smpl_vert_segmentation.json deleted file mode 100644 index b3244cce450e13f1095a1c3af676f4c8fdea5633..0000000000000000000000000000000000000000 --- a/lib /common /smpl_vert_segmentation.json +++ /dev/null @@ -1,7440 +0,0 @@ -{ - "rightHand": [ - 5442, - 5443, - 5444, - 5445, - 5446, - 5447, - 5448, - 5449, - 5450, - 5451, - 5452, - 5453, - 5454, - 5455, - 5456, - 5457, - 5458, - 5459, - 5460, - 5461, - 5462, - 5463, - 5464, - 5465, - 5466, - 5467, - 5468, - 5469, - 5470, - 5471, - 5472, - 5473, - 5474, - 5475, - 5476, - 5477, - 5478, - 5479, - 5480, - 5481, - 5482, - 5483, - 5484, - 5485, - 5486, - 5487, - 5492, - 5493, - 5494, - 5495, - 5496, - 5497, - 5502, - 5503, - 5504, - 5505, - 5506, - 5507, - 5508, - 5509, - 5510, - 5511, - 5512, - 5513, - 5514, - 5515, - 5516, - 5517, - 5518, - 5519, - 5520, - 5521, - 5522, - 5523, - 5524, - 5525, - 5526, - 5527, - 5530, - 5531, - 5532, - 5533, - 5534, - 5535, - 5536, - 5537, - 5538, - 5539, - 5540, - 5541, - 5542, - 5543, - 5544, - 5545, - 5546, - 5547, - 5548, - 5549, - 5550, - 5551, - 5552, - 5553, - 5554, - 5555, - 5556, - 5557, - 5558, - 5559, - 5560, - 5561, - 5562, - 5569, - 5571, - 5574, - 5575, - 5576, - 5577, - 5578, - 5579, - 5580, - 5581, - 5582, - 5583, - 5588, - 5589, - 5592, - 5593, - 5594, - 5595, - 5596, - 5597, - 5598, - 5599, - 5600, - 5601, - 5602, - 5603, - 5604, - 5605, - 5610, - 5611, - 5612, - 5613, - 5614, - 5621, - 5622, - 5625, - 5631, - 5632, - 5633, - 5634, - 5635, - 5636, - 5637, - 5638, - 5639, - 5640, - 5641, - 5643, - 5644, - 5645, - 5646, - 5649, - 5650, - 5652, - 5653, - 5654, - 5655, - 5656, - 5657, - 5658, - 5659, - 5660, - 5661, - 5662, - 5663, - 5664, - 5667, - 5670, - 5671, - 5672, - 5673, - 5674, - 5675, - 5682, - 5683, - 5684, - 5685, - 5686, - 5687, - 5688, - 5689, - 5690, - 5692, - 5695, - 5697, - 5698, - 5699, - 5700, - 5701, - 5707, - 5708, - 5709, - 5710, - 5711, - 5712, - 5713, - 5714, - 5715, - 5716, - 5717, - 5718, - 5719, - 5720, - 5721, - 5723, - 5724, - 5725, - 5726, - 5727, - 5728, - 5729, - 5730, - 5731, - 5732, - 5735, - 5736, - 5737, - 5738, - 5739, - 5740, - 5745, - 5746, - 5748, - 5749, - 5750, - 5751, - 5752, - 6056, - 6057, - 6066, - 6067, - 6158, - 6159, - 6160, - 6161, - 6162, - 6163, - 6164, - 6165, - 6166, - 6167, - 6168, - 6169, - 6170, - 6171, - 6172, - 6173, - 6174, - 6175, - 6176, - 6177, - 6178, - 6179, - 6180, - 6181, - 6182, - 6183, - 6184, - 6185, - 6186, - 6187, - 6188, - 6189, - 6190, - 6191, - 6192, - 6193, - 6194, - 6195, - 6196, - 6197, - 6198, - 6199, - 6200, - 6201, - 6202, - 6203, - 6204, - 6205, - 6206, - 6207, - 6208, - 6209, - 6210, - 6211, - 6212, - 6213, - 6214, - 6215, - 6216, - 6217, - 6218, - 6219, - 6220, - 6221, - 6222, - 6223, - 6224, - 6225, - 6226, - 6227, - 6228, - 6229, - 6230, - 6231, - 6232, - 6233, - 6234, - 6235, - 6236, - 6237, - 6238, - 6239 - ], - "rightUpLeg": [ - 4320, - 4321, - 4323, - 4324, - 4333, - 4334, - 4335, - 4336, - 4337, - 4338, - 4339, - 4340, - 4356, - 4357, - 4358, - 4359, - 4360, - 4361, - 4362, - 4363, - 4364, - 4365, - 4366, - 4367, - 4383, - 4384, - 4385, - 4386, - 4387, - 4388, - 4389, - 4390, - 4391, - 4392, - 4393, - 4394, - 4395, - 4396, - 4397, - 4398, - 4399, - 4400, - 4401, - 4419, - 4420, - 4421, - 4422, - 4430, - 4431, - 4432, - 4433, - 4434, - 4435, - 4436, - 4437, - 4438, - 4439, - 4440, - 4441, - 4442, - 4443, - 4444, - 4445, - 4446, - 4447, - 4448, - 4449, - 4450, - 4451, - 4452, - 4453, - 4454, - 4455, - 4456, - 4457, - 4458, - 4459, - 4460, - 4461, - 4462, - 4463, - 4464, - 4465, - 4466, - 4467, - 4468, - 4469, - 4470, - 4471, - 4472, - 4473, - 4474, - 4475, - 4476, - 4477, - 4478, - 4479, - 4480, - 4481, - 4482, - 4483, - 4484, - 4485, - 4486, - 4487, - 4488, - 4489, - 4490, - 4491, - 4492, - 4493, - 4494, - 4495, - 4496, - 4497, - 4498, - 4499, - 4500, - 4501, - 4502, - 4503, - 4504, - 4505, - 4506, - 4507, - 4508, - 4509, - 4510, - 4511, - 4512, - 4513, - 4514, - 4515, - 4516, - 4517, - 4518, - 4519, - 4520, - 4521, - 4522, - 4523, - 4524, - 4525, - 4526, - 4527, - 4528, - 4529, - 4530, - 4531, - 4532, - 4623, - 4624, - 4625, - 4626, - 4627, - 4628, - 4629, - 4630, - 4631, - 4632, - 4633, - 4634, - 4645, - 4646, - 4647, - 4648, - 4649, - 4650, - 4651, - 4652, - 4653, - 4654, - 4655, - 4656, - 4657, - 4658, - 4659, - 4660, - 4670, - 4671, - 4672, - 4673, - 4704, - 4705, - 4706, - 4707, - 4708, - 4709, - 4710, - 4711, - 4712, - 4713, - 4745, - 4746, - 4757, - 4758, - 4759, - 4760, - 4801, - 4802, - 4829, - 4834, - 4835, - 4836, - 4837, - 4838, - 4839, - 4840, - 4841, - 4924, - 4925, - 4926, - 4928, - 4929, - 4930, - 4931, - 4932, - 4933, - 4934, - 4935, - 4936, - 4948, - 4949, - 4950, - 4951, - 4952, - 4970, - 4971, - 4972, - 4973, - 4983, - 4984, - 4985, - 4986, - 4987, - 4988, - 4989, - 4990, - 4991, - 4992, - 4993, - 5004, - 5005, - 6546, - 6547, - 6548, - 6549, - 6552, - 6553, - 6554, - 6555, - 6556, - 6873, - 6877 - ], - "leftArm": [ - 626, - 627, - 628, - 629, - 634, - 635, - 680, - 681, - 716, - 717, - 718, - 719, - 769, - 770, - 771, - 772, - 773, - 774, - 775, - 776, - 777, - 778, - 779, - 780, - 784, - 785, - 786, - 787, - 788, - 789, - 790, - 791, - 792, - 793, - 1231, - 1232, - 1233, - 1234, - 1258, - 1259, - 1260, - 1261, - 1271, - 1281, - 1282, - 1310, - 1311, - 1314, - 1315, - 1340, - 1341, - 1342, - 1343, - 1355, - 1356, - 1357, - 1358, - 1376, - 1377, - 1378, - 1379, - 1380, - 1381, - 1382, - 1383, - 1384, - 1385, - 1386, - 1387, - 1388, - 1389, - 1390, - 1391, - 1392, - 1393, - 1394, - 1395, - 1396, - 1397, - 1398, - 1399, - 1400, - 1402, - 1403, - 1405, - 1406, - 1407, - 1408, - 1409, - 1410, - 1411, - 1412, - 1413, - 1414, - 1415, - 1416, - 1428, - 1429, - 1430, - 1431, - 1432, - 1433, - 1438, - 1439, - 1440, - 1441, - 1442, - 1443, - 1444, - 1445, - 1502, - 1505, - 1506, - 1507, - 1508, - 1509, - 1510, - 1538, - 1541, - 1542, - 1543, - 1545, - 1619, - 1620, - 1621, - 1622, - 1631, - 1632, - 1633, - 1634, - 1635, - 1636, - 1637, - 1638, - 1639, - 1640, - 1641, - 1642, - 1645, - 1646, - 1647, - 1648, - 1649, - 1650, - 1651, - 1652, - 1653, - 1654, - 1655, - 1656, - 1658, - 1659, - 1661, - 1662, - 1664, - 1666, - 1667, - 1668, - 1669, - 1670, - 1671, - 1672, - 1673, - 1674, - 1675, - 1676, - 1677, - 1678, - 1679, - 1680, - 1681, - 1682, - 1683, - 1684, - 1696, - 1697, - 1698, - 1703, - 1704, - 1705, - 1706, - 1707, - 1708, - 1709, - 1710, - 1711, - 1712, - 1713, - 1714, - 1715, - 1716, - 1717, - 1718, - 1719, - 1720, - 1725, - 1731, - 1732, - 1733, - 1734, - 1735, - 1737, - 1739, - 1740, - 1745, - 1746, - 1747, - 1748, - 1749, - 1751, - 1761, - 1830, - 1831, - 1844, - 1845, - 1846, - 1850, - 1851, - 1854, - 1855, - 1858, - 1860, - 1865, - 1866, - 1867, - 1869, - 1870, - 1871, - 1874, - 1875, - 1876, - 1877, - 1878, - 1882, - 1883, - 1888, - 1889, - 1892, - 1900, - 1901, - 1902, - 1903, - 1904, - 1909, - 2819, - 2820, - 2821, - 2822, - 2895, - 2896, - 2897, - 2898, - 2899, - 2900, - 2901, - 2902, - 2903, - 2945, - 2946, - 2974, - 2975, - 2976, - 2977, - 2978, - 2979, - 2980, - 2981, - 2982, - 2983, - 2984, - 2985, - 2986, - 2987, - 2988, - 2989, - 2990, - 2991, - 2992, - 2993, - 2994, - 2995, - 2996, - 3002, - 3013 - ], - "leftLeg": [ - 995, - 998, - 999, - 1002, - 1004, - 1005, - 1008, - 1010, - 1012, - 1015, - 1016, - 1018, - 1019, - 1043, - 1044, - 1047, - 1048, - 1049, - 1050, - 1051, - 1052, - 1053, - 1054, - 1055, - 1056, - 1057, - 1058, - 1059, - 1060, - 1061, - 1062, - 1063, - 1064, - 1065, - 1066, - 1067, - 1068, - 1069, - 1070, - 1071, - 1072, - 1073, - 1074, - 1075, - 1076, - 1077, - 1078, - 1079, - 1080, - 1081, - 1082, - 1083, - 1084, - 1085, - 1086, - 1087, - 1088, - 1089, - 1090, - 1091, - 1092, - 1093, - 1094, - 1095, - 1096, - 1097, - 1098, - 1099, - 1100, - 1101, - 1102, - 1103, - 1104, - 1105, - 1106, - 1107, - 1108, - 1109, - 1110, - 1111, - 1112, - 1113, - 1114, - 1115, - 1116, - 1117, - 1118, - 1119, - 1120, - 1121, - 1122, - 1123, - 1124, - 1125, - 1126, - 1127, - 1128, - 1129, - 1130, - 1131, - 1132, - 1133, - 1134, - 1135, - 1136, - 1148, - 1149, - 1150, - 1151, - 1152, - 1153, - 1154, - 1155, - 1156, - 1157, - 1158, - 1175, - 1176, - 1177, - 1178, - 1179, - 1180, - 1181, - 1182, - 1183, - 1369, - 1370, - 1371, - 1372, - 1373, - 1374, - 1375, - 1464, - 1465, - 1466, - 1467, - 1468, - 1469, - 1470, - 1471, - 1472, - 1473, - 1474, - 1522, - 1523, - 1524, - 1525, - 1526, - 1527, - 1528, - 1529, - 1530, - 1531, - 1532, - 3174, - 3175, - 3176, - 3177, - 3178, - 3179, - 3180, - 3181, - 3182, - 3183, - 3184, - 3185, - 3186, - 3187, - 3188, - 3189, - 3190, - 3191, - 3192, - 3193, - 3194, - 3195, - 3196, - 3197, - 3198, - 3199, - 3200, - 3201, - 3202, - 3203, - 3204, - 3205, - 3206, - 3207, - 3208, - 3209, - 3210, - 3319, - 3320, - 3321, - 3322, - 3323, - 3324, - 3325, - 3326, - 3327, - 3328, - 3329, - 3330, - 3331, - 3332, - 3333, - 3334, - 3335, - 3432, - 3433, - 3434, - 3435, - 3436, - 3469, - 3472, - 3473, - 3474 - ], - "leftToeBase": [ - 3211, - 3212, - 3213, - 3214, - 3215, - 3216, - 3217, - 3218, - 3219, - 3220, - 3221, - 3222, - 3223, - 3224, - 3225, - 3226, - 3227, - 3228, - 3229, - 3230, - 3231, - 3232, - 3233, - 3234, - 3235, - 3236, - 3237, - 3238, - 3239, - 3240, - 3241, - 3242, - 3243, - 3244, - 3245, - 3246, - 3247, - 3248, - 3249, - 3250, - 3251, - 3252, - 3253, - 3254, - 3255, - 3256, - 3257, - 3258, - 3259, - 3260, - 3261, - 3262, - 3263, - 3264, - 3265, - 3266, - 3267, - 3268, - 3269, - 3270, - 3271, - 3272, - 3273, - 3274, - 3275, - 3276, - 3277, - 3278, - 3279, - 3280, - 3281, - 3282, - 3283, - 3284, - 3285, - 3286, - 3287, - 3288, - 3289, - 3290, - 3291, - 3292, - 3293, - 3294, - 3295, - 3296, - 3297, - 3298, - 3299, - 3300, - 3301, - 3302, - 3303, - 3304, - 3305, - 3306, - 3307, - 3308, - 3309, - 3310, - 3311, - 3312, - 3313, - 3314, - 3315, - 3316, - 3317, - 3318, - 3336, - 3337, - 3340, - 3342, - 3344, - 3346, - 3348, - 3350, - 3352, - 3354, - 3357, - 3358, - 3360, - 3362 - ], - "leftFoot": [ - 3327, - 3328, - 3329, - 3330, - 3331, - 3332, - 3333, - 3334, - 3335, - 3336, - 3337, - 3338, - 3339, - 3340, - 3341, - 3342, - 3343, - 3344, - 3345, - 3346, - 3347, - 3348, - 3349, - 3350, - 3351, - 3352, - 3353, - 3354, - 3355, - 3356, - 3357, - 3358, - 3359, - 3360, - 3361, - 3362, - 3363, - 3364, - 3365, - 3366, - 3367, - 3368, - 3369, - 3370, - 3371, - 3372, - 3373, - 3374, - 3375, - 3376, - 3377, - 3378, - 3379, - 3380, - 3381, - 3382, - 3383, - 3384, - 3385, - 3386, - 3387, - 3388, - 3389, - 3390, - 3391, - 3392, - 3393, - 3394, - 3395, - 3396, - 3397, - 3398, - 3399, - 3400, - 3401, - 3402, - 3403, - 3404, - 3405, - 3406, - 3407, - 3408, - 3409, - 3410, - 3411, - 3412, - 3413, - 3414, - 3415, - 3416, - 3417, - 3418, - 3419, - 3420, - 3421, - 3422, - 3423, - 3424, - 3425, - 3426, - 3427, - 3428, - 3429, - 3430, - 3431, - 3432, - 3433, - 3434, - 3435, - 3436, - 3437, - 3438, - 3439, - 3440, - 3441, - 3442, - 3443, - 3444, - 3445, - 3446, - 3447, - 3448, - 3449, - 3450, - 3451, - 3452, - 3453, - 3454, - 3455, - 3456, - 3457, - 3458, - 3459, - 3460, - 3461, - 3462, - 3463, - 3464, - 3465, - 3466, - 3467, - 3468, - 3469 - ], - "spine1": [ - 598, - 599, - 600, - 601, - 610, - 611, - 612, - 613, - 614, - 615, - 616, - 617, - 618, - 619, - 620, - 621, - 642, - 645, - 646, - 647, - 652, - 653, - 658, - 659, - 660, - 661, - 668, - 669, - 670, - 671, - 684, - 685, - 686, - 687, - 688, - 689, - 690, - 691, - 692, - 722, - 723, - 724, - 725, - 736, - 750, - 751, - 761, - 764, - 766, - 767, - 794, - 795, - 891, - 892, - 893, - 894, - 925, - 926, - 927, - 928, - 929, - 940, - 941, - 942, - 943, - 1190, - 1191, - 1192, - 1193, - 1194, - 1195, - 1196, - 1197, - 1200, - 1201, - 1202, - 1212, - 1236, - 1252, - 1253, - 1254, - 1255, - 1268, - 1269, - 1270, - 1329, - 1330, - 1348, - 1349, - 1351, - 1420, - 1421, - 1423, - 1424, - 1425, - 1426, - 1436, - 1437, - 1756, - 1757, - 1758, - 2839, - 2840, - 2841, - 2842, - 2843, - 2844, - 2845, - 2846, - 2847, - 2848, - 2849, - 2850, - 2851, - 2870, - 2871, - 2883, - 2906, - 2908, - 3014, - 3017, - 3025, - 3030, - 3033, - 3034, - 3037, - 3039, - 3040, - 3041, - 3042, - 3043, - 3044, - 3076, - 3077, - 3079, - 3480, - 3505, - 3511, - 4086, - 4087, - 4088, - 4089, - 4098, - 4099, - 4100, - 4101, - 4102, - 4103, - 4104, - 4105, - 4106, - 4107, - 4108, - 4109, - 4130, - 4131, - 4134, - 4135, - 4140, - 4141, - 4146, - 4147, - 4148, - 4149, - 4156, - 4157, - 4158, - 4159, - 4172, - 4173, - 4174, - 4175, - 4176, - 4177, - 4178, - 4179, - 4180, - 4210, - 4211, - 4212, - 4213, - 4225, - 4239, - 4240, - 4249, - 4250, - 4255, - 4256, - 4282, - 4283, - 4377, - 4378, - 4379, - 4380, - 4411, - 4412, - 4413, - 4414, - 4415, - 4426, - 4427, - 4428, - 4429, - 4676, - 4677, - 4678, - 4679, - 4680, - 4681, - 4682, - 4683, - 4686, - 4687, - 4688, - 4695, - 4719, - 4735, - 4736, - 4737, - 4740, - 4751, - 4752, - 4753, - 4824, - 4825, - 4828, - 4893, - 4894, - 4895, - 4897, - 4898, - 4899, - 4908, - 4909, - 5223, - 5224, - 5225, - 6300, - 6301, - 6302, - 6303, - 6304, - 6305, - 6306, - 6307, - 6308, - 6309, - 6310, - 6311, - 6312, - 6331, - 6332, - 6342, - 6366, - 6367, - 6475, - 6477, - 6478, - 6481, - 6482, - 6485, - 6487, - 6488, - 6489, - 6490, - 6491, - 6878 - ], - "spine2": [ - 570, - 571, - 572, - 573, - 584, - 585, - 586, - 587, - 588, - 589, - 590, - 591, - 592, - 593, - 594, - 595, - 596, - 597, - 602, - 603, - 604, - 605, - 606, - 607, - 608, - 609, - 622, - 623, - 624, - 625, - 638, - 639, - 640, - 641, - 643, - 644, - 648, - 649, - 650, - 651, - 666, - 667, - 672, - 673, - 674, - 675, - 680, - 681, - 682, - 683, - 693, - 694, - 695, - 696, - 697, - 698, - 699, - 700, - 701, - 702, - 703, - 704, - 713, - 714, - 715, - 716, - 717, - 726, - 727, - 728, - 729, - 730, - 731, - 732, - 733, - 735, - 737, - 738, - 739, - 740, - 741, - 742, - 743, - 744, - 745, - 746, - 747, - 748, - 749, - 752, - 753, - 754, - 755, - 756, - 757, - 758, - 759, - 760, - 762, - 763, - 803, - 804, - 805, - 806, - 811, - 812, - 813, - 814, - 817, - 818, - 819, - 820, - 821, - 824, - 825, - 826, - 827, - 828, - 895, - 896, - 930, - 931, - 1198, - 1199, - 1213, - 1214, - 1215, - 1216, - 1217, - 1218, - 1219, - 1220, - 1235, - 1237, - 1256, - 1257, - 1271, - 1272, - 1273, - 1279, - 1280, - 1283, - 1284, - 1285, - 1286, - 1287, - 1288, - 1289, - 1290, - 1291, - 1292, - 1293, - 1294, - 1295, - 1296, - 1297, - 1298, - 1299, - 1300, - 1301, - 1302, - 1303, - 1304, - 1305, - 1306, - 1307, - 1308, - 1309, - 1312, - 1313, - 1319, - 1320, - 1346, - 1347, - 1350, - 1352, - 1401, - 1417, - 1418, - 1419, - 1422, - 1427, - 1434, - 1435, - 1503, - 1504, - 1536, - 1537, - 1544, - 1545, - 1753, - 1754, - 1755, - 1759, - 1760, - 1761, - 1762, - 1763, - 1808, - 1809, - 1810, - 1811, - 1816, - 1817, - 1818, - 1819, - 1820, - 1834, - 1835, - 1836, - 1837, - 1838, - 1839, - 1868, - 1879, - 1880, - 2812, - 2813, - 2852, - 2853, - 2854, - 2855, - 2856, - 2857, - 2858, - 2859, - 2860, - 2861, - 2862, - 2863, - 2864, - 2865, - 2866, - 2867, - 2868, - 2869, - 2872, - 2875, - 2876, - 2877, - 2878, - 2881, - 2882, - 2884, - 2885, - 2886, - 2904, - 2905, - 2907, - 2931, - 2932, - 2933, - 2934, - 2935, - 2936, - 2937, - 2941, - 2950, - 2951, - 2952, - 2953, - 2954, - 2955, - 2956, - 2957, - 2958, - 2959, - 2960, - 2961, - 2962, - 2963, - 2964, - 2965, - 2966, - 2967, - 2968, - 2969, - 2970, - 2971, - 2972, - 2973, - 2997, - 2998, - 3006, - 3007, - 3012, - 3015, - 3026, - 3027, - 3028, - 3029, - 3031, - 3032, - 3035, - 3036, - 3038, - 3059, - 3060, - 3061, - 3062, - 3063, - 3064, - 3065, - 3066, - 3067, - 3073, - 3074, - 3075, - 3078, - 3168, - 3169, - 3171, - 3470, - 3471, - 3482, - 3483, - 3495, - 3496, - 3497, - 3498, - 3506, - 3508, - 4058, - 4059, - 4060, - 4061, - 4072, - 4073, - 4074, - 4075, - 4076, - 4077, - 4078, - 4079, - 4080, - 4081, - 4082, - 4083, - 4084, - 4085, - 4090, - 4091, - 4092, - 4093, - 4094, - 4095, - 4096, - 4097, - 4110, - 4111, - 4112, - 4113, - 4126, - 4127, - 4128, - 4129, - 4132, - 4133, - 4136, - 4137, - 4138, - 4139, - 4154, - 4155, - 4160, - 4161, - 4162, - 4163, - 4168, - 4169, - 4170, - 4171, - 4181, - 4182, - 4183, - 4184, - 4185, - 4186, - 4187, - 4188, - 4189, - 4190, - 4191, - 4192, - 4201, - 4202, - 4203, - 4204, - 4207, - 4214, - 4215, - 4216, - 4217, - 4218, - 4219, - 4220, - 4221, - 4223, - 4224, - 4226, - 4227, - 4228, - 4229, - 4230, - 4231, - 4232, - 4233, - 4234, - 4235, - 4236, - 4237, - 4238, - 4241, - 4242, - 4243, - 4244, - 4245, - 4246, - 4247, - 4248, - 4251, - 4252, - 4291, - 4292, - 4293, - 4294, - 4299, - 4300, - 4301, - 4302, - 4305, - 4306, - 4307, - 4308, - 4309, - 4312, - 4313, - 4314, - 4315, - 4381, - 4382, - 4416, - 4417, - 4684, - 4685, - 4696, - 4697, - 4698, - 4699, - 4700, - 4701, - 4702, - 4703, - 4718, - 4720, - 4738, - 4739, - 4754, - 4755, - 4756, - 4761, - 4762, - 4765, - 4766, - 4767, - 4768, - 4769, - 4770, - 4771, - 4772, - 4773, - 4774, - 4775, - 4776, - 4777, - 4778, - 4779, - 4780, - 4781, - 4782, - 4783, - 4784, - 4785, - 4786, - 4787, - 4788, - 4789, - 4792, - 4793, - 4799, - 4800, - 4822, - 4823, - 4826, - 4827, - 4874, - 4890, - 4891, - 4892, - 4896, - 4900, - 4907, - 4910, - 4975, - 4976, - 5007, - 5008, - 5013, - 5014, - 5222, - 5226, - 5227, - 5228, - 5229, - 5230, - 5269, - 5270, - 5271, - 5272, - 5277, - 5278, - 5279, - 5280, - 5281, - 5295, - 5296, - 5297, - 5298, - 5299, - 5300, - 5329, - 5340, - 5341, - 6273, - 6274, - 6313, - 6314, - 6315, - 6316, - 6317, - 6318, - 6319, - 6320, - 6321, - 6322, - 6323, - 6324, - 6325, - 6326, - 6327, - 6328, - 6329, - 6330, - 6333, - 6336, - 6337, - 6340, - 6341, - 6343, - 6344, - 6345, - 6363, - 6364, - 6365, - 6390, - 6391, - 6392, - 6393, - 6394, - 6395, - 6396, - 6398, - 6409, - 6410, - 6411, - 6412, - 6413, - 6414, - 6415, - 6416, - 6417, - 6418, - 6419, - 6420, - 6421, - 6422, - 6423, - 6424, - 6425, - 6426, - 6427, - 6428, - 6429, - 6430, - 6431, - 6432, - 6456, - 6457, - 6465, - 6466, - 6476, - 6479, - 6480, - 6483, - 6484, - 6486, - 6496, - 6497, - 6498, - 6499, - 6500, - 6501, - 6502, - 6503, - 6879 - ], - "leftShoulder": [ - 591, - 604, - 605, - 606, - 609, - 634, - 635, - 636, - 637, - 674, - 706, - 707, - 708, - 709, - 710, - 711, - 712, - 713, - 715, - 717, - 730, - 733, - 734, - 735, - 781, - 782, - 783, - 1238, - 1239, - 1240, - 1241, - 1242, - 1243, - 1244, - 1245, - 1290, - 1291, - 1294, - 1316, - 1317, - 1318, - 1401, - 1402, - 1403, - 1404, - 1509, - 1535, - 1545, - 1808, - 1810, - 1811, - 1812, - 1813, - 1814, - 1815, - 1818, - 1819, - 1821, - 1822, - 1823, - 1824, - 1825, - 1826, - 1827, - 1828, - 1829, - 1830, - 1831, - 1832, - 1833, - 1837, - 1840, - 1841, - 1842, - 1843, - 1844, - 1845, - 1846, - 1847, - 1848, - 1849, - 1850, - 1851, - 1852, - 1853, - 1854, - 1855, - 1856, - 1857, - 1858, - 1859, - 1861, - 1862, - 1863, - 1864, - 1872, - 1873, - 1880, - 1881, - 1884, - 1885, - 1886, - 1887, - 1890, - 1891, - 1893, - 1894, - 1895, - 1896, - 1897, - 1898, - 1899, - 2879, - 2880, - 2881, - 2886, - 2887, - 2888, - 2889, - 2890, - 2891, - 2892, - 2893, - 2894, - 2903, - 2938, - 2939, - 2940, - 2941, - 2942, - 2943, - 2944, - 2945, - 2946, - 2947, - 2948, - 2949, - 2965, - 2967, - 2969, - 2999, - 3000, - 3001, - 3002, - 3003, - 3004, - 3005, - 3008, - 3009, - 3010, - 3011 - ], - "rightShoulder": [ - 4077, - 4091, - 4092, - 4094, - 4095, - 4122, - 4123, - 4124, - 4125, - 4162, - 4194, - 4195, - 4196, - 4197, - 4198, - 4199, - 4200, - 4201, - 4203, - 4207, - 4218, - 4219, - 4222, - 4223, - 4269, - 4270, - 4271, - 4721, - 4722, - 4723, - 4724, - 4725, - 4726, - 4727, - 4728, - 4773, - 4774, - 4778, - 4796, - 4797, - 4798, - 4874, - 4875, - 4876, - 4877, - 4982, - 5006, - 5014, - 5269, - 5271, - 5272, - 5273, - 5274, - 5275, - 5276, - 5279, - 5281, - 5282, - 5283, - 5284, - 5285, - 5286, - 5287, - 5288, - 5289, - 5290, - 5291, - 5292, - 5293, - 5294, - 5298, - 5301, - 5302, - 5303, - 5304, - 5305, - 5306, - 5307, - 5308, - 5309, - 5310, - 5311, - 5312, - 5313, - 5314, - 5315, - 5316, - 5317, - 5318, - 5319, - 5320, - 5322, - 5323, - 5324, - 5325, - 5333, - 5334, - 5341, - 5342, - 5345, - 5346, - 5347, - 5348, - 5351, - 5352, - 5354, - 5355, - 5356, - 5357, - 5358, - 5359, - 5360, - 6338, - 6339, - 6340, - 6345, - 6346, - 6347, - 6348, - 6349, - 6350, - 6351, - 6352, - 6353, - 6362, - 6397, - 6398, - 6399, - 6400, - 6401, - 6402, - 6403, - 6404, - 6405, - 6406, - 6407, - 6408, - 6424, - 6425, - 6428, - 6458, - 6459, - 6460, - 6461, - 6462, - 6463, - 6464, - 6467, - 6468, - 6469, - 6470 - ], - "rightFoot": [ - 6727, - 6728, - 6729, - 6730, - 6731, - 6732, - 6733, - 6734, - 6735, - 6736, - 6737, - 6738, - 6739, - 6740, - 6741, - 6742, - 6743, - 6744, - 6745, - 6746, - 6747, - 6748, - 6749, - 6750, - 6751, - 6752, - 6753, - 6754, - 6755, - 6756, - 6757, - 6758, - 6759, - 6760, - 6761, - 6762, - 6763, - 6764, - 6765, - 6766, - 6767, - 6768, - 6769, - 6770, - 6771, - 6772, - 6773, - 6774, - 6775, - 6776, - 6777, - 6778, - 6779, - 6780, - 6781, - 6782, - 6783, - 6784, - 6785, - 6786, - 6787, - 6788, - 6789, - 6790, - 6791, - 6792, - 6793, - 6794, - 6795, - 6796, - 6797, - 6798, - 6799, - 6800, - 6801, - 6802, - 6803, - 6804, - 6805, - 6806, - 6807, - 6808, - 6809, - 6810, - 6811, - 6812, - 6813, - 6814, - 6815, - 6816, - 6817, - 6818, - 6819, - 6820, - 6821, - 6822, - 6823, - 6824, - 6825, - 6826, - 6827, - 6828, - 6829, - 6830, - 6831, - 6832, - 6833, - 6834, - 6835, - 6836, - 6837, - 6838, - 6839, - 6840, - 6841, - 6842, - 6843, - 6844, - 6845, - 6846, - 6847, - 6848, - 6849, - 6850, - 6851, - 6852, - 6853, - 6854, - 6855, - 6856, - 6857, - 6858, - 6859, - 6860, - 6861, - 6862, - 6863, - 6864, - 6865, - 6866, - 6867, - 6868, - 6869 - ], - "head": [ - 0, - 1, - 2, - 3, - 4, - 5, - 6, - 7, - 8, - 9, - 10, - 11, - 12, - 13, - 14, - 15, - 16, - 17, - 18, - 19, - 20, - 21, - 22, - 23, - 24, - 25, - 26, - 27, - 28, - 29, - 30, - 31, - 32, - 33, - 34, - 35, - 36, - 37, - 38, - 39, - 40, - 41, - 42, - 43, - 44, - 45, - 46, - 47, - 48, - 49, - 50, - 51, - 52, - 53, - 54, - 55, - 56, - 57, - 58, - 59, - 60, - 61, - 62, - 63, - 64, - 65, - 66, - 67, - 68, - 69, - 70, - 71, - 72, - 73, - 74, - 75, - 76, - 77, - 78, - 79, - 80, - 81, - 82, - 83, - 84, - 85, - 86, - 87, - 88, - 89, - 90, - 91, - 92, - 93, - 94, - 95, - 96, - 97, - 98, - 99, - 100, - 101, - 102, - 103, - 104, - 105, - 106, - 107, - 108, - 109, - 110, - 111, - 112, - 113, - 114, - 115, - 116, - 117, - 118, - 119, - 120, - 121, - 122, - 123, - 124, - 125, - 126, - 127, - 128, - 129, - 130, - 131, - 132, - 133, - 134, - 135, - 136, - 137, - 138, - 139, - 140, - 141, - 142, - 143, - 144, - 145, - 146, - 147, - 148, - 149, - 154, - 155, - 156, - 157, - 158, - 159, - 160, - 161, - 162, - 163, - 164, - 165, - 166, - 167, - 168, - 169, - 170, - 171, - 172, - 173, - 176, - 177, - 178, - 179, - 180, - 181, - 182, - 183, - 184, - 185, - 186, - 187, - 188, - 189, - 190, - 191, - 192, - 193, - 194, - 195, - 196, - 197, - 198, - 199, - 200, - 201, - 202, - 203, - 204, - 205, - 220, - 221, - 225, - 226, - 227, - 228, - 229, - 230, - 231, - 232, - 233, - 234, - 235, - 236, - 237, - 238, - 239, - 240, - 241, - 242, - 243, - 244, - 245, - 246, - 247, - 248, - 249, - 250, - 251, - 252, - 253, - 254, - 255, - 258, - 259, - 260, - 261, - 262, - 263, - 264, - 265, - 266, - 267, - 268, - 269, - 270, - 271, - 272, - 273, - 274, - 275, - 276, - 277, - 278, - 279, - 280, - 281, - 282, - 283, - 286, - 287, - 288, - 289, - 290, - 291, - 292, - 293, - 294, - 295, - 303, - 304, - 306, - 307, - 310, - 311, - 312, - 313, - 314, - 315, - 316, - 317, - 318, - 319, - 320, - 321, - 322, - 323, - 324, - 325, - 326, - 327, - 328, - 329, - 330, - 331, - 332, - 335, - 336, - 337, - 338, - 339, - 340, - 341, - 342, - 343, - 344, - 345, - 346, - 347, - 348, - 349, - 350, - 351, - 352, - 353, - 354, - 355, - 356, - 357, - 358, - 359, - 360, - 361, - 362, - 363, - 364, - 365, - 366, - 367, - 368, - 369, - 370, - 371, - 372, - 373, - 374, - 375, - 376, - 377, - 378, - 379, - 380, - 381, - 382, - 383, - 384, - 385, - 386, - 387, - 388, - 389, - 390, - 391, - 392, - 393, - 394, - 395, - 396, - 397, - 398, - 399, - 400, - 401, - 402, - 403, - 404, - 405, - 406, - 407, - 408, - 409, - 410, - 411, - 412, - 413, - 414, - 415, - 416, - 417, - 418, - 419, - 420, - 421, - 422, - 427, - 428, - 429, - 430, - 431, - 432, - 433, - 434, - 435, - 436, - 437, - 438, - 439, - 442, - 443, - 444, - 445, - 446, - 447, - 448, - 449, - 450, - 454, - 455, - 456, - 457, - 458, - 459, - 461, - 462, - 463, - 464, - 465, - 466, - 467, - 468, - 469, - 470, - 471, - 472, - 473, - 474, - 475, - 476, - 477, - 478, - 479, - 480, - 481, - 482, - 483, - 484, - 485, - 486, - 487, - 488, - 489, - 490, - 491, - 492, - 493, - 494, - 495, - 496, - 497, - 498, - 499, - 500, - 501, - 502, - 503, - 504, - 505, - 506, - 507, - 508, - 509, - 510, - 511, - 512, - 513, - 514, - 515, - 516, - 517, - 518, - 519, - 520, - 521, - 522, - 523, - 524, - 525, - 526, - 527, - 528, - 529, - 530, - 531, - 532, - 533, - 534, - 535, - 536, - 537, - 538, - 539, - 540, - 541, - 542, - 543, - 544, - 545, - 546, - 547, - 548, - 549, - 550, - 551, - 552, - 553, - 554, - 555, - 556, - 557, - 558, - 559, - 560, - 561, - 562, - 563, - 564, - 565, - 566, - 567, - 568, - 569, - 574, - 575, - 576, - 577, - 578, - 579, - 580, - 581, - 582, - 583, - 1764, - 1765, - 1766, - 1770, - 1771, - 1772, - 1773, - 1774, - 1775, - 1776, - 1777, - 1778, - 1905, - 1906, - 1907, - 1908, - 2779, - 2780, - 2781, - 2782, - 2783, - 2784, - 2785, - 2786, - 2787, - 2788, - 2789, - 2790, - 2791, - 2792, - 2793, - 2794, - 2795, - 2796, - 2797, - 2798, - 2799, - 2800, - 2801, - 2802, - 2803, - 2804, - 2805, - 2806, - 2807, - 2808, - 2809, - 2810, - 2811, - 2814, - 2815, - 2816, - 2817, - 2818, - 3045, - 3046, - 3047, - 3048, - 3051, - 3052, - 3053, - 3054, - 3055, - 3056, - 3058, - 3069, - 3070, - 3071, - 3072, - 3161, - 3162, - 3163, - 3165, - 3166, - 3167, - 3485, - 3486, - 3487, - 3488, - 3489, - 3490, - 3491, - 3492, - 3493, - 3494, - 3499, - 3512, - 3513, - 3514, - 3515, - 3516, - 3517, - 3518, - 3519, - 3520, - 3521, - 3522, - 3523, - 3524, - 3525, - 3526, - 3527, - 3528, - 3529, - 3530, - 3531, - 3532, - 3533, - 3534, - 3535, - 3536, - 3537, - 3538, - 3539, - 3540, - 3541, - 3542, - 3543, - 3544, - 3545, - 3546, - 3547, - 3548, - 3549, - 3550, - 3551, - 3552, - 3553, - 3554, - 3555, - 3556, - 3557, - 3558, - 3559, - 3560, - 3561, - 3562, - 3563, - 3564, - 3565, - 3566, - 3567, - 3568, - 3569, - 3570, - 3571, - 3572, - 3573, - 3574, - 3575, - 3576, - 3577, - 3578, - 3579, - 3580, - 3581, - 3582, - 3583, - 3584, - 3585, - 3586, - 3587, - 3588, - 3589, - 3590, - 3591, - 3592, - 3593, - 3594, - 3595, - 3596, - 3597, - 3598, - 3599, - 3600, - 3601, - 3602, - 3603, - 3604, - 3605, - 3606, - 3607, - 3608, - 3609, - 3610, - 3611, - 3612, - 3613, - 3614, - 3615, - 3616, - 3617, - 3618, - 3619, - 3620, - 3621, - 3622, - 3623, - 3624, - 3625, - 3626, - 3627, - 3628, - 3629, - 3630, - 3631, - 3632, - 3633, - 3634, - 3635, - 3636, - 3637, - 3638, - 3639, - 3640, - 3641, - 3642, - 3643, - 3644, - 3645, - 3646, - 3647, - 3648, - 3649, - 3650, - 3651, - 3652, - 3653, - 3654, - 3655, - 3656, - 3657, - 3658, - 3659, - 3660, - 3661, - 3666, - 3667, - 3668, - 3669, - 3670, - 3671, - 3672, - 3673, - 3674, - 3675, - 3676, - 3677, - 3678, - 3679, - 3680, - 3681, - 3682, - 3683, - 3684, - 3685, - 3688, - 3689, - 3690, - 3691, - 3692, - 3693, - 3694, - 3695, - 3696, - 3697, - 3698, - 3699, - 3700, - 3701, - 3702, - 3703, - 3704, - 3705, - 3706, - 3707, - 3708, - 3709, - 3710, - 3711, - 3712, - 3713, - 3714, - 3715, - 3716, - 3717, - 3732, - 3733, - 3737, - 3738, - 3739, - 3740, - 3741, - 3742, - 3743, - 3744, - 3745, - 3746, - 3747, - 3748, - 3749, - 3750, - 3751, - 3752, - 3753, - 3754, - 3755, - 3756, - 3757, - 3758, - 3759, - 3760, - 3761, - 3762, - 3763, - 3764, - 3765, - 3766, - 3767, - 3770, - 3771, - 3772, - 3773, - 3774, - 3775, - 3776, - 3777, - 3778, - 3779, - 3780, - 3781, - 3782, - 3783, - 3784, - 3785, - 3786, - 3787, - 3788, - 3789, - 3790, - 3791, - 3792, - 3793, - 3794, - 3795, - 3798, - 3799, - 3800, - 3801, - 3802, - 3803, - 3804, - 3805, - 3806, - 3807, - 3815, - 3816, - 3819, - 3820, - 3821, - 3822, - 3823, - 3824, - 3825, - 3826, - 3827, - 3828, - 3829, - 3830, - 3831, - 3832, - 3833, - 3834, - 3835, - 3836, - 3837, - 3838, - 3841, - 3842, - 3843, - 3844, - 3845, - 3846, - 3847, - 3848, - 3849, - 3850, - 3851, - 3852, - 3853, - 3854, - 3855, - 3856, - 3857, - 3858, - 3859, - 3860, - 3861, - 3862, - 3863, - 3864, - 3865, - 3866, - 3867, - 3868, - 3869, - 3870, - 3871, - 3872, - 3873, - 3874, - 3875, - 3876, - 3877, - 3878, - 3879, - 3880, - 3881, - 3882, - 3883, - 3884, - 3885, - 3886, - 3887, - 3888, - 3889, - 3890, - 3891, - 3892, - 3893, - 3894, - 3895, - 3896, - 3897, - 3898, - 3899, - 3900, - 3901, - 3902, - 3903, - 3904, - 3905, - 3906, - 3907, - 3908, - 3909, - 3910, - 3911, - 3912, - 3913, - 3914, - 3915, - 3916, - 3917, - 3922, - 3923, - 3924, - 3925, - 3926, - 3927, - 3928, - 3929, - 3930, - 3931, - 3932, - 3933, - 3936, - 3937, - 3938, - 3939, - 3940, - 3941, - 3945, - 3946, - 3947, - 3948, - 3949, - 3950, - 3951, - 3952, - 3953, - 3954, - 3955, - 3956, - 3957, - 3958, - 3959, - 3960, - 3961, - 3962, - 3963, - 3964, - 3965, - 3966, - 3967, - 3968, - 3969, - 3970, - 3971, - 3972, - 3973, - 3974, - 3975, - 3976, - 3977, - 3978, - 3979, - 3980, - 3981, - 3982, - 3983, - 3984, - 3985, - 3986, - 3987, - 3988, - 3989, - 3990, - 3991, - 3992, - 3993, - 3994, - 3995, - 3996, - 3997, - 3998, - 3999, - 4000, - 4001, - 4002, - 4003, - 4004, - 4005, - 4006, - 4007, - 4008, - 4009, - 4010, - 4011, - 4012, - 4013, - 4014, - 4015, - 4016, - 4017, - 4018, - 4019, - 4020, - 4021, - 4022, - 4023, - 4024, - 4025, - 4026, - 4027, - 4028, - 4029, - 4030, - 4031, - 4032, - 4033, - 4034, - 4035, - 4036, - 4037, - 4038, - 4039, - 4040, - 4041, - 4042, - 4043, - 4044, - 4045, - 4046, - 4047, - 4048, - 4049, - 4050, - 4051, - 4052, - 4053, - 4054, - 4055, - 4056, - 4057, - 4062, - 4063, - 4064, - 4065, - 4066, - 4067, - 4068, - 4069, - 4070, - 4071, - 5231, - 5232, - 5233, - 5235, - 5236, - 5237, - 5238, - 5239, - 5240, - 5241, - 5242, - 5243, - 5366, - 5367, - 5368, - 5369, - 6240, - 6241, - 6242, - 6243, - 6244, - 6245, - 6246, - 6247, - 6248, - 6249, - 6250, - 6251, - 6252, - 6253, - 6254, - 6255, - 6256, - 6257, - 6258, - 6259, - 6260, - 6261, - 6262, - 6263, - 6264, - 6265, - 6266, - 6267, - 6268, - 6269, - 6270, - 6271, - 6272, - 6275, - 6276, - 6277, - 6278, - 6279, - 6492, - 6493, - 6494, - 6495, - 6880, - 6881, - 6882, - 6883, - 6884, - 6885, - 6886, - 6887, - 6888, - 6889 - ], - "rightArm": [ - 4114, - 4115, - 4116, - 4117, - 4122, - 4125, - 4168, - 4171, - 4204, - 4205, - 4206, - 4207, - 4257, - 4258, - 4259, - 4260, - 4261, - 4262, - 4263, - 4264, - 4265, - 4266, - 4267, - 4268, - 4272, - 4273, - 4274, - 4275, - 4276, - 4277, - 4278, - 4279, - 4280, - 4281, - 4714, - 4715, - 4716, - 4717, - 4741, - 4742, - 4743, - 4744, - 4756, - 4763, - 4764, - 4790, - 4791, - 4794, - 4795, - 4816, - 4817, - 4818, - 4819, - 4830, - 4831, - 4832, - 4833, - 4849, - 4850, - 4851, - 4852, - 4853, - 4854, - 4855, - 4856, - 4857, - 4858, - 4859, - 4860, - 4861, - 4862, - 4863, - 4864, - 4865, - 4866, - 4867, - 4868, - 4869, - 4870, - 4871, - 4872, - 4873, - 4876, - 4877, - 4878, - 4879, - 4880, - 4881, - 4882, - 4883, - 4884, - 4885, - 4886, - 4887, - 4888, - 4889, - 4901, - 4902, - 4903, - 4904, - 4905, - 4906, - 4911, - 4912, - 4913, - 4914, - 4915, - 4916, - 4917, - 4918, - 4974, - 4977, - 4978, - 4979, - 4980, - 4981, - 4982, - 5009, - 5010, - 5011, - 5012, - 5014, - 5088, - 5089, - 5090, - 5091, - 5100, - 5101, - 5102, - 5103, - 5104, - 5105, - 5106, - 5107, - 5108, - 5109, - 5110, - 5111, - 5114, - 5115, - 5116, - 5117, - 5118, - 5119, - 5120, - 5121, - 5122, - 5123, - 5124, - 5125, - 5128, - 5129, - 5130, - 5131, - 5134, - 5135, - 5136, - 5137, - 5138, - 5139, - 5140, - 5141, - 5142, - 5143, - 5144, - 5145, - 5146, - 5147, - 5148, - 5149, - 5150, - 5151, - 5152, - 5153, - 5165, - 5166, - 5167, - 5172, - 5173, - 5174, - 5175, - 5176, - 5177, - 5178, - 5179, - 5180, - 5181, - 5182, - 5183, - 5184, - 5185, - 5186, - 5187, - 5188, - 5189, - 5194, - 5200, - 5201, - 5202, - 5203, - 5204, - 5206, - 5208, - 5209, - 5214, - 5215, - 5216, - 5217, - 5218, - 5220, - 5229, - 5292, - 5293, - 5303, - 5306, - 5309, - 5311, - 5314, - 5315, - 5318, - 5319, - 5321, - 5326, - 5327, - 5328, - 5330, - 5331, - 5332, - 5335, - 5336, - 5337, - 5338, - 5339, - 5343, - 5344, - 5349, - 5350, - 5353, - 5361, - 5362, - 5363, - 5364, - 5365, - 5370, - 6280, - 6281, - 6282, - 6283, - 6354, - 6355, - 6356, - 6357, - 6358, - 6359, - 6360, - 6361, - 6362, - 6404, - 6405, - 6433, - 6434, - 6435, - 6436, - 6437, - 6438, - 6439, - 6440, - 6441, - 6442, - 6443, - 6444, - 6445, - 6446, - 6447, - 6448, - 6449, - 6450, - 6451, - 6452, - 6453, - 6454, - 6455, - 6461, - 6471 - ], - "leftHandIndex1": [ - 2027, - 2028, - 2029, - 2030, - 2037, - 2038, - 2039, - 2040, - 2057, - 2067, - 2068, - 2123, - 2124, - 2125, - 2126, - 2127, - 2128, - 2129, - 2130, - 2132, - 2145, - 2146, - 2152, - 2153, - 2154, - 2156, - 2157, - 2158, - 2159, - 2160, - 2161, - 2162, - 2163, - 2164, - 2165, - 2166, - 2167, - 2168, - 2169, - 2177, - 2178, - 2179, - 2181, - 2186, - 2187, - 2190, - 2191, - 2204, - 2205, - 2215, - 2216, - 2217, - 2218, - 2219, - 2220, - 2232, - 2233, - 2245, - 2246, - 2247, - 2258, - 2259, - 2261, - 2262, - 2263, - 2269, - 2270, - 2272, - 2273, - 2274, - 2276, - 2277, - 2280, - 2281, - 2282, - 2283, - 2291, - 2292, - 2293, - 2294, - 2295, - 2296, - 2297, - 2298, - 2299, - 2300, - 2301, - 2302, - 2303, - 2304, - 2305, - 2306, - 2307, - 2308, - 2309, - 2310, - 2311, - 2312, - 2313, - 2314, - 2315, - 2316, - 2317, - 2318, - 2319, - 2320, - 2321, - 2322, - 2323, - 2324, - 2325, - 2326, - 2327, - 2328, - 2329, - 2330, - 2331, - 2332, - 2333, - 2334, - 2335, - 2336, - 2337, - 2338, - 2339, - 2340, - 2341, - 2342, - 2343, - 2344, - 2345, - 2346, - 2347, - 2348, - 2349, - 2350, - 2351, - 2352, - 2353, - 2354, - 2355, - 2356, - 2357, - 2358, - 2359, - 2360, - 2361, - 2362, - 2363, - 2364, - 2365, - 2366, - 2367, - 2368, - 2369, - 2370, - 2371, - 2372, - 2373, - 2374, - 2375, - 2376, - 2377, - 2378, - 2379, - 2380, - 2381, - 2382, - 2383, - 2384, - 2385, - 2386, - 2387, - 2388, - 2389, - 2390, - 2391, - 2392, - 2393, - 2394, - 2395, - 2396, - 2397, - 2398, - 2399, - 2400, - 2401, - 2402, - 2403, - 2404, - 2405, - 2406, - 2407, - 2408, - 2409, - 2410, - 2411, - 2412, - 2413, - 2414, - 2415, - 2416, - 2417, - 2418, - 2419, - 2420, - 2421, - 2422, - 2423, - 2424, - 2425, - 2426, - 2427, - 2428, - 2429, - 2430, - 2431, - 2432, - 2433, - 2434, - 2435, - 2436, - 2437, - 2438, - 2439, - 2440, - 2441, - 2442, - 2443, - 2444, - 2445, - 2446, - 2447, - 2448, - 2449, - 2450, - 2451, - 2452, - 2453, - 2454, - 2455, - 2456, - 2457, - 2458, - 2459, - 2460, - 2461, - 2462, - 2463, - 2464, - 2465, - 2466, - 2467, - 2468, - 2469, - 2470, - 2471, - 2472, - 2473, - 2474, - 2475, - 2476, - 2477, - 2478, - 2479, - 2480, - 2481, - 2482, - 2483, - 2484, - 2485, - 2486, - 2487, - 2488, - 2489, - 2490, - 2491, - 2492, - 2493, - 2494, - 2495, - 2496, - 2497, - 2498, - 2499, - 2500, - 2501, - 2502, - 2503, - 2504, - 2505, - 2506, - 2507, - 2508, - 2509, - 2510, - 2511, - 2512, - 2513, - 2514, - 2515, - 2516, - 2517, - 2518, - 2519, - 2520, - 2521, - 2522, - 2523, - 2524, - 2525, - 2526, - 2527, - 2528, - 2529, - 2530, - 2531, - 2532, - 2533, - 2534, - 2535, - 2536, - 2537, - 2538, - 2539, - 2540, - 2541, - 2542, - 2543, - 2544, - 2545, - 2546, - 2547, - 2548, - 2549, - 2550, - 2551, - 2552, - 2553, - 2554, - 2555, - 2556, - 2557, - 2558, - 2559, - 2560, - 2561, - 2562, - 2563, - 2564, - 2565, - 2566, - 2567, - 2568, - 2569, - 2570, - 2571, - 2572, - 2573, - 2574, - 2575, - 2576, - 2577, - 2578, - 2579, - 2580, - 2581, - 2582, - 2583, - 2584, - 2585, - 2586, - 2587, - 2588, - 2589, - 2590, - 2591, - 2592, - 2593, - 2594, - 2596, - 2597, - 2599, - 2600, - 2601, - 2602, - 2603, - 2604, - 2606, - 2607, - 2609, - 2610, - 2611, - 2612, - 2613, - 2614, - 2615, - 2616, - 2617, - 2618, - 2619, - 2620, - 2621, - 2622, - 2623, - 2624, - 2625, - 2626, - 2627, - 2628, - 2629, - 2630, - 2631, - 2632, - 2633, - 2634, - 2635, - 2636, - 2637, - 2638, - 2639, - 2640, - 2641, - 2642, - 2643, - 2644, - 2645, - 2646, - 2647, - 2648, - 2649, - 2650, - 2651, - 2652, - 2653, - 2654, - 2655, - 2656, - 2657, - 2658, - 2659, - 2660, - 2661, - 2662, - 2663, - 2664, - 2665, - 2666, - 2667, - 2668, - 2669, - 2670, - 2671, - 2672, - 2673, - 2674, - 2675, - 2676, - 2677, - 2678, - 2679, - 2680, - 2681, - 2682, - 2683, - 2684, - 2685, - 2686, - 2687, - 2688, - 2689, - 2690, - 2691, - 2692, - 2693, - 2694, - 2695, - 2696 - ], - "rightLeg": [ - 4481, - 4482, - 4485, - 4486, - 4491, - 4492, - 4493, - 4495, - 4498, - 4500, - 4501, - 4505, - 4506, - 4529, - 4532, - 4533, - 4534, - 4535, - 4536, - 4537, - 4538, - 4539, - 4540, - 4541, - 4542, - 4543, - 4544, - 4545, - 4546, - 4547, - 4548, - 4549, - 4550, - 4551, - 4552, - 4553, - 4554, - 4555, - 4556, - 4557, - 4558, - 4559, - 4560, - 4561, - 4562, - 4563, - 4564, - 4565, - 4566, - 4567, - 4568, - 4569, - 4570, - 4571, - 4572, - 4573, - 4574, - 4575, - 4576, - 4577, - 4578, - 4579, - 4580, - 4581, - 4582, - 4583, - 4584, - 4585, - 4586, - 4587, - 4588, - 4589, - 4590, - 4591, - 4592, - 4593, - 4594, - 4595, - 4596, - 4597, - 4598, - 4599, - 4600, - 4601, - 4602, - 4603, - 4604, - 4605, - 4606, - 4607, - 4608, - 4609, - 4610, - 4611, - 4612, - 4613, - 4614, - 4615, - 4616, - 4617, - 4618, - 4619, - 4620, - 4621, - 4622, - 4634, - 4635, - 4636, - 4637, - 4638, - 4639, - 4640, - 4641, - 4642, - 4643, - 4644, - 4661, - 4662, - 4663, - 4664, - 4665, - 4666, - 4667, - 4668, - 4669, - 4842, - 4843, - 4844, - 4845, - 4846, - 4847, - 4848, - 4937, - 4938, - 4939, - 4940, - 4941, - 4942, - 4943, - 4944, - 4945, - 4946, - 4947, - 4993, - 4994, - 4995, - 4996, - 4997, - 4998, - 4999, - 5000, - 5001, - 5002, - 5003, - 6574, - 6575, - 6576, - 6577, - 6578, - 6579, - 6580, - 6581, - 6582, - 6583, - 6584, - 6585, - 6586, - 6587, - 6588, - 6589, - 6590, - 6591, - 6592, - 6593, - 6594, - 6595, - 6596, - 6597, - 6598, - 6599, - 6600, - 6601, - 6602, - 6603, - 6604, - 6605, - 6606, - 6607, - 6608, - 6609, - 6610, - 6719, - 6720, - 6721, - 6722, - 6723, - 6724, - 6725, - 6726, - 6727, - 6728, - 6729, - 6730, - 6731, - 6732, - 6733, - 6734, - 6735, - 6832, - 6833, - 6834, - 6835, - 6836, - 6869, - 6870, - 6871, - 6872 - ], - "rightHandIndex1": [ - 5488, - 5489, - 5490, - 5491, - 5498, - 5499, - 5500, - 5501, - 5518, - 5528, - 5529, - 5584, - 5585, - 5586, - 5587, - 5588, - 5589, - 5590, - 5591, - 5592, - 5606, - 5607, - 5613, - 5615, - 5616, - 5617, - 5618, - 5619, - 5620, - 5621, - 5622, - 5623, - 5624, - 5625, - 5626, - 5627, - 5628, - 5629, - 5630, - 5638, - 5639, - 5640, - 5642, - 5647, - 5648, - 5650, - 5651, - 5665, - 5666, - 5676, - 5677, - 5678, - 5679, - 5680, - 5681, - 5693, - 5694, - 5706, - 5707, - 5708, - 5719, - 5721, - 5722, - 5723, - 5724, - 5730, - 5731, - 5733, - 5734, - 5735, - 5737, - 5738, - 5741, - 5742, - 5743, - 5744, - 5752, - 5753, - 5754, - 5755, - 5756, - 5757, - 5758, - 5759, - 5760, - 5761, - 5762, - 5763, - 5764, - 5765, - 5766, - 5767, - 5768, - 5769, - 5770, - 5771, - 5772, - 5773, - 5774, - 5775, - 5776, - 5777, - 5778, - 5779, - 5780, - 5781, - 5782, - 5783, - 5784, - 5785, - 5786, - 5787, - 5788, - 5789, - 5790, - 5791, - 5792, - 5793, - 5794, - 5795, - 5796, - 5797, - 5798, - 5799, - 5800, - 5801, - 5802, - 5803, - 5804, - 5805, - 5806, - 5807, - 5808, - 5809, - 5810, - 5811, - 5812, - 5813, - 5814, - 5815, - 5816, - 5817, - 5818, - 5819, - 5820, - 5821, - 5822, - 5823, - 5824, - 5825, - 5826, - 5827, - 5828, - 5829, - 5830, - 5831, - 5832, - 5833, - 5834, - 5835, - 5836, - 5837, - 5838, - 5839, - 5840, - 5841, - 5842, - 5843, - 5844, - 5845, - 5846, - 5847, - 5848, - 5849, - 5850, - 5851, - 5852, - 5853, - 5854, - 5855, - 5856, - 5857, - 5858, - 5859, - 5860, - 5861, - 5862, - 5863, - 5864, - 5865, - 5866, - 5867, - 5868, - 5869, - 5870, - 5871, - 5872, - 5873, - 5874, - 5875, - 5876, - 5877, - 5878, - 5879, - 5880, - 5881, - 5882, - 5883, - 5884, - 5885, - 5886, - 5887, - 5888, - 5889, - 5890, - 5891, - 5892, - 5893, - 5894, - 5895, - 5896, - 5897, - 5898, - 5899, - 5900, - 5901, - 5902, - 5903, - 5904, - 5905, - 5906, - 5907, - 5908, - 5909, - 5910, - 5911, - 5912, - 5913, - 5914, - 5915, - 5916, - 5917, - 5918, - 5919, - 5920, - 5921, - 5922, - 5923, - 5924, - 5925, - 5926, - 5927, - 5928, - 5929, - 5930, - 5931, - 5932, - 5933, - 5934, - 5935, - 5936, - 5937, - 5938, - 5939, - 5940, - 5941, - 5942, - 5943, - 5944, - 5945, - 5946, - 5947, - 5948, - 5949, - 5950, - 5951, - 5952, - 5953, - 5954, - 5955, - 5956, - 5957, - 5958, - 5959, - 5960, - 5961, - 5962, - 5963, - 5964, - 5965, - 5966, - 5967, - 5968, - 5969, - 5970, - 5971, - 5972, - 5973, - 5974, - 5975, - 5976, - 5977, - 5978, - 5979, - 5980, - 5981, - 5982, - 5983, - 5984, - 5985, - 5986, - 5987, - 5988, - 5989, - 5990, - 5991, - 5992, - 5993, - 5994, - 5995, - 5996, - 5997, - 5998, - 5999, - 6000, - 6001, - 6002, - 6003, - 6004, - 6005, - 6006, - 6007, - 6008, - 6009, - 6010, - 6011, - 6012, - 6013, - 6014, - 6015, - 6016, - 6017, - 6018, - 6019, - 6020, - 6021, - 6022, - 6023, - 6024, - 6025, - 6026, - 6027, - 6028, - 6029, - 6030, - 6031, - 6032, - 6033, - 6034, - 6035, - 6036, - 6037, - 6038, - 6039, - 6040, - 6041, - 6042, - 6043, - 6044, - 6045, - 6046, - 6047, - 6048, - 6049, - 6050, - 6051, - 6052, - 6053, - 6054, - 6055, - 6058, - 6059, - 6060, - 6061, - 6062, - 6063, - 6064, - 6065, - 6068, - 6069, - 6070, - 6071, - 6072, - 6073, - 6074, - 6075, - 6076, - 6077, - 6078, - 6079, - 6080, - 6081, - 6082, - 6083, - 6084, - 6085, - 6086, - 6087, - 6088, - 6089, - 6090, - 6091, - 6092, - 6093, - 6094, - 6095, - 6096, - 6097, - 6098, - 6099, - 6100, - 6101, - 6102, - 6103, - 6104, - 6105, - 6106, - 6107, - 6108, - 6109, - 6110, - 6111, - 6112, - 6113, - 6114, - 6115, - 6116, - 6117, - 6118, - 6119, - 6120, - 6121, - 6122, - 6123, - 6124, - 6125, - 6126, - 6127, - 6128, - 6129, - 6130, - 6131, - 6132, - 6133, - 6134, - 6135, - 6136, - 6137, - 6138, - 6139, - 6140, - 6141, - 6142, - 6143, - 6144, - 6145, - 6146, - 6147, - 6148, - 6149, - 6150, - 6151, - 6152, - 6153, - 6154, - 6155, - 6156, - 6157 - ], - "leftForeArm": [ - 1546, - 1547, - 1548, - 1549, - 1550, - 1551, - 1552, - 1553, - 1554, - 1555, - 1556, - 1557, - 1558, - 1559, - 1560, - 1561, - 1562, - 1563, - 1564, - 1565, - 1566, - 1567, - 1568, - 1569, - 1570, - 1571, - 1572, - 1573, - 1574, - 1575, - 1576, - 1577, - 1578, - 1579, - 1580, - 1581, - 1582, - 1583, - 1584, - 1585, - 1586, - 1587, - 1588, - 1589, - 1590, - 1591, - 1592, - 1593, - 1594, - 1595, - 1596, - 1597, - 1598, - 1599, - 1600, - 1601, - 1602, - 1603, - 1604, - 1605, - 1606, - 1607, - 1608, - 1609, - 1610, - 1611, - 1612, - 1613, - 1614, - 1615, - 1616, - 1617, - 1618, - 1620, - 1621, - 1623, - 1624, - 1625, - 1626, - 1627, - 1628, - 1629, - 1630, - 1643, - 1644, - 1646, - 1647, - 1650, - 1651, - 1654, - 1655, - 1657, - 1658, - 1659, - 1660, - 1661, - 1662, - 1663, - 1664, - 1665, - 1666, - 1685, - 1686, - 1687, - 1688, - 1689, - 1690, - 1691, - 1692, - 1693, - 1694, - 1695, - 1699, - 1700, - 1701, - 1702, - 1721, - 1722, - 1723, - 1724, - 1725, - 1726, - 1727, - 1728, - 1729, - 1730, - 1732, - 1736, - 1738, - 1741, - 1742, - 1743, - 1744, - 1750, - 1752, - 1900, - 1909, - 1910, - 1911, - 1912, - 1913, - 1914, - 1915, - 1916, - 1917, - 1918, - 1919, - 1920, - 1921, - 1922, - 1923, - 1924, - 1925, - 1926, - 1927, - 1928, - 1929, - 1930, - 1931, - 1932, - 1933, - 1934, - 1935, - 1936, - 1937, - 1938, - 1939, - 1940, - 1941, - 1942, - 1943, - 1944, - 1945, - 1946, - 1947, - 1948, - 1949, - 1950, - 1951, - 1952, - 1953, - 1954, - 1955, - 1956, - 1957, - 1958, - 1959, - 1960, - 1961, - 1962, - 1963, - 1964, - 1965, - 1966, - 1967, - 1968, - 1969, - 1970, - 1971, - 1972, - 1973, - 1974, - 1975, - 1976, - 1977, - 1978, - 1979, - 1980, - 2019, - 2059, - 2060, - 2073, - 2089, - 2098, - 2099, - 2100, - 2101, - 2102, - 2103, - 2104, - 2105, - 2106, - 2107, - 2108, - 2109, - 2110, - 2111, - 2112, - 2147, - 2148, - 2206, - 2207, - 2208, - 2209, - 2228, - 2230, - 2234, - 2235, - 2241, - 2242, - 2243, - 2244, - 2279, - 2286, - 2873, - 2874 - ], - "rightForeArm": [ - 5015, - 5016, - 5017, - 5018, - 5019, - 5020, - 5021, - 5022, - 5023, - 5024, - 5025, - 5026, - 5027, - 5028, - 5029, - 5030, - 5031, - 5032, - 5033, - 5034, - 5035, - 5036, - 5037, - 5038, - 5039, - 5040, - 5041, - 5042, - 5043, - 5044, - 5045, - 5046, - 5047, - 5048, - 5049, - 5050, - 5051, - 5052, - 5053, - 5054, - 5055, - 5056, - 5057, - 5058, - 5059, - 5060, - 5061, - 5062, - 5063, - 5064, - 5065, - 5066, - 5067, - 5068, - 5069, - 5070, - 5071, - 5072, - 5073, - 5074, - 5075, - 5076, - 5077, - 5078, - 5079, - 5080, - 5081, - 5082, - 5083, - 5084, - 5085, - 5086, - 5087, - 5090, - 5091, - 5092, - 5093, - 5094, - 5095, - 5096, - 5097, - 5098, - 5099, - 5112, - 5113, - 5116, - 5117, - 5120, - 5121, - 5124, - 5125, - 5126, - 5127, - 5128, - 5129, - 5130, - 5131, - 5132, - 5133, - 5134, - 5135, - 5154, - 5155, - 5156, - 5157, - 5158, - 5159, - 5160, - 5161, - 5162, - 5163, - 5164, - 5168, - 5169, - 5170, - 5171, - 5190, - 5191, - 5192, - 5193, - 5194, - 5195, - 5196, - 5197, - 5198, - 5199, - 5202, - 5205, - 5207, - 5210, - 5211, - 5212, - 5213, - 5219, - 5221, - 5361, - 5370, - 5371, - 5372, - 5373, - 5374, - 5375, - 5376, - 5377, - 5378, - 5379, - 5380, - 5381, - 5382, - 5383, - 5384, - 5385, - 5386, - 5387, - 5388, - 5389, - 5390, - 5391, - 5392, - 5393, - 5394, - 5395, - 5396, - 5397, - 5398, - 5399, - 5400, - 5401, - 5402, - 5403, - 5404, - 5405, - 5406, - 5407, - 5408, - 5409, - 5410, - 5411, - 5412, - 5413, - 5414, - 5415, - 5416, - 5417, - 5418, - 5419, - 5420, - 5421, - 5422, - 5423, - 5424, - 5425, - 5426, - 5427, - 5428, - 5429, - 5430, - 5431, - 5432, - 5433, - 5434, - 5435, - 5436, - 5437, - 5438, - 5439, - 5440, - 5441, - 5480, - 5520, - 5521, - 5534, - 5550, - 5559, - 5560, - 5561, - 5562, - 5563, - 5564, - 5565, - 5566, - 5567, - 5568, - 5569, - 5570, - 5571, - 5572, - 5573, - 5608, - 5609, - 5667, - 5668, - 5669, - 5670, - 5689, - 5691, - 5695, - 5696, - 5702, - 5703, - 5704, - 5705, - 5740, - 5747, - 6334, - 6335 - ], - "neck": [ - 148, - 150, - 151, - 152, - 153, - 172, - 174, - 175, - 201, - 202, - 204, - 205, - 206, - 207, - 208, - 209, - 210, - 211, - 212, - 213, - 214, - 215, - 216, - 217, - 218, - 219, - 222, - 223, - 224, - 225, - 256, - 257, - 284, - 285, - 295, - 296, - 297, - 298, - 299, - 300, - 301, - 302, - 303, - 304, - 305, - 306, - 307, - 308, - 309, - 333, - 334, - 423, - 424, - 425, - 426, - 440, - 441, - 451, - 452, - 453, - 460, - 461, - 571, - 572, - 824, - 825, - 826, - 827, - 828, - 829, - 1279, - 1280, - 1312, - 1313, - 1319, - 1320, - 1331, - 3049, - 3050, - 3057, - 3058, - 3059, - 3068, - 3164, - 3661, - 3662, - 3663, - 3664, - 3665, - 3685, - 3686, - 3687, - 3714, - 3715, - 3716, - 3717, - 3718, - 3719, - 3720, - 3721, - 3722, - 3723, - 3724, - 3725, - 3726, - 3727, - 3728, - 3729, - 3730, - 3731, - 3734, - 3735, - 3736, - 3737, - 3768, - 3769, - 3796, - 3797, - 3807, - 3808, - 3809, - 3810, - 3811, - 3812, - 3813, - 3814, - 3815, - 3816, - 3817, - 3818, - 3819, - 3839, - 3840, - 3918, - 3919, - 3920, - 3921, - 3934, - 3935, - 3942, - 3943, - 3944, - 3950, - 4060, - 4061, - 4312, - 4313, - 4314, - 4315, - 4761, - 4762, - 4792, - 4793, - 4799, - 4800, - 4807 - ], - "rightToeBase": [ - 6611, - 6612, - 6613, - 6614, - 6615, - 6616, - 6617, - 6618, - 6619, - 6620, - 6621, - 6622, - 6623, - 6624, - 6625, - 6626, - 6627, - 6628, - 6629, - 6630, - 6631, - 6632, - 6633, - 6634, - 6635, - 6636, - 6637, - 6638, - 6639, - 6640, - 6641, - 6642, - 6643, - 6644, - 6645, - 6646, - 6647, - 6648, - 6649, - 6650, - 6651, - 6652, - 6653, - 6654, - 6655, - 6656, - 6657, - 6658, - 6659, - 6660, - 6661, - 6662, - 6663, - 6664, - 6665, - 6666, - 6667, - 6668, - 6669, - 6670, - 6671, - 6672, - 6673, - 6674, - 6675, - 6676, - 6677, - 6678, - 6679, - 6680, - 6681, - 6682, - 6683, - 6684, - 6685, - 6686, - 6687, - 6688, - 6689, - 6690, - 6691, - 6692, - 6693, - 6694, - 6695, - 6696, - 6697, - 6698, - 6699, - 6700, - 6701, - 6702, - 6703, - 6704, - 6705, - 6706, - 6707, - 6708, - 6709, - 6710, - 6711, - 6712, - 6713, - 6714, - 6715, - 6716, - 6717, - 6718, - 6736, - 6739, - 6741, - 6743, - 6745, - 6747, - 6749, - 6750, - 6752, - 6754, - 6757, - 6758, - 6760, - 6762 - ], - "spine": [ - 616, - 617, - 630, - 631, - 632, - 633, - 654, - 655, - 656, - 657, - 662, - 663, - 664, - 665, - 720, - 721, - 765, - 766, - 767, - 768, - 796, - 797, - 798, - 799, - 889, - 890, - 916, - 917, - 918, - 919, - 921, - 922, - 923, - 924, - 925, - 926, - 1188, - 1189, - 1211, - 1212, - 1248, - 1249, - 1250, - 1251, - 1264, - 1265, - 1266, - 1267, - 1323, - 1324, - 1325, - 1326, - 1327, - 1328, - 1332, - 1333, - 1334, - 1335, - 1336, - 1344, - 1345, - 1481, - 1482, - 1483, - 1484, - 1485, - 1486, - 1487, - 1488, - 1489, - 1490, - 1491, - 1492, - 1493, - 1494, - 1495, - 1496, - 1767, - 2823, - 2824, - 2825, - 2826, - 2827, - 2828, - 2829, - 2830, - 2831, - 2832, - 2833, - 2834, - 2835, - 2836, - 2837, - 2838, - 2839, - 2840, - 2841, - 2842, - 2843, - 2844, - 2845, - 2847, - 2848, - 2851, - 3016, - 3017, - 3018, - 3019, - 3020, - 3023, - 3024, - 3124, - 3173, - 3476, - 3477, - 3478, - 3480, - 3500, - 3501, - 3502, - 3504, - 3509, - 3511, - 4103, - 4104, - 4118, - 4119, - 4120, - 4121, - 4142, - 4143, - 4144, - 4145, - 4150, - 4151, - 4152, - 4153, - 4208, - 4209, - 4253, - 4254, - 4255, - 4256, - 4284, - 4285, - 4286, - 4287, - 4375, - 4376, - 4402, - 4403, - 4405, - 4406, - 4407, - 4408, - 4409, - 4410, - 4411, - 4412, - 4674, - 4675, - 4694, - 4695, - 4731, - 4732, - 4733, - 4734, - 4747, - 4748, - 4749, - 4750, - 4803, - 4804, - 4805, - 4806, - 4808, - 4809, - 4810, - 4811, - 4812, - 4820, - 4821, - 4953, - 4954, - 4955, - 4956, - 4957, - 4958, - 4959, - 4960, - 4961, - 4962, - 4963, - 4964, - 4965, - 4966, - 4967, - 4968, - 5234, - 6284, - 6285, - 6286, - 6287, - 6288, - 6289, - 6290, - 6291, - 6292, - 6293, - 6294, - 6295, - 6296, - 6297, - 6298, - 6299, - 6300, - 6301, - 6302, - 6303, - 6304, - 6305, - 6306, - 6308, - 6309, - 6312, - 6472, - 6473, - 6474, - 6545, - 6874, - 6875, - 6876, - 6878 - ], - "leftUpLeg": [ - 833, - 834, - 838, - 839, - 847, - 848, - 849, - 850, - 851, - 852, - 853, - 854, - 870, - 871, - 872, - 873, - 874, - 875, - 876, - 877, - 878, - 879, - 880, - 881, - 897, - 898, - 899, - 900, - 901, - 902, - 903, - 904, - 905, - 906, - 907, - 908, - 909, - 910, - 911, - 912, - 913, - 914, - 915, - 933, - 934, - 935, - 936, - 944, - 945, - 946, - 947, - 948, - 949, - 950, - 951, - 952, - 953, - 954, - 955, - 956, - 957, - 958, - 959, - 960, - 961, - 962, - 963, - 964, - 965, - 966, - 967, - 968, - 969, - 970, - 971, - 972, - 973, - 974, - 975, - 976, - 977, - 978, - 979, - 980, - 981, - 982, - 983, - 984, - 985, - 986, - 987, - 988, - 989, - 990, - 991, - 992, - 993, - 994, - 995, - 996, - 997, - 998, - 999, - 1000, - 1001, - 1002, - 1003, - 1004, - 1005, - 1006, - 1007, - 1008, - 1009, - 1010, - 1011, - 1012, - 1013, - 1014, - 1015, - 1016, - 1017, - 1018, - 1019, - 1020, - 1021, - 1022, - 1023, - 1024, - 1025, - 1026, - 1027, - 1028, - 1029, - 1030, - 1031, - 1032, - 1033, - 1034, - 1035, - 1036, - 1037, - 1038, - 1039, - 1040, - 1041, - 1042, - 1043, - 1044, - 1045, - 1046, - 1137, - 1138, - 1139, - 1140, - 1141, - 1142, - 1143, - 1144, - 1145, - 1146, - 1147, - 1148, - 1159, - 1160, - 1161, - 1162, - 1163, - 1164, - 1165, - 1166, - 1167, - 1168, - 1169, - 1170, - 1171, - 1172, - 1173, - 1174, - 1184, - 1185, - 1186, - 1187, - 1221, - 1222, - 1223, - 1224, - 1225, - 1226, - 1227, - 1228, - 1229, - 1230, - 1262, - 1263, - 1274, - 1275, - 1276, - 1277, - 1321, - 1322, - 1354, - 1359, - 1360, - 1361, - 1362, - 1365, - 1366, - 1367, - 1368, - 1451, - 1452, - 1453, - 1455, - 1456, - 1457, - 1458, - 1459, - 1460, - 1461, - 1462, - 1463, - 1475, - 1477, - 1478, - 1479, - 1480, - 1498, - 1499, - 1500, - 1501, - 1511, - 1512, - 1513, - 1514, - 1516, - 1517, - 1518, - 1519, - 1520, - 1521, - 1522, - 1533, - 1534, - 3125, - 3126, - 3127, - 3128, - 3131, - 3132, - 3133, - 3134, - 3135, - 3475, - 3479 - ], - "leftHand": [ - 1981, - 1982, - 1983, - 1984, - 1985, - 1986, - 1987, - 1988, - 1989, - 1990, - 1991, - 1992, - 1993, - 1994, - 1995, - 1996, - 1997, - 1998, - 1999, - 2000, - 2001, - 2002, - 2003, - 2004, - 2005, - 2006, - 2007, - 2008, - 2009, - 2010, - 2011, - 2012, - 2013, - 2014, - 2015, - 2016, - 2017, - 2018, - 2019, - 2020, - 2021, - 2022, - 2023, - 2024, - 2025, - 2026, - 2031, - 2032, - 2033, - 2034, - 2035, - 2036, - 2041, - 2042, - 2043, - 2044, - 2045, - 2046, - 2047, - 2048, - 2049, - 2050, - 2051, - 2052, - 2053, - 2054, - 2055, - 2056, - 2057, - 2058, - 2059, - 2060, - 2061, - 2062, - 2063, - 2064, - 2065, - 2066, - 2069, - 2070, - 2071, - 2072, - 2073, - 2074, - 2075, - 2076, - 2077, - 2078, - 2079, - 2080, - 2081, - 2082, - 2083, - 2084, - 2085, - 2086, - 2087, - 2088, - 2089, - 2090, - 2091, - 2092, - 2093, - 2094, - 2095, - 2096, - 2097, - 2098, - 2099, - 2100, - 2101, - 2107, - 2111, - 2113, - 2114, - 2115, - 2116, - 2117, - 2118, - 2119, - 2120, - 2121, - 2122, - 2127, - 2130, - 2131, - 2132, - 2133, - 2134, - 2135, - 2136, - 2137, - 2138, - 2139, - 2140, - 2141, - 2142, - 2143, - 2144, - 2149, - 2150, - 2151, - 2152, - 2155, - 2160, - 2163, - 2164, - 2170, - 2171, - 2172, - 2173, - 2174, - 2175, - 2176, - 2177, - 2178, - 2179, - 2180, - 2182, - 2183, - 2184, - 2185, - 2188, - 2189, - 2191, - 2192, - 2193, - 2194, - 2195, - 2196, - 2197, - 2198, - 2199, - 2200, - 2201, - 2202, - 2203, - 2207, - 2209, - 2210, - 2211, - 2212, - 2213, - 2214, - 2221, - 2222, - 2223, - 2224, - 2225, - 2226, - 2227, - 2228, - 2229, - 2231, - 2234, - 2236, - 2237, - 2238, - 2239, - 2240, - 2246, - 2247, - 2248, - 2249, - 2250, - 2251, - 2252, - 2253, - 2254, - 2255, - 2256, - 2257, - 2258, - 2259, - 2260, - 2262, - 2263, - 2264, - 2265, - 2266, - 2267, - 2268, - 2269, - 2270, - 2271, - 2274, - 2275, - 2276, - 2277, - 2278, - 2279, - 2284, - 2285, - 2287, - 2288, - 2289, - 2290, - 2293, - 2595, - 2598, - 2605, - 2608, - 2697, - 2698, - 2699, - 2700, - 2701, - 2702, - 2703, - 2704, - 2705, - 2706, - 2707, - 2708, - 2709, - 2710, - 2711, - 2712, - 2713, - 2714, - 2715, - 2716, - 2717, - 2718, - 2719, - 2720, - 2721, - 2722, - 2723, - 2724, - 2725, - 2726, - 2727, - 2728, - 2729, - 2730, - 2731, - 2732, - 2733, - 2734, - 2735, - 2736, - 2737, - 2738, - 2739, - 2740, - 2741, - 2742, - 2743, - 2744, - 2745, - 2746, - 2747, - 2748, - 2749, - 2750, - 2751, - 2752, - 2753, - 2754, - 2755, - 2756, - 2757, - 2758, - 2759, - 2760, - 2761, - 2762, - 2763, - 2764, - 2765, - 2766, - 2767, - 2768, - 2769, - 2770, - 2771, - 2772, - 2773, - 2774, - 2775, - 2776, - 2777, - 2778 - ], - "hips": [ - 631, - 632, - 654, - 657, - 662, - 665, - 676, - 677, - 678, - 679, - 705, - 720, - 796, - 799, - 800, - 801, - 802, - 807, - 808, - 809, - 810, - 815, - 816, - 822, - 823, - 830, - 831, - 832, - 833, - 834, - 835, - 836, - 837, - 838, - 839, - 840, - 841, - 842, - 843, - 844, - 845, - 846, - 855, - 856, - 857, - 858, - 859, - 860, - 861, - 862, - 863, - 864, - 865, - 866, - 867, - 868, - 869, - 871, - 878, - 881, - 882, - 883, - 884, - 885, - 886, - 887, - 888, - 889, - 890, - 912, - 915, - 916, - 917, - 918, - 919, - 920, - 932, - 937, - 938, - 939, - 1163, - 1166, - 1203, - 1204, - 1205, - 1206, - 1207, - 1208, - 1209, - 1210, - 1246, - 1247, - 1262, - 1263, - 1276, - 1277, - 1278, - 1321, - 1336, - 1337, - 1338, - 1339, - 1353, - 1354, - 1361, - 1362, - 1363, - 1364, - 1446, - 1447, - 1448, - 1449, - 1450, - 1454, - 1476, - 1497, - 1511, - 1513, - 1514, - 1515, - 1533, - 1534, - 1539, - 1540, - 1768, - 1769, - 1779, - 1780, - 1781, - 1782, - 1783, - 1784, - 1785, - 1786, - 1787, - 1788, - 1789, - 1790, - 1791, - 1792, - 1793, - 1794, - 1795, - 1796, - 1797, - 1798, - 1799, - 1800, - 1801, - 1802, - 1803, - 1804, - 1805, - 1806, - 1807, - 2909, - 2910, - 2911, - 2912, - 2913, - 2914, - 2915, - 2916, - 2917, - 2918, - 2919, - 2920, - 2921, - 2922, - 2923, - 2924, - 2925, - 2926, - 2927, - 2928, - 2929, - 2930, - 3018, - 3019, - 3021, - 3022, - 3080, - 3081, - 3082, - 3083, - 3084, - 3085, - 3086, - 3087, - 3088, - 3089, - 3090, - 3091, - 3092, - 3093, - 3094, - 3095, - 3096, - 3097, - 3098, - 3099, - 3100, - 3101, - 3102, - 3103, - 3104, - 3105, - 3106, - 3107, - 3108, - 3109, - 3110, - 3111, - 3112, - 3113, - 3114, - 3115, - 3116, - 3117, - 3118, - 3119, - 3120, - 3121, - 3122, - 3123, - 3124, - 3128, - 3129, - 3130, - 3136, - 3137, - 3138, - 3139, - 3140, - 3141, - 3142, - 3143, - 3144, - 3145, - 3146, - 3147, - 3148, - 3149, - 3150, - 3151, - 3152, - 3153, - 3154, - 3155, - 3156, - 3157, - 3158, - 3159, - 3160, - 3170, - 3172, - 3481, - 3484, - 3500, - 3502, - 3503, - 3507, - 3510, - 4120, - 4121, - 4142, - 4143, - 4150, - 4151, - 4164, - 4165, - 4166, - 4167, - 4193, - 4208, - 4284, - 4285, - 4288, - 4289, - 4290, - 4295, - 4296, - 4297, - 4298, - 4303, - 4304, - 4310, - 4311, - 4316, - 4317, - 4318, - 4319, - 4320, - 4321, - 4322, - 4323, - 4324, - 4325, - 4326, - 4327, - 4328, - 4329, - 4330, - 4331, - 4332, - 4341, - 4342, - 4343, - 4344, - 4345, - 4346, - 4347, - 4348, - 4349, - 4350, - 4351, - 4352, - 4353, - 4354, - 4355, - 4356, - 4364, - 4365, - 4368, - 4369, - 4370, - 4371, - 4372, - 4373, - 4374, - 4375, - 4376, - 4398, - 4399, - 4402, - 4403, - 4404, - 4405, - 4406, - 4418, - 4423, - 4424, - 4425, - 4649, - 4650, - 4689, - 4690, - 4691, - 4692, - 4693, - 4729, - 4730, - 4745, - 4746, - 4759, - 4760, - 4801, - 4812, - 4813, - 4814, - 4815, - 4829, - 4836, - 4837, - 4919, - 4920, - 4921, - 4922, - 4923, - 4927, - 4969, - 4983, - 4984, - 4986, - 5004, - 5005, - 5244, - 5245, - 5246, - 5247, - 5248, - 5249, - 5250, - 5251, - 5252, - 5253, - 5254, - 5255, - 5256, - 5257, - 5258, - 5259, - 5260, - 5261, - 5262, - 5263, - 5264, - 5265, - 5266, - 5267, - 5268, - 6368, - 6369, - 6370, - 6371, - 6372, - 6373, - 6374, - 6375, - 6376, - 6377, - 6378, - 6379, - 6380, - 6381, - 6382, - 6383, - 6384, - 6385, - 6386, - 6387, - 6388, - 6389, - 6473, - 6474, - 6504, - 6505, - 6506, - 6507, - 6508, - 6509, - 6510, - 6511, - 6512, - 6513, - 6514, - 6515, - 6516, - 6517, - 6518, - 6519, - 6520, - 6521, - 6522, - 6523, - 6524, - 6525, - 6526, - 6527, - 6528, - 6529, - 6530, - 6531, - 6532, - 6533, - 6534, - 6535, - 6536, - 6537, - 6538, - 6539, - 6540, - 6541, - 6542, - 6543, - 6544, - 6545, - 6549, - 6550, - 6551, - 6557, - 6558, - 6559, - 6560, - 6561, - 6562, - 6563, - 6564, - 6565, - 6566, - 6567, - 6568, - 6569, - 6570, - 6571, - 6572, - 6573 - ] -} \ No newline at end of file diff --git a/lib /common /train_util.py b/lib /common /train_util.py deleted file mode 100644 index 24af5a05f70588a65cb26658ddd3ea87430e8c44..0000000000000000000000000000000000000000 --- a/lib /common /train_util.py +++ /dev/null @@ -1,594 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import yaml -import os.path as osp -import torch -import numpy as np -import torch.nn.functional as F -from ..dataset.mesh_util import * -from ..net.geometry import orthogonal -from pytorch3d.renderer.mesh import rasterize_meshes -from .render_utils import Pytorch3dRasterizer -from pytorch3d.structures import Meshes -import cv2 -from PIL import Image -from tqdm import tqdm -import os -from termcolor import colored - - - - -def reshape_sample_tensor(sample_tensor, num_views): - if num_views == 1: - return sample_tensor - # Need to repeat sample_tensor along the batch dim num_views times - sample_tensor = sample_tensor.unsqueeze(dim=1) - sample_tensor = sample_tensor.repeat(1, num_views, 1, 1) - sample_tensor = sample_tensor.view( - sample_tensor.shape[0] * sample_tensor.shape[1], - sample_tensor.shape[2], sample_tensor.shape[3]) - return sample_tensor - - -def gen_mesh_eval(opt, net, cuda, data, resolution=None): - resolution = opt.resolution if resolution is None else resolution - image_tensor = data['img'].to(device=cuda) - calib_tensor = data['calib'].to(device=cuda) - - net.filter(image_tensor) - - b_min = data['b_min'] - b_max = data['b_max'] - try: - verts, faces, _, _ = reconstruction_faster(net, - cuda, - calib_tensor, - resolution, - b_min, - b_max, - use_octree=False) - - except Exception as e: - print(e) - print('Can not create marching cubes at this time.') - verts, faces = None, None - return verts, faces - - -def gen_mesh(opt, net, cuda, data, save_path, resolution=None): - resolution = opt.resolution if resolution is None else resolution - image_tensor = data['img'].to(device=cuda) - calib_tensor = data['calib'].to(device=cuda) - - net.filter(image_tensor) - - b_min = data['b_min'] - b_max = data['b_max'] - try: - save_img_path = save_path[:-4] + '.png' - save_img_list = [] - for v in range(image_tensor.shape[0]): - save_img = (np.transpose(image_tensor[v].detach().cpu().numpy(), - (1, 2, 0)) * 0.5 + - 0.5)[:, :, ::-1] * 255.0 - save_img_list.append(save_img) - save_img = np.concatenate(save_img_list, axis=1) - Image.fromarray(np.uint8(save_img[:, :, ::-1])).save(save_img_path) - - verts, faces, _, _ = reconstruction_faster(net, cuda, calib_tensor, - resolution, b_min, b_max) - verts_tensor = torch.from_numpy( - verts.T).unsqueeze(0).to(device=cuda).float() - xyz_tensor = net.projection(verts_tensor, calib_tensor[:1]) - uv = xyz_tensor[:, :2, :] - color = netG.index(image_tensor[:1], uv).detach().cpu().numpy()[0].T - color = color * 0.5 + 0.5 - save_obj_mesh_with_color(save_path, verts, faces, color) - except Exception as e: - print(e) - print('Can not create marching cubes at this time.') - verts, faces, color = None, None, None - return verts, faces, color - - -def gen_mesh_color(opt, netG, netC, cuda, data, save_path, use_octree=True): - image_tensor = data['img'].to(device=cuda) - calib_tensor = data['calib'].to(device=cuda) - - netG.filter(image_tensor) - netC.filter(image_tensor) - netC.attach(netG.get_im_feat()) - - b_min = data['b_min'] - b_max = data['b_max'] - try: - save_img_path = save_path[:-4] + '.png' - save_img_list = [] - for v in range(image_tensor.shape[0]): - save_img = (np.transpose(image_tensor[v].detach().cpu().numpy(), - (1, 2, 0)) * 0.5 + - 0.5)[:, :, ::-1] * 255.0 - save_img_list.append(save_img) - save_img = np.concatenate(save_img_list, axis=1) - Image.fromarray(np.uint8(save_img[:, :, ::-1])).save(save_img_path) - - verts, faces, _, _ = reconstruction_faster(netG, - cuda, - calib_tensor, - opt.resolution, - b_min, - b_max, - use_octree=use_octree) - - # Now Getting colors - verts_tensor = torch.from_numpy( - verts.T).unsqueeze(0).to(device=cuda).float() - verts_tensor = reshape_sample_tensor(verts_tensor, opt.num_views) - color = np.zeros(verts.shape) - interval = 10000 - for i in range(len(color) // interval): - left = i * interval - right = i * interval + interval - if i == len(color) // interval - 1: - right = -1 - netC.query(verts_tensor[:, :, left:right], calib_tensor) - rgb = netC.get_preds()[0].detach().cpu().numpy() * 0.5 + 0.5 - color[left:right] = rgb.T - - save_obj_mesh_with_color(save_path, verts, faces, color) - except Exception as e: - print(e) - print('Can not create marching cubes at this time.') - verts, faces, color = None, None, None - return verts, faces, color - - -def adjust_learning_rate(optimizer, epoch, lr, schedule, gamma): - """Sets the learning rate to the initial LR decayed by schedule""" - if epoch in schedule: - lr *= gamma - for param_group in optimizer.param_groups: - param_group['lr'] = lr - return lr - - -def compute_acc(pred, gt, thresh=0.5): - ''' - return: - IOU, precision, and recall - ''' - with torch.no_grad(): - vol_pred = pred > thresh - vol_gt = gt > thresh - - union = vol_pred | vol_gt - inter = vol_pred & vol_gt - - true_pos = inter.sum().float() - - union = union.sum().float() - if union == 0: - union = 1 - vol_pred = vol_pred.sum().float() - if vol_pred == 0: - vol_pred = 1 - vol_gt = vol_gt.sum().float() - if vol_gt == 0: - vol_gt = 1 - return true_pos / union, true_pos / vol_pred, true_pos / vol_gt - - -# def calc_metrics(opt, net, cuda, dataset, num_tests, -# resolution=128, sampled_points=1000, use_kaolin=True): -# if num_tests > len(dataset): -# num_tests = len(dataset) -# with torch.no_grad(): -# chamfer_arr, p2s_arr = [], [] -# for idx in tqdm(range(num_tests)): -# data = dataset[idx * len(dataset) // num_tests] - -# verts, faces = gen_mesh_eval(opt, net, cuda, data, resolution) -# if verts is None: -# continue - -# mesh_gt = trimesh.load(data['mesh_path']) -# mesh_gt = mesh_gt.split(only_watertight=False) -# comp_num = [mesh.vertices.shape[0] for mesh in mesh_gt] -# mesh_gt = mesh_gt[comp_num.index(max(comp_num))] - -# mesh_pred = trimesh.Trimesh(verts, faces) - -# gt_surface_pts, _ = trimesh.sample.sample_surface_even( -# mesh_gt, sampled_points) -# pred_surface_pts, _ = trimesh.sample.sample_surface_even( -# mesh_pred, sampled_points) - -# if use_kaolin and has_kaolin: -# kal_mesh_gt = kal.rep.TriangleMesh.from_tensors( -# torch.tensor(mesh_gt.vertices).float().to(device=cuda), -# torch.tensor(mesh_gt.faces).long().to(device=cuda)) -# kal_mesh_pred = kal.rep.TriangleMesh.from_tensors( -# torch.tensor(mesh_pred.vertices).float().to(device=cuda), -# torch.tensor(mesh_pred.faces).long().to(device=cuda)) - -# kal_distance_0 = kal.metrics.mesh.point_to_surface( -# torch.tensor(pred_surface_pts).float().to(device=cuda), kal_mesh_gt) -# kal_distance_1 = kal.metrics.mesh.point_to_surface( -# torch.tensor(gt_surface_pts).float().to(device=cuda), kal_mesh_pred) - -# dist_gt_pred = torch.sqrt(kal_distance_0).cpu().numpy() -# dist_pred_gt = torch.sqrt(kal_distance_1).cpu().numpy() -# else: -# try: -# _, dist_pred_gt, _ = trimesh.proximity.closest_point(mesh_pred, gt_surface_pts) -# _, dist_gt_pred, _ = trimesh.proximity.closest_point(mesh_gt, pred_surface_pts) -# except Exception as e: -# print (e) -# continue - -# chamfer_dist = 0.5 * (dist_pred_gt.mean() + dist_gt_pred.mean()) -# p2s_dist = dist_pred_gt.mean() - -# chamfer_arr.append(chamfer_dist) -# p2s_arr.append(p2s_dist) - -# return np.average(chamfer_arr), np.average(p2s_arr) - - -def calc_error(opt, net, cuda, dataset, num_tests): - if num_tests > len(dataset): - num_tests = len(dataset) - with torch.no_grad(): - erorr_arr, IOU_arr, prec_arr, recall_arr = [], [], [], [] - for idx in tqdm(range(num_tests)): - data = dataset[idx * len(dataset) // num_tests] - # retrieve the data - image_tensor = data['img'].to(device=cuda) - calib_tensor = data['calib'].to(device=cuda) - sample_tensor = data['samples'].to(device=cuda).unsqueeze(0) - if opt.num_views > 1: - sample_tensor = reshape_sample_tensor(sample_tensor, - opt.num_views) - label_tensor = data['labels'].to(device=cuda).unsqueeze(0) - - res, error = net.forward(image_tensor, - sample_tensor, - calib_tensor, - labels=label_tensor) - - IOU, prec, recall = compute_acc(res, label_tensor) - - # print( - # '{0}/{1} | Error: {2:06f} IOU: {3:06f} prec: {4:06f} recall: {5:06f}' - # .format(idx, num_tests, error.item(), IOU.item(), prec.item(), recall.item())) - erorr_arr.append(error.item()) - IOU_arr.append(IOU.item()) - prec_arr.append(prec.item()) - recall_arr.append(recall.item()) - - return np.average(erorr_arr), np.average(IOU_arr), np.average( - prec_arr), np.average(recall_arr) - - -def calc_error_color(opt, netG, netC, cuda, dataset, num_tests): - if num_tests > len(dataset): - num_tests = len(dataset) - with torch.no_grad(): - error_color_arr = [] - - for idx in tqdm(range(num_tests)): - data = dataset[idx * len(dataset) // num_tests] - # retrieve the data - image_tensor = data['img'].to(device=cuda) - calib_tensor = data['calib'].to(device=cuda) - color_sample_tensor = data['color_samples'].to( - device=cuda).unsqueeze(0) - - if opt.num_views > 1: - color_sample_tensor = reshape_sample_tensor( - color_sample_tensor, opt.num_views) - - rgb_tensor = data['rgbs'].to(device=cuda).unsqueeze(0) - - netG.filter(image_tensor) - _, errorC = netC.forward(image_tensor, - netG.get_im_feat(), - color_sample_tensor, - calib_tensor, - labels=rgb_tensor) - - # print('{0}/{1} | Error inout: {2:06f} | Error color: {3:06f}' - # .format(idx, num_tests, errorG.item(), errorC.item())) - error_color_arr.append(errorC.item()) - - return np.average(error_color_arr) - - -# pytorch lightning training related fucntions - - -def query_func(opt, netG, features, points, proj_matrix=None): - ''' - - points: size of (bz, N, 3) - - proj_matrix: size of (bz, 4, 4) - return: size of (bz, 1, N) - ''' - assert len(points) == 1 - samples = points.repeat(opt.num_views, 1, 1) - samples = samples.permute(0, 2, 1) # [bz, 3, N] - - # view specific query - if proj_matrix is not None: - samples = orthogonal(samples, proj_matrix) - - calib_tensor = torch.stack([torch.eye(4).float()], dim=0).type_as(samples) - - preds = netG.query(features=features, - points=samples, - calibs=calib_tensor, - regressor=netG.if_regressor) - - if type(preds) is list: - preds = preds[0] - - return preds - - -def isin(ar1, ar2): - return (ar1[..., None] == ar2).any(-1) - - -def in1d(ar1, ar2): - mask = ar2.new_zeros((max(ar1.max(), ar2.max()) + 1, ), dtype=torch.bool) - mask[ar2.unique()] = True - return mask[ar1] - - -def get_visibility(xy, z, faces): - """get the visibility of vertices - Args: - xy (torch.tensor): [N,2] - z (torch.tensor): [N,1] - faces (torch.tensor): [N,3] - size (int): resolution of rendered image - """ - - xyz = torch.cat((xy, -z), dim=1) - xyz = (xyz + 1.0) / 2.0 - faces = faces.long() - - rasterizer = Pytorch3dRasterizer(image_size=2**12) - meshes_screen = Meshes(verts=xyz[None, ...], faces=faces[None, ...]) - raster_settings = rasterizer.raster_settings - - pix_to_face, zbuf, bary_coords, dists = rasterize_meshes( - meshes_screen, - image_size=raster_settings.image_size, - blur_radius=raster_settings.blur_radius, - faces_per_pixel=raster_settings.faces_per_pixel, - bin_size=raster_settings.bin_size, - max_faces_per_bin=raster_settings.max_faces_per_bin, - perspective_correct=raster_settings.perspective_correct, - cull_backfaces=raster_settings.cull_backfaces, - ) - - vis_vertices_id = torch.unique(faces[torch.unique(pix_to_face), :]) - vis_mask = torch.zeros(size=(z.shape[0], 1)) - vis_mask[vis_vertices_id] = 1.0 - - # print("------------------------\n") - # print(f"keep points : {vis_mask.sum()/len(vis_mask)}") - - return vis_mask - - -def batch_mean(res, key): - # recursive mean for multilevel dicts - return torch.stack([ - x[key] if isinstance(x, dict) else batch_mean(x, key) for x in res - ]).mean() - - -def tf_log_convert(log_dict): - new_log_dict = log_dict.copy() - for k, v in log_dict.items(): - new_log_dict[k.replace("_", "/")] = v - del new_log_dict[k] - - return new_log_dict - - -def bar_log_convert(log_dict, name=None, rot=None): - from decimal import Decimal - - new_log_dict = {} - - if name is not None: - new_log_dict['name'] = name[0] - if rot is not None: - new_log_dict['rot'] = rot[0] - - for k, v in log_dict.items(): - color = "yellow" - if 'loss' in k: - color = "red" - k = k.replace("loss", "L") - elif 'acc' in k: - color = "green" - k = k.replace("acc", "A") - elif 'iou' in k: - color = "green" - k = k.replace("iou", "I") - elif 'prec' in k: - color = "green" - k = k.replace("prec", "P") - elif 'recall' in k: - color = "green" - k = k.replace("recall", "R") - - if 'lr' not in k: - new_log_dict[colored(k.split("_")[1], - color)] = colored(f"{v:.3f}", color) - else: - new_log_dict[colored(k.split("_")[1], - color)] = colored(f"{Decimal(str(v)):.1E}", - color) - - if 'loss' in new_log_dict.keys(): - del new_log_dict['loss'] - - return new_log_dict - - -def accumulate(outputs, rot_num, split): - - hparam_log_dict = {} - - metrics = outputs[0].keys() - datasets = split.keys() - - for dataset in datasets: - for metric in metrics: - keyword = f"hparam/{dataset}-{metric}" - if keyword not in hparam_log_dict.keys(): - hparam_log_dict[keyword] = 0 - for idx in range(split[dataset][0] * rot_num, - split[dataset][1] * rot_num): - hparam_log_dict[keyword] += outputs[idx][metric] - hparam_log_dict[keyword] /= (split[dataset][1] - - split[dataset][0]) * rot_num - - print(colored(hparam_log_dict, "green")) - - return hparam_log_dict - - -def calc_error_N(outputs, targets): - """calculate the error of normal (IGR) - Args: - outputs (torch.tensor): [B, 3, N] - target (torch.tensor): [B, N, 3] - # manifold loss and grad_loss in IGR paper - grad_loss = ((nonmnfld_grad.norm(2, dim=-1) - 1) ** 2).mean() - normals_loss = ((mnfld_grad - normals).abs()).norm(2, dim=1).mean() - Returns: - torch.tensor: error of valid normals on the surface - """ - # outputs = torch.tanh(-outputs.permute(0,2,1).reshape(-1,3)) - outputs = -outputs.permute(0, 2, 1).reshape(-1, 1) - targets = targets.reshape(-1, 3)[:, 2:3] - with_normals = targets.sum(dim=1).abs() > 0.0 - - # eikonal loss - grad_loss = ((outputs[with_normals].norm(2, dim=-1) - 1)**2).mean() - # normals loss - normal_loss = (outputs - targets)[with_normals].abs().norm(2, dim=1).mean() - - return grad_loss * 0.0 + normal_loss - - -def calc_knn_acc(preds, carn_verts, labels, pick_num): - """calculate knn accuracy - Args: - preds (torch.tensor): [B, 3, N] - carn_verts (torch.tensor): [SMPLX_V_num, 3] - labels (torch.tensor): [B, N_knn, N] - """ - N_knn_full = labels.shape[1] - preds = preds.permute(0, 2, 1).reshape(-1, 3) - labels = labels.permute(0, 2, 1).reshape(-1, N_knn_full) # [BxN, num_knn] - labels = labels[:, :pick_num] - - dist = torch.cdist(preds, carn_verts, p=2) # [BxN, SMPL_V_num] - knn = dist.topk(k=pick_num, dim=1, largest=False)[1] # [BxN, num_knn] - cat_mat = torch.sort(torch.cat((knn, labels), dim=1))[0] - bool_col = torch.zeros_like(cat_mat)[:, 0] - for i in range(pick_num * 2 - 1): - bool_col += cat_mat[:, i] == cat_mat[:, i + 1] - acc = (bool_col > 0).sum() / len(bool_col) - - return acc - - -def calc_acc_seg(output, target, num_multiseg): - from pytorch_lightning.metrics import Accuracy - return Accuracy()(output.reshape(-1, num_multiseg).cpu(), - target.flatten().cpu()) - - -def add_watermark(imgs, titles): - - # Write some Text - - font = cv2.FONT_HERSHEY_SIMPLEX - bottomLeftCornerOfText = (350, 50) - bottomRightCornerOfText = (800, 50) - fontScale = 1 - fontColor = (1.0, 1.0, 1.0) - lineType = 2 - - for i in range(len(imgs)): - - title = titles[i + 1] - cv2.putText(imgs[i], title, bottomLeftCornerOfText, font, fontScale, - fontColor, lineType) - - if i == 0: - cv2.putText(imgs[i], str(titles[i][0]), bottomRightCornerOfText, - font, fontScale, fontColor, lineType) - - result = np.concatenate(imgs, axis=0).transpose(2, 0, 1) - - return result - - -def make_test_gif(img_dir): - - if img_dir is not None and len(os.listdir(img_dir)) > 0: - for dataset in os.listdir(img_dir): - for subject in sorted(os.listdir(osp.join(img_dir, dataset))): - img_lst = [] - im1 = None - for file in sorted( - os.listdir(osp.join(img_dir, dataset, subject))): - if file[-3:] not in ['obj', 'gif']: - img_path = os.path.join(img_dir, dataset, subject, - file) - if im1 == None: - im1 = Image.open(img_path) - else: - img_lst.append(Image.open(img_path)) - - print(os.path.join(img_dir, dataset, subject, "out.gif")) - im1.save(os.path.join(img_dir, dataset, subject, "out.gif"), - save_all=True, - append_images=img_lst, - duration=500, - loop=0) - - -def export_cfg(logger, cfg): - - cfg_export_file = osp.join(logger.save_dir, logger.name, - f"version_{logger.version}", "cfg.yaml") - - if not osp.exists(cfg_export_file): - os.makedirs(osp.dirname(cfg_export_file), exist_ok=True) - with open(cfg_export_file, "w+") as file: - _ = yaml.dump(cfg, file) diff --git a/lib /dataset / Evaluator.py b/lib /dataset / Evaluator.py deleted file mode 100644 index 68746d1b76c6328f7e4a904bb4d6fb1f5aa1de29..0000000000000000000000000000000000000000 --- a/lib /dataset / Evaluator.py +++ /dev/null @@ -1,262 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - - -from lib.renderer.gl.normal_render import NormalRender -from lib.dataset.mesh_util import projection -from lib.common.render import Render -from PIL import Image -import numpy as np -import torch -from torch import nn -import trimesh -import os.path as osp -from PIL import Image - - -class Evaluator: - - _normal_render = None - - @staticmethod - def init_gl(): - Evaluator._normal_render = NormalRender(width=512, height=512) - - def __init__(self, device): - self.device = device - self.render = Render(size=512, device=self.device) - self.error_term = nn.MSELoss() - - self.offset = 0.0 - self.scale_factor = None - - def set_mesh(self, result_dict, scale_factor=1.0, offset=0.0): - - for key in result_dict.keys(): - if torch.is_tensor(result_dict[key]): - result_dict[key] = result_dict[key].detach().cpu().numpy() - - for k, v in result_dict.items(): - setattr(self, k, v) - - self.scale_factor = scale_factor - self.offset = offset - - def _render_normal(self, mesh, deg, norms=None): - view_mat = np.identity(4) - rz = deg / 180.0 * np.pi - model_mat = np.identity(4) - model_mat[:3, :3] = self._normal_render.euler_to_rot_mat(0, rz, 0) - model_mat[1, 3] = self.offset - view_mat[2, 2] *= -1 - - self._normal_render.set_matrices(view_mat, model_mat) - if norms is None: - norms = mesh.vertex_normals - self._normal_render.set_normal_mesh(self.scale_factor * mesh.vertices, - mesh.faces, norms, mesh.faces) - self._normal_render.draw() - normal_img = self._normal_render.get_color() - return normal_img - - def render_mesh_list(self, mesh_lst): - - self.offset = 0.0 - self.scale_factor = 1.0 - - full_list = [] - for mesh in mesh_lst: - row_lst = [] - for deg in np.arange(0, 360, 90): - normal = self._render_normal(mesh, deg) - row_lst.append(normal) - full_list.append(np.concatenate(row_lst, axis=1)) - - res_array = np.concatenate(full_list, axis=0) - - return res_array - - def _get_reproj_normal_error(self, deg): - - tgt_normal = self._render_normal(self.tgt_mesh, deg) - src_normal = self._render_normal(self.src_mesh, deg) - error = (((src_normal[:, :, :3] - - tgt_normal[:, :, :3])**2).sum(axis=2).mean(axis=(0, 1))) - - return error, [src_normal, tgt_normal] - - def render_normal(self, verts, faces): - - verts = verts[0].detach().cpu().numpy() - faces = faces[0].detach().cpu().numpy() - - mesh_F = trimesh.Trimesh(verts * np.array([1.0, -1.0, 1.0]), faces) - mesh_B = trimesh.Trimesh(verts * np.array([1.0, -1.0, -1.0]), faces) - - self.scale_factor = 1.0 - - normal_F = self._render_normal(mesh_F, 0) - normal_B = self._render_normal(mesh_B, - 0, - norms=mesh_B.vertex_normals * - np.array([-1.0, -1.0, 1.0])) - - mask = normal_F[:, :, 3:4] - normal_F = (torch.as_tensor(2.0 * (normal_F - 0.5) * mask).permute( - 2, 0, 1)[:3, :, :].float().unsqueeze(0).to(self.device)) - normal_B = (torch.as_tensor(2.0 * (normal_B - 0.5) * mask).permute( - 2, 0, 1)[:3, :, :].float().unsqueeze(0).to(self.device)) - - return {"T_normal_F": normal_F, "T_normal_B": normal_B} - - def calculate_normal_consist( - self, - frontal=True, - back=True, - left=True, - right=True, - save_demo_img=None, - return_demo=False, - ): - - # reproj error - # if save_demo_img is not None, save a visualization at the given path (etc, "./test.png") - if self._normal_render is None: - print( - "In order to use normal render, " - "you have to call init_gl() before initialing any evaluator objects." - ) - return -1 - - side_cnt = 0 - total_error = 0 - demo_list = [] - - if frontal: - side_cnt += 1 - error, normal_lst = self._get_reproj_normal_error(0) - total_error += error - demo_list.append(np.concatenate(normal_lst, axis=0)) - if back: - side_cnt += 1 - error, normal_lst = self._get_reproj_normal_error(180) - total_error += error - demo_list.append(np.concatenate(normal_lst, axis=0)) - if left: - side_cnt += 1 - error, normal_lst = self._get_reproj_normal_error(90) - total_error += error - demo_list.append(np.concatenate(normal_lst, axis=0)) - if right: - side_cnt += 1 - error, normal_lst = self._get_reproj_normal_error(270) - total_error += error - demo_list.append(np.concatenate(normal_lst, axis=0)) - if save_demo_img is not None: - res_array = np.concatenate(demo_list, axis=1) - res_img = Image.fromarray((res_array * 255).astype(np.uint8)) - res_img.save(save_demo_img) - - if return_demo: - res_array = np.concatenate(demo_list, axis=1) - return res_array - else: - return total_error - - def space_transfer(self): - - # convert from GT to SDF - self.verts_pr -= self.recon_size / 2.0 - self.verts_pr /= self.recon_size / 2.0 - - self.verts_gt = projection(self.verts_gt, self.calib) - self.verts_gt[:, 1] *= -1 - - self.tgt_mesh = trimesh.Trimesh(self.verts_gt, self.faces_gt) - self.src_mesh = trimesh.Trimesh(self.verts_pr, self.faces_pr) - - # (self.tgt_mesh+self.src_mesh).show() - - def export_mesh(self, dir, name): - self.tgt_mesh.visual.vertex_colors = np.array([255, 0, 0]) - self.src_mesh.visual.vertex_colors = np.array([0, 255, 0]) - - (self.tgt_mesh + self.src_mesh).export( - osp.join(dir, f"{name}_gt_pr.obj")) - - def calculate_chamfer_p2s(self, sampled_points=1000): - """calculate the geometry metrics [chamfer, p2s, chamfer_H, p2s_H] - Args: - verts_gt (torch.cuda.tensor): [N, 3] - faces_gt (torch.cuda.tensor): [M, 3] - verts_pr (torch.cuda.tensor): [N', 3] - faces_pr (torch.cuda.tensor): [M', 3] - sampled_points (int, optional): use smaller number for faster testing. Defaults to 1000. - Returns: - tuple: chamfer, p2s, chamfer_H, p2s_H - """ - - gt_surface_pts, _ = trimesh.sample.sample_surface_even( - self.tgt_mesh, sampled_points) - pred_surface_pts, _ = trimesh.sample.sample_surface_even( - self.src_mesh, sampled_points) - - _, dist_pred_gt, _ = trimesh.proximity.closest_point( - self.src_mesh, gt_surface_pts) - _, dist_gt_pred, _ = trimesh.proximity.closest_point( - self.tgt_mesh, pred_surface_pts) - - dist_pred_gt[np.isnan(dist_pred_gt)] = 0 - dist_gt_pred[np.isnan(dist_gt_pred)] = 0 - chamfer_dist = 0.5 * (dist_pred_gt.mean() + - dist_gt_pred.mean()).item() * 100 - p2s_dist = dist_pred_gt.mean().item() * 100 - - return chamfer_dist, p2s_dist - - def calc_acc(self, output, target, thres=0.5, use_sdf=False): - - # # remove the surface points with thres - # non_surf_ids = (target != thres) - # output = output[non_surf_ids] - # target = target[non_surf_ids] - - with torch.no_grad(): - output = output.masked_fill(output < thres, 0.0) - output = output.masked_fill(output > thres, 1.0) - - if use_sdf: - target = target.masked_fill(target < thres, 0.0) - target = target.masked_fill(target > thres, 1.0) - - acc = output.eq(target).float().mean() - - # iou, precison, recall - output = output > thres - target = target > thres - - union = output | target - inter = output & target - - _max = torch.tensor(1.0).to(output.device) - - union = max(union.sum().float(), _max) - true_pos = max(inter.sum().float(), _max) - vol_pred = max(output.sum().float(), _max) - vol_gt = max(target.sum().float(), _max) - - return acc, true_pos / union, true_pos / vol_pred, true_pos / vol_gt \ No newline at end of file diff --git a/lib /dataset / PIFuDataModule.py b/lib /dataset / PIFuDataModule.py deleted file mode 100644 index 453238cf616d14ee0034253c91f661d3a98f6319..0000000000000000000000000000000000000000 --- a/lib /dataset / PIFuDataModule.py +++ /dev/null @@ -1,71 +0,0 @@ -import numpy as np -from torch.utils.data import DataLoader -from .PIFuDataset import PIFuDataset -import pytorch_lightning as pl - - -class PIFuDataModule(pl.LightningDataModule): - def __init__(self, cfg): - super(PIFuDataModule, self).__init__() - self.cfg = cfg - self.overfit = self.cfg.overfit - - if self.overfit: - self.batch_size = 1 - else: - self.batch_size = self.cfg.batch_size - - self.data_size = {} - - def prepare_data(self): - - pass - - @staticmethod - def worker_init_fn(worker_id): - np.random.seed(np.random.get_state()[1][0] + worker_id) - - def setup(self, stage): - - if stage == 'fit': - self.train_dataset = PIFuDataset(cfg=self.cfg, split="train") - self.val_dataset = PIFuDataset(cfg=self.cfg, split="val") - self.data_size = {'train': len(self.train_dataset), - 'val': len(self.val_dataset)} - - if stage == 'test': - self.test_dataset = PIFuDataset(cfg=self.cfg, split="test") - - def train_dataloader(self): - - train_data_loader = DataLoader( - self.train_dataset, - batch_size=self.batch_size, shuffle=True, - num_workers=self.cfg.num_threads, pin_memory=True, - worker_init_fn=self.worker_init_fn) - - return train_data_loader - - def val_dataloader(self): - - if self.overfit: - current_dataset = self.train_dataset - else: - current_dataset = self.val_dataset - - val_data_loader = DataLoader( - current_dataset, - batch_size=1, shuffle=False, - num_workers=self.cfg.num_threads, pin_memory=True, - worker_init_fn=self.worker_init_fn) - - return val_data_loader - - def test_dataloader(self): - - test_data_loader = DataLoader( - self.test_dataset, - batch_size=1, shuffle=False, - num_workers=self.cfg.num_threads, pin_memory=True) - - return test_data_loader \ No newline at end of file diff --git a/lib /dataset / hoppeMesh.py b/lib /dataset / hoppeMesh.py deleted file mode 100644 index 924d54b0676a2b61823c884d361906c765c9d14c..0000000000000000000000000000000000000000 --- a/lib /dataset / hoppeMesh.py +++ /dev/null @@ -1,116 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import numpy as np -from scipy.spatial import cKDTree -import trimesh - -import logging - -logging.getLogger("trimesh").setLevel(logging.ERROR) - - -def save_obj_mesh(mesh_path, verts, faces): - file = open(mesh_path, 'w') - for v in verts: - file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2])) - for f in faces: - f_plus = f + 1 - file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2])) - file.close() - - -def save_obj_mesh_with_color(mesh_path, verts, faces, colors): - file = open(mesh_path, 'w') - - for idx, v in enumerate(verts): - c = colors[idx] - file.write('v %.4f %.4f %.4f %.4f %.4f %.4f\n' % - (v[0], v[1], v[2], c[0], c[1], c[2])) - for f in faces: - f_plus = f + 1 - file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2])) - file.close() - - -def save_ply(mesh_path, points, rgb): - ''' - Save the visualization of sampling to a ply file. - Red points represent positive predictions. - Green points represent negative predictions. - :param mesh_path: File name to save - :param points: [N, 3] array of points - :param rgb: [N, 3] array of rgb values in the range [0~1] - :return: - ''' - to_save = np.concatenate([points, rgb * 255], axis=-1) - return np.savetxt( - mesh_path, - to_save, - fmt='%.6f %.6f %.6f %d %d %d', - comments='', - header=( - 'ply\nformat ascii 1.0\nelement vertex {:d}\n' + - 'property float x\nproperty float y\nproperty float z\n' + - 'property uchar red\nproperty uchar green\nproperty uchar blue\n' + - 'end_header').format(points.shape[0])) - - -class HoppeMesh: - def __init__(self, verts, faces, vert_normals, face_normals): - ''' - The HoppeSDF calculates signed distance towards a predefined oriented point cloud - http://hhoppe.com/recon.pdf - For clean and high-resolution pcl data, this is the fastest and accurate approximation of sdf - :param points: pts - :param normals: normals - ''' - self.verts = verts # [n, 3] - self.faces = faces # [m, 3] - self.vert_normals = vert_normals # [n, 3] - self.face_normals = face_normals # [m, 3] - - self.kd_tree = cKDTree(self.verts) - self.len = len(self.verts) - - def query(self, points): - dists, idx = self.kd_tree.query(points, n_jobs=1) - # FIXME: because the eyebows are removed, cKDTree around eyebows - # are not accurate. Cause a few false-inside labels here. - dirs = points - self.verts[idx] - signs = (dirs * self.vert_normals[idx]).sum(axis=1) - signs = (signs > 0) * 2 - 1 - return signs * dists - - def contains(self, points): - - labels = trimesh.Trimesh(vertices=self.verts, - faces=self.faces).contains(points) - return labels - - def export(self, path): - if self.colors is not None: - save_obj_mesh_with_color(path, self.verts, self.faces, - self.colors[:, 0:3] / 255.0) - else: - save_obj_mesh(path, self.verts, self.faces) - - def export_ply(self, path): - save_ply(path, self.verts, self.colors[:, 0:3] / 255.0) - - def triangles(self): - return self.verts[self.faces] # [n, 3, 3] \ No newline at end of file diff --git a/lib /dataset /NormalDataset.py b/lib /dataset /NormalDataset.py deleted file mode 100644 index 4e71cd3ed344867c80cf842c369e06cfccedf4c7..0000000000000000000000000000000000000000 --- a/lib /dataset /NormalDataset.py +++ /dev/null @@ -1,212 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import os.path as osp -import numpy as np -from PIL import Image -import torchvision.transforms as transforms - - -class NormalDataset(): - def __init__(self, cfg, split='train'): - - self.split = split - self.root = cfg.root - self.overfit = cfg.overfit - - self.opt = cfg.dataset - self.datasets = self.opt.types - self.input_size = self.opt.input_size - self.set_splits = self.opt.set_splits - self.scales = self.opt.scales - self.pifu = self.opt.pifu - - # input data types and dimensions - self.in_nml = [item[0] for item in cfg.net.in_nml] - self.in_nml_dim = [item[1] for item in cfg.net.in_nml] - self.in_total = self.in_nml + ['normal_F', 'normal_B'] - self.in_total_dim = self.in_nml_dim + [3, 3] - - if self.split != 'train': - self.rotations = range(0, 360, 120) - else: - self.rotations = np.arange(0, 360, 360 / - self.opt.rotation_num).astype(np.int) - - self.datasets_dict = {} - for dataset_id, dataset in enumerate(self.datasets): - dataset_dir = osp.join(self.root, dataset, "smplx") - self.datasets_dict[dataset] = { - "subjects": - np.loadtxt(osp.join(self.root, dataset, "all.txt"), dtype=str), - "path": - dataset_dir, - "scale": - self.scales[dataset_id] - } - - self.subject_list = self.get_subject_list(split) - - # PIL to tensor - self.image_to_tensor = transforms.Compose([ - transforms.Resize(self.input_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ]) - - # PIL to tensor - self.mask_to_tensor = transforms.Compose([ - transforms.Resize(self.input_size), - transforms.ToTensor(), - transforms.Normalize((0.0, ), (1.0, )) - ]) - - def get_subject_list(self, split): - - subject_list = [] - - for dataset in self.datasets: - - if self.pifu: - txt = osp.join(self.root, dataset, f'{split}_pifu.txt') - else: - txt = osp.join(self.root, dataset, f'{split}.txt') - - if osp.exists(txt): - print(f"load from {txt}") - subject_list += sorted(np.loadtxt(txt, dtype=str).tolist()) - - if self.pifu: - miss_pifu = sorted( - np.loadtxt(osp.join(self.root, dataset, - "miss_pifu.txt"), - dtype=str).tolist()) - subject_list = [ - subject for subject in subject_list - if subject not in miss_pifu - ] - subject_list = [ - "renderpeople/" + subject for subject in subject_list - ] - - else: - train_txt = osp.join(self.root, dataset, 'train.txt') - val_txt = osp.join(self.root, dataset, 'val.txt') - test_txt = osp.join(self.root, dataset, 'test.txt') - - print( - f"generate lists of [train, val, test] \n {train_txt} \n {val_txt} \n {test_txt} \n" - ) - - split_txt = osp.join(self.root, dataset, f'{split}.txt') - - subjects = self.datasets_dict[dataset]['subjects'] - train_split = int(len(subjects) * self.set_splits[0]) - val_split = int( - len(subjects) * self.set_splits[1]) + train_split - - with open(train_txt, "w") as f: - f.write("\n".join(dataset + "/" + item - for item in subjects[:train_split])) - with open(val_txt, "w") as f: - f.write("\n".join( - dataset + "/" + item - for item in subjects[train_split:val_split])) - with open(test_txt, "w") as f: - f.write("\n".join(dataset + "/" + item - for item in subjects[val_split:])) - - subject_list += sorted( - np.loadtxt(split_txt, dtype=str).tolist()) - - bug_list = sorted( - np.loadtxt(osp.join(self.root, 'bug.txt'), dtype=str).tolist()) - - subject_list = [ - subject for subject in subject_list if (subject not in bug_list) - ] - - return subject_list - - def __len__(self): - return len(self.subject_list) * len(self.rotations) - - def __getitem__(self, index): - - # only pick the first data if overfitting - if self.overfit: - index = 0 - - rid = index % len(self.rotations) - mid = index // len(self.rotations) - - rotation = self.rotations[rid] - - # choose specific test sets - subject = self.subject_list[mid] - - subject_render = "/".join( - [subject.split("/")[0] + "_12views", - subject.split("/")[1]]) - - # setup paths - data_dict = { - 'dataset': - subject.split("/")[0], - 'subject': - subject, - 'rotation': - rotation, - 'image_path': - osp.join(self.root, subject_render, 'render', - f'{rotation:03d}.png') - } - - # image/normal/depth loader - for name, channel in zip(self.in_total, self.in_total_dim): - - if name != 'image': - data_dict.update({ - f'{name}_path': - osp.join(self.root, subject_render, name, - f'{rotation:03d}.png') - }) - data_dict.update({ - name: - self.imagepath2tensor(data_dict[f'{name}_path'], - channel, - inv='depth_B' in name) - }) - - path_keys = [ - key for key in data_dict.keys() if '_path' in key or '_dir' in key - ] - for key in path_keys: - del data_dict[key] - - return data_dict - - def imagepath2tensor(self, path, channel=3, inv=False): - - rgba = Image.open(path).convert('RGBA') - mask = rgba.split()[-1] - image = rgba.convert('RGB') - image = self.image_to_tensor(image) - mask = self.mask_to_tensor(mask) - image = (image * mask)[:channel] - - return (image * (0.5 - inv) * 2.0).float() \ No newline at end of file diff --git a/lib /dataset /NormalModule.py b/lib /dataset /NormalModule.py deleted file mode 100644 index ea25725bc8d553084c63baf8900a1b3035b00f10..0000000000000000000000000000000000000000 --- a/lib /dataset /NormalModule.py +++ /dev/null @@ -1,94 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import numpy as np -from torch.utils.data import DataLoader -from .NormalDataset import NormalDataset - -# pytorch lightning related libs -import pytorch_lightning as pl - - -class NormalModule(pl.LightningDataModule): - def __init__(self, cfg): - super(NormalModule, self).__init__() - self.cfg = cfg - self.overfit = self.cfg.overfit - - if self.overfit: - self.batch_size = 1 - else: - self.batch_size = self.cfg.batch_size - - self.data_size = {} - - def prepare_data(self): - - pass - - @staticmethod - def worker_init_fn(worker_id): - np.random.seed(np.random.get_state()[1][0] + worker_id) - - def setup(self, stage): - - if stage == 'fit' or stage is None: - self.train_dataset = NormalDataset(cfg=self.cfg, split="train") - self.val_dataset = NormalDataset(cfg=self.cfg, split="val") - self.data_size = { - 'train': len(self.train_dataset), - 'val': len(self.val_dataset) - } - - if stage == 'test' or stage is None: - self.test_dataset = NormalDataset(cfg=self.cfg, split="test") - - def train_dataloader(self): - - train_data_loader = DataLoader(self.train_dataset, - batch_size=self.batch_size, - shuffle=not self.overfit, - num_workers=self.cfg.num_threads, - pin_memory=True, - worker_init_fn=self.worker_init_fn) - - return train_data_loader - - def val_dataloader(self): - - if self.overfit: - current_dataset = self.train_dataset - else: - current_dataset = self.val_dataset - - val_data_loader = DataLoader(current_dataset, - batch_size=self.batch_size, - shuffle=False, - num_workers=self.cfg.num_threads, - pin_memory=True) - - return val_data_loader - - def test_dataloader(self): - - test_data_loader = DataLoader(self.test_dataset, - batch_size=1, - shuffle=False, - num_workers=self.cfg.num_threads, - pin_memory=True) - - return test_data_loader \ No newline at end of file diff --git a/lib /dataset /PIFuDataset.py b/lib /dataset /PIFuDataset.py deleted file mode 100644 index 9bf240c1f1a35b45c4bd38aed6c1827bdd09e56d..0000000000000000000000000000000000000000 --- a/lib /dataset /PIFuDataset.py +++ /dev/null @@ -1,589 +0,0 @@ -from lib.renderer.mesh import load_fit_body -from lib.dataset.hoppeMesh import HoppeMesh -from lib.dataset.body_model import TetraSMPLModel -from lib.common.render import Render -from lib.dataset.mesh_util import SMPLX, projection, cal_sdf_batch, get_visibility -from lib.pare.pare.utils.geometry import rotation_matrix_to_angle_axis -from termcolor import colored -import os.path as osp -import numpy as np -from PIL import Image -import random -import os -import trimesh -import torch -from kaolin.ops.mesh import check_sign -import torchvision.transforms as transforms -from huggingface_hub import hf_hub_download, cached_download - - -class PIFuDataset(): - def __init__(self, cfg, split='train', vis=False): - - self.split = split - self.root = cfg.root - self.bsize = cfg.batch_size - self.overfit = cfg.overfit - - # for debug, only used in visualize_sampling3D - self.vis = vis - - self.opt = cfg.dataset - self.datasets = self.opt.types - self.input_size = self.opt.input_size - self.scales = self.opt.scales - self.workers = cfg.num_threads - self.prior_type = cfg.net.prior_type - - self.noise_type = self.opt.noise_type - self.noise_scale = self.opt.noise_scale - - noise_joints = [4, 5, 7, 8, 13, 14, 16, 17, 18, 19, 20, 21] - - self.noise_smpl_idx = [] - self.noise_smplx_idx = [] - - for idx in noise_joints: - self.noise_smpl_idx.append(idx * 3) - self.noise_smpl_idx.append(idx * 3 + 1) - self.noise_smpl_idx.append(idx * 3 + 2) - - self.noise_smplx_idx.append((idx-1) * 3) - self.noise_smplx_idx.append((idx-1) * 3 + 1) - self.noise_smplx_idx.append((idx-1) * 3 + 2) - - self.use_sdf = cfg.sdf - self.sdf_clip = cfg.sdf_clip - - # [(feat_name, channel_num),...] - self.in_geo = [item[0] for item in cfg.net.in_geo] - self.in_nml = [item[0] for item in cfg.net.in_nml] - - self.in_geo_dim = [item[1] for item in cfg.net.in_geo] - self.in_nml_dim = [item[1] for item in cfg.net.in_nml] - - self.in_total = self.in_geo + self.in_nml - self.in_total_dim = self.in_geo_dim + self.in_nml_dim - - if self.split == 'train': - self.rotations = np.arange( - 0, 360, 360 / self.opt.rotation_num).astype(np.int32) - else: - self.rotations = range(0, 360, 120) - - self.datasets_dict = {} - - for dataset_id, dataset in enumerate(self.datasets): - - mesh_dir = None - smplx_dir = None - - dataset_dir = osp.join(self.root, dataset) - - if dataset in ['thuman2']: - mesh_dir = osp.join(dataset_dir, "scans") - smplx_dir = osp.join(dataset_dir, "fits") - smpl_dir = osp.join(dataset_dir, "smpl") - - self.datasets_dict[dataset] = { - "subjects": np.loadtxt(osp.join(dataset_dir, "all.txt"), dtype=str), - "smplx_dir": smplx_dir, - "smpl_dir": smpl_dir, - "mesh_dir": mesh_dir, - "scale": self.scales[dataset_id] - } - - self.subject_list = self.get_subject_list(split) - self.smplx = SMPLX() - - # PIL to tensor - self.image_to_tensor = transforms.Compose([ - transforms.Resize(self.input_size), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ]) - - # PIL to tensor - self.mask_to_tensor = transforms.Compose([ - transforms.Resize(self.input_size), - transforms.ToTensor(), - transforms.Normalize((0.0, ), (1.0, )) - ]) - - self.device = torch.device(f"cuda:{cfg.gpus[0]}") - self.render = Render(size=512, device=self.device) - - def render_normal(self, verts, faces): - - # render optimized mesh (normal, T_normal, image [-1,1]) - self.render.load_meshes(verts, faces) - return self.render.get_rgb_image() - - def get_subject_list(self, split): - - subject_list = [] - - for dataset in self.datasets: - - split_txt = osp.join(self.root, dataset, f'{split}.txt') - - if osp.exists(split_txt): - print(f"load from {split_txt}") - subject_list += np.loadtxt(split_txt, dtype=str).tolist() - else: - full_txt = osp.join(self.root, dataset, 'all.txt') - print(f"split {full_txt} into train/val/test") - - full_lst = np.loadtxt(full_txt, dtype=str) - full_lst = [dataset+"/"+item for item in full_lst] - [train_lst, test_lst, val_lst] = np.split( - full_lst, [500, 500+5, ]) - - np.savetxt(full_txt.replace( - "all", "train"), train_lst, fmt="%s") - np.savetxt(full_txt.replace("all", "test"), test_lst, fmt="%s") - np.savetxt(full_txt.replace("all", "val"), val_lst, fmt="%s") - - print(f"load from {split_txt}") - subject_list += np.loadtxt(split_txt, dtype=str).tolist() - - if self.split != 'test': - subject_list += subject_list[:self.bsize - - len(subject_list) % self.bsize] - print(colored(f"total: {len(subject_list)}", "yellow")) - random.shuffle(subject_list) - - # subject_list = ["thuman2/0008"] - return subject_list - - def __len__(self): - return len(self.subject_list) * len(self.rotations) - - def __getitem__(self, index): - - # only pick the first data if overfitting - if self.overfit: - index = 0 - - rid = index % len(self.rotations) - mid = index // len(self.rotations) - - rotation = self.rotations[rid] - subject = self.subject_list[mid].split("/")[1] - dataset = self.subject_list[mid].split("/")[0] - render_folder = "/".join([dataset + - f"_{self.opt.rotation_num}views", subject]) - - # setup paths - data_dict = { - 'dataset': dataset, - 'subject': subject, - 'rotation': rotation, - 'scale': self.datasets_dict[dataset]["scale"], - 'mesh_path': osp.join(self.datasets_dict[dataset]["mesh_dir"], f"{subject}/{subject}.obj"), - 'smplx_path': osp.join(self.datasets_dict[dataset]["smplx_dir"], f"{subject}/smplx_param.pkl"), - 'smpl_path': osp.join(self.datasets_dict[dataset]["smpl_dir"], f"{subject}.pkl"), - 'calib_path': osp.join(self.root, render_folder, 'calib', f'{rotation:03d}.txt'), - 'vis_path': osp.join(self.root, render_folder, 'vis', f'{rotation:03d}.pt'), - 'image_path': osp.join(self.root, render_folder, 'render', f'{rotation:03d}.png') - } - - # load training data - data_dict.update(self.load_calib(data_dict)) - - # image/normal/depth loader - for name, channel in zip(self.in_total, self.in_total_dim): - - if f'{name}_path' not in data_dict.keys(): - data_dict.update({ - f'{name}_path': osp.join(self.root, render_folder, name, f'{rotation:03d}.png') - }) - - # tensor update - data_dict.update({ - name: self.imagepath2tensor( - data_dict[f'{name}_path'], channel, inv=False) - }) - - data_dict.update(self.load_mesh(data_dict)) - data_dict.update(self.get_sampling_geo( - data_dict, is_valid=self.split == "val", is_sdf=self.use_sdf)) - data_dict.update(self.load_smpl(data_dict, self.vis)) - - if self.prior_type == 'pamir': - data_dict.update(self.load_smpl_voxel(data_dict)) - - if (self.split != 'test') and (not self.vis): - - del data_dict['verts'] - del data_dict['faces'] - - if not self.vis: - del data_dict['mesh'] - - path_keys = [ - key for key in data_dict.keys() if '_path' in key or '_dir' in key - ] - for key in path_keys: - del data_dict[key] - - return data_dict - - def imagepath2tensor(self, path, channel=3, inv=False): - - rgba = Image.open(path).convert('RGBA') - mask = rgba.split()[-1] - image = rgba.convert('RGB') - image = self.image_to_tensor(image) - mask = self.mask_to_tensor(mask) - image = (image * mask)[:channel] - - return (image * (0.5 - inv) * 2.0).float() - - def load_calib(self, data_dict): - calib_data = np.loadtxt(data_dict['calib_path'], dtype=float) - extrinsic = calib_data[:4, :4] - intrinsic = calib_data[4:8, :4] - calib_mat = np.matmul(intrinsic, extrinsic) - calib_mat = torch.from_numpy(calib_mat).float() - return {'calib': calib_mat} - - def load_mesh(self, data_dict): - mesh_path = data_dict['mesh_path'] - scale = data_dict['scale'] - - mesh_ori = trimesh.load(mesh_path, - skip_materials=True, - process=False, - maintain_order=True) - verts = mesh_ori.vertices * scale - faces = mesh_ori.faces - - vert_normals = np.array(mesh_ori.vertex_normals) - face_normals = np.array(mesh_ori.face_normals) - - mesh = HoppeMesh(verts, faces, vert_normals, face_normals) - - return { - 'mesh': mesh, - 'verts': torch.as_tensor(mesh.verts).float(), - 'faces': torch.as_tensor(mesh.faces).long() - } - - def add_noise(self, - beta_num, - smpl_pose, - smpl_betas, - noise_type, - noise_scale, - type, - hashcode): - - np.random.seed(hashcode) - - if type == 'smplx': - noise_idx = self.noise_smplx_idx - else: - noise_idx = self.noise_smpl_idx - - if 'beta' in noise_type and noise_scale[noise_type.index("beta")] > 0.0: - smpl_betas += (np.random.rand(beta_num) - - 0.5) * 2.0 * noise_scale[noise_type.index("beta")] - smpl_betas = smpl_betas.astype(np.float32) - - if 'pose' in noise_type and noise_scale[noise_type.index("pose")] > 0.0: - smpl_pose[noise_idx] += ( - np.random.rand(len(noise_idx)) - - 0.5) * 2.0 * np.pi * noise_scale[noise_type.index("pose")] - smpl_pose = smpl_pose.astype(np.float32) - if type == 'smplx': - return torch.as_tensor(smpl_pose[None, ...]), torch.as_tensor(smpl_betas[None, ...]) - else: - return smpl_pose, smpl_betas - - def compute_smpl_verts(self, data_dict, noise_type=None, noise_scale=None): - - dataset = data_dict['dataset'] - smplx_dict = {} - - smplx_param = np.load(data_dict['smplx_path'], allow_pickle=True) - smplx_pose = smplx_param["body_pose"] # [1,63] - smplx_betas = smplx_param["betas"] # [1,10] - smplx_pose, smplx_betas = self.add_noise( - smplx_betas.shape[1], - smplx_pose[0], - smplx_betas[0], - noise_type, - noise_scale, - type='smplx', - hashcode=(hash(f"{data_dict['subject']}_{data_dict['rotation']}")) % (10**8)) - - smplx_out, _ = load_fit_body(fitted_path=data_dict['smplx_path'], - scale=self.datasets_dict[dataset]['scale'], - smpl_type='smplx', - smpl_gender='male', - noise_dict=dict(betas=smplx_betas, body_pose=smplx_pose)) - - smplx_dict.update({"type": "smplx", - "gender": 'male', - "body_pose": torch.as_tensor(smplx_pose), - "betas": torch.as_tensor(smplx_betas)}) - - return smplx_out.vertices, smplx_dict - - def compute_voxel_verts(self, - data_dict, - noise_type=None, - noise_scale=None): - - smpl_param = np.load(data_dict['smpl_path'], allow_pickle=True) - smplx_param = np.load(data_dict['smplx_path'], allow_pickle=True) - - smpl_pose = rotation_matrix_to_angle_axis( - torch.as_tensor(smpl_param['full_pose'][0])).numpy() - smpl_betas = smpl_param["betas"] - - smpl_path = cached_download(osp.join(self.smplx.model_dir, "smpl/SMPL_MALE.pkl"), use_auth_token=os.environ['ICON']) - tetra_path = cached_download(osp.join(self.smplx.tedra_dir, - "tetra_male_adult_smpl.npz"), use_auth_token=os.environ['ICON']) - - smpl_model = TetraSMPLModel(smpl_path, tetra_path, 'adult') - - smpl_pose, smpl_betas = self.add_noise( - smpl_model.beta_shape[0], - smpl_pose.flatten(), - smpl_betas[0], - noise_type, - noise_scale, - type='smpl', - hashcode=(hash(f"{data_dict['subject']}_{data_dict['rotation']}")) % (10**8)) - - smpl_model.set_params(pose=smpl_pose.reshape(-1, 3), - beta=smpl_betas, - trans=smpl_param["transl"]) - - verts = (np.concatenate([smpl_model.verts, smpl_model.verts_added], - axis=0) * smplx_param["scale"] + smplx_param["translation"] - ) * self.datasets_dict[data_dict['dataset']]['scale'] - faces = np.loadtxt(cached_download(osp.join(self.smplx.tedra_dir, "tetrahedrons_male_adult.txt"), use_auth_token=os.environ['ICON']), - dtype=np.int32) - 1 - - pad_v_num = int(8000 - verts.shape[0]) - pad_f_num = int(25100 - faces.shape[0]) - - verts = np.pad(verts, ((0, pad_v_num), (0, 0)), - mode='constant', - constant_values=0.0).astype(np.float32) - faces = np.pad(faces, ((0, pad_f_num), (0, 0)), - mode='constant', - constant_values=0.0).astype(np.int32) - - - return verts, faces, pad_v_num, pad_f_num - - def load_smpl(self, data_dict, vis=False): - - smplx_verts, smplx_dict = self.compute_smpl_verts( - data_dict, self.noise_type, - self.noise_scale) # compute using smpl model - - smplx_verts = projection(smplx_verts, data_dict['calib']).float() - smplx_faces = torch.as_tensor(self.smplx.faces).long() - smplx_vis = torch.load(data_dict['vis_path']).float() - smplx_cmap = torch.as_tensor( - np.load(self.smplx.cmap_vert_path)).float() - - # get smpl_signs - query_points = projection(data_dict['samples_geo'], - data_dict['calib']).float() - - pts_signs = 2.0 * (check_sign(smplx_verts.unsqueeze(0), - smplx_faces, - query_points.unsqueeze(0)).float() - 0.5).squeeze(0) - - return_dict = { - 'smpl_verts': smplx_verts, - 'smpl_faces': smplx_faces, - 'smpl_vis': smplx_vis, - 'smpl_cmap': smplx_cmap, - 'pts_signs': pts_signs - } - if smplx_dict is not None: - return_dict.update(smplx_dict) - - if vis: - - (xy, z) = torch.as_tensor(smplx_verts).to( - self.device).split([2, 1], dim=1) - smplx_vis = get_visibility(xy, z, torch.as_tensor( - smplx_faces).to(self.device).long()) - - T_normal_F, T_normal_B = self.render_normal( - (smplx_verts*torch.tensor([1.0, -1.0, 1.0])).to(self.device), - smplx_faces.to(self.device)) - - return_dict.update({"T_normal_F": T_normal_F.squeeze(0), - "T_normal_B": T_normal_B.squeeze(0)}) - query_points = projection(data_dict['samples_geo'], - data_dict['calib']).float() - - smplx_sdf, smplx_norm, smplx_cmap, smplx_vis = cal_sdf_batch( - smplx_verts.unsqueeze(0).to(self.device), - smplx_faces.unsqueeze(0).to(self.device), - smplx_cmap.unsqueeze(0).to(self.device), - smplx_vis.unsqueeze(0).to(self.device), - query_points.unsqueeze(0).contiguous().to(self.device)) - - return_dict.update({ - 'smpl_feat': - torch.cat( - (smplx_sdf[0].detach().cpu(), - smplx_cmap[0].detach().cpu(), - smplx_norm[0].detach().cpu(), - smplx_vis[0].detach().cpu()), - dim=1) - }) - - return return_dict - - def load_smpl_voxel(self, data_dict): - - smpl_verts, smpl_faces, pad_v_num, pad_f_num = self.compute_voxel_verts( - data_dict, self.noise_type, - self.noise_scale) # compute using smpl model - smpl_verts = projection(smpl_verts, data_dict['calib']) - - smpl_verts *= 0.5 - - return { - 'voxel_verts': smpl_verts, - 'voxel_faces': smpl_faces, - 'pad_v_num': pad_v_num, - 'pad_f_num': pad_f_num - } - - def get_sampling_geo(self, data_dict, is_valid=False, is_sdf=False): - - mesh = data_dict['mesh'] - calib = data_dict['calib'] - - # Samples are around the true surface with an offset - n_samples_surface = 4 * self.opt.num_sample_geo - vert_ids = np.arange(mesh.verts.shape[0]) - thickness_sample_ratio = np.ones_like(vert_ids).astype(np.float32) - - thickness_sample_ratio /= thickness_sample_ratio.sum() - - samples_surface_ids = np.random.choice(vert_ids, - n_samples_surface, - replace=True, - p=thickness_sample_ratio) - - samples_normal_ids = np.random.choice(vert_ids, - self.opt.num_sample_geo // 2, - replace=False, - p=thickness_sample_ratio) - - surf_samples = mesh.verts[samples_normal_ids, :] - surf_normals = mesh.vert_normals[samples_normal_ids, :] - - samples_surface = mesh.verts[samples_surface_ids, :] - - # Sampling offsets are random noise with constant scale (15cm - 20cm) - offset = np.random.normal(scale=self.opt.sigma_geo, - size=(n_samples_surface, 1)) - samples_surface += mesh.vert_normals[samples_surface_ids, :] * offset - - # Uniform samples in [-1, 1] - calib_inv = np.linalg.inv(calib) - n_samples_space = self.opt.num_sample_geo // 4 - samples_space_img = 2.0 * np.random.rand(n_samples_space, 3) - 1.0 - samples_space = projection(samples_space_img, calib_inv) - - # z-ray direction samples - if self.opt.zray_type and not is_valid: - n_samples_rayz = self.opt.ray_sample_num - samples_surface_cube = projection(samples_surface, calib) - samples_surface_cube_repeat = np.repeat(samples_surface_cube, - n_samples_rayz, - axis=0) - - thickness_repeat = np.repeat(0.5 * - np.ones_like(samples_surface_ids), - n_samples_rayz, - axis=0) - - noise_repeat = np.random.normal(scale=0.40, - size=(n_samples_surface * - n_samples_rayz, )) - samples_surface_cube_repeat[:, - -1] += thickness_repeat * noise_repeat - samples_surface_rayz = projection(samples_surface_cube_repeat, - calib_inv) - - samples = np.concatenate( - [samples_surface, samples_space, samples_surface_rayz], 0) - else: - samples = np.concatenate([samples_surface, samples_space], 0) - - np.random.shuffle(samples) - - # labels: in->1.0; out->0.0. - if is_sdf: - sdfs = mesh.get_sdf(samples) - inside_samples = samples[sdfs < 0] - outside_samples = samples[sdfs >= 0] - - inside_sdfs = sdfs[sdfs < 0] - outside_sdfs = sdfs[sdfs >= 0] - else: - inside = mesh.contains(samples) - inside_samples = samples[inside >= 0.5] - outside_samples = samples[inside < 0.5] - - nin = inside_samples.shape[0] - - if nin > self.opt.num_sample_geo // 2: - inside_samples = inside_samples[:self.opt.num_sample_geo // 2] - outside_samples = outside_samples[:self.opt.num_sample_geo // 2] - if is_sdf: - inside_sdfs = inside_sdfs[:self.opt.num_sample_geo // 2] - outside_sdfs = outside_sdfs[:self.opt.num_sample_geo // 2] - else: - outside_samples = outside_samples[:(self.opt.num_sample_geo - nin)] - if is_sdf: - outside_sdfs = outside_sdfs[:(self.opt.num_sample_geo - nin)] - - if is_sdf: - samples = np.concatenate( - [inside_samples, outside_samples, surf_samples], 0) - - labels = np.concatenate([ - inside_sdfs, outside_sdfs, 0.0 * np.ones(surf_samples.shape[0]) - ]) - - normals = np.zeros_like(samples) - normals[-self.opt.num_sample_geo // 2:, :] = surf_normals - - # convert sdf from [-14, 130] to [0, 1] - # outside: 0, inside: 1 - # Note: Marching cubes is defined on occupancy space (inside=1.0, outside=0.0) - - labels = -labels.clip(min=-self.sdf_clip, max=self.sdf_clip) - labels += self.sdf_clip - labels /= (self.sdf_clip * 2) - - else: - samples = np.concatenate([inside_samples, outside_samples]) - labels = np.concatenate([ - np.ones(inside_samples.shape[0]), - np.zeros(outside_samples.shape[0]) - ]) - - normals = np.zeros_like(samples) - - samples = torch.from_numpy(samples).float() - labels = torch.from_numpy(labels).float() - normals = torch.from_numpy(normals).float() - - return {'samples_geo': samples, 'labels_geo': labels} \ No newline at end of file diff --git a/lib /dataset /TestDataset.py b/lib /dataset /TestDataset.py deleted file mode 100644 index 2ea83cc7f96edfec1eae31f660c23d204a653e38..0000000000000000000000000000000000000000 --- a/lib /dataset /TestDataset.py +++ /dev/null @@ -1,254 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import os - -import lib.smplx as smplx -from lib.pymaf.utils.geometry import rotation_matrix_to_angle_axis, batch_rodrigues -from lib.pymaf.utils.imutils import process_image -from lib.pymaf.core import path_config -from lib.pymaf.models import pymaf_net -from lib.common.config import cfg -from lib.common.render import Render -from lib.dataset.body_model import TetraSMPLModel -from lib.dataset.mesh_util import get_visibility, SMPLX -import os.path as osp -import torch -import numpy as np -import random -from termcolor import colored -from PIL import ImageFile -from huggingface_hub import cached_download - -ImageFile.LOAD_TRUNCATED_IMAGES = True - - -class TestDataset(): - def __init__(self, cfg, device): - - random.seed(1993) - - self.image_path = cfg['image_path'] - self.seg_dir = cfg['seg_dir'] - self.has_det = cfg['has_det'] - self.hps_type = cfg['hps_type'] - self.smpl_type = 'smpl' if cfg['hps_type'] != 'pixie' else 'smplx' - self.smpl_gender = 'neutral' - - self.device = device - - self.subject_list = [self.image_path] - - # smpl related - self.smpl_data = SMPLX() - - self.get_smpl_model = lambda smpl_type, smpl_gender: smplx.create( - model_path=self.smpl_data.model_dir, - gender=smpl_gender, - model_type=smpl_type, - ext='npz') - - # Load SMPL model - self.smpl_model = self.get_smpl_model( - self.smpl_type, self.smpl_gender).to(self.device) - self.faces = self.smpl_model.faces - - self.hps = pymaf_net(path_config.SMPL_MEAN_PARAMS, - pretrained=True).to(self.device) - self.hps.load_state_dict(torch.load( - path_config.CHECKPOINT_FILE)['model'], - strict=True) - self.hps.eval() - - print(colored(f"Using {self.hps_type} as HPS Estimator\n", "green")) - - self.render = Render(size=512, device=device) - - def __len__(self): - return len(self.subject_list) - - def compute_vis_cmap(self, smpl_verts, smpl_faces): - - (xy, z) = torch.as_tensor(smpl_verts).split([2, 1], dim=1) - smpl_vis = get_visibility(xy, -z, torch.as_tensor(smpl_faces).long()) - if self.smpl_type == 'smpl': - smplx_ind = self.smpl_data.smpl2smplx(np.arange(smpl_vis.shape[0])) - else: - smplx_ind = np.arange(smpl_vis.shape[0]) - smpl_cmap = self.smpl_data.get_smpl_mat(smplx_ind) - - return { - 'smpl_vis': smpl_vis.unsqueeze(0).to(self.device), - 'smpl_cmap': smpl_cmap.unsqueeze(0).to(self.device), - 'smpl_verts': smpl_verts.unsqueeze(0) - } - - def compute_voxel_verts(self, body_pose, global_orient, betas, trans, - scale): - - smpl_path = cached_download(osp.join(self.smpl_data.model_dir, "smpl/SMPL_NEUTRAL.pkl"), use_auth_token=os.environ['ICON']) - tetra_path = cached_download(osp.join(self.smpl_data.tedra_dir, - 'tetra_neutral_adult_smpl.npz'), use_auth_token=os.environ['ICON']) - smpl_model = TetraSMPLModel(smpl_path, tetra_path, 'adult') - - pose = torch.cat([global_orient[0], body_pose[0]], dim=0) - smpl_model.set_params(rotation_matrix_to_angle_axis(pose), - beta=betas[0]) - - verts = np.concatenate( - [smpl_model.verts, smpl_model.verts_added], - axis=0) * scale.item() + trans.detach().cpu().numpy() - faces = np.loadtxt(cached_download(osp.join(self.smpl_data.tedra_dir, - 'tetrahedrons_neutral_adult.txt'), use_auth_token=os.environ['ICON']), - dtype=np.int32) - 1 - - pad_v_num = int(8000 - verts.shape[0]) - pad_f_num = int(25100 - faces.shape[0]) - - verts = np.pad(verts, ((0, pad_v_num), (0, 0)), - mode='constant', - constant_values=0.0).astype(np.float32) * 0.5 - faces = np.pad(faces, ((0, pad_f_num), (0, 0)), - mode='constant', - constant_values=0.0).astype(np.int32) - - verts[:, 2] *= -1.0 - - voxel_dict = { - 'voxel_verts': - torch.from_numpy(verts).to(self.device).unsqueeze(0).float(), - 'voxel_faces': - torch.from_numpy(faces).to(self.device).unsqueeze(0).long(), - 'pad_v_num': - torch.tensor(pad_v_num).to(self.device).unsqueeze(0).long(), - 'pad_f_num': - torch.tensor(pad_f_num).to(self.device).unsqueeze(0).long() - } - - return voxel_dict - - def __getitem__(self, index): - - img_path = self.subject_list[index] - img_name = img_path.split("/")[-1].rsplit(".", 1)[0] - - if self.seg_dir is None: - img_icon, img_hps, img_ori, img_mask, uncrop_param = process_image( - img_path, self.hps_type, 512, self.device) - - data_dict = { - 'name': img_name, - 'image': img_icon.to(self.device).unsqueeze(0), - 'ori_image': img_ori, - 'mask': img_mask, - 'uncrop_param': uncrop_param - } - - else: - img_icon, img_hps, img_ori, img_mask, uncrop_param, segmentations = process_image( - img_path, self.hps_type, 512, self.device, - seg_path=os.path.join(self.seg_dir, f'{img_name}.json')) - data_dict = { - 'name': img_name, - 'image': img_icon.to(self.device).unsqueeze(0), - 'ori_image': img_ori, - 'mask': img_mask, - 'uncrop_param': uncrop_param, - 'segmentations': segmentations - } - - with torch.no_grad(): - # import ipdb; ipdb.set_trace() - preds_dict = self.hps.forward(img_hps) - - data_dict['smpl_faces'] = torch.Tensor( - self.faces.astype(np.int16)).long().unsqueeze(0).to( - self.device) - - if self.hps_type == 'pymaf': - output = preds_dict['smpl_out'][-1] - scale, tranX, tranY = output['theta'][0, :3] - data_dict['betas'] = output['pred_shape'] - data_dict['body_pose'] = output['rotmat'][:, 1:] - data_dict['global_orient'] = output['rotmat'][:, 0:1] - data_dict['smpl_verts'] = output['verts'] - - elif self.hps_type == 'pare': - data_dict['body_pose'] = preds_dict['pred_pose'][:, 1:] - data_dict['global_orient'] = preds_dict['pred_pose'][:, 0:1] - data_dict['betas'] = preds_dict['pred_shape'] - data_dict['smpl_verts'] = preds_dict['smpl_vertices'] - scale, tranX, tranY = preds_dict['pred_cam'][0, :3] - - elif self.hps_type == 'pixie': - data_dict.update(preds_dict) - data_dict['body_pose'] = preds_dict['body_pose'] - data_dict['global_orient'] = preds_dict['global_pose'] - data_dict['betas'] = preds_dict['shape'] - data_dict['smpl_verts'] = preds_dict['vertices'] - scale, tranX, tranY = preds_dict['cam'][0, :3] - - elif self.hps_type == 'hybrik': - data_dict['body_pose'] = preds_dict['pred_theta_mats'][:, 1:] - data_dict['global_orient'] = preds_dict['pred_theta_mats'][:, [0]] - data_dict['betas'] = preds_dict['pred_shape'] - data_dict['smpl_verts'] = preds_dict['pred_vertices'] - scale, tranX, tranY = preds_dict['pred_camera'][0, :3] - scale = scale * 2 - - elif self.hps_type == 'bev': - data_dict['betas'] = torch.from_numpy(preds_dict['smpl_betas'])[ - [0], :10].to(self.device).float() - pred_thetas = batch_rodrigues(torch.from_numpy( - preds_dict['smpl_thetas'][0]).reshape(-1, 3)).float() - data_dict['body_pose'] = pred_thetas[1:][None].to(self.device) - data_dict['global_orient'] = pred_thetas[[0]][None].to(self.device) - data_dict['smpl_verts'] = torch.from_numpy( - preds_dict['verts'][[0]]).to(self.device).float() - tranX = preds_dict['cam_trans'][0, 0] - tranY = preds_dict['cam'][0, 1] + 0.28 - scale = preds_dict['cam'][0, 0] * 1.1 - - data_dict['scale'] = scale - data_dict['trans'] = torch.tensor( - [tranX, tranY, 0.0]).to(self.device).float() - - # data_dict info (key-shape): - # scale, tranX, tranY - tensor.float - # betas - [1,10] / [1, 200] - # body_pose - [1, 23, 3, 3] / [1, 21, 3, 3] - # global_orient - [1, 1, 3, 3] - # smpl_verts - [1, 6890, 3] / [1, 10475, 3] - - # from rot_mat to rot_6d for better optimization - N_body = data_dict["body_pose"].shape[1] - data_dict["body_pose"] = data_dict["body_pose"][:, :, :, :2].reshape(1, N_body,-1) - data_dict["global_orient"] = data_dict["global_orient"][:, :, :, :2].reshape(1, 1,-1) - - return data_dict - - def render_normal(self, verts, faces): - - # render optimized mesh (normal, T_normal, image [-1,1]) - self.render.load_meshes(verts, faces) - return self.render.get_rgb_image() - - def render_depth(self, verts, faces): - - # render optimized mesh (normal, T_normal, image [-1,1]) - self.render.load_meshes(verts, faces) - return self.render.get_depth_map(cam_ids=[0, 2]) \ No newline at end of file diff --git a/lib /dataset /__init__.py b/lib /dataset /__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/lib /dataset /body_model.py b/lib /dataset /body_model.py deleted file mode 100644 index 5f46fe7db22cdd61f372094c6d89bf12b5a8d14e..0000000000000000000000000000000000000000 --- a/lib /dataset /body_model.py +++ /dev/null @@ -1,454 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import numpy as np -import pickle -import torch -import os - - -class SMPLModel(): - def __init__(self, model_path, age): - """ - SMPL model. - Parameter: - --------- - model_path: Path to the SMPL model parameters, pre-processed by - `preprocess.py`. - """ - with open(model_path, 'rb') as f: - params = pickle.load(f, encoding='latin1') - - self.J_regressor = params['J_regressor'] - self.weights = np.asarray(params['weights']) - self.posedirs = np.asarray(params['posedirs']) - self.v_template = np.asarray(params['v_template']) - self.shapedirs = np.asarray(params['shapedirs']) - self.faces = np.asarray(params['f']) - self.kintree_table = np.asarray(params['kintree_table']) - - self.pose_shape = [24, 3] - self.beta_shape = [10] - self.trans_shape = [3] - - if age == 'kid': - v_template_smil = np.load( - os.path.join(os.path.dirname(model_path), - "smpl/smpl_kid_template.npy")) - v_template_smil -= np.mean(v_template_smil, axis=0) - v_template_diff = np.expand_dims(v_template_smil - self.v_template, - axis=2) - self.shapedirs = np.concatenate( - (self.shapedirs[:, :, :self.beta_shape[0]], v_template_diff), - axis=2) - self.beta_shape[0] += 1 - - id_to_col = { - self.kintree_table[1, i]: i - for i in range(self.kintree_table.shape[1]) - } - self.parent = { - i: id_to_col[self.kintree_table[0, i]] - for i in range(1, self.kintree_table.shape[1]) - } - - self.pose = np.zeros(self.pose_shape) - self.beta = np.zeros(self.beta_shape) - self.trans = np.zeros(self.trans_shape) - - self.verts = None - self.J = None - self.R = None - self.G = None - - self.update() - - def set_params(self, pose=None, beta=None, trans=None): - """ - Set pose, shape, and/or translation parameters of SMPL model. Verices of the - model will be updated and returned. - Prameters: - --------- - pose: Also known as 'theta', a [24,3] matrix indicating child joint rotation - relative to parent joint. For root joint it's global orientation. - Represented in a axis-angle format. - beta: Parameter for model shape. A vector of shape [10]. Coefficients for - PCA component. Only 10 components were released by MPI. - trans: Global translation of shape [3]. - Return: - ------ - Updated vertices. - """ - if pose is not None: - self.pose = pose - if beta is not None: - self.beta = beta - if trans is not None: - self.trans = trans - self.update() - return self.verts - - def update(self): - """ - Called automatically when parameters are updated. - """ - # how beta affect body shape - v_shaped = self.shapedirs.dot(self.beta) + self.v_template - # joints location - self.J = self.J_regressor.dot(v_shaped) - pose_cube = self.pose.reshape((-1, 1, 3)) - # rotation matrix for each joint - self.R = self.rodrigues(pose_cube) - I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), - (self.R.shape[0] - 1, 3, 3)) - lrotmin = (self.R[1:] - I_cube).ravel() - # how pose affect body shape in zero pose - v_posed = v_shaped + self.posedirs.dot(lrotmin) - # world transformation of each joint - G = np.empty((self.kintree_table.shape[1], 4, 4)) - G[0] = self.with_zeros( - np.hstack((self.R[0], self.J[0, :].reshape([3, 1])))) - for i in range(1, self.kintree_table.shape[1]): - G[i] = G[self.parent[i]].dot( - self.with_zeros( - np.hstack([ - self.R[i], - ((self.J[i, :] - self.J[self.parent[i], :]).reshape( - [3, 1])) - ]))) - # remove the transformation due to the rest pose - G = G - self.pack( - np.matmul( - G, - np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1]))) - # transformation of each vertex - T = np.tensordot(self.weights, G, axes=[[1], [0]]) - rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1]))) - v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1, - 4])[:, :3] - self.verts = v + self.trans.reshape([1, 3]) - self.G = G - - def rodrigues(self, r): - """ - Rodrigues' rotation formula that turns axis-angle vector into rotation - matrix in a batch-ed manner. - Parameter: - ---------- - r: Axis-angle rotation vector of shape [batch_size, 1, 3]. - Return: - ------- - Rotation matrix of shape [batch_size, 3, 3]. - """ - theta = np.linalg.norm(r, axis=(1, 2), keepdims=True) - # avoid zero divide - theta = np.maximum(theta, np.finfo(np.float64).tiny) - r_hat = r / theta - cos = np.cos(theta) - z_stick = np.zeros(theta.shape[0]) - m = np.dstack([ - z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick, - -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick - ]).reshape([-1, 3, 3]) - i_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), - [theta.shape[0], 3, 3]) - A = np.transpose(r_hat, axes=[0, 2, 1]) - B = r_hat - dot = np.matmul(A, B) - R = cos * i_cube + (1 - cos) * dot + np.sin(theta) * m - return R - - def with_zeros(self, x): - """ - Append a [0, 0, 0, 1] vector to a [3, 4] matrix. - Parameter: - --------- - x: Matrix to be appended. - Return: - ------ - Matrix after appending of shape [4,4] - """ - return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]]))) - - def pack(self, x): - """ - Append zero matrices of shape [4, 3] to vectors of [4, 1] shape in a batched - manner. - Parameter: - ---------- - x: Matrices to be appended of shape [batch_size, 4, 1] - Return: - ------ - Matrix of shape [batch_size, 4, 4] after appending. - """ - return np.dstack((np.zeros((x.shape[0], 4, 3)), x)) - - def save_to_obj(self, path): - """ - Save the SMPL model into .obj file. - Parameter: - --------- - path: Path to save. - """ - with open(path, 'w') as fp: - for v in self.verts: - fp.write('v %f %f %f\n' % (v[0], v[1], v[2])) - for f in self.faces + 1: - fp.write('f %d %d %d\n' % (f[0], f[1], f[2])) - - -class TetraSMPLModel(): - def __init__(self, - model_path, - model_addition_path, - age='adult', - v_template=None): - """ - SMPL model. - Parameter: - --------- - model_path: Path to the SMPL model parameters, pre-processed by - `preprocess.py`. - """ - with open(model_path, 'rb') as f: - params = pickle.load(f, encoding='latin1') - - self.J_regressor = params['J_regressor'] - self.weights = np.asarray(params['weights']) - self.posedirs = np.asarray(params['posedirs']) - - if v_template is not None: - self.v_template = v_template - else: - self.v_template = np.asarray(params['v_template']) - - self.shapedirs = np.asarray(params['shapedirs']) - self.faces = np.asarray(params['f']) - self.kintree_table = np.asarray(params['kintree_table']) - - params_added = np.load(model_addition_path) - self.v_template_added = params_added['v_template_added'] - self.weights_added = params_added['weights_added'] - self.shapedirs_added = params_added['shapedirs_added'] - self.posedirs_added = params_added['posedirs_added'] - self.tetrahedrons = params_added['tetrahedrons'] - - id_to_col = { - self.kintree_table[1, i]: i - for i in range(self.kintree_table.shape[1]) - } - self.parent = { - i: id_to_col[self.kintree_table[0, i]] - for i in range(1, self.kintree_table.shape[1]) - } - - self.pose_shape = [24, 3] - self.beta_shape = [10] - self.trans_shape = [3] - - if age == 'kid': - v_template_smil = np.load( - os.path.join(os.path.dirname(model_path), - "smpl/smpl_kid_template.npy")) - v_template_smil -= np.mean(v_template_smil, axis=0) - v_template_diff = np.expand_dims(v_template_smil - self.v_template, - axis=2) - self.shapedirs = np.concatenate( - (self.shapedirs[:, :, :self.beta_shape[0]], v_template_diff), - axis=2) - self.beta_shape[0] += 1 - - self.pose = np.zeros(self.pose_shape) - self.beta = np.zeros(self.beta_shape) - self.trans = np.zeros(self.trans_shape) - - self.verts = None - self.verts_added = None - self.J = None - self.R = None - self.G = None - - self.update() - - def set_params(self, pose=None, beta=None, trans=None): - """ - Set pose, shape, and/or translation parameters of SMPL model. Verices of the - model will be updated and returned. - Prameters: - --------- - pose: Also known as 'theta', a [24,3] matrix indicating child joint rotation - relative to parent joint. For root joint it's global orientation. - Represented in a axis-angle format. - beta: Parameter for model shape. A vector of shape [10]. Coefficients for - PCA component. Only 10 components were released by MPI. - trans: Global translation of shape [3]. - Return: - ------ - Updated vertices. - """ - - if torch.is_tensor(pose): - pose = pose.detach().cpu().numpy() - if torch.is_tensor(beta): - beta = beta.detach().cpu().numpy() - - if pose is not None: - self.pose = pose - if beta is not None: - self.beta = beta - if trans is not None: - self.trans = trans - self.update() - return self.verts - - def update(self): - """ - Called automatically when parameters are updated. - """ - # how beta affect body shape - v_shaped = self.shapedirs.dot(self.beta) + self.v_template - v_shaped_added = self.shapedirs_added.dot( - self.beta) + self.v_template_added - # joints location - self.J = self.J_regressor.dot(v_shaped) - pose_cube = self.pose.reshape((-1, 1, 3)) - # rotation matrix for each joint - self.R = self.rodrigues(pose_cube) - I_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), - (self.R.shape[0] - 1, 3, 3)) - lrotmin = (self.R[1:] - I_cube).ravel() - # how pose affect body shape in zero pose - v_posed = v_shaped + self.posedirs.dot(lrotmin) - v_posed_added = v_shaped_added + self.posedirs_added.dot(lrotmin) - # world transformation of each joint - G = np.empty((self.kintree_table.shape[1], 4, 4)) - G[0] = self.with_zeros( - np.hstack((self.R[0], self.J[0, :].reshape([3, 1])))) - for i in range(1, self.kintree_table.shape[1]): - G[i] = G[self.parent[i]].dot( - self.with_zeros( - np.hstack([ - self.R[i], - ((self.J[i, :] - self.J[self.parent[i], :]).reshape( - [3, 1])) - ]))) - # remove the transformation due to the rest pose - G = G - self.pack( - np.matmul( - G, - np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1]))) - self.G = G - # transformation of each vertex - T = np.tensordot(self.weights, G, axes=[[1], [0]]) - rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1]))) - v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1, - 4])[:, :3] - self.verts = v + self.trans.reshape([1, 3]) - T_added = np.tensordot(self.weights_added, G, axes=[[1], [0]]) - rest_shape_added_h = np.hstack( - (v_posed_added, np.ones([v_posed_added.shape[0], 1]))) - v_added = np.matmul(T_added, - rest_shape_added_h.reshape([-1, 4, - 1])).reshape([-1, 4 - ])[:, :3] - self.verts_added = v_added + self.trans.reshape([1, 3]) - - def rodrigues(self, r): - """ - Rodrigues' rotation formula that turns axis-angle vector into rotation - matrix in a batch-ed manner. - Parameter: - ---------- - r: Axis-angle rotation vector of shape [batch_size, 1, 3]. - Return: - ------- - Rotation matrix of shape [batch_size, 3, 3]. - """ - theta = np.linalg.norm(r, axis=(1, 2), keepdims=True) - # avoid zero divide - theta = np.maximum(theta, np.finfo(np.float64).tiny) - r_hat = r / theta - cos = np.cos(theta) - z_stick = np.zeros(theta.shape[0]) - m = np.dstack([ - z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick, - -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick - ]).reshape([-1, 3, 3]) - i_cube = np.broadcast_to(np.expand_dims(np.eye(3), axis=0), - [theta.shape[0], 3, 3]) - A = np.transpose(r_hat, axes=[0, 2, 1]) - B = r_hat - dot = np.matmul(A, B) - R = cos * i_cube + (1 - cos) * dot + np.sin(theta) * m - return R - - def with_zeros(self, x): - """ - Append a [0, 0, 0, 1] vector to a [3, 4] matrix. - Parameter: - --------- - x: Matrix to be appended. - Return: - ------ - Matrix after appending of shape [4,4] - """ - return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]]))) - - def pack(self, x): - """ - Append zero matrices of shape [4, 3] to vectors of [4, 1] shape in a batched - manner. - Parameter: - ---------- - x: Matrices to be appended of shape [batch_size, 4, 1] - Return: - ------ - Matrix of shape [batch_size, 4, 4] after appending. - """ - return np.dstack((np.zeros((x.shape[0], 4, 3)), x)) - - def save_mesh_to_obj(self, path): - """ - Save the SMPL model into .obj file. - Parameter: - --------- - path: Path to save. - """ - with open(path, 'w') as fp: - for v in self.verts: - fp.write('v %f %f %f\n' % (v[0], v[1], v[2])) - for f in self.faces + 1: - fp.write('f %d %d %d\n' % (f[0], f[1], f[2])) - - def save_tetrahedron_to_obj(self, path): - """ - Save the tetrahedron SMPL model into .obj file. - Parameter: - --------- - path: Path to save. - """ - - with open(path, 'w') as fp: - for v in self.verts: - fp.write('v %f %f %f 1 0 0\n' % (v[0], v[1], v[2])) - for va in self.verts_added: - fp.write('v %f %f %f 0 0 1\n' % (va[0], va[1], va[2])) - for t in self.tetrahedrons + 1: - fp.write('f %d %d %d\n' % (t[0], t[2], t[1])) - fp.write('f %d %d %d\n' % (t[0], t[3], t[2])) - fp.write('f %d %d %d\n' % (t[0], t[1], t[3])) - fp.write('f %d %d %d\n' % (t[1], t[2], t[3])) \ No newline at end of file diff --git a/lib /dataset /mesh_util.py b/lib /dataset /mesh_util.py deleted file mode 100644 index 6109d53323630252adbae092482dae1d26c331be..0000000000000000000000000000000000000000 --- a/lib /dataset /mesh_util.py +++ /dev/null @@ -1,905 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import numpy as np -import cv2 -import pymeshlab -import torch -import torchvision -import trimesh -from pytorch3d.io import load_obj -from termcolor import colored -from scipy.spatial import cKDTree - -from pytorch3d.structures import Meshes -import torch.nn.functional as F - -import os -from lib.pymaf.utils.imutils import uncrop -from lib.common.render_utils import Pytorch3dRasterizer, face_vertices - -from pytorch3d.renderer.mesh import rasterize_meshes -from PIL import Image, ImageFont, ImageDraw -from kaolin.ops.mesh import check_sign -from kaolin.metrics.trianglemesh import point_to_mesh_distance - -from pytorch3d.loss import ( - mesh_laplacian_smoothing, - mesh_normal_consistency -) - -from huggingface_hub import hf_hub_download, hf_hub_url, cached_download - -def rot6d_to_rotmat(x): - """Convert 6D rotation representation to 3x3 rotation matrix. - Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019 - Input: - (B,6) Batch of 6-D rotation representations - Output: - (B,3,3) Batch of corresponding rotation matrices - """ - x = x.view(-1, 3, 2) - a1 = x[:, :, 0] - a2 = x[:, :, 1] - b1 = F.normalize(a1) - b2 = F.normalize(a2 - torch.einsum("bi,bi->b", b1, a2).unsqueeze(-1) * b1) - b3 = torch.cross(b1, b2) - return torch.stack((b1, b2, b3), dim=-1) - - -def tensor2variable(tensor, device): - # [1,23,3,3] - return torch.tensor(tensor, device=device, requires_grad=True) - - -def normal_loss(vec1, vec2): - - # vec1_mask = vec1.sum(dim=1) != 0.0 - # vec2_mask = vec2.sum(dim=1) != 0.0 - # union_mask = vec1_mask * vec2_mask - vec_sim = torch.nn.CosineSimilarity(dim=1, eps=1e-6)(vec1, vec2) - # vec_diff = ((vec_sim-1.0)**2)[union_mask].mean() - vec_diff = ((vec_sim-1.0)**2).mean() - - return vec_diff - - -class GMoF(torch.nn.Module): - def __init__(self, rho=1): - super(GMoF, self).__init__() - self.rho = rho - - def extra_repr(self): - return 'rho = {}'.format(self.rho) - - def forward(self, residual): - dist = torch.div(residual, residual + self.rho ** 2) - return self.rho ** 2 * dist - - -def mesh_edge_loss(meshes, target_length: float = 0.0): - """ - Computes mesh edge length regularization loss averaged across all meshes - in a batch. Each mesh contributes equally to the final loss, regardless of - the number of edges per mesh in the batch by weighting each mesh with the - inverse number of edges. For example, if mesh 3 (out of N) has only E=4 - edges, then the loss for each edge in mesh 3 should be multiplied by 1/E to - contribute to the final loss. - Args: - meshes: Meshes object with a batch of meshes. - target_length: Resting value for the edge length. - Returns: - loss: Average loss across the batch. Returns 0 if meshes contains - no meshes or all empty meshes. - """ - if meshes.isempty(): - return torch.tensor( - [0.0], dtype=torch.float32, device=meshes.device, requires_grad=True - ) - - N = len(meshes) - edges_packed = meshes.edges_packed() # (sum(E_n), 3) - verts_packed = meshes.verts_packed() # (sum(V_n), 3) - edge_to_mesh_idx = meshes.edges_packed_to_mesh_idx() # (sum(E_n), ) - num_edges_per_mesh = meshes.num_edges_per_mesh() # N - - # Determine the weight for each edge based on the number of edges in the - # mesh it corresponds to. - # TODO (nikhilar) Find a faster way of computing the weights for each edge - # as this is currently a bottleneck for meshes with a large number of faces. - weights = num_edges_per_mesh.gather(0, edge_to_mesh_idx) - weights = 1.0 / weights.float() - - verts_edges = verts_packed[edges_packed] - v0, v1 = verts_edges.unbind(1) - loss = ((v0 - v1).norm(dim=1, p=2) - target_length) ** 2.0 - loss_vertex = loss * weights - # loss_outlier = torch.topk(loss, 100)[0].mean() - # loss_all = (loss_vertex.sum() + loss_outlier.mean()) / N - loss_all = loss_vertex.sum() / N - - return loss_all - - -def remesh(obj_path, perc, device): - - ms = pymeshlab.MeshSet() - ms.load_new_mesh(obj_path) - ms.laplacian_smooth() - ms.remeshing_isotropic_explicit_remeshing( - targetlen=pymeshlab.Percentage(perc), adaptive=True) - ms.save_current_mesh(obj_path.replace("recon", "remesh")) - polished_mesh = trimesh.load_mesh(obj_path.replace("recon", "remesh")) - verts_pr = torch.tensor(polished_mesh.vertices).float().unsqueeze(0).to(device) - faces_pr = torch.tensor(polished_mesh.faces).long().unsqueeze(0).to(device) - - return verts_pr, faces_pr - - -def possion(mesh, obj_path): - - mesh.export(obj_path) - ms = pymeshlab.MeshSet() - ms.load_new_mesh(obj_path) - ms.surface_reconstruction_screened_poisson(depth=10) - ms.set_current_mesh(1) - ms.save_current_mesh(obj_path) - - return trimesh.load(obj_path) - - -def get_mask(tensor, dim): - - mask = torch.abs(tensor).sum(dim=dim, keepdims=True) > 0.0 - mask = mask.type_as(tensor) - - return mask - - -def blend_rgb_norm(rgb, norm, mask): - - # [0,0,0] or [127,127,127] should be marked as mask - final = rgb * (1-mask) + norm * (mask) - - return final.astype(np.uint8) - - -def unwrap(image, data): - - img_uncrop = uncrop(np.array(Image.fromarray(image).resize(data['uncrop_param']['box_shape'][:2])), - data['uncrop_param']['center'], - data['uncrop_param']['scale'], - data['uncrop_param']['crop_shape']) - - img_orig = cv2.warpAffine(img_uncrop, - np.linalg.inv(data['uncrop_param']['M'])[:2, :], - data['uncrop_param']['ori_shape'][::-1][1:], - flags=cv2.INTER_CUBIC) - - return img_orig - - -# Losses to smooth / regularize the mesh shape -def update_mesh_shape_prior_losses(mesh, losses): - - # and (b) the edge length of the predicted mesh - losses["edge"]['value'] = mesh_edge_loss(mesh) - # mesh normal consistency - losses["nc"]['value'] = mesh_normal_consistency(mesh) - # mesh laplacian smoothing - losses["laplacian"]['value'] = mesh_laplacian_smoothing( - mesh, method="uniform") - - -def rename(old_dict, old_name, new_name): - new_dict = {} - for key, value in zip(old_dict.keys(), old_dict.values()): - new_key = key if key != old_name else new_name - new_dict[new_key] = old_dict[key] - return new_dict - - -def load_checkpoint(model, cfg): - - model_dict = model.state_dict() - main_dict = {} - normal_dict = {} - - device = torch.device(f"cuda:{cfg['test_gpus'][0]}") - - main_dict = torch.load(cached_download(cfg.resume_path, use_auth_token=os.environ['ICON']), - map_location=device)['state_dict'] - - main_dict = { - k: v - for k, v in main_dict.items() - if k in model_dict and v.shape == model_dict[k].shape and ( - 'reconEngine' not in k) and ("normal_filter" not in k) and ( - 'voxelization' not in k) - } - print(colored(f"Resume MLP weights from {cfg.resume_path}", 'green')) - - normal_dict = torch.load(cached_download(cfg.normal_path, use_auth_token=os.environ['ICON']), - map_location=device)['state_dict'] - - for key in normal_dict.keys(): - normal_dict = rename(normal_dict, key, - key.replace("netG", "netG.normal_filter")) - - normal_dict = { - k: v - for k, v in normal_dict.items() - if k in model_dict and v.shape == model_dict[k].shape - } - print(colored(f"Resume normal model from {cfg.normal_path}", 'green')) - - model_dict.update(main_dict) - model_dict.update(normal_dict) - model.load_state_dict(model_dict) - - model.netG = model.netG.to(device) - model.reconEngine = model.reconEngine.to(device) - - model.netG.training = False - model.netG.eval() - - del main_dict - del normal_dict - del model_dict - - return model - - -def read_smpl_constants(folder): - """Load smpl vertex code""" - smpl_vtx_std = np.loadtxt(cached_download(os.path.join(folder, 'vertices.txt'), use_auth_token=os.environ['ICON'])) - min_x = np.min(smpl_vtx_std[:, 0]) - max_x = np.max(smpl_vtx_std[:, 0]) - min_y = np.min(smpl_vtx_std[:, 1]) - max_y = np.max(smpl_vtx_std[:, 1]) - min_z = np.min(smpl_vtx_std[:, 2]) - max_z = np.max(smpl_vtx_std[:, 2]) - - smpl_vtx_std[:, 0] = (smpl_vtx_std[:, 0] - min_x) / (max_x - min_x) - smpl_vtx_std[:, 1] = (smpl_vtx_std[:, 1] - min_y) / (max_y - min_y) - smpl_vtx_std[:, 2] = (smpl_vtx_std[:, 2] - min_z) / (max_z - min_z) - smpl_vertex_code = np.float32(np.copy(smpl_vtx_std)) - """Load smpl faces & tetrahedrons""" - smpl_faces = np.loadtxt(cached_download(os.path.join(folder, 'faces.txt'), use_auth_token=os.environ['ICON']), - dtype=np.int32) - 1 - smpl_face_code = (smpl_vertex_code[smpl_faces[:, 0]] + - smpl_vertex_code[smpl_faces[:, 1]] + - smpl_vertex_code[smpl_faces[:, 2]]) / 3.0 - smpl_tetras = np.loadtxt(cached_download(os.path.join(folder, 'tetrahedrons.txt'), use_auth_token=os.environ['ICON']), - dtype=np.int32) - 1 - - return smpl_vertex_code, smpl_face_code, smpl_faces, smpl_tetras - - -def feat_select(feat, select): - - # feat [B, featx2, N] - # select [B, 1, N] - # return [B, feat, N] - - dim = feat.shape[1] // 2 - idx = torch.tile((1-select), (1, dim, 1))*dim + \ - torch.arange(0, dim).unsqueeze(0).unsqueeze(2).type_as(select) - feat_select = torch.gather(feat, 1, idx.long()) - - return feat_select - - -def get_visibility(xy, z, faces): - """get the visibility of vertices - Args: - xy (torch.tensor): [N,2] - z (torch.tensor): [N,1] - faces (torch.tensor): [N,3] - size (int): resolution of rendered image - """ - - xyz = torch.cat((xy, -z), dim=1) - xyz = (xyz + 1.0) / 2.0 - faces = faces.long() - - rasterizer = Pytorch3dRasterizer(image_size=2**12) - meshes_screen = Meshes(verts=xyz[None, ...], faces=faces[None, ...]) - raster_settings = rasterizer.raster_settings - - pix_to_face, zbuf, bary_coords, dists = rasterize_meshes( - meshes_screen, - image_size=raster_settings.image_size, - blur_radius=raster_settings.blur_radius, - faces_per_pixel=raster_settings.faces_per_pixel, - bin_size=raster_settings.bin_size, - max_faces_per_bin=raster_settings.max_faces_per_bin, - perspective_correct=raster_settings.perspective_correct, - cull_backfaces=raster_settings.cull_backfaces, - ) - - vis_vertices_id = torch.unique(faces[torch.unique(pix_to_face), :]) - vis_mask = torch.zeros(size=(z.shape[0], 1)) - vis_mask[vis_vertices_id] = 1.0 - - # print("------------------------\n") - # print(f"keep points : {vis_mask.sum()/len(vis_mask)}") - - return vis_mask - - -def barycentric_coordinates_of_projection(points, vertices): - ''' https://github.com/MPI-IS/mesh/blob/master/mesh/geometry/barycentric_coordinates_of_projection.py - ''' - """Given a point, gives projected coords of that point to a triangle - in barycentric coordinates. - See - **Heidrich**, Computing the Barycentric Coordinates of a Projected Point, JGT 05 - at http://www.cs.ubc.ca/~heidrich/Papers/JGT.05.pdf - - :param p: point to project. [B, 3] - :param v0: first vertex of triangles. [B, 3] - :returns: barycentric coordinates of ``p``'s projection in triangle defined by ``q``, ``u``, ``v`` - vectorized so ``p``, ``q``, ``u``, ``v`` can all be ``3xN`` - """ - #(p, q, u, v) - v0, v1, v2 = vertices[:, 0], vertices[:, 1], vertices[:, 2] - p = points - - q = v0 - u = v1 - v0 - v = v2 - v0 - n = torch.cross(u, v) - s = torch.sum(n * n, dim=1) - # If the triangle edges are collinear, cross-product is zero, - # which makes "s" 0, which gives us divide by zero. So we - # make the arbitrary choice to set s to epsv (=numpy.spacing(1)), - # the closest thing to zero - s[s == 0] = 1e-6 - oneOver4ASquared = 1.0 / s - w = p - q - b2 = torch.sum(torch.cross(u, w) * n, dim=1) * oneOver4ASquared - b1 = torch.sum(torch.cross(w, v) * n, dim=1) * oneOver4ASquared - weights = torch.stack((1 - b1 - b2, b1, b2), dim=-1) - # check barycenric weights - # p_n = v0*weights[:,0:1] + v1*weights[:,1:2] + v2*weights[:,2:3] - return weights - - -def cal_sdf_batch(verts, faces, cmaps, vis, points): - - # verts [B, N_vert, 3] - # faces [B, N_face, 3] - # triangles [B, N_face, 3, 3] - # points [B, N_point, 3] - # cmaps [B, N_vert, 3] - - Bsize = points.shape[0] - - normals = Meshes(verts, faces).verts_normals_padded() - - triangles = face_vertices(verts, faces) - normals = face_vertices(normals, faces) - cmaps = face_vertices(cmaps, faces) - vis = face_vertices(vis, faces) - - residues, pts_ind, _ = point_to_mesh_distance(points, triangles) - closest_triangles = torch.gather( - triangles, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3) - closest_normals = torch.gather( - normals, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3) - closest_cmaps = torch.gather( - cmaps, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 3)).view(-1, 3, 3) - closest_vis = torch.gather( - vis, 1, pts_ind[:, :, None, None].expand(-1, -1, 3, 1)).view(-1, 3, 1) - bary_weights = barycentric_coordinates_of_projection( - points.view(-1, 3), closest_triangles) - - pts_cmap = (closest_cmaps*bary_weights[:, :, None]).sum(1).unsqueeze(0).clamp_(min=0.0, max=1.0) - pts_vis = (closest_vis*bary_weights[:, - :, None]).sum(1).unsqueeze(0).ge(1e-1) - pts_norm = (closest_normals*bary_weights[:, :, None]).sum( - 1).unsqueeze(0) * torch.tensor([-1.0, 1.0, -1.0]).type_as(normals) - pts_norm = F.normalize(pts_norm, dim=2) - pts_dist = torch.sqrt(residues) / torch.sqrt(torch.tensor(3)) - - pts_signs = 2.0 * (check_sign(verts, faces[0], points).float() - 0.5) - pts_sdf = (pts_dist * pts_signs).unsqueeze(-1) - - return pts_sdf.view(Bsize, -1, 1), pts_norm.view(Bsize, -1, 3), pts_cmap.view(Bsize, -1, 3), pts_vis.view(Bsize, -1, 1) - - -def orthogonal(points, calibrations, transforms=None): - ''' - Compute the orthogonal projections of 3D points into the image plane by given projection matrix - :param points: [B, 3, N] Tensor of 3D points - :param calibrations: [B, 3, 4] Tensor of projection matrix - :param transforms: [B, 2, 3] Tensor of image transform matrix - :return: xyz: [B, 3, N] Tensor of xyz coordinates in the image plane - ''' - rot = calibrations[:, :3, :3] - trans = calibrations[:, :3, 3:4] - pts = torch.baddbmm(trans, rot, points) # [B, 3, N] - if transforms is not None: - scale = transforms[:2, :2] - shift = transforms[:2, 2:3] - pts[:, :2, :] = torch.baddbmm(shift, scale, pts[:, :2, :]) - return pts - - -def projection(points, calib, format='numpy'): - if format == 'tensor': - return torch.mm(calib[:3, :3], points.T).T + calib[:3, 3] - else: - return np.matmul(calib[:3, :3], points.T).T + calib[:3, 3] - - -def load_calib(calib_path): - calib_data = np.loadtxt(calib_path, dtype=float) - extrinsic = calib_data[:4, :4] - intrinsic = calib_data[4:8, :4] - calib_mat = np.matmul(intrinsic, extrinsic) - calib_mat = torch.from_numpy(calib_mat).float() - return calib_mat - - -def load_obj_mesh_for_Hoppe(mesh_file): - vertex_data = [] - face_data = [] - - if isinstance(mesh_file, str): - f = open(mesh_file, "r") - else: - f = mesh_file - for line in f: - if isinstance(line, bytes): - line = line.decode("utf-8") - if line.startswith('#'): - continue - values = line.split() - if not values: - continue - - if values[0] == 'v': - v = list(map(float, values[1:4])) - vertex_data.append(v) - - elif values[0] == 'f': - # quad mesh - if len(values) > 4: - f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) - face_data.append(f) - f = list( - map(lambda x: int(x.split('/')[0]), - [values[3], values[4], values[1]])) - face_data.append(f) - # tri mesh - else: - f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) - face_data.append(f) - - vertices = np.array(vertex_data) - faces = np.array(face_data) - faces[faces > 0] -= 1 - - normals, _ = compute_normal(vertices, faces) - - return vertices, normals, faces - - -def load_obj_mesh_with_color(mesh_file): - vertex_data = [] - color_data = [] - face_data = [] - - if isinstance(mesh_file, str): - f = open(mesh_file, "r") - else: - f = mesh_file - for line in f: - if isinstance(line, bytes): - line = line.decode("utf-8") - if line.startswith('#'): - continue - values = line.split() - if not values: - continue - - if values[0] == 'v': - v = list(map(float, values[1:4])) - vertex_data.append(v) - c = list(map(float, values[4:7])) - color_data.append(c) - - elif values[0] == 'f': - # quad mesh - if len(values) > 4: - f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) - face_data.append(f) - f = list( - map(lambda x: int(x.split('/')[0]), - [values[3], values[4], values[1]])) - face_data.append(f) - # tri mesh - else: - f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) - face_data.append(f) - - vertices = np.array(vertex_data) - colors = np.array(color_data) - faces = np.array(face_data) - faces[faces > 0] -= 1 - - return vertices, colors, faces - - -def load_obj_mesh(mesh_file, with_normal=False, with_texture=False): - vertex_data = [] - norm_data = [] - uv_data = [] - - face_data = [] - face_norm_data = [] - face_uv_data = [] - - if isinstance(mesh_file, str): - f = open(mesh_file, "r") - else: - f = mesh_file - for line in f: - if isinstance(line, bytes): - line = line.decode("utf-8") - if line.startswith('#'): - continue - values = line.split() - if not values: - continue - - if values[0] == 'v': - v = list(map(float, values[1:4])) - vertex_data.append(v) - elif values[0] == 'vn': - vn = list(map(float, values[1:4])) - norm_data.append(vn) - elif values[0] == 'vt': - vt = list(map(float, values[1:3])) - uv_data.append(vt) - - elif values[0] == 'f': - # quad mesh - if len(values) > 4: - f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) - face_data.append(f) - f = list( - map(lambda x: int(x.split('/')[0]), - [values[3], values[4], values[1]])) - face_data.append(f) - # tri mesh - else: - f = list(map(lambda x: int(x.split('/')[0]), values[1:4])) - face_data.append(f) - - # deal with texture - if len(values[1].split('/')) >= 2: - # quad mesh - if len(values) > 4: - f = list(map(lambda x: int(x.split('/')[1]), values[1:4])) - face_uv_data.append(f) - f = list( - map(lambda x: int(x.split('/')[1]), - [values[3], values[4], values[1]])) - face_uv_data.append(f) - # tri mesh - elif len(values[1].split('/')[1]) != 0: - f = list(map(lambda x: int(x.split('/')[1]), values[1:4])) - face_uv_data.append(f) - # deal with normal - if len(values[1].split('/')) == 3: - # quad mesh - if len(values) > 4: - f = list(map(lambda x: int(x.split('/')[2]), values[1:4])) - face_norm_data.append(f) - f = list( - map(lambda x: int(x.split('/')[2]), - [values[3], values[4], values[1]])) - face_norm_data.append(f) - # tri mesh - elif len(values[1].split('/')[2]) != 0: - f = list(map(lambda x: int(x.split('/')[2]), values[1:4])) - face_norm_data.append(f) - - vertices = np.array(vertex_data) - faces = np.array(face_data) - faces[faces > 0] -= 1 - - if with_texture and with_normal: - uvs = np.array(uv_data) - face_uvs = np.array(face_uv_data) - face_uvs[face_uvs > 0] -= 1 - norms = np.array(norm_data) - if norms.shape[0] == 0: - norms, _ = compute_normal(vertices, faces) - face_normals = faces - else: - norms = normalize_v3(norms) - face_normals = np.array(face_norm_data) - face_normals[face_normals > 0] -= 1 - return vertices, faces, norms, face_normals, uvs, face_uvs - - if with_texture: - uvs = np.array(uv_data) - face_uvs = np.array(face_uv_data) - 1 - return vertices, faces, uvs, face_uvs - - if with_normal: - norms = np.array(norm_data) - norms = normalize_v3(norms) - face_normals = np.array(face_norm_data) - 1 - return vertices, faces, norms, face_normals - - return vertices, faces - - -def normalize_v3(arr): - ''' Normalize a numpy array of 3 component vectors shape=(n,3) ''' - lens = np.sqrt(arr[:, 0]**2 + arr[:, 1]**2 + arr[:, 2]**2) - eps = 0.00000001 - lens[lens < eps] = eps - arr[:, 0] /= lens - arr[:, 1] /= lens - arr[:, 2] /= lens - return arr - - -def compute_normal(vertices, faces): - # Create a zeroed array with the same type and shape as our vertices i.e., per vertex normal - vert_norms = np.zeros(vertices.shape, dtype=vertices.dtype) - # Create an indexed view into the vertex array using the array of three indices for triangles - tris = vertices[faces] - # Calculate the normal for all the triangles, by taking the cross product of the vectors v1-v0, and v2-v0 in each triangle - face_norms = np.cross(tris[::, 1] - tris[::, 0], tris[::, 2] - tris[::, 0]) - # n is now an array of normals per triangle. The length of each normal is dependent the vertices, - # we need to normalize these, so that our next step weights each normal equally. - normalize_v3(face_norms) - # now we have a normalized array of normals, one per triangle, i.e., per triangle normals. - # But instead of one per triangle (i.e., flat shading), we add to each vertex in that triangle, - # the triangles' normal. Multiple triangles would then contribute to every vertex, so we need to normalize again afterwards. - # The cool part, we can actually add the normals through an indexed view of our (zeroed) per vertex normal array - vert_norms[faces[:, 0]] += face_norms - vert_norms[faces[:, 1]] += face_norms - vert_norms[faces[:, 2]] += face_norms - normalize_v3(vert_norms) - - return vert_norms, face_norms - - -def save_obj_mesh(mesh_path, verts, faces): - file = open(mesh_path, 'w') - for v in verts: - file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2])) - for f in faces: - f_plus = f + 1 - file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2])) - file.close() - - -def save_obj_mesh_with_color(mesh_path, verts, faces, colors): - file = open(mesh_path, 'w') - - for idx, v in enumerate(verts): - c = colors[idx] - file.write('v %.4f %.4f %.4f %.4f %.4f %.4f\n' % - (v[0], v[1], v[2], c[0], c[1], c[2])) - for f in faces: - f_plus = f + 1 - file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2])) - file.close() - - -def calculate_mIoU(outputs, labels): - - SMOOTH = 1e-6 - - outputs = outputs.int() - labels = labels.int() - - intersection = ( - outputs - & labels).float().sum() # Will be zero if Truth=0 or Prediction=0 - union = (outputs | labels).float().sum() # Will be zzero if both are 0 - - iou = (intersection + SMOOTH) / (union + SMOOTH - ) # We smooth our devision to avoid 0/0 - - thresholded = torch.clamp( - 20 * (iou - 0.5), 0, - 10).ceil() / 10 # This is equal to comparing with thresolds - - return thresholded.mean().detach().cpu().numpy( - ) # Or thresholded.mean() if you are interested in average across the batch - - -def mask_filter(mask, number=1000): - """only keep {number} True items within a mask - Args: - mask (bool array): [N, ] - number (int, optional): total True item. Defaults to 1000. - """ - true_ids = np.where(mask)[0] - keep_ids = np.random.choice(true_ids, size=number) - filter_mask = np.isin(np.arange(len(mask)), keep_ids) - - return filter_mask - - -def query_mesh(path): - - verts, faces_idx, _ = load_obj(path) - - return verts, faces_idx.verts_idx - - -def add_alpha(colors, alpha=0.7): - - colors_pad = np.pad(colors, ((0, 0), (0, 1)), - mode='constant', - constant_values=alpha) - - return colors_pad - - -def get_optim_grid_image(per_loop_lst, loss=None, nrow=4, type='smpl'): - - font_path = os.path.join(os.path.dirname(__file__), "tbfo.ttf") - font = ImageFont.truetype(font_path, 30) - grid_img = torchvision.utils.make_grid(torch.cat(per_loop_lst, dim=0), - nrow=nrow) - grid_img = Image.fromarray( - ((grid_img.permute(1, 2, 0).detach().cpu().numpy() + 1.0) * 0.5 * - 255.0).astype(np.uint8)) - - # add text - draw = ImageDraw.Draw(grid_img) - grid_size = 512 - if loss is not None: - draw.text((10, 5), f"error: {loss:.3f}", (255, 0, 0), font=font) - - if type == 'smpl': - for col_id, col_txt in enumerate( - ['image', 'smpl-norm(render)', 'cloth-norm(pred)', 'diff-norm', 'diff-mask']): - draw.text((10+(col_id*grid_size), 5), - col_txt, (255, 0, 0), font=font) - elif type == 'cloth': - for col_id, col_txt in enumerate( - ['image', 'cloth-norm(recon)', 'cloth-norm(pred)', 'diff-norm']): - draw.text((10+(col_id*grid_size), 5), - col_txt, (255, 0, 0), font=font) - for col_id, col_txt in enumerate( - ['0', '90', '180', '270']): - draw.text((10+(col_id*grid_size), grid_size*2+5), - col_txt, (255, 0, 0), font=font) - else: - print(f"{type} should be 'smpl' or 'cloth'") - - grid_img = grid_img.resize((grid_img.size[0], grid_img.size[1]), - Image.ANTIALIAS) - - return grid_img - - -def clean_mesh(verts, faces): - - device = verts.device - - mesh_lst = trimesh.Trimesh(verts.detach().cpu().numpy(), - faces.detach().cpu().numpy()) - mesh_lst = mesh_lst.split(only_watertight=False) - comp_num = [mesh.vertices.shape[0] for mesh in mesh_lst] - mesh_clean = mesh_lst[comp_num.index(max(comp_num))] - - final_verts = torch.as_tensor(mesh_clean.vertices).float().to(device) - final_faces = torch.as_tensor(mesh_clean.faces).int().to(device) - - return final_verts, final_faces - - -def merge_mesh(verts_A, faces_A, verts_B, faces_B, color=False): - - sep_mesh = trimesh.Trimesh(np.concatenate([verts_A, verts_B], axis=0), - np.concatenate( - [faces_A, faces_B + faces_A.max() + 1], - axis=0), - maintain_order=True, - process=False) - if color: - colors = np.ones_like(sep_mesh.vertices) - colors[:verts_A.shape[0]] *= np.array([255.0, 0.0, 0.0]) - colors[verts_A.shape[0]:] *= np.array([0.0, 255.0, 0.0]) - sep_mesh.visual.vertex_colors = colors - - # union_mesh = trimesh.boolean.union([trimesh.Trimesh(verts_A, faces_A), - # trimesh.Trimesh(verts_B, faces_B)], engine='blender') - - return sep_mesh - - -def mesh_move(mesh_lst, step, scale=1.0): - - trans = np.array([1.0, 0.0, 0.0]) * step - - resize_matrix = trimesh.transformations.scale_and_translate( - scale=(scale), translate=trans) - - results = [] - - for mesh in mesh_lst: - mesh.apply_transform(resize_matrix) - results.append(mesh) - - return results - - -class SMPLX(): - def __init__(self): - - REPO_ID = "Yuliang/SMPL" - - self.smpl_verts_path = hf_hub_download(REPO_ID, filename='smpl_data/smpl_verts.npy', use_auth_token=os.environ['ICON']) - self.smplx_verts_path = hf_hub_download(REPO_ID, filename='smpl_data/smplx_verts.npy', use_auth_token=os.environ['ICON']) - self.faces_path = hf_hub_download(REPO_ID, filename='smpl_data/smplx_faces.npy', use_auth_token=os.environ['ICON']) - self.cmap_vert_path = hf_hub_download(REPO_ID, filename='smpl_data/smplx_cmap.npy', use_auth_token=os.environ['ICON']) - - self.faces = np.load(self.faces_path) - self.verts = np.load(self.smplx_verts_path) - self.smpl_verts = np.load(self.smpl_verts_path) - - self.model_dir = hf_hub_url(REPO_ID, filename='models') - self.tedra_dir = hf_hub_url(REPO_ID, filename='tedra_data') - - def get_smpl_mat(self, vert_ids): - - mat = torch.as_tensor(np.load(self.cmap_vert_path)).float() - return mat[vert_ids, :] - - def smpl2smplx(self, vert_ids=None): - """convert vert_ids in smpl to vert_ids in smplx - Args: - vert_ids ([int.array]): [n, knn_num] - """ - smplx_tree = cKDTree(self.verts, leafsize=1) - _, ind = smplx_tree.query(self.smpl_verts, k=1) # ind: [smpl_num, 1] - - if vert_ids is not None: - smplx_vert_ids = ind[vert_ids] - else: - smplx_vert_ids = ind - - return smplx_vert_ids - - def smplx2smpl(self, vert_ids=None): - """convert vert_ids in smplx to vert_ids in smpl - Args: - vert_ids ([int.array]): [n, knn_num] - """ - smpl_tree = cKDTree(self.smpl_verts, leafsize=1) - _, ind = smpl_tree.query(self.verts, k=1) # ind: [smplx_num, 1] - if vert_ids is not None: - smpl_vert_ids = ind[vert_ids] - else: - smpl_vert_ids = ind - - return smpl_vert_ids diff --git a/lib /dataset /tbfo.ttf b/lib /dataset /tbfo.ttf deleted file mode 100644 index 6cc76fcd568a5a42edd71272a19b15214de0b0d5..0000000000000000000000000000000000000000 Binary files a/lib /dataset /tbfo.ttf and /dev/null differ diff --git a/lib /net / BasePIFuNet.py b/lib /net / BasePIFuNet.py deleted file mode 100644 index 674cd61e958f16db78823005630fcb9e4bb80657..0000000000000000000000000000000000000000 --- a/lib /net / BasePIFuNet.py +++ /dev/null @@ -1,84 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import torch.nn as nn -import pytorch_lightning as pl - -from .geometry import index, orthogonal, perspective - - -class BasePIFuNet(pl.LightningModule): - def __init__( - self, - projection_mode='orthogonal', - error_term=nn.MSELoss(), - ): - """ - :param projection_mode: - Either orthogonal or perspective. - It will call the corresponding function for projection. - :param error_term: - nn Loss between the predicted [B, Res, N] and the label [B, Res, N] - """ - super(BasePIFuNet, self).__init__() - self.name = 'base' - - self.error_term = error_term - - self.index = index - self.projection = orthogonal if projection_mode == 'orthogonal' else perspective - - def forward(self, points, images, calibs, transforms=None): - ''' - :param points: [B, 3, N] world space coordinates of points - :param images: [B, C, H, W] input images - :param calibs: [B, 3, 4] calibration matrices for each image - :param transforms: Optional [B, 2, 3] image space coordinate transforms - :return: [B, Res, N] predictions for each point - ''' - features = self.filter(images) - preds = self.query(features, points, calibs, transforms) - return preds - - def filter(self, images): - ''' - Filter the input images - store all intermediate features. - :param images: [B, C, H, W] input images - ''' - return None - - def query(self, features, points, calibs, transforms=None): - ''' - Given 3D points, query the network predictions for each point. - Image features should be pre-computed before this call. - store all intermediate features. - query() function may behave differently during training/testing. - :param points: [B, 3, N] world space coordinates of points - :param calibs: [B, 3, 4] calibration matrices for each image - :param transforms: Optional [B, 2, 3] image space coordinate transforms - :param labels: Optional [B, Res, N] gt labeling - :return: [B, Res, N] predictions for each point - ''' - return None - - def get_error(self, preds, labels): - ''' - Get the network loss from the last query - :return: loss term - ''' - return self.error_term(preds, labels) \ No newline at end of file diff --git a/lib /net / VE.py b/lib /net / VE.py deleted file mode 100644 index dfcef12c42cf935db00b26a114c3d3f0642de5a8..0000000000000000000000000000000000000000 --- a/lib /net / VE.py +++ /dev/null @@ -1,183 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - - -import torch.nn as nn -import pytorch_lightning as pl - - -class BaseNetwork(pl.LightningModule): - def __init__(self): - super(BaseNetwork, self).__init__() - - def init_weights(self, init_type='xavier', gain=0.02): - ''' - initializes network's weights - init_type: normal | xavier | kaiming | orthogonal - https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39 - ''' - def init_func(m): - classname = m.__class__.__name__ - if hasattr(m, 'weight') and (classname.find('Conv') != -1 - or classname.find('Linear') != -1): - if init_type == 'normal': - nn.init.normal_(m.weight.data, 0.0, gain) - elif init_type == 'xavier': - nn.init.xavier_normal_(m.weight.data, gain=gain) - elif init_type == 'kaiming': - nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') - elif init_type == 'orthogonal': - nn.init.orthogonal_(m.weight.data, gain=gain) - - if hasattr(m, 'bias') and m.bias is not None: - nn.init.constant_(m.bias.data, 0.0) - - elif classname.find('BatchNorm2d') != -1: - nn.init.normal_(m.weight.data, 1.0, gain) - nn.init.constant_(m.bias.data, 0.0) - - self.apply(init_func) - - -class Residual3D(BaseNetwork): - def __init__(self, numIn, numOut): - super(Residual3D, self).__init__() - self.numIn = numIn - self.numOut = numOut - self.with_bias = True - # self.bn = nn.GroupNorm(4, self.numIn) - self.bn = nn.BatchNorm3d(self.numIn) - self.relu = nn.ReLU(inplace=True) - self.conv1 = nn.Conv3d(self.numIn, - self.numOut, - bias=self.with_bias, - kernel_size=3, - stride=1, - padding=2, - dilation=2) - # self.bn1 = nn.GroupNorm(4, self.numOut) - self.bn1 = nn.BatchNorm3d(self.numOut) - self.conv2 = nn.Conv3d(self.numOut, - self.numOut, - bias=self.with_bias, - kernel_size=3, - stride=1, - padding=1) - # self.bn2 = nn.GroupNorm(4, self.numOut) - self.bn2 = nn.BatchNorm3d(self.numOut) - self.conv3 = nn.Conv3d(self.numOut, - self.numOut, - bias=self.with_bias, - kernel_size=3, - stride=1, - padding=1) - - if self.numIn != self.numOut: - self.conv4 = nn.Conv3d(self.numIn, - self.numOut, - bias=self.with_bias, - kernel_size=1) - self.init_weights() - - def forward(self, x): - residual = x - # out = self.bn(x) - # out = self.relu(out) - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - out = self.conv2(out) - out = self.bn2(out) - # out = self.conv3(out) - # out = self.relu(out) - - if self.numIn != self.numOut: - residual = self.conv4(x) - - return out + residual - - -class VolumeEncoder(BaseNetwork): - """CycleGan Encoder""" - - def __init__(self, num_in=3, num_out=32, num_stacks=2): - super(VolumeEncoder, self).__init__() - self.num_in = num_in - self.num_out = num_out - self.num_inter = 8 - self.num_stacks = num_stacks - self.with_bias = True - - self.relu = nn.ReLU(inplace=True) - self.conv1 = nn.Conv3d(self.num_in, - self.num_inter, - bias=self.with_bias, - kernel_size=5, - stride=2, - padding=4, - dilation=2) - # self.bn1 = nn.GroupNorm(4, self.num_inter) - self.bn1 = nn.BatchNorm3d(self.num_inter) - self.conv2 = nn.Conv3d(self.num_inter, - self.num_out, - bias=self.with_bias, - kernel_size=5, - stride=2, - padding=4, - dilation=2) - # self.bn2 = nn.GroupNorm(4, self.num_out) - self.bn2 = nn.BatchNorm3d(self.num_out) - - self.conv_out1 = nn.Conv3d(self.num_out, - self.num_out, - bias=self.with_bias, - kernel_size=3, - stride=1, - padding=1, - dilation=1) - self.conv_out2 = nn.Conv3d(self.num_out, - self.num_out, - bias=self.with_bias, - kernel_size=3, - stride=1, - padding=1, - dilation=1) - - for idx in range(self.num_stacks): - self.add_module("res" + str(idx), - Residual3D(self.num_out, self.num_out)) - - self.init_weights() - - def forward(self, x, intermediate_output=True): - 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_lst = [] - for idx in range(self.num_stacks): - out = self._modules["res" + str(idx)](out) - out_lst.append(out) - - if intermediate_output: - return out_lst - else: - return [out_lst[-1]] \ No newline at end of file diff --git a/lib /net / __init__.py b/lib /net / __init__.py deleted file mode 100644 index d300c0941f58c2172bab5dcdadb309e80d3c2b4d..0000000000000000000000000000000000000000 --- a/lib /net / __init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .BasePIFuNet import BasePIFuNet -from .HGPIFuNet import HGPIFuNet -from .NormalNet import NormalNet -from .VE import VolumeEncoder \ No newline at end of file diff --git a/lib /net / geometry.py b/lib /net / geometry.py deleted file mode 100644 index 27958c565cc933ad9114b078094727231c97d49d..0000000000000000000000000000000000000000 --- a/lib /net / geometry.py +++ /dev/null @@ -1,82 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -import torch - - -def index(feat, uv): - ''' - :param feat: [B, C, H, W] image features - :param uv: [B, 2, N] uv coordinates in the image plane, range [0, 1] - :return: [B, C, N] image features at the uv coordinates - ''' - uv = uv.transpose(1, 2) # [B, N, 2] - - (B, N, _) = uv.shape - C = feat.shape[1] - - if uv.shape[-1] == 3: - # uv = uv[:,:,[2,1,0]] - # uv = uv * torch.tensor([1.0,-1.0,1.0]).type_as(uv)[None,None,...] - uv = uv.unsqueeze(2).unsqueeze(3) # [B, N, 1, 1, 3] - else: - uv = uv.unsqueeze(2) # [B, N, 1, 2] - - # NOTE: for newer PyTorch, it seems that training results are degraded due to implementation diff in F.grid_sample - # for old versions, simply remove the aligned_corners argument. - samples = torch.nn.functional.grid_sample( - feat, uv, align_corners=True) # [B, C, N, 1] - return samples.view(B, C, N) # [B, C, N] - - -def orthogonal(points, calibrations, transforms=None): - ''' - Compute the orthogonal projections of 3D points into the image plane by given projection matrix - :param points: [B, 3, N] Tensor of 3D points - :param calibrations: [B, 3, 4] Tensor of projection matrix - :param transforms: [B, 2, 3] Tensor of image transform matrix - :return: xyz: [B, 3, N] Tensor of xyz coordinates in the image plane - ''' - rot = calibrations[:, :3, :3] - trans = calibrations[:, :3, 3:4] - pts = torch.baddbmm(trans, rot, points) # [B, 3, N] - if transforms is not None: - scale = transforms[:2, :2] - shift = transforms[:2, 2:3] - pts[:, :2, :] = torch.baddbmm(shift, scale, pts[:, :2, :]) - return pts - - -def perspective(points, calibrations, transforms=None): - ''' - Compute the perspective projections of 3D points into the image plane by given projection matrix - :param points: [Bx3xN] Tensor of 3D points - :param calibrations: [Bx3x4] Tensor of projection matrix - :param transforms: [Bx2x3] Tensor of image transform matrix - :return: xy: [Bx2xN] Tensor of xy coordinates in the image plane - ''' - rot = calibrations[:, :3, :3] - trans = calibrations[:, :3, 3:4] - homo = torch.baddbmm(trans, rot, points) # [B, 3, N] - xy = homo[:, :2, :] / homo[:, 2:3, :] - if transforms is not None: - scale = transforms[:2, :2] - shift = transforms[:2, 2:3] - xy = torch.baddbmm(shift, scale, xy) - - xyz = torch.cat([xy, homo[:, 2:3, :]], 1) - return xyz \ No newline at end of file diff --git a/lib /net /FBNet.py b/lib /net /FBNet.py deleted file mode 100644 index dab9b7cae8869cb0d907cbb5cfac79b5cb016644..0000000000000000000000000000000000000000 --- a/lib /net /FBNet.py +++ /dev/null @@ -1,384 +0,0 @@ -''' -Copyright (C) 2019 NVIDIA Corporation. Ting-Chun Wang, Ming-Yu Liu, Jun-Yan Zhu. -BSD License. All rights reserved. -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: -* Redistributions of source code must retain the above copyright notice, this - list of conditions and the following disclaimer. -* Redistributions in binary form must reproduce the above copyright notice, - this list of conditions and the following disclaimer in the documentation - and/or other materials provided with the distribution. -THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING ALL -IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR ANY PARTICULAR PURPOSE. -IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL -DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, -WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING -OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. -''' -import torch -import torch.nn as nn -import functools -import numpy as np -import pytorch_lightning as pl - - -############################################################################### -# Functions -############################################################################### -def weights_init(m): - classname = m.__class__.__name__ - if classname.find('Conv') != -1: - m.weight.data.normal_(0.0, 0.02) - elif classname.find('BatchNorm2d') != -1: - m.weight.data.normal_(1.0, 0.02) - m.bias.data.fill_(0) - - -def get_norm_layer(norm_type='instance'): - if norm_type == 'batch': - norm_layer = functools.partial(nn.BatchNorm2d, affine=True) - elif norm_type == 'instance': - norm_layer = functools.partial(nn.InstanceNorm2d, affine=False) - else: - raise NotImplementedError('normalization layer [%s] is not found' % - norm_type) - return norm_layer - - -def define_G(input_nc, - output_nc, - ngf, - netG, - n_downsample_global=3, - n_blocks_global=9, - n_local_enhancers=1, - n_blocks_local=3, - norm='instance', - gpu_ids=[], - last_op=nn.Tanh()): - norm_layer = get_norm_layer(norm_type=norm) - if netG == 'global': - netG = GlobalGenerator(input_nc, - output_nc, - ngf, - n_downsample_global, - n_blocks_global, - norm_layer, - last_op=last_op) - elif netG == 'local': - netG = LocalEnhancer(input_nc, output_nc, ngf, n_downsample_global, - n_blocks_global, n_local_enhancers, - n_blocks_local, norm_layer) - elif netG == 'encoder': - netG = Encoder(input_nc, output_nc, ngf, n_downsample_global, - norm_layer) - else: - raise ('generator not implemented!') - # print(netG) - if len(gpu_ids) > 0: - assert (torch.cuda.is_available()) - device=torch.device(f"cuda:{gpu_ids[0]}") - netG = netG.to(device) - netG.apply(weights_init) - return netG - - -def print_network(net): - if isinstance(net, list): - net = net[0] - num_params = 0 - for param in net.parameters(): - num_params += param.numel() - print(net) - print('Total number of parameters: %d' % num_params) - - -############################################################################## -# Generator -############################################################################## -class LocalEnhancer(pl.LightningModule): - def __init__(self, - input_nc, - output_nc, - ngf=32, - n_downsample_global=3, - n_blocks_global=9, - n_local_enhancers=1, - n_blocks_local=3, - norm_layer=nn.BatchNorm2d, - padding_type='reflect'): - super(LocalEnhancer, self).__init__() - self.n_local_enhancers = n_local_enhancers - - ###### global generator model ##### - ngf_global = ngf * (2**n_local_enhancers) - model_global = GlobalGenerator(input_nc, output_nc, ngf_global, - n_downsample_global, n_blocks_global, - norm_layer).model - model_global = [model_global[i] for i in range(len(model_global) - 3) - ] # get rid of final convolution layers - self.model = nn.Sequential(*model_global) - - ###### local enhancer layers ##### - for n in range(1, n_local_enhancers + 1): - # downsample - ngf_global = ngf * (2**(n_local_enhancers - n)) - model_downsample = [ - nn.ReflectionPad2d(3), - nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0), - norm_layer(ngf_global), - nn.ReLU(True), - nn.Conv2d(ngf_global, - ngf_global * 2, - kernel_size=3, - stride=2, - padding=1), - norm_layer(ngf_global * 2), - nn.ReLU(True) - ] - # residual blocks - model_upsample = [] - for i in range(n_blocks_local): - model_upsample += [ - ResnetBlock(ngf_global * 2, - padding_type=padding_type, - norm_layer=norm_layer) - ] - - # upsample - model_upsample += [ - nn.ConvTranspose2d(ngf_global * 2, - ngf_global, - kernel_size=3, - stride=2, - padding=1, - output_padding=1), - norm_layer(ngf_global), - nn.ReLU(True) - ] - - # final convolution - if n == n_local_enhancers: - model_upsample += [ - nn.ReflectionPad2d(3), - nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), - nn.Tanh() - ] - - setattr(self, 'model' + str(n) + '_1', - nn.Sequential(*model_downsample)) - setattr(self, 'model' + str(n) + '_2', - nn.Sequential(*model_upsample)) - - self.downsample = nn.AvgPool2d(3, - stride=2, - padding=[1, 1], - count_include_pad=False) - - def forward(self, input): - # create input pyramid - input_downsampled = [input] - for i in range(self.n_local_enhancers): - input_downsampled.append(self.downsample(input_downsampled[-1])) - - # output at coarest level - output_prev = self.model(input_downsampled[-1]) - # build up one layer at a time - for n_local_enhancers in range(1, self.n_local_enhancers + 1): - model_downsample = getattr(self, - 'model' + str(n_local_enhancers) + '_1') - model_upsample = getattr(self, - 'model' + str(n_local_enhancers) + '_2') - input_i = input_downsampled[self.n_local_enhancers - - n_local_enhancers] - output_prev = model_upsample( - model_downsample(input_i) + output_prev) - return output_prev - - -class GlobalGenerator(pl.LightningModule): - def __init__(self, - input_nc, - output_nc, - ngf=64, - n_downsampling=3, - n_blocks=9, - norm_layer=nn.BatchNorm2d, - padding_type='reflect', - last_op=nn.Tanh()): - assert (n_blocks >= 0) - super(GlobalGenerator, self).__init__() - activation = nn.ReLU(True) - - model = [ - nn.ReflectionPad2d(3), - nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), - norm_layer(ngf), activation - ] - # downsample - for i in range(n_downsampling): - mult = 2**i - model += [ - nn.Conv2d(ngf * mult, - ngf * mult * 2, - kernel_size=3, - stride=2, - padding=1), - norm_layer(ngf * mult * 2), activation - ] - - # resnet blocks - mult = 2**n_downsampling - for i in range(n_blocks): - model += [ - ResnetBlock(ngf * mult, - padding_type=padding_type, - activation=activation, - norm_layer=norm_layer) - ] - - # upsample - for i in range(n_downsampling): - mult = 2**(n_downsampling - i) - model += [ - nn.ConvTranspose2d(ngf * mult, - int(ngf * mult / 2), - kernel_size=3, - stride=2, - padding=1, - output_padding=1), - norm_layer(int(ngf * mult / 2)), activation - ] - model += [ - nn.ReflectionPad2d(3), - nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0) - ] - if last_op is not None: - model += [last_op] - self.model = nn.Sequential(*model) - - def forward(self, input): - return self.model(input) - - -# Define a resnet block -class ResnetBlock(pl.LightningModule): - def __init__(self, - dim, - padding_type, - norm_layer, - activation=nn.ReLU(True), - use_dropout=False): - super(ResnetBlock, self).__init__() - self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, - activation, use_dropout) - - def build_conv_block(self, dim, padding_type, norm_layer, activation, - use_dropout): - conv_block = [] - p = 0 - if padding_type == 'reflect': - conv_block += [nn.ReflectionPad2d(1)] - elif padding_type == 'replicate': - conv_block += [nn.ReplicationPad2d(1)] - elif padding_type == 'zero': - p = 1 - else: - raise NotImplementedError('padding [%s] is not implemented' % - padding_type) - - conv_block += [ - nn.Conv2d(dim, dim, kernel_size=3, padding=p), - norm_layer(dim), activation - ] - if use_dropout: - conv_block += [nn.Dropout(0.5)] - - p = 0 - if padding_type == 'reflect': - conv_block += [nn.ReflectionPad2d(1)] - elif padding_type == 'replicate': - conv_block += [nn.ReplicationPad2d(1)] - elif padding_type == 'zero': - p = 1 - else: - raise NotImplementedError('padding [%s] is not implemented' % - padding_type) - conv_block += [ - nn.Conv2d(dim, dim, kernel_size=3, padding=p), - norm_layer(dim) - ] - - return nn.Sequential(*conv_block) - - def forward(self, x): - out = x + self.conv_block(x) - return out - - -class Encoder(pl.LightningModule): - def __init__(self, - input_nc, - output_nc, - ngf=32, - n_downsampling=4, - norm_layer=nn.BatchNorm2d): - super(Encoder, self).__init__() - self.output_nc = output_nc - - model = [ - nn.ReflectionPad2d(3), - nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), - norm_layer(ngf), - nn.ReLU(True) - ] - # downsample - for i in range(n_downsampling): - mult = 2**i - model += [ - nn.Conv2d(ngf * mult, - ngf * mult * 2, - kernel_size=3, - stride=2, - padding=1), - norm_layer(ngf * mult * 2), - nn.ReLU(True) - ] - - # upsample - for i in range(n_downsampling): - mult = 2**(n_downsampling - i) - model += [ - nn.ConvTranspose2d(ngf * mult, - int(ngf * mult / 2), - kernel_size=3, - stride=2, - padding=1, - output_padding=1), - norm_layer(int(ngf * mult / 2)), - nn.ReLU(True) - ] - - model += [ - nn.ReflectionPad2d(3), - nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), - nn.Tanh() - ] - self.model = nn.Sequential(*model) - - def forward(self, input, inst): - outputs = self.model(input) - - # instance-wise average pooling - outputs_mean = outputs.clone() - inst_list = np.unique(inst.cpu().numpy().astype(int)) - for i in inst_list: - for b in range(input.size()[0]): - indices = (inst[b:b + 1] == int(i)).nonzero() # n x 4 - for j in range(self.output_nc): - output_ins = outputs[indices[:, 0] + b, indices[:, 1] + j, - indices[:, 2], indices[:, 3]] - mean_feat = torch.mean(output_ins).expand_as(output_ins) - outputs_mean[indices[:, 0] + b, indices[:, 1] + j, - indices[:, 2], indices[:, 3]] = mean_feat - return outputs_mean \ No newline at end of file diff --git a/lib /net /HGFilters.py b/lib /net /HGFilters.py deleted file mode 100644 index 46b977807608861add77085f4ed48bf3585e928f..0000000000000000000000000000000000000000 --- a/lib /net /HGFilters.py +++ /dev/null @@ -1,197 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from lib.net.net_util import * -import torch.nn as nn -import torch.nn.functional as F - - -class HourGlass(nn.Module): - def __init__(self, num_modules, depth, num_features, opt): - super(HourGlass, self).__init__() - self.num_modules = num_modules - self.depth = depth - self.features = num_features - self.opt = opt - - self._generate_network(self.depth) - - def _generate_network(self, level): - self.add_module('b1_' + str(level), - ConvBlock(self.features, self.features, self.opt)) - - self.add_module('b2_' + str(level), - ConvBlock(self.features, self.features, self.opt)) - - if level > 1: - self._generate_network(level - 1) - else: - self.add_module('b2_plus_' + str(level), - ConvBlock(self.features, self.features, self.opt)) - - self.add_module('b3_' + str(level), - ConvBlock(self.features, self.features, self.opt)) - - def _forward(self, level, inp): - # Upper branch - up1 = inp - up1 = self._modules['b1_' + str(level)](up1) - - # Lower branch - low1 = F.avg_pool2d(inp, 2, stride=2) - low1 = self._modules['b2_' + str(level)](low1) - - if level > 1: - low2 = self._forward(level - 1, low1) - else: - low2 = low1 - low2 = self._modules['b2_plus_' + str(level)](low2) - - low3 = low2 - low3 = self._modules['b3_' + str(level)](low3) - - # NOTE: for newer PyTorch (1.3~), it seems that training results are degraded due to implementation diff in F.grid_sample - # if the pretrained model behaves weirdly, switch with the commented line. - # NOTE: I also found that "bicubic" works better. - up2 = F.interpolate(low3, - scale_factor=2, - mode='bicubic', - align_corners=True) - # up2 = F.interpolate(low3, scale_factor=2, mode='nearest) - - return up1 + up2 - - def forward(self, x): - return self._forward(self.depth, x) - - -class HGFilter(nn.Module): - def __init__(self, opt, num_modules, in_dim): - super(HGFilter, self).__init__() - self.num_modules = num_modules - - self.opt = opt - [k, s, d, p] = self.opt.conv1 - - # self.conv1 = nn.Conv2d(in_dim, 64, kernel_size=7, stride=2, padding=3) - self.conv1 = nn.Conv2d(in_dim, - 64, - kernel_size=k, - stride=s, - dilation=d, - padding=p) - - if self.opt.norm == 'batch': - self.bn1 = nn.BatchNorm2d(64) - elif self.opt.norm == 'group': - self.bn1 = nn.GroupNorm(32, 64) - - if self.opt.hg_down == 'conv64': - self.conv2 = ConvBlock(64, 64, self.opt) - self.down_conv2 = nn.Conv2d(64, - 128, - kernel_size=3, - stride=2, - padding=1) - elif self.opt.hg_down == 'conv128': - self.conv2 = ConvBlock(64, 128, self.opt) - self.down_conv2 = nn.Conv2d(128, - 128, - kernel_size=3, - stride=2, - padding=1) - elif self.opt.hg_down == 'ave_pool': - self.conv2 = ConvBlock(64, 128, self.opt) - else: - raise NameError('Unknown Fan Filter setting!') - - self.conv3 = ConvBlock(128, 128, self.opt) - self.conv4 = ConvBlock(128, 256, self.opt) - - # Stacking part - for hg_module in range(self.num_modules): - self.add_module('m' + str(hg_module), - HourGlass(1, opt.num_hourglass, 256, self.opt)) - - self.add_module('top_m_' + str(hg_module), - ConvBlock(256, 256, self.opt)) - self.add_module( - 'conv_last' + str(hg_module), - nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)) - if self.opt.norm == 'batch': - self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256)) - elif self.opt.norm == 'group': - self.add_module('bn_end' + str(hg_module), - nn.GroupNorm(32, 256)) - - self.add_module( - 'l' + str(hg_module), - nn.Conv2d(256, - opt.hourglass_dim, - kernel_size=1, - stride=1, - padding=0)) - - if hg_module < self.num_modules - 1: - self.add_module( - 'bl' + str(hg_module), - nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)) - self.add_module( - 'al' + str(hg_module), - nn.Conv2d(opt.hourglass_dim, - 256, - kernel_size=1, - stride=1, - padding=0)) - - def forward(self, x): - x = F.relu(self.bn1(self.conv1(x)), True) - tmpx = x - if self.opt.hg_down == 'ave_pool': - x = F.avg_pool2d(self.conv2(x), 2, stride=2) - elif self.opt.hg_down in ['conv64', 'conv128']: - x = self.conv2(x) - x = self.down_conv2(x) - else: - raise NameError('Unknown Fan Filter setting!') - - x = self.conv3(x) - x = self.conv4(x) - - previous = x - - outputs = [] - for i in range(self.num_modules): - hg = self._modules['m' + str(i)](previous) - - ll = hg - ll = self._modules['top_m_' + str(i)](ll) - - ll = F.relu( - self._modules['bn_end' + str(i)]( - self._modules['conv_last' + str(i)](ll)), True) - - # Predict heatmaps - tmp_out = self._modules['l' + str(i)](ll) - outputs.append(tmp_out) - - if i < self.num_modules - 1: - ll = self._modules['bl' + str(i)](ll) - tmp_out_ = self._modules['al' + str(i)](tmp_out) - previous = previous + ll + tmp_out_ - - return outputs \ No newline at end of file diff --git a/lib /net /HGPIFuNet.py b/lib /net /HGPIFuNet.py deleted file mode 100644 index 0e8f00fe1011a6293e0f37ed977d92e9fd1791ef..0000000000000000000000000000000000000000 --- a/lib /net /HGPIFuNet.py +++ /dev/null @@ -1,401 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from lib.net.voxelize import Voxelization -from lib.dataset.mesh_util import cal_sdf_batch, feat_select, read_smpl_constants -from lib.net.NormalNet import NormalNet -from lib.net.MLP import MLP -from lib.dataset.mesh_util import SMPLX -from lib.net.VE import VolumeEncoder -from lib.net.HGFilters import * -from termcolor import colored -from lib.net.BasePIFuNet import BasePIFuNet -import torch.nn as nn -import torch - - -maskout = False - - -class HGPIFuNet(BasePIFuNet): - ''' - HG PIFu network uses Hourglass stacks as the image filter. - It does the following: - 1. Compute image feature stacks and store it in self.im_feat_list - self.im_feat_list[-1] is the last stack (output stack) - 2. Calculate calibration - 3. If training, it index on every intermediate stacks, - If testing, it index on the last stack. - 4. Classification. - 5. During training, error is calculated on all stacks. - ''' - - def __init__(self, - cfg, - projection_mode='orthogonal', - error_term=nn.MSELoss()): - - super(HGPIFuNet, self).__init__(projection_mode=projection_mode, - error_term=error_term) - - self.l1_loss = nn.SmoothL1Loss() - self.opt = cfg.net - self.root = cfg.root - self.overfit = cfg.overfit - - channels_IF = self.opt.mlp_dim - - self.use_filter = self.opt.use_filter - self.prior_type = self.opt.prior_type - self.smpl_feats = self.opt.smpl_feats - - self.smpl_dim = self.opt.smpl_dim - self.voxel_dim = self.opt.voxel_dim - self.hourglass_dim = self.opt.hourglass_dim - self.sdf_clip = cfg.sdf_clip / 100.0 - - self.in_geo = [item[0] for item in self.opt.in_geo] - self.in_nml = [item[0] for item in self.opt.in_nml] - - self.in_geo_dim = sum([item[1] for item in self.opt.in_geo]) - self.in_nml_dim = sum([item[1] for item in self.opt.in_nml]) - - self.in_total = self.in_geo + self.in_nml - self.smpl_feat_dict = None - self.smplx_data = SMPLX() - - if self.prior_type == 'icon': - if 'image' in self.in_geo: - self.channels_filter = [[0, 1, 2, 3, 4, 5], [0, 1, 2, 6, 7, 8]] - else: - self.channels_filter = [[0, 1, 2], [3, 4, 5]] - - else: - if 'image' in self.in_geo: - self.channels_filter = [[0, 1, 2, 3, 4, 5, 6, 7, 8]] - else: - self.channels_filter = [[0, 1, 2, 3, 4, 5]] - - channels_IF[0] = self.hourglass_dim if self.use_filter else len( - self.channels_filter[0]) - - if self.prior_type == 'icon' and 'vis' not in self.smpl_feats: - if self.use_filter: - channels_IF[0] += self.hourglass_dim - else: - channels_IF[0] += len(self.channels_filter[0]) - - if self.prior_type == 'icon': - channels_IF[0] += self.smpl_dim - elif self.prior_type == 'pamir': - channels_IF[0] += self.voxel_dim - smpl_vertex_code, smpl_face_code, smpl_faces, smpl_tetras = read_smpl_constants( - self.smplx_data.tedra_dir) - self.voxelization = Voxelization( - smpl_vertex_code, - smpl_face_code, - smpl_faces, - smpl_tetras, - volume_res=128, - sigma=0.05, - smooth_kernel_size=7, - batch_size=cfg.batch_size, - device=torch.device(f"cuda:{cfg.gpus[0]}")) - self.ve = VolumeEncoder(3, self.voxel_dim, self.opt.num_stack) - else: - channels_IF[0] += 1 - - self.icon_keys = ["smpl_verts", "smpl_faces", "smpl_vis", "smpl_cmap"] - self.pamir_keys = [ - "voxel_verts", "voxel_faces", "pad_v_num", "pad_f_num" - ] - - self.if_regressor = MLP( - filter_channels=channels_IF, - name='if', - res_layers=self.opt.res_layers, - norm=self.opt.norm_mlp, - last_op=nn.Sigmoid() if not cfg.test_mode else None) - - # network - if self.use_filter: - if self.opt.gtype == "HGPIFuNet": - self.F_filter = HGFilter(self.opt, self.opt.num_stack, - len(self.channels_filter[0])) - else: - print( - colored(f"Backbone {self.opt.gtype} is unimplemented", - 'green')) - - summary_log = f"{self.prior_type.upper()}:\n" + \ - f"w/ Global Image Encoder: {self.use_filter}\n" + \ - f"Image Features used by MLP: {self.in_geo}\n" - - if self.prior_type == "icon": - summary_log += f"Geometry Features used by MLP: {self.smpl_feats}\n" - summary_log += f"Dim of Image Features (local): 6\n" - summary_log += f"Dim of Geometry Features (ICON): {self.smpl_dim}\n" - elif self.prior_type == "pamir": - summary_log += f"Dim of Image Features (global): {self.hourglass_dim}\n" - summary_log += f"Dim of Geometry Features (PaMIR): {self.voxel_dim}\n" - else: - summary_log += f"Dim of Image Features (global): {self.hourglass_dim}\n" - summary_log += f"Dim of Geometry Features (PIFu): 1 (z-value)\n" - - summary_log += f"Dim of MLP's first layer: {channels_IF[0]}\n" - - print(colored(summary_log, "yellow")) - - self.normal_filter = NormalNet(cfg) - init_net(self) - - def get_normal(self, in_tensor_dict): - - # insert normal features - if (not self.training) and (not self.overfit): - # print(colored("infer normal","blue")) - with torch.no_grad(): - feat_lst = [] - if "image" in self.in_geo: - feat_lst.append( - in_tensor_dict['image']) # [1, 3, 512, 512] - if 'normal_F' in self.in_geo and 'normal_B' in self.in_geo: - if 'normal_F' not in in_tensor_dict.keys( - ) or 'normal_B' not in in_tensor_dict.keys(): - (nmlF, nmlB) = self.normal_filter(in_tensor_dict) - else: - nmlF = in_tensor_dict['normal_F'] - nmlB = in_tensor_dict['normal_B'] - feat_lst.append(nmlF) # [1, 3, 512, 512] - feat_lst.append(nmlB) # [1, 3, 512, 512] - in_filter = torch.cat(feat_lst, dim=1) - - else: - in_filter = torch.cat([in_tensor_dict[key] for key in self.in_geo], - dim=1) - - return in_filter - - def get_mask(self, in_filter, size=128): - - mask = F.interpolate(in_filter[:, self.channels_filter[0]], - size=(size, size), - mode="bilinear", - align_corners=True).abs().sum(dim=1, - keepdim=True) != 0.0 - - return mask - - def filter(self, in_tensor_dict, return_inter=False): - ''' - Filter the input images - store all intermediate features. - :param images: [B, C, H, W] input images - ''' - - in_filter = self.get_normal(in_tensor_dict) - - features_G = [] - - if self.prior_type == 'icon': - if self.use_filter: - features_F = self.F_filter(in_filter[:, - self.channels_filter[0]] - ) # [(B,hg_dim,128,128) * 4] - features_B = self.F_filter(in_filter[:, - self.channels_filter[1]] - ) # [(B,hg_dim,128,128) * 4] - else: - features_F = [in_filter[:, self.channels_filter[0]]] - features_B = [in_filter[:, self.channels_filter[1]]] - for idx in range(len(features_F)): - features_G.append( - torch.cat([features_F[idx], features_B[idx]], dim=1)) - else: - if self.use_filter: - features_G = self.F_filter(in_filter[:, - self.channels_filter[0]]) - else: - features_G = [in_filter[:, self.channels_filter[0]]] - - if self.prior_type == 'icon': - self.smpl_feat_dict = { - k: in_tensor_dict[k] - for k in self.icon_keys - } - elif self.prior_type == "pamir": - self.smpl_feat_dict = { - k: in_tensor_dict[k] - for k in self.pamir_keys - } - else: - pass - # print(colored("use z rather than icon or pamir", "green")) - - # If it is not in training, only produce the last im_feat - if not self.training: - features_out = [features_G[-1]] - else: - features_out = features_G - - if maskout: - features_out_mask = [] - for feat in features_out: - features_out_mask.append( - feat * self.get_mask(in_filter, size=feat.shape[2])) - features_out = features_out_mask - - if return_inter: - return features_out, in_filter - else: - return features_out - - def query(self, features, points, calibs, transforms=None, regressor=None): - - xyz = self.projection(points, calibs, transforms) - - (xy, z) = xyz.split([2, 1], dim=1) - - in_cube = (xyz > -1.0) & (xyz < 1.0) - in_cube = in_cube.all(dim=1, keepdim=True).detach().float() - - preds_list = [] - - if self.prior_type == 'icon': - - # smpl_verts [B, N_vert, 3] - # smpl_faces [B, N_face, 3] - # points [B, 3, N] - - smpl_sdf, smpl_norm, smpl_cmap, smpl_vis = cal_sdf_batch( - self.smpl_feat_dict['smpl_verts'], - self.smpl_feat_dict['smpl_faces'], - self.smpl_feat_dict['smpl_cmap'], - self.smpl_feat_dict['smpl_vis'], - xyz.permute(0, 2, 1).contiguous()) - - # smpl_sdf [B, N, 1] - # smpl_norm [B, N, 3] - # smpl_cmap [B, N, 3] - # smpl_vis [B, N, 1] - - feat_lst = [smpl_sdf] - if 'cmap' in self.smpl_feats: - feat_lst.append(smpl_cmap) - if 'norm' in self.smpl_feats: - feat_lst.append(smpl_norm) - if 'vis' in self.smpl_feats: - feat_lst.append(smpl_vis) - - smpl_feat = torch.cat(feat_lst, dim=2).permute(0, 2, 1) - vol_feats = features - - elif self.prior_type == "pamir": - - voxel_verts = self.smpl_feat_dict[ - 'voxel_verts'][:, :-self.smpl_feat_dict['pad_v_num'][0], :] - voxel_faces = self.smpl_feat_dict[ - 'voxel_faces'][:, :-self.smpl_feat_dict['pad_f_num'][0], :] - - self.voxelization.update_param( - batch_size=voxel_faces.shape[0], - smpl_tetra=voxel_faces[0].detach().cpu().numpy()) - vol = self.voxelization(voxel_verts) # vol ~ [0,1] - vol_feats = self.ve(vol, intermediate_output=self.training) - else: - vol_feats = features - - for im_feat, vol_feat in zip(features, vol_feats): - - # [B, Feat_i + z, N] - # normal feature choice by smpl_vis - if self.prior_type == 'icon': - if 'vis' in self.smpl_feats: - point_local_feat = feat_select(self.index(im_feat, xy), - smpl_feat[:, [-1], :]) - if maskout: - normal_mask = torch.tile( - point_local_feat.sum(dim=1, keepdims=True) == 0.0, - (1, smpl_feat.shape[1], 1)) - normal_mask[:, 1:, :] = False - smpl_feat[normal_mask] = -1.0 - point_feat_list = [point_local_feat, smpl_feat[:, :-1, :]] - else: - point_local_feat = self.index(im_feat, xy) - point_feat_list = [point_local_feat, smpl_feat[:, :, :]] - - elif self.prior_type == 'pamir': - # im_feat [B, hg_dim, 128, 128] - # vol_feat [B, vol_dim, 32, 32, 32] - point_feat_list = [ - self.index(im_feat, xy), - self.index(vol_feat, xyz) - ] - - else: - point_feat_list = [self.index(im_feat, xy), z] - - point_feat = torch.cat(point_feat_list, 1) - - # out of image plane is always set to 0 - preds = regressor(point_feat) - preds = in_cube * preds - - preds_list.append(preds) - - return preds_list - - def get_error(self, preds_if_list, labels): - """calcaulate error - Args: - preds_list (list): list of torch.tensor(B, 3, N) - labels (torch.tensor): (B, N_knn, N) - Returns: - torch.tensor: error - """ - error_if = 0 - - for pred_id in range(len(preds_if_list)): - pred_if = preds_if_list[pred_id] - error_if += self.error_term(pred_if, labels) - - error_if /= len(preds_if_list) - - return error_if - - def forward(self, in_tensor_dict): - """ - sample_tensor [B, 3, N] - calib_tensor [B, 4, 4] - label_tensor [B, 1, N] - smpl_feat_tensor [B, 59, N] - """ - - sample_tensor = in_tensor_dict['sample'] - calib_tensor = in_tensor_dict['calib'] - label_tensor = in_tensor_dict['label'] - - in_feat = self.filter(in_tensor_dict) - - preds_if_list = self.query(in_feat, - sample_tensor, - calib_tensor, - regressor=self.if_regressor) - - error = self.get_error(preds_if_list, label_tensor) - - return preds_if_list[-1], error \ No newline at end of file diff --git a/lib /net /MLP.py b/lib /net /MLP.py deleted file mode 100644 index de9628f7ab35de92e9241cf09a12f33aa96b82a7..0000000000000000000000000000000000000000 --- a/lib /net /MLP.py +++ /dev/null @@ -1,72 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. - -import torch -import torch.nn as nn -import pytorch_lightning as pl - - -class MLP(pl.LightningModule): - def __init__(self, - filter_channels, - name=None, - res_layers=[], - norm='group', - last_op=None): - - super(MLP, self).__init__() - - self.filters = nn.ModuleList() - self.norms = nn.ModuleList() - self.res_layers = res_layers - self.norm = norm - self.last_op = last_op - self.name = name - self.activate = nn.LeakyReLU(inplace=True) - - for l in range(0, len(filter_channels) - 1): - if l in self.res_layers: - self.filters.append( - nn.Conv1d(filter_channels[l] + filter_channels[0], - filter_channels[l + 1], 1)) - else: - self.filters.append( - nn.Conv1d(filter_channels[l], filter_channels[l + 1], 1)) - - if l != len(filter_channels) - 2: - if norm == 'group': - self.norms.append(nn.GroupNorm(32, filter_channels[l + 1])) - elif norm == 'batch': - self.norms.append(nn.BatchNorm1d(filter_channels[l + 1])) - elif norm == 'instance': - self.norms.append(nn.InstanceNorm1d(filter_channels[l + - 1])) - elif norm == 'weight': - self.filters[l] = nn.utils.weight_norm(self.filters[l], - name='weight') - # print(self.filters[l].weight_g.size(), - # self.filters[l].weight_v.size()) - - def forward(self, feature): - ''' - feature may include multiple view inputs - args: - feature: [B, C_in, N] - return: - [B, C_out, N] prediction - ''' - y = feature - tmpy = feature - - for i, f in enumerate(self.filters): - - y = f(y if i not in self.res_layers else torch.cat([y, tmpy], 1)) - if i != len(self.filters) - 1: - if self.norm not in ['batch', 'group', 'instance']: - y = self.activate(y) - else: - y = self.activate(self.norms[i](y)) - - if self.last_op is not None: - y = self.last_op(y) - - return y \ No newline at end of file diff --git a/lib /net /NormalNet.py b/lib /net /NormalNet.py deleted file mode 100644 index 022cde1050fe29d007ab1a5d8913a74d1af81e38..0000000000000000000000000000000000000000 --- a/lib /net /NormalNet.py +++ /dev/null @@ -1,121 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from lib.net.FBNet import define_G -from lib.net.net_util import init_net, VGGLoss -from lib.net.HGFilters import * -from lib.net.BasePIFuNet import BasePIFuNet -import torch -import torch.nn as nn - - -class NormalNet(BasePIFuNet): - ''' - HG PIFu network uses Hourglass stacks as the image filter. - It does the following: - 1. Compute image feature stacks and store it in self.im_feat_list - self.im_feat_list[-1] is the last stack (output stack) - 2. Calculate calibration - 3. If training, it index on every intermediate stacks, - If testing, it index on the last stack. - 4. Classification. - 5. During training, error is calculated on all stacks. - ''' - - def __init__(self, cfg, error_term=nn.SmoothL1Loss()): - - super(NormalNet, self).__init__(error_term=error_term) - - self.l1_loss = nn.SmoothL1Loss() - - self.opt = cfg.net - - if self.training: - self.vgg_loss = [VGGLoss()] - - self.in_nmlF = [ - item[0] for item in self.opt.in_nml - if '_F' in item[0] or item[0] == 'image' - ] - self.in_nmlB = [ - item[0] for item in self.opt.in_nml - if '_B' in item[0] or item[0] == 'image' - ] - self.in_nmlF_dim = sum([ - item[1] for item in self.opt.in_nml - if '_F' in item[0] or item[0] == 'image' - ]) - self.in_nmlB_dim = sum([ - item[1] for item in self.opt.in_nml - if '_B' in item[0] or item[0] == 'image' - ]) - - self.netF = define_G(self.in_nmlF_dim, 3, 64, "global", 4, 9, 1, 3, - "instance") - self.netB = define_G(self.in_nmlB_dim, 3, 64, "global", 4, 9, 1, 3, - "instance") - - init_net(self) - - def forward(self, in_tensor): - - inF_list = [] - inB_list = [] - - for name in self.in_nmlF: - inF_list.append(in_tensor[name]) - for name in self.in_nmlB: - inB_list.append(in_tensor[name]) - - nmlF = self.netF(torch.cat(inF_list, dim=1)) - nmlB = self.netB(torch.cat(inB_list, dim=1)) - - # ||normal|| == 1 - nmlF /= torch.norm(nmlF, dim=1) - nmlB /= torch.norm(nmlB, dim=1) - - # output: float_arr [-1,1] with [B, C, H, W] - - mask = (in_tensor['image'].abs().sum(dim=1, keepdim=True) != - 0.0).detach().float() - - nmlF = nmlF * mask - nmlB = nmlB * mask - - return nmlF, nmlB - - def get_norm_error(self, prd_F, prd_B, tgt): - """calculate normal loss - Args: - pred (torch.tensor): [B, 6, 512, 512] - tagt (torch.tensor): [B, 6, 512, 512] - """ - - tgt_F, tgt_B = tgt['normal_F'], tgt['normal_B'] - - l1_F_loss = self.l1_loss(prd_F, tgt_F) - l1_B_loss = self.l1_loss(prd_B, tgt_B) - - with torch.no_grad(): - vgg_F_loss = self.vgg_loss[0](prd_F, tgt_F) - vgg_B_loss = self.vgg_loss[0](prd_B, tgt_B) - - total_loss = [ - 5.0 * l1_F_loss + vgg_F_loss, 5.0 * l1_B_loss + vgg_B_loss - ] - - return total_loss \ No newline at end of file diff --git a/lib /net /local_affine.py b/lib /net /local_affine.py deleted file mode 100644 index 5399842502c8b1adf8f60cf53e38b6daa7f35068..0000000000000000000000000000000000000000 --- a/lib /net /local_affine.py +++ /dev/null @@ -1,57 +0,0 @@ -# Copyright 2021 by Haozhe Wu, Tsinghua University, Department of Computer Science and Technology. -# All rights reserved. -# This file is part of the pytorch-nicp, -# and is released under the "MIT License Agreement". Please see the LICENSE -# file that should have been included as part of this package. - -import torch -import torch.nn as nn -import torch.sparse as sp - -# reference: https://github.com/wuhaozhe/pytorch-nicp -class LocalAffine(nn.Module): - def __init__(self, num_points, batch_size=1, edges=None): - ''' - specify the number of points, the number of points should be constant across the batch - and the edges torch.Longtensor() with shape N * 2 - the local affine operator supports batch operation - batch size must be constant - add additional pooling on top of w matrix - ''' - super(LocalAffine, self).__init__() - self.A = nn.Parameter(torch.eye(3).unsqueeze( - 0).unsqueeze(0).repeat(batch_size, num_points, 1, 1)) - self.b = nn.Parameter(torch.zeros(3).unsqueeze(0).unsqueeze( - 0).unsqueeze(3).repeat(batch_size, num_points, 1, 1)) - self.edges = edges - self.num_points = num_points - - def stiffness(self): - ''' - calculate the stiffness of local affine transformation - f norm get infinity gradient when w is zero matrix, - ''' - if self.edges is None: - raise Exception("edges cannot be none when calculate stiff") - idx1 = self.edges[:, 0] - idx2 = self.edges[:, 1] - affine_weight = torch.cat((self.A, self.b), dim=3) - w1 = torch.index_select(affine_weight, dim=1, index=idx1) - w2 = torch.index_select(affine_weight, dim=1, index=idx2) - w_diff = (w1 - w2) ** 2 - w_rigid = (torch.linalg.det(self.A) - 1.0) ** 2 - return w_diff, w_rigid - - def forward(self, x, return_stiff=False): - ''' - x should have shape of B * N * 3 - ''' - x = x.unsqueeze(3) - out_x = torch.matmul(self.A, x) - out_x = out_x + self.b - out_x.squeeze_(3) - if return_stiff: - stiffness, rigid = self.stiffness() - return out_x, stiffness, rigid - else: - return out_x \ No newline at end of file diff --git a/lib /net /net_util.py b/lib /net /net_util.py deleted file mode 100644 index 0b794eaa3995d32e098299ea3a2f7a4893c8d3d4..0000000000000000000000000000000000000000 --- a/lib /net /net_util.py +++ /dev/null @@ -1,323 +0,0 @@ - -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from torchvision import models -import torch -from torch.nn import init -import torch.nn as nn -import torch.nn.functional as F -import functools -from torch.autograd import grad - - -def gradient(inputs, outputs): - d_points = torch.ones_like(outputs, - requires_grad=False, - device=outputs.device) - points_grad = grad(outputs=outputs, - inputs=inputs, - grad_outputs=d_points, - create_graph=True, - retain_graph=True, - only_inputs=True, - allow_unused=True)[0] - return points_grad - - -# def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False): -# "3x3 convolution with padding" -# return nn.Conv2d(in_planes, out_planes, kernel_size=3, -# stride=strd, padding=padding, bias=bias) - - -def conv3x3(in_planes, - out_planes, - kernel=3, - strd=1, - dilation=1, - padding=1, - bias=False): - "3x3 convolution with padding" - return nn.Conv2d(in_planes, - out_planes, - kernel_size=kernel, - dilation=dilation, - stride=strd, - padding=padding, - bias=bias) - - -def conv1x1(in_planes, out_planes, stride=1): - """1x1 convolution""" - return nn.Conv2d(in_planes, - out_planes, - kernel_size=1, - stride=stride, - bias=False) - - -def init_weights(net, init_type='normal', init_gain=0.02): - """Initialize network weights. - Parameters: - net (network) -- network to be initialized - init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal - init_gain (float) -- scaling factor for normal, xavier and orthogonal. - We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might - work better for some applications. Feel free to try yourself. - """ - def init_func(m): # define the initialization function - classname = m.__class__.__name__ - if hasattr(m, 'weight') and (classname.find('Conv') != -1 - or classname.find('Linear') != -1): - if init_type == 'normal': - init.normal_(m.weight.data, 0.0, init_gain) - elif init_type == 'xavier': - init.xavier_normal_(m.weight.data, gain=init_gain) - elif init_type == 'kaiming': - init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') - elif init_type == 'orthogonal': - init.orthogonal_(m.weight.data, gain=init_gain) - else: - raise NotImplementedError( - 'initialization method [%s] is not implemented' % - init_type) - if hasattr(m, 'bias') and m.bias is not None: - init.constant_(m.bias.data, 0.0) - elif classname.find( - 'BatchNorm2d' - ) != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies. - init.normal_(m.weight.data, 1.0, init_gain) - init.constant_(m.bias.data, 0.0) - - # print('initialize network with %s' % init_type) - net.apply(init_func) # apply the initialization function - - -def init_net(net, init_type='xavier', init_gain=0.02, gpu_ids=[]): - """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights - Parameters: - net (network) -- the network to be initialized - init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal - gain (float) -- scaling factor for normal, xavier and orthogonal. - gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 - Return an initialized network. - """ - if len(gpu_ids) > 0: - assert (torch.cuda.is_available()) - net = torch.nn.DataParallel(net) # multi-GPUs - init_weights(net, init_type, init_gain=init_gain) - return net - - -def imageSpaceRotation(xy, rot): - ''' - args: - xy: (B, 2, N) input - rot: (B, 2) x,y axis rotation angles - rotation center will be always image center (other rotation center can be represented by additional z translation) - ''' - disp = rot.unsqueeze(2).sin().expand_as(xy) - return (disp * xy).sum(dim=1) - - -def cal_gradient_penalty(netD, - real_data, - fake_data, - device, - type='mixed', - constant=1.0, - lambda_gp=10.0): - """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028 - Arguments: - netD (network) -- discriminator network - real_data (tensor array) -- real images - fake_data (tensor array) -- generated images from the generator - device (str) -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') - type (str) -- if we mix real and fake data or not [real | fake | mixed]. - constant (float) -- the constant used in formula ( | |gradient||_2 - constant)^2 - lambda_gp (float) -- weight for this loss - Returns the gradient penalty loss - """ - if lambda_gp > 0.0: - # either use real images, fake images, or a linear interpolation of two. - if type == 'real': - interpolatesv = real_data - elif type == 'fake': - interpolatesv = fake_data - elif type == 'mixed': - alpha = torch.rand(real_data.shape[0], 1) - alpha = alpha.expand( - real_data.shape[0], - real_data.nelement() // - real_data.shape[0]).contiguous().view(*real_data.shape) - alpha = alpha.to(device) - interpolatesv = alpha * real_data + ((1 - alpha) * fake_data) - else: - raise NotImplementedError('{} not implemented'.format(type)) - interpolatesv.requires_grad_(True) - disc_interpolates = netD(interpolatesv) - gradients = torch.autograd.grad( - outputs=disc_interpolates, - inputs=interpolatesv, - grad_outputs=torch.ones(disc_interpolates.size()).to(device), - create_graph=True, - retain_graph=True, - only_inputs=True) - gradients = gradients[0].view(real_data.size(0), -1) # flat the data - gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** - 2).mean() * lambda_gp # added eps - return gradient_penalty, gradients - else: - return 0.0, None - - -def get_norm_layer(norm_type='instance'): - """Return a normalization layer - Parameters: - norm_type (str) -- the name of the normalization layer: batch | instance | none - For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev). - For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics. - """ - if norm_type == 'batch': - norm_layer = functools.partial(nn.BatchNorm2d, - affine=True, - track_running_stats=True) - elif norm_type == 'instance': - norm_layer = functools.partial(nn.InstanceNorm2d, - affine=False, - track_running_stats=False) - elif norm_type == 'group': - norm_layer = functools.partial(nn.GroupNorm, 32) - elif norm_type == 'none': - norm_layer = None - else: - raise NotImplementedError('normalization layer [%s] is not found' % - norm_type) - return norm_layer - - -class Flatten(nn.Module): - def forward(self, input): - return input.view(input.size(0), -1) - - -class ConvBlock(nn.Module): - def __init__(self, in_planes, out_planes, opt): - super(ConvBlock, self).__init__() - [k, s, d, p] = opt.conv3x3 - self.conv1 = conv3x3(in_planes, int(out_planes / 2), k, s, d, p) - self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4), k, s, d, - p) - self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4), k, s, d, - p) - - if opt.norm == 'batch': - self.bn1 = nn.BatchNorm2d(in_planes) - self.bn2 = nn.BatchNorm2d(int(out_planes / 2)) - self.bn3 = nn.BatchNorm2d(int(out_planes / 4)) - self.bn4 = nn.BatchNorm2d(in_planes) - elif opt.norm == 'group': - self.bn1 = nn.GroupNorm(32, in_planes) - self.bn2 = nn.GroupNorm(32, int(out_planes / 2)) - self.bn3 = nn.GroupNorm(32, int(out_planes / 4)) - self.bn4 = nn.GroupNorm(32, in_planes) - - if in_planes != out_planes: - self.downsample = nn.Sequential( - self.bn4, - nn.ReLU(True), - nn.Conv2d(in_planes, - out_planes, - kernel_size=1, - stride=1, - bias=False), - ) - else: - self.downsample = None - - def forward(self, x): - residual = x - - out1 = self.bn1(x) - out1 = F.relu(out1, True) - out1 = self.conv1(out1) - - out2 = self.bn2(out1) - out2 = F.relu(out2, True) - out2 = self.conv2(out2) - - out3 = self.bn3(out2) - out3 = F.relu(out3, True) - out3 = self.conv3(out3) - - out3 = torch.cat((out1, out2, out3), 1) - - if self.downsample is not None: - residual = self.downsample(residual) - - out3 += residual - - return out3 - - -class Vgg19(torch.nn.Module): - def __init__(self, requires_grad=False): - super(Vgg19, self).__init__() - vgg_pretrained_features = models.vgg19(pretrained=True).features - self.slice1 = torch.nn.Sequential() - self.slice2 = torch.nn.Sequential() - self.slice3 = torch.nn.Sequential() - self.slice4 = torch.nn.Sequential() - self.slice5 = torch.nn.Sequential() - for x in range(2): - self.slice1.add_module(str(x), vgg_pretrained_features[x]) - for x in range(2, 7): - self.slice2.add_module(str(x), vgg_pretrained_features[x]) - for x in range(7, 12): - self.slice3.add_module(str(x), vgg_pretrained_features[x]) - for x in range(12, 21): - self.slice4.add_module(str(x), vgg_pretrained_features[x]) - for x in range(21, 30): - self.slice5.add_module(str(x), vgg_pretrained_features[x]) - if not requires_grad: - for param in self.parameters(): - param.requires_grad = False - - def forward(self, X): - h_relu1 = self.slice1(X) - h_relu2 = self.slice2(h_relu1) - h_relu3 = self.slice3(h_relu2) - h_relu4 = self.slice4(h_relu3) - h_relu5 = self.slice5(h_relu4) - out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5] - return out - - -class VGGLoss(nn.Module): - def __init__(self): - super(VGGLoss, self).__init__() - self.vgg = Vgg19() - self.criterion = nn.L1Loss() - self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0] - - def forward(self, x, y): - x_vgg, y_vgg = self.vgg(x), self.vgg(y) - loss = 0 - for i in range(len(x_vgg)): - loss += self.weights[i] * self.criterion(x_vgg[i], - y_vgg[i].detach()) - return loss diff --git a/lib /net /voxelize.py b/lib /net /voxelize.py deleted file mode 100644 index bf2936bd646693a2af1bffac766020057119f6d5..0000000000000000000000000000000000000000 --- a/lib /net /voxelize.py +++ /dev/null @@ -1,184 +0,0 @@ -from __future__ import division, print_function -import torch -import torch.nn as nn -import torch.nn.functional as F -import numpy as np -from torch.autograd import Function - -import voxelize_cuda - - -class VoxelizationFunction(Function): - """ - Definition of differentiable voxelization function - Currently implemented only for cuda Tensors - """ - @staticmethod - def forward(ctx, smpl_vertices, smpl_face_center, smpl_face_normal, - smpl_vertex_code, smpl_face_code, smpl_tetrahedrons, - volume_res, sigma, smooth_kernel_size): - """ - forward pass - Output format: (batch_size, z_dims, y_dims, x_dims, channel_num) - """ - assert (smpl_vertices.size()[1] == smpl_vertex_code.size()[1]) - assert (smpl_face_center.size()[1] == smpl_face_normal.size()[1]) - assert (smpl_face_center.size()[1] == smpl_face_code.size()[1]) - ctx.batch_size = smpl_vertices.size()[0] - ctx.volume_res = volume_res - ctx.sigma = sigma - ctx.smooth_kernel_size = smooth_kernel_size - ctx.smpl_vertex_num = smpl_vertices.size()[1] - ctx.device = smpl_vertices.device - - smpl_vertices = smpl_vertices.contiguous() - smpl_face_center = smpl_face_center.contiguous() - smpl_face_normal = smpl_face_normal.contiguous() - smpl_vertex_code = smpl_vertex_code.contiguous() - smpl_face_code = smpl_face_code.contiguous() - smpl_tetrahedrons = smpl_tetrahedrons.contiguous() - - occ_volume = torch.cuda.FloatTensor(ctx.batch_size, ctx.volume_res, - ctx.volume_res, - ctx.volume_res).fill_(0.0) - semantic_volume = torch.cuda.FloatTensor(ctx.batch_size, - ctx.volume_res, - ctx.volume_res, - ctx.volume_res, 3).fill_(0.0) - weight_sum_volume = torch.cuda.FloatTensor(ctx.batch_size, - ctx.volume_res, - ctx.volume_res, - ctx.volume_res).fill_(1e-3) - - # occ_volume [B, volume_res, volume_res, volume_res] - # semantic_volume [B, volume_res, volume_res, volume_res, 3] - # weight_sum_volume [B, volume_res, volume_res, volume_res] - - occ_volume, semantic_volume, weight_sum_volume = voxelize_cuda.forward_semantic_voxelization( - smpl_vertices, smpl_vertex_code, smpl_tetrahedrons, occ_volume, - semantic_volume, weight_sum_volume, sigma) - - return semantic_volume - - -class Voxelization(nn.Module): - """ - Wrapper around the autograd function VoxelizationFunction - """ - - def __init__(self, smpl_vertex_code, smpl_face_code, smpl_face_indices, - smpl_tetraderon_indices, volume_res, sigma, - smooth_kernel_size, batch_size, device): - super(Voxelization, self).__init__() - assert (len(smpl_face_indices.shape) == 2) - assert (len(smpl_tetraderon_indices.shape) == 2) - assert (smpl_face_indices.shape[1] == 3) - assert (smpl_tetraderon_indices.shape[1] == 4) - - self.volume_res = volume_res - self.sigma = sigma - self.smooth_kernel_size = smooth_kernel_size - self.batch_size = batch_size - self.device = device - - self.smpl_vertex_code = smpl_vertex_code - self.smpl_face_code = smpl_face_code - self.smpl_face_indices = smpl_face_indices - self.smpl_tetraderon_indices = smpl_tetraderon_indices - - def update_param(self, batch_size, smpl_tetra): - - self.batch_size = batch_size - self.smpl_tetraderon_indices = smpl_tetra - - smpl_vertex_code_batch = np.tile(self.smpl_vertex_code, - (self.batch_size, 1, 1)) - smpl_face_code_batch = np.tile(self.smpl_face_code, - (self.batch_size, 1, 1)) - smpl_face_indices_batch = np.tile(self.smpl_face_indices, - (self.batch_size, 1, 1)) - smpl_tetraderon_indices_batch = np.tile(self.smpl_tetraderon_indices, - (self.batch_size, 1, 1)) - - smpl_vertex_code_batch = torch.from_numpy( - smpl_vertex_code_batch).contiguous().to(self.device) - smpl_face_code_batch = torch.from_numpy( - smpl_face_code_batch).contiguous().to(self.device) - smpl_face_indices_batch = torch.from_numpy( - smpl_face_indices_batch).contiguous().to(self.device) - smpl_tetraderon_indices_batch = torch.from_numpy( - smpl_tetraderon_indices_batch).contiguous().to(self.device) - - self.register_buffer('smpl_vertex_code_batch', smpl_vertex_code_batch) - self.register_buffer('smpl_face_code_batch', smpl_face_code_batch) - self.register_buffer('smpl_face_indices_batch', - smpl_face_indices_batch) - self.register_buffer('smpl_tetraderon_indices_batch', - smpl_tetraderon_indices_batch) - - def forward(self, smpl_vertices): - """ - Generate semantic volumes from SMPL vertices - """ - assert (smpl_vertices.size()[0] == self.batch_size) - self.check_input(smpl_vertices) - smpl_faces = self.vertices_to_faces(smpl_vertices) - smpl_tetrahedrons = self.vertices_to_tetrahedrons(smpl_vertices) - smpl_face_center = self.calc_face_centers(smpl_faces) - smpl_face_normal = self.calc_face_normals(smpl_faces) - smpl_surface_vertex_num = self.smpl_vertex_code_batch.size()[1] - smpl_vertices_surface = smpl_vertices[:, :smpl_surface_vertex_num, :] - vol = VoxelizationFunction.apply(smpl_vertices_surface, - smpl_face_center, smpl_face_normal, - self.smpl_vertex_code_batch, - self.smpl_face_code_batch, - smpl_tetrahedrons, self.volume_res, - self.sigma, self.smooth_kernel_size) - return vol.permute((0, 4, 1, 2, 3)) # (bzyxc --> bcdhw) - - def vertices_to_faces(self, vertices): - assert (vertices.ndimension() == 3) - bs, nv = vertices.shape[:2] - device = vertices.device - face = self.smpl_face_indices_batch + ( - torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None] - vertices_ = vertices.reshape((bs * nv, 3)) - return vertices_[face.long()] - - def vertices_to_tetrahedrons(self, vertices): - assert (vertices.ndimension() == 3) - bs, nv = vertices.shape[:2] - device = vertices.device - tets = self.smpl_tetraderon_indices_batch + ( - torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None] - vertices_ = vertices.reshape((bs * nv, 3)) - return vertices_[tets.long()] - - def calc_face_centers(self, face_verts): - assert len(face_verts.shape) == 4 - assert face_verts.shape[2] == 3 - assert face_verts.shape[3] == 3 - bs, nf = face_verts.shape[:2] - face_centers = (face_verts[:, :, 0, :] + face_verts[:, :, 1, :] + - face_verts[:, :, 2, :]) / 3.0 - face_centers = face_centers.reshape((bs, nf, 3)) - return face_centers - - def calc_face_normals(self, face_verts): - assert len(face_verts.shape) == 4 - assert face_verts.shape[2] == 3 - assert face_verts.shape[3] == 3 - bs, nf = face_verts.shape[:2] - face_verts = face_verts.reshape((bs * nf, 3, 3)) - v10 = face_verts[:, 0] - face_verts[:, 1] - v12 = face_verts[:, 2] - face_verts[:, 1] - normals = F.normalize(torch.cross(v10, v12), eps=1e-5) - normals = normals.reshape((bs, nf, 3)) - return normals - - def check_input(self, x): - if x.device == 'cpu': - raise TypeError('Voxelization module supports only cuda tensors') - if x.type() != 'torch.cuda.FloatTensor': - raise TypeError( - 'Voxelization module supports only float32 tensors') \ No newline at end of file diff --git a/lib /smplx / .gitignore b/lib /smplx / .gitignore deleted file mode 100644 index f1bf708c5b86d8e13496eba43dae45dcfec7bcc8..0000000000000000000000000000000000000000 --- a/lib /smplx / .gitignore +++ /dev/null @@ -1,114 +0,0 @@ -#### joe made this: http://goel.io/joe - -#####=== Python ===##### - -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$py.class - -# C extensions -*.so - -# Distribution / packaging -.Python -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -*.egg-info/ -.installed.cfg -*.egg -MANIFEST - -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -.hypothesis/ -.pytest_cache/ - -# Translations -*.mo -*.pot - -# Django stuff: -*.log -local_settings.py -db.sqlite3 - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - -# Sphinx documentation -docs/_build/ - -# PyBuilder -target/ - -# Jupyter Notebook -.ipynb_checkpoints - -# pyenv -.python-version - -# celery beat schedule file -celerybeat-schedule - -# SageMath parsed files -*.sage.py - -# Environments -.env -.venv -env/ -venv/ -ENV/ -env.bak/ -venv.bak/ - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site - -# mypy -.mypy_cache/ -models/ -output/ -outputs/ -transfer_data/ -torch-trust-ncg/ -build/ \ No newline at end of file diff --git a/lib /smplx / __init__.py b/lib /smplx / __init__.py deleted file mode 100644 index 9bd031b5cd03d61226709fc324b66b8403e6f08b..0000000000000000000000000000000000000000 --- a/lib /smplx / __init__.py +++ /dev/null @@ -1,30 +0,0 @@ -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from .body_models import ( - create, - SMPL, - SMPLH, - SMPLX, - MANO, - FLAME, - build_layer, - SMPLLayer, - SMPLHLayer, - SMPLXLayer, - MANOLayer, - FLAMELayer, -) \ No newline at end of file diff --git a/lib /smplx / vertex_ids.py b/lib /smplx / vertex_ids.py deleted file mode 100644 index 5c77e31f290bad38c2815208c645e0211be09df2..0000000000000000000000000000000000000000 --- a/lib /smplx / vertex_ids.py +++ /dev/null @@ -1,77 +0,0 @@ -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from __future__ import print_function -from __future__ import absolute_import -from __future__ import division - -# Joint name to vertex mapping. SMPL/SMPL-H/SMPL-X vertices that correspond to -# MSCOCO and OpenPose joints -vertex_ids = { - 'smplh': { - 'nose': 332, - 'reye': 6260, - 'leye': 2800, - 'rear': 4071, - 'lear': 583, - 'rthumb': 6191, - 'rindex': 5782, - 'rmiddle': 5905, - 'rring': 6016, - 'rpinky': 6133, - 'lthumb': 2746, - 'lindex': 2319, - 'lmiddle': 2445, - 'lring': 2556, - 'lpinky': 2673, - 'LBigToe': 3216, - 'LSmallToe': 3226, - 'LHeel': 3387, - 'RBigToe': 6617, - 'RSmallToe': 6624, - 'RHeel': 6787 - }, - 'smplx': { - 'nose': 9120, - 'reye': 9929, - 'leye': 9448, - 'rear': 616, - 'lear': 6, - 'rthumb': 8079, - 'rindex': 7669, - 'rmiddle': 7794, - 'rring': 7905, - 'rpinky': 8022, - 'lthumb': 5361, - 'lindex': 4933, - 'lmiddle': 5058, - 'lring': 5169, - 'lpinky': 5286, - 'LBigToe': 5770, - 'LSmallToe': 5780, - 'LHeel': 8846, - 'RBigToe': 8463, - 'RSmallToe': 8474, - 'RHeel': 8635 - }, - 'mano': { - 'thumb': 744, - 'index': 320, - 'middle': 443, - 'ring': 554, - 'pinky': 671, - } -} \ No newline at end of file diff --git a/lib /smplx /README.md b/lib /smplx /README.md deleted file mode 100644 index ef5d84a7abdab40b9a2f7602f9111d2c40288411..0000000000000000000000000000000000000000 --- a/lib /smplx /README.md +++ /dev/null @@ -1,32 +0,0 @@ - -Citation -Depending on which model is loaded for your project, i.e. SMPL-X or SMPL+H or SMPL, please cite the most relevant work below, listed in the same order: - -@inproceedings{SMPL-X:2019, - title = {Expressive Body Capture: 3D Hands, Face, and Body from a Single Image}, - author = {Pavlakos, Georgios and Choutas, Vasileios and Ghorbani, Nima and Bolkart, Timo and Osman, Ahmed A. A. and Tzionas, Dimitrios and Black, Michael J.}, - booktitle = {Proceedings IEEE Conf. on Computer Vision and Pattern Recognition (CVPR)}, - year = {2019} -} -@article{MANO:SIGGRAPHASIA:2017, - title = {Embodied Hands: Modeling and Capturing Hands and Bodies Together}, - author = {Romero, Javier and Tzionas, Dimitrios and Black, Michael J.}, - journal = {ACM Transactions on Graphics, (Proc. SIGGRAPH Asia)}, - volume = {36}, - number = {6}, - series = {245:1--245:17}, - month = nov, - year = {2017}, - month_numeric = {11} - } -@article{SMPL:2015, - author = {Loper, Matthew and Mahmood, Naureen and Romero, Javier and Pons-Moll, Gerard and Black, Michael J.}, - title = {{SMPL}: A Skinned Multi-Person Linear Model}, - journal = {ACM Transactions on Graphics, (Proc. SIGGRAPH Asia)}, - month = oct, - number = {6}, - pages = {248:1--248:16}, - publisher = {ACM}, - volume = {34}, - year = {2015} -} diff --git a/lib /smplx /body_models.py b/lib /smplx /body_models.py deleted file mode 100644 index d0896a7368ce7368d500d7333585fabe00193cc8..0000000000000000000000000000000000000000 --- a/lib /smplx /body_models.py +++ /dev/null @@ -1,2392 +0,0 @@ -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from typing import Optional, Dict, Union -import os -import os.path as osp -import pickle - -import numpy as np -from termcolor import colored - -import torch -import torch.nn as nn -from collections import namedtuple -from huggingface_hub import cached_download - -import logging -logging.getLogger("smplx").setLevel(logging.ERROR) - -from .lbs import ( - lbs, vertices2landmarks, find_dynamic_lmk_idx_and_bcoords) - -from .vertex_ids import vertex_ids as VERTEX_IDS -from .utils import ( - Struct, to_np, to_tensor, Tensor, Array, - SMPLOutput, - SMPLHOutput, - SMPLXOutput, - MANOOutput, - FLAMEOutput, - find_joint_kin_chain) -from .vertex_joint_selector import VertexJointSelector - -ModelOutput = namedtuple('ModelOutput', - ['vertices', 'joints', 'full_pose', 'betas', - 'global_orient', - 'body_pose', 'expression', - 'left_hand_pose', 'right_hand_pose', - 'jaw_pose']) -ModelOutput.__new__.__defaults__ = (None,) * len(ModelOutput._fields) - -class SMPL(nn.Module): - - NUM_JOINTS = 23 - NUM_BODY_JOINTS = 23 - SHAPE_SPACE_DIM = 300 - - def __init__( - self, model_path: str, - kid_template_path: str = '', - data_struct: Optional[Struct] = None, - create_betas: bool = True, - betas: Optional[Tensor] = None, - num_betas: int = 10, - create_global_orient: bool = True, - global_orient: Optional[Tensor] = None, - create_body_pose: bool = True, - body_pose: Optional[Tensor] = None, - create_transl: bool = True, - transl: Optional[Tensor] = None, - dtype=torch.float32, - batch_size: int = 1, - joint_mapper=None, - gender: str = 'neutral', - age: str = 'adult', - vertex_ids: Dict[str, int] = None, - v_template: Optional[Union[Tensor, Array]] = None, - v_personal: Optional[Union[Tensor, Array]] = None, - **kwargs - ) -> None: - ''' SMPL model constructor - Parameters - ---------- - model_path: str - The path to the folder or to the file where the model - parameters are stored - data_struct: Strct - A struct object. If given, then the parameters of the model are - read from the object. Otherwise, the model tries to read the - parameters from the given `model_path`. (default = None) - create_global_orient: bool, optional - Flag for creating a member variable for the global orientation - of the body. (default = True) - global_orient: torch.tensor, optional, Bx3 - The default value for the global orientation variable. - (default = None) - create_body_pose: bool, optional - Flag for creating a member variable for the pose of the body. - (default = True) - body_pose: torch.tensor, optional, Bx(Body Joints * 3) - The default value for the body pose variable. - (default = None) - num_betas: int, optional - Number of shape components to use - (default = 10). - create_betas: bool, optional - Flag for creating a member variable for the shape space - (default = True). - betas: torch.tensor, optional, Bx10 - The default value for the shape member variable. - (default = None) - create_transl: bool, optional - Flag for creating a member variable for the translation - of the body. (default = True) - transl: torch.tensor, optional, Bx3 - The default value for the transl variable. - (default = None) - dtype: torch.dtype, optional - The data type for the created variables - batch_size: int, optional - The batch size used for creating the member variables - joint_mapper: object, optional - An object that re-maps the joints. Useful if one wants to - re-order the SMPL joints to some other convention (e.g. MSCOCO) - (default = None) - gender: str, optional - Which gender to load - vertex_ids: dict, optional - A dictionary containing the indices of the extra vertices that - will be selected - ''' - - self.gender = gender - self.age = age - - if data_struct is None: - model_fn = 'SMPL_{}.{ext}'.format(gender.upper(), ext='pkl') - smpl_path = cached_download(os.path.join(model_path, model_fn), use_auth_token=os.environ['ICON']) - - with open(smpl_path, 'rb') as smpl_file: - data_struct = Struct(**pickle.load(smpl_file, - encoding='latin1')) - - super(SMPL, self).__init__() - self.batch_size = batch_size - shapedirs = data_struct.shapedirs - if (shapedirs.shape[-1] < self.SHAPE_SPACE_DIM): - # print(f'WARNING: You are using a {self.name()} model, with only' - # ' 10 shape coefficients.') - num_betas = min(num_betas, 10) - else: - num_betas = min(num_betas, self.SHAPE_SPACE_DIM) - - if self.age=='kid': - v_template_smil = np.load(kid_template_path) - v_template_smil -= np.mean(v_template_smil, axis=0) - v_template_diff = np.expand_dims(v_template_smil - data_struct.v_template, axis=2) - shapedirs = np.concatenate((shapedirs[:, :, :num_betas], v_template_diff), axis=2) - num_betas = num_betas + 1 - - self._num_betas = num_betas - shapedirs = shapedirs[:, :, :num_betas] - # The shape components - self.register_buffer( - 'shapedirs', - to_tensor(to_np(shapedirs), dtype=dtype)) - - if vertex_ids is None: - # SMPL and SMPL-H share the same topology, so any extra joints can - # be drawn from the same place - vertex_ids = VERTEX_IDS['smplh'] - - self.dtype = dtype - - self.joint_mapper = joint_mapper - - self.vertex_joint_selector = VertexJointSelector( - vertex_ids=vertex_ids, **kwargs) - - self.faces = data_struct.f - self.register_buffer('faces_tensor', - to_tensor(to_np(self.faces, dtype=np.int64), - dtype=torch.long)) - - if create_betas: - if betas is None: - default_betas = torch.zeros( - [batch_size, self.num_betas], dtype=dtype) - else: - if torch.is_tensor(betas): - default_betas = betas.clone().detach() - else: - default_betas = torch.tensor(betas, dtype=dtype) - - self.register_parameter( - 'betas', nn.Parameter(default_betas, requires_grad=True)) - - # The tensor that contains the global rotation of the model - # It is separated from the pose of the joints in case we wish to - # optimize only over one of them - if create_global_orient: - if global_orient is None: - default_global_orient = torch.zeros( - [batch_size, 3], dtype=dtype) - else: - if torch.is_tensor(global_orient): - default_global_orient = global_orient.clone().detach() - else: - default_global_orient = torch.tensor( - global_orient, dtype=dtype) - - global_orient = nn.Parameter(default_global_orient, - requires_grad=True) - self.register_parameter('global_orient', global_orient) - - if create_body_pose: - if body_pose is None: - default_body_pose = torch.zeros( - [batch_size, self.NUM_BODY_JOINTS * 3], dtype=dtype) - else: - if torch.is_tensor(body_pose): - default_body_pose = body_pose.clone().detach() - else: - default_body_pose = torch.tensor(body_pose, - dtype=dtype) - self.register_parameter( - 'body_pose', - nn.Parameter(default_body_pose, requires_grad=True)) - - if create_transl: - if transl is None: - default_transl = torch.zeros([batch_size, 3], - dtype=dtype, - requires_grad=True) - else: - default_transl = torch.tensor(transl, dtype=dtype) - self.register_parameter( - 'transl', nn.Parameter(default_transl, requires_grad=True)) - - if v_template is None: - v_template = data_struct.v_template - - if not torch.is_tensor(v_template): - v_template = to_tensor(to_np(v_template), dtype=dtype) - - if v_personal is not None: - v_personal = to_tensor(to_np(v_personal), dtype=dtype) - v_template += v_personal - - # The vertices of the template model - self.register_buffer('v_template', v_template) - - j_regressor = to_tensor(to_np( - data_struct.J_regressor), dtype=dtype) - self.register_buffer('J_regressor', j_regressor) - - # Pose blend shape basis: 6890 x 3 x 207, reshaped to 6890*3 x 207 - num_pose_basis = data_struct.posedirs.shape[-1] - # 207 x 20670 - posedirs = np.reshape(data_struct.posedirs, [-1, num_pose_basis]).T - self.register_buffer('posedirs', - to_tensor(to_np(posedirs), dtype=dtype)) - - # indices of parents for each joints - parents = to_tensor(to_np(data_struct.kintree_table[0])).long() - parents[0] = -1 - self.register_buffer('parents', parents) - - self.register_buffer( - 'lbs_weights', to_tensor(to_np(data_struct.weights), dtype=dtype)) - - @property - def num_betas(self): - return self._num_betas - - @property - def num_expression_coeffs(self): - return 0 - - def create_mean_pose(self, data_struct) -> Tensor: - pass - - def name(self) -> str: - return 'SMPL' - - @torch.no_grad() - def reset_params(self, **params_dict) -> None: - for param_name, param in self.named_parameters(): - if param_name in params_dict: - param[:] = torch.tensor(params_dict[param_name]) - else: - param.fill_(0) - - def get_num_verts(self) -> int: - return self.v_template.shape[0] - - def get_num_faces(self) -> int: - return self.faces.shape[0] - - def extra_repr(self) -> str: - msg = [ - f'Gender: {self.gender.upper()}', - f'Number of joints: {self.J_regressor.shape[0]}', - f'Betas: {self.num_betas}', - ] - return '\n'.join(msg) - - def forward( - self, - betas: Optional[Tensor] = None, - body_pose: Optional[Tensor] = None, - global_orient: Optional[Tensor] = None, - transl: Optional[Tensor] = None, - return_verts=True, - return_full_pose: bool = False, - pose2rot: bool = True, - **kwargs - ) -> SMPLOutput: - ''' Forward pass for the SMPL model - Parameters - ---------- - global_orient: torch.tensor, optional, shape Bx3 - If given, ignore the member variable and use it as the global - rotation of the body. Useful if someone wishes to predicts this - with an external model. (default=None) - betas: torch.tensor, optional, shape BxN_b - If given, ignore the member variable `betas` and use it - instead. For example, it can used if shape parameters - `betas` are predicted from some external model. - (default=None) - body_pose: torch.tensor, optional, shape Bx(J*3) - If given, ignore the member variable `body_pose` and use it - instead. For example, it can used if someone predicts the - pose of the body joints are predicted from some external model. - It should be a tensor that contains joint rotations in - axis-angle format. (default=None) - transl: torch.tensor, optional, shape Bx3 - If given, ignore the member variable `transl` and use it - instead. For example, it can used if the translation - `transl` is predicted from some external model. - (default=None) - return_verts: bool, optional - Return the vertices. (default=True) - return_full_pose: bool, optional - Returns the full axis-angle pose vector (default=False) - Returns - ------- - ''' - # If no shape and pose parameters are passed along, then use the - # ones from the module - global_orient = (global_orient if global_orient is not None else - self.global_orient) - body_pose = body_pose if body_pose is not None else self.body_pose - betas = betas if betas is not None else self.betas - - apply_trans = transl is not None or hasattr(self, 'transl') - if transl is None and hasattr(self, 'transl'): - transl = self.transl - - full_pose = torch.cat([global_orient, body_pose], dim=1) - - batch_size = max(betas.shape[0], global_orient.shape[0], - body_pose.shape[0]) - - if betas.shape[0] != batch_size: - num_repeats = int(batch_size / betas.shape[0]) - betas = betas.expand(num_repeats, -1) - - vertices, joints = lbs(betas, full_pose, self.v_template, - self.shapedirs, self.posedirs, - self.J_regressor, self.parents, - self.lbs_weights, pose2rot=pose2rot) - - joints = self.vertex_joint_selector(vertices, joints) - # Map the joints to the current dataset - if self.joint_mapper is not None: - joints = self.joint_mapper(joints) - - if apply_trans: - joints += transl.unsqueeze(dim=1) - vertices += transl.unsqueeze(dim=1) - - output = SMPLOutput(vertices=vertices if return_verts else None, - global_orient=global_orient, - body_pose=body_pose, - joints=joints, - betas=betas, - full_pose=full_pose if return_full_pose else None) - - return output - - -class SMPLLayer(SMPL): - def __init__( - self, - *args, - **kwargs - ) -> None: - # Just create a SMPL module without any member variables - super(SMPLLayer, self).__init__( - create_body_pose=False, - create_betas=False, - create_global_orient=False, - create_transl=False, - *args, - **kwargs, - ) - - def forward( - self, - betas: Optional[Tensor] = None, - body_pose: Optional[Tensor] = None, - global_orient: Optional[Tensor] = None, - transl: Optional[Tensor] = None, - return_verts=True, - return_full_pose: bool = False, - pose2rot: bool = True, - **kwargs - ) -> SMPLOutput: - ''' Forward pass for the SMPL model - Parameters - ---------- - global_orient: torch.tensor, optional, shape Bx3x3 - Global rotation of the body. Useful if someone wishes to - predicts this with an external model. It is expected to be in - rotation matrix format. (default=None) - betas: torch.tensor, optional, shape BxN_b - Shape parameters. For example, it can used if shape parameters - `betas` are predicted from some external model. - (default=None) - body_pose: torch.tensor, optional, shape BxJx3x3 - Body pose. For example, it can used if someone predicts the - pose of the body joints are predicted from some external model. - It should be a tensor that contains joint rotations in - rotation matrix format. (default=None) - transl: torch.tensor, optional, shape Bx3 - Translation vector of the body. - For example, it can used if the translation - `transl` is predicted from some external model. - (default=None) - return_verts: bool, optional - Return the vertices. (default=True) - return_full_pose: bool, optional - Returns the full axis-angle pose vector (default=False) - Returns - ------- - ''' - model_vars = [betas, global_orient, body_pose, transl] - batch_size = 1 - for var in model_vars: - if var is None: - continue - batch_size = max(batch_size, len(var)) - device, dtype = self.shapedirs.device, self.shapedirs.dtype - if global_orient is None: - global_orient = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, -1, -1, -1).contiguous() - if body_pose is None: - body_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand( - batch_size, self.NUM_BODY_JOINTS, -1, -1).contiguous() - if betas is None: - betas = torch.zeros([batch_size, self.num_betas], - dtype=dtype, device=device) - if transl is None: - transl = torch.zeros([batch_size, 3], dtype=dtype, device=device) - full_pose = torch.cat( - [global_orient.reshape(-1, 1, 3, 3), - body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3, 3)], - dim=1) - - vertices, joints = lbs(betas, full_pose, self.v_template, - self.shapedirs, self.posedirs, - self.J_regressor, self.parents, - self.lbs_weights, - pose2rot=False) - - joints = self.vertex_joint_selector(vertices, joints) - # Map the joints to the current dataset - if self.joint_mapper is not None: - joints = self.joint_mapper(joints) - - if transl is not None: - joints += transl.unsqueeze(dim=1) - vertices += transl.unsqueeze(dim=1) - - output = SMPLOutput(vertices=vertices if return_verts else None, - global_orient=global_orient, - body_pose=body_pose, - joints=joints, - betas=betas, - full_pose=full_pose if return_full_pose else None) - - return output - - -class SMPLH(SMPL): - - # The hand joints are replaced by MANO - NUM_BODY_JOINTS = SMPL.NUM_JOINTS - 2 - NUM_HAND_JOINTS = 15 - NUM_JOINTS = NUM_BODY_JOINTS + 2 * NUM_HAND_JOINTS - - def __init__( - self, model_path, - kid_template_path: str = '', - data_struct: Optional[Struct] = None, - create_left_hand_pose: bool = True, - left_hand_pose: Optional[Tensor] = None, - create_right_hand_pose: bool = True, - right_hand_pose: Optional[Tensor] = None, - use_pca: bool = True, - num_pca_comps: int = 6, - flat_hand_mean: bool = False, - batch_size: int = 1, - gender: str = 'neutral', - age: str = 'adult', - dtype=torch.float32, - vertex_ids=None, - use_compressed: bool = True, - ext: str = 'pkl', - **kwargs - ) -> None: - ''' SMPLH model constructor - Parameters - ---------- - model_path: str - The path to the folder or to the file where the model - parameters are stored - data_struct: Strct - A struct object. If given, then the parameters of the model are - read from the object. Otherwise, the model tries to read the - parameters from the given `model_path`. (default = None) - create_left_hand_pose: bool, optional - Flag for creating a member variable for the pose of the left - hand. (default = True) - left_hand_pose: torch.tensor, optional, BxP - The default value for the left hand pose member variable. - (default = None) - create_right_hand_pose: bool, optional - Flag for creating a member variable for the pose of the right - hand. (default = True) - right_hand_pose: torch.tensor, optional, BxP - The default value for the right hand pose member variable. - (default = None) - num_pca_comps: int, optional - The number of PCA components to use for each hand. - (default = 6) - flat_hand_mean: bool, optional - If False, then the pose of the hand is initialized to False. - batch_size: int, optional - The batch size used for creating the member variables - gender: str, optional - Which gender to load - dtype: torch.dtype, optional - The data type for the created variables - vertex_ids: dict, optional - A dictionary containing the indices of the extra vertices that - will be selected - ''' - - self.num_pca_comps = num_pca_comps - # If no data structure is passed, then load the data from the given - # model folder - if data_struct is None: - # Load the model - if osp.isdir(model_path): - model_fn = 'SMPLH_{}.{ext}'.format(gender.upper(), ext=ext) - smplh_path = os.path.join(model_path, model_fn) - else: - smplh_path = model_path - assert osp.exists(smplh_path), 'Path {} does not exist!'.format( - smplh_path) - - if ext == 'pkl': - with open(smplh_path, 'rb') as smplh_file: - model_data = pickle.load(smplh_file, encoding='latin1') - elif ext == 'npz': - model_data = np.load(smplh_path, allow_pickle=True) - else: - raise ValueError('Unknown extension: {}'.format(ext)) - data_struct = Struct(**model_data) - - if vertex_ids is None: - vertex_ids = VERTEX_IDS['smplh'] - - super(SMPLH, self).__init__( - model_path=model_path, - kid_template_path=kid_template_path, - data_struct=data_struct, - batch_size=batch_size, vertex_ids=vertex_ids, gender=gender, age=age, - use_compressed=use_compressed, dtype=dtype, ext=ext, **kwargs) - - self.use_pca = use_pca - self.num_pca_comps = num_pca_comps - self.flat_hand_mean = flat_hand_mean - - left_hand_components = data_struct.hands_componentsl[:num_pca_comps] - right_hand_components = data_struct.hands_componentsr[:num_pca_comps] - - self.np_left_hand_components = left_hand_components - self.np_right_hand_components = right_hand_components - if self.use_pca: - self.register_buffer( - 'left_hand_components', - torch.tensor(left_hand_components, dtype=dtype)) - self.register_buffer( - 'right_hand_components', - torch.tensor(right_hand_components, dtype=dtype)) - - if self.flat_hand_mean: - left_hand_mean = np.zeros_like(data_struct.hands_meanl) - else: - left_hand_mean = data_struct.hands_meanl - - if self.flat_hand_mean: - right_hand_mean = np.zeros_like(data_struct.hands_meanr) - else: - right_hand_mean = data_struct.hands_meanr - - self.register_buffer('left_hand_mean', - to_tensor(left_hand_mean, dtype=self.dtype)) - self.register_buffer('right_hand_mean', - to_tensor(right_hand_mean, dtype=self.dtype)) - - # Create the buffers for the pose of the left hand - hand_pose_dim = num_pca_comps if use_pca else 3 * self.NUM_HAND_JOINTS - if create_left_hand_pose: - if left_hand_pose is None: - default_lhand_pose = torch.zeros([batch_size, hand_pose_dim], - dtype=dtype) - else: - default_lhand_pose = torch.tensor(left_hand_pose, dtype=dtype) - - left_hand_pose_param = nn.Parameter(default_lhand_pose, - requires_grad=True) - self.register_parameter('left_hand_pose', - left_hand_pose_param) - - if create_right_hand_pose: - if right_hand_pose is None: - default_rhand_pose = torch.zeros([batch_size, hand_pose_dim], - dtype=dtype) - else: - default_rhand_pose = torch.tensor(right_hand_pose, dtype=dtype) - - right_hand_pose_param = nn.Parameter(default_rhand_pose, - requires_grad=True) - self.register_parameter('right_hand_pose', - right_hand_pose_param) - - # Create the buffer for the mean pose. - pose_mean_tensor = self.create_mean_pose( - data_struct, flat_hand_mean=flat_hand_mean) - if not torch.is_tensor(pose_mean_tensor): - pose_mean_tensor = torch.tensor(pose_mean_tensor, dtype=dtype) - self.register_buffer('pose_mean', pose_mean_tensor) - - def create_mean_pose(self, data_struct, flat_hand_mean=False): - # Create the array for the mean pose. If flat_hand is false, then use - # the mean that is given by the data, rather than the flat open hand - global_orient_mean = torch.zeros([3], dtype=self.dtype) - body_pose_mean = torch.zeros([self.NUM_BODY_JOINTS * 3], - dtype=self.dtype) - - pose_mean = torch.cat([global_orient_mean, body_pose_mean, - self.left_hand_mean, - self.right_hand_mean], dim=0) - return pose_mean - - def name(self) -> str: - return 'SMPL+H' - - def extra_repr(self): - msg = super(SMPLH, self).extra_repr() - msg = [msg] - if self.use_pca: - msg.append(f'Number of PCA components: {self.num_pca_comps}') - msg.append(f'Flat hand mean: {self.flat_hand_mean}') - return '\n'.join(msg) - - def forward( - self, - betas: Optional[Tensor] = None, - global_orient: Optional[Tensor] = None, - body_pose: Optional[Tensor] = None, - left_hand_pose: Optional[Tensor] = None, - right_hand_pose: Optional[Tensor] = None, - transl: Optional[Tensor] = None, - return_verts: bool = True, - return_full_pose: bool = False, - pose2rot: bool = True, - **kwargs - ) -> SMPLHOutput: - ''' - ''' - - # If no shape and pose parameters are passed along, then use the - # ones from the module - global_orient = (global_orient if global_orient is not None else - self.global_orient) - body_pose = body_pose if body_pose is not None else self.body_pose - betas = betas if betas is not None else self.betas - left_hand_pose = (left_hand_pose if left_hand_pose is not None else - self.left_hand_pose) - right_hand_pose = (right_hand_pose if right_hand_pose is not None else - self.right_hand_pose) - - apply_trans = transl is not None or hasattr(self, 'transl') - if transl is None: - if hasattr(self, 'transl'): - transl = self.transl - - if self.use_pca: - left_hand_pose = torch.einsum( - 'bi,ij->bj', [left_hand_pose, self.left_hand_components]) - right_hand_pose = torch.einsum( - 'bi,ij->bj', [right_hand_pose, self.right_hand_components]) - - full_pose = torch.cat([global_orient, body_pose, - left_hand_pose, - right_hand_pose], dim=1) - - full_pose += self.pose_mean - - vertices, joints = lbs(betas, full_pose, self.v_template, - self.shapedirs, self.posedirs, - self.J_regressor, self.parents, - self.lbs_weights, pose2rot=pose2rot) - - # Add any extra joints that might be needed - joints = self.vertex_joint_selector(vertices, joints) - if self.joint_mapper is not None: - joints = self.joint_mapper(joints) - - if apply_trans: - joints += transl.unsqueeze(dim=1) - vertices += transl.unsqueeze(dim=1) - - output = SMPLHOutput(vertices=vertices if return_verts else None, - joints=joints, - betas=betas, - global_orient=global_orient, - body_pose=body_pose, - left_hand_pose=left_hand_pose, - right_hand_pose=right_hand_pose, - full_pose=full_pose if return_full_pose else None) - - return output - - -class SMPLHLayer(SMPLH): - - def __init__( - self, *args, **kwargs - ) -> None: - ''' SMPL+H as a layer model constructor - ''' - super(SMPLHLayer, self).__init__( - create_global_orient=False, - create_body_pose=False, - create_left_hand_pose=False, - create_right_hand_pose=False, - create_betas=False, - create_transl=False, - *args, - **kwargs) - - def forward( - self, - betas: Optional[Tensor] = None, - global_orient: Optional[Tensor] = None, - body_pose: Optional[Tensor] = None, - left_hand_pose: Optional[Tensor] = None, - right_hand_pose: Optional[Tensor] = None, - transl: Optional[Tensor] = None, - return_verts: bool = True, - return_full_pose: bool = False, - pose2rot: bool = True, - **kwargs - ) -> SMPLHOutput: - ''' Forward pass for the SMPL+H model - Parameters - ---------- - global_orient: torch.tensor, optional, shape Bx3x3 - Global rotation of the body. Useful if someone wishes to - predicts this with an external model. It is expected to be in - rotation matrix format. (default=None) - betas: torch.tensor, optional, shape BxN_b - Shape parameters. For example, it can used if shape parameters - `betas` are predicted from some external model. - (default=None) - body_pose: torch.tensor, optional, shape BxJx3x3 - If given, ignore the member variable `body_pose` and use it - instead. For example, it can used if someone predicts the - pose of the body joints are predicted from some external model. - It should be a tensor that contains joint rotations in - rotation matrix format. (default=None) - left_hand_pose: torch.tensor, optional, shape Bx15x3x3 - If given, contains the pose of the left hand. - It should be a tensor that contains joint rotations in - rotation matrix format. (default=None) - right_hand_pose: torch.tensor, optional, shape Bx15x3x3 - If given, contains the pose of the right hand. - It should be a tensor that contains joint rotations in - rotation matrix format. (default=None) - transl: torch.tensor, optional, shape Bx3 - Translation vector of the body. - For example, it can used if the translation - `transl` is predicted from some external model. - (default=None) - return_verts: bool, optional - Return the vertices. (default=True) - return_full_pose: bool, optional - Returns the full axis-angle pose vector (default=False) - Returns - ------- - ''' - model_vars = [betas, global_orient, body_pose, transl, left_hand_pose, - right_hand_pose] - batch_size = 1 - for var in model_vars: - if var is None: - continue - batch_size = max(batch_size, len(var)) - device, dtype = self.shapedirs.device, self.shapedirs.dtype - if global_orient is None: - global_orient = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, -1, -1, -1).contiguous() - if body_pose is None: - body_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, 21, -1, -1).contiguous() - if left_hand_pose is None: - left_hand_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, 15, -1, -1).contiguous() - if right_hand_pose is None: - right_hand_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, 15, -1, -1).contiguous() - if betas is None: - betas = torch.zeros([batch_size, self.num_betas], - dtype=dtype, device=device) - if transl is None: - transl = torch.zeros([batch_size, 3], dtype=dtype, device=device) - - # Concatenate all pose vectors - full_pose = torch.cat( - [global_orient.reshape(-1, 1, 3, 3), - body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3, 3), - left_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3), - right_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3)], - dim=1) - - vertices, joints = lbs(betas, full_pose, self.v_template, - self.shapedirs, self.posedirs, - self.J_regressor, self.parents, - self.lbs_weights, pose2rot=False) - - # Add any extra joints that might be needed - joints = self.vertex_joint_selector(vertices, joints) - if self.joint_mapper is not None: - joints = self.joint_mapper(joints) - - if transl is not None: - joints += transl.unsqueeze(dim=1) - vertices += transl.unsqueeze(dim=1) - - output = SMPLHOutput(vertices=vertices if return_verts else None, - joints=joints, - betas=betas, - global_orient=global_orient, - body_pose=body_pose, - left_hand_pose=left_hand_pose, - right_hand_pose=right_hand_pose, - full_pose=full_pose if return_full_pose else None) - - return output - - -class SMPLX(SMPLH): - ''' - SMPL-X (SMPL eXpressive) is a unified body model, with shape parameters - trained jointly for the face, hands and body. - SMPL-X uses standard vertex based linear blend skinning with learned - corrective blend shapes, has N=10475 vertices and K=54 joints, - which includes joints for the neck, jaw, eyeballs and fingers. - ''' - - NUM_BODY_JOINTS = SMPLH.NUM_BODY_JOINTS # 21 - NUM_HAND_JOINTS = 15 - NUM_FACE_JOINTS = 3 - NUM_JOINTS = NUM_BODY_JOINTS + 2 * NUM_HAND_JOINTS + NUM_FACE_JOINTS - EXPRESSION_SPACE_DIM = 100 - NECK_IDX = 12 - - def __init__( - self, model_path: str, - kid_template_path: str = '', - num_expression_coeffs: int = 10, - create_expression: bool = True, - expression: Optional[Tensor] = None, - create_jaw_pose: bool = True, - jaw_pose: Optional[Tensor] = None, - create_leye_pose: bool = True, - leye_pose: Optional[Tensor] = None, - create_reye_pose=True, - reye_pose: Optional[Tensor] = None, - use_face_contour: bool = False, - batch_size: int = 1, - gender: str = 'neutral', - age: str = 'adult', - dtype=torch.float32, - ext: str = 'npz', - **kwargs - ) -> None: - ''' SMPLX model constructor - Parameters - ---------- - model_path: str - The path to the folder or to the file where the model - parameters are stored - num_expression_coeffs: int, optional - Number of expression components to use - (default = 10). - create_expression: bool, optional - Flag for creating a member variable for the expression space - (default = True). - expression: torch.tensor, optional, Bx10 - The default value for the expression member variable. - (default = None) - create_jaw_pose: bool, optional - Flag for creating a member variable for the jaw pose. - (default = False) - jaw_pose: torch.tensor, optional, Bx3 - The default value for the jaw pose variable. - (default = None) - create_leye_pose: bool, optional - Flag for creating a member variable for the left eye pose. - (default = False) - leye_pose: torch.tensor, optional, Bx10 - The default value for the left eye pose variable. - (default = None) - create_reye_pose: bool, optional - Flag for creating a member variable for the right eye pose. - (default = False) - reye_pose: torch.tensor, optional, Bx10 - The default value for the right eye pose variable. - (default = None) - use_face_contour: bool, optional - Whether to compute the keypoints that form the facial contour - batch_size: int, optional - The batch size used for creating the member variables - gender: str, optional - Which gender to load - dtype: torch.dtype - The data type for the created variables - ''' - - # Load the model - if osp.isdir(model_path): - model_fn = 'SMPLX_{}.{ext}'.format(gender.upper(), ext=ext) - smplx_path = os.path.join(model_path, model_fn) - else: - smplx_path = model_path - assert osp.exists(smplx_path), 'Path {} does not exist!'.format( - smplx_path) - - if ext == 'pkl': - with open(smplx_path, 'rb') as smplx_file: - model_data = pickle.load(smplx_file, encoding='latin1') - elif ext == 'npz': - model_data = np.load(smplx_path, allow_pickle=True) - else: - raise ValueError('Unknown extension: {}'.format(ext)) - - # print(colored(f"Use SMPL-X: {smplx_path}", "green")) - - data_struct = Struct(**model_data) - - super(SMPLX, self).__init__( - model_path=model_path, - kid_template_path=kid_template_path, - data_struct=data_struct, - dtype=dtype, - batch_size=batch_size, - vertex_ids=VERTEX_IDS['smplx'], - gender=gender, age=age, ext=ext, - **kwargs) - - lmk_faces_idx = data_struct.lmk_faces_idx - self.register_buffer('lmk_faces_idx', - torch.tensor(lmk_faces_idx, dtype=torch.long)) - lmk_bary_coords = data_struct.lmk_bary_coords - self.register_buffer('lmk_bary_coords', - torch.tensor(lmk_bary_coords, dtype=dtype)) - - self.use_face_contour = use_face_contour - if self.use_face_contour: - dynamic_lmk_faces_idx = data_struct.dynamic_lmk_faces_idx - dynamic_lmk_faces_idx = torch.tensor( - dynamic_lmk_faces_idx, - dtype=torch.long) - self.register_buffer('dynamic_lmk_faces_idx', - dynamic_lmk_faces_idx) - - dynamic_lmk_bary_coords = data_struct.dynamic_lmk_bary_coords - dynamic_lmk_bary_coords = torch.tensor( - dynamic_lmk_bary_coords, dtype=dtype) - self.register_buffer('dynamic_lmk_bary_coords', - dynamic_lmk_bary_coords) - - neck_kin_chain = find_joint_kin_chain(self.NECK_IDX, self.parents) - self.register_buffer( - 'neck_kin_chain', - torch.tensor(neck_kin_chain, dtype=torch.long)) - - if create_jaw_pose: - if jaw_pose is None: - default_jaw_pose = torch.zeros([batch_size, 3], dtype=dtype) - else: - default_jaw_pose = torch.tensor(jaw_pose, dtype=dtype) - jaw_pose_param = nn.Parameter(default_jaw_pose, - requires_grad=True) - self.register_parameter('jaw_pose', jaw_pose_param) - - if create_leye_pose: - if leye_pose is None: - default_leye_pose = torch.zeros([batch_size, 3], dtype=dtype) - else: - default_leye_pose = torch.tensor(leye_pose, dtype=dtype) - leye_pose_param = nn.Parameter(default_leye_pose, - requires_grad=True) - self.register_parameter('leye_pose', leye_pose_param) - - if create_reye_pose: - if reye_pose is None: - default_reye_pose = torch.zeros([batch_size, 3], dtype=dtype) - else: - default_reye_pose = torch.tensor(reye_pose, dtype=dtype) - reye_pose_param = nn.Parameter(default_reye_pose, - requires_grad=True) - self.register_parameter('reye_pose', reye_pose_param) - - shapedirs = data_struct.shapedirs - if len(shapedirs.shape) < 3: - shapedirs = shapedirs[:, :, None] - if (shapedirs.shape[-1] < self.SHAPE_SPACE_DIM + - self.EXPRESSION_SPACE_DIM): - # print(f'WARNING: You are using a {self.name()} model, with only' - # ' 10 shape and 10 expression coefficients.') - expr_start_idx = 10 - expr_end_idx = 20 - num_expression_coeffs = min(num_expression_coeffs, 10) - else: - expr_start_idx = self.SHAPE_SPACE_DIM - expr_end_idx = self.SHAPE_SPACE_DIM + num_expression_coeffs - num_expression_coeffs = min( - num_expression_coeffs, self.EXPRESSION_SPACE_DIM) - - self._num_expression_coeffs = num_expression_coeffs - - expr_dirs = shapedirs[:, :, expr_start_idx:expr_end_idx] - self.register_buffer( - 'expr_dirs', to_tensor(to_np(expr_dirs), dtype=dtype)) - - if create_expression: - if expression is None: - default_expression = torch.zeros( - [batch_size, self.num_expression_coeffs], dtype=dtype) - else: - default_expression = torch.tensor(expression, dtype=dtype) - expression_param = nn.Parameter(default_expression, - requires_grad=True) - self.register_parameter('expression', expression_param) - - def name(self) -> str: - return 'SMPL-X' - - @property - def num_expression_coeffs(self): - return self._num_expression_coeffs - - def create_mean_pose(self, data_struct, flat_hand_mean=False): - # Create the array for the mean pose. If flat_hand is false, then use - # the mean that is given by the data, rather than the flat open hand - global_orient_mean = torch.zeros([3], dtype=self.dtype) - body_pose_mean = torch.zeros([self.NUM_BODY_JOINTS * 3], - dtype=self.dtype) - jaw_pose_mean = torch.zeros([3], dtype=self.dtype) - leye_pose_mean = torch.zeros([3], dtype=self.dtype) - reye_pose_mean = torch.zeros([3], dtype=self.dtype) - - pose_mean = np.concatenate([global_orient_mean, body_pose_mean, - jaw_pose_mean, - leye_pose_mean, reye_pose_mean, - self.left_hand_mean, self.right_hand_mean], - axis=0) - - return pose_mean - - def extra_repr(self): - msg = super(SMPLX, self).extra_repr() - msg = [ - msg, - f'Number of Expression Coefficients: {self.num_expression_coeffs}' - ] - return '\n'.join(msg) - - def forward( - self, - betas: Optional[Tensor] = None, - global_orient: Optional[Tensor] = None, - body_pose: Optional[Tensor] = None, - left_hand_pose: Optional[Tensor] = None, - right_hand_pose: Optional[Tensor] = None, - transl: Optional[Tensor] = None, - expression: Optional[Tensor] = None, - jaw_pose: Optional[Tensor] = None, - leye_pose: Optional[Tensor] = None, - reye_pose: Optional[Tensor] = None, - return_verts: bool = True, - return_full_pose: bool = False, - pose2rot: bool = True, - return_joint_transformation: bool = False, - return_vertex_transformation: bool = False, - **kwargs - ) -> SMPLXOutput: - ''' - Forward pass for the SMPLX model - Parameters - ---------- - global_orient: torch.tensor, optional, shape Bx3 - If given, ignore the member variable and use it as the global - rotation of the body. Useful if someone wishes to predicts this - with an external model. (default=None) - betas: torch.tensor, optional, shape BxN_b - If given, ignore the member variable `betas` and use it - instead. For example, it can used if shape parameters - `betas` are predicted from some external model. - (default=None) - expression: torch.tensor, optional, shape BxN_e - If given, ignore the member variable `expression` and use it - instead. For example, it can used if expression parameters - `expression` are predicted from some external model. - body_pose: torch.tensor, optional, shape Bx(J*3) - If given, ignore the member variable `body_pose` and use it - instead. For example, it can used if someone predicts the - pose of the body joints are predicted from some external model. - It should be a tensor that contains joint rotations in - axis-angle format. (default=None) - left_hand_pose: torch.tensor, optional, shape BxP - If given, ignore the member variable `left_hand_pose` and - use this instead. It should either contain PCA coefficients or - joint rotations in axis-angle format. - right_hand_pose: torch.tensor, optional, shape BxP - If given, ignore the member variable `right_hand_pose` and - use this instead. It should either contain PCA coefficients or - joint rotations in axis-angle format. - jaw_pose: torch.tensor, optional, shape Bx3 - If given, ignore the member variable `jaw_pose` and - use this instead. It should either joint rotations in - axis-angle format. - transl: torch.tensor, optional, shape Bx3 - If given, ignore the member variable `transl` and use it - instead. For example, it can used if the translation - `transl` is predicted from some external model. - (default=None) - return_verts: bool, optional - Return the vertices. (default=True) - return_full_pose: bool, optional - Returns the full axis-angle pose vector (default=False) - Returns - ------- - output: ModelOutput - A named tuple of type `ModelOutput` - ''' - - # If no shape and pose parameters are passed along, then use the - # ones from the module - global_orient = (global_orient if global_orient is not None else - self.global_orient) - body_pose = body_pose if body_pose is not None else self.body_pose - betas = betas if betas is not None else self.betas - - left_hand_pose = (left_hand_pose if left_hand_pose is not None else - self.left_hand_pose) - right_hand_pose = (right_hand_pose if right_hand_pose is not None else - self.right_hand_pose) - jaw_pose = jaw_pose if jaw_pose is not None else self.jaw_pose - leye_pose = leye_pose if leye_pose is not None else self.leye_pose - reye_pose = reye_pose if reye_pose is not None else self.reye_pose - expression = expression if expression is not None else self.expression - - apply_trans = transl is not None or hasattr(self, 'transl') - if transl is None: - if hasattr(self, 'transl'): - transl = self.transl - - if self.use_pca: - left_hand_pose = torch.einsum('bi,ij->bj', [left_hand_pose, self.left_hand_components]) - right_hand_pose = torch.einsum( - 'bi,ij->bj', [right_hand_pose, self.right_hand_components]) - - full_pose = torch.cat([global_orient, body_pose, - jaw_pose, leye_pose, reye_pose, - left_hand_pose, - right_hand_pose], dim=1) - - # Add the mean pose of the model. Does not affect the body, only the - # hands when flat_hand_mean == False - full_pose += self.pose_mean - - batch_size = max(betas.shape[0], global_orient.shape[0], - body_pose.shape[0]) - # Concatenate the shape and expression coefficients - scale = int(batch_size / betas.shape[0]) - if scale > 1: - betas = betas.expand(scale, -1) - shape_components = torch.cat([betas, expression], dim=-1) - - shapedirs = torch.cat([self.shapedirs, self.expr_dirs], dim=-1) - - if return_joint_transformation or return_vertex_transformation: - vertices, joints, joint_transformation, vertex_transformation = lbs(shape_components, full_pose, self.v_template, - shapedirs, self.posedirs, - self.J_regressor, self.parents, - self.lbs_weights, pose2rot=pose2rot, return_transformation=True - ) - else: - vertices, joints = lbs(shape_components, full_pose, self.v_template, - shapedirs, self.posedirs, - self.J_regressor, self.parents, - self.lbs_weights, pose2rot=pose2rot, - ) - - lmk_faces_idx = self.lmk_faces_idx.unsqueeze( - dim=0).expand(batch_size, -1).contiguous() - lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat( - self.batch_size, 1, 1) - if self.use_face_contour: - lmk_idx_and_bcoords = find_dynamic_lmk_idx_and_bcoords( - vertices, full_pose, self.dynamic_lmk_faces_idx, - self.dynamic_lmk_bary_coords, - self.neck_kin_chain, - pose2rot=True, - ) - dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords - - lmk_faces_idx = torch.cat([lmk_faces_idx, - dyn_lmk_faces_idx], 1) - lmk_bary_coords = torch.cat( - [lmk_bary_coords.expand(batch_size, -1, -1), - dyn_lmk_bary_coords], 1) - - landmarks = vertices2landmarks(vertices, self.faces_tensor, - lmk_faces_idx, - lmk_bary_coords) - - # Add any extra joints that might be needed - joints = self.vertex_joint_selector(vertices, joints) - # Add the landmarks to the joints - joints = torch.cat([joints, landmarks], dim=1) - # Map the joints to the current dataset - - if self.joint_mapper is not None: - joints = self.joint_mapper(joints=joints, vertices=vertices) - - if apply_trans: - joints += transl.unsqueeze(dim=1) - vertices += transl.unsqueeze(dim=1) - - output = SMPLXOutput(vertices=vertices if return_verts else None, - joints=joints, - betas=betas, - expression=expression, - global_orient=global_orient, - body_pose=body_pose, - left_hand_pose=left_hand_pose, - right_hand_pose=right_hand_pose, - jaw_pose=jaw_pose, - full_pose=full_pose if return_full_pose else None, - joint_transformation=joint_transformation if return_joint_transformation else None, - vertex_transformation=vertex_transformation if return_vertex_transformation else None) - return output - - -class SMPLXLayer(SMPLX): - def __init__( - self, - *args, - **kwargs - ) -> None: - # Just create a SMPLX module without any member variables - super(SMPLXLayer, self).__init__( - create_global_orient=False, - create_body_pose=False, - create_left_hand_pose=False, - create_right_hand_pose=False, - create_jaw_pose=False, - create_leye_pose=False, - create_reye_pose=False, - create_betas=False, - create_expression=False, - create_transl=False, - *args, **kwargs, - ) - - def forward( - self, - betas: Optional[Tensor] = None, - global_orient: Optional[Tensor] = None, - body_pose: Optional[Tensor] = None, - left_hand_pose: Optional[Tensor] = None, - right_hand_pose: Optional[Tensor] = None, - transl: Optional[Tensor] = None, - expression: Optional[Tensor] = None, - jaw_pose: Optional[Tensor] = None, - leye_pose: Optional[Tensor] = None, - reye_pose: Optional[Tensor] = None, - return_verts: bool = True, - return_full_pose: bool = False, - **kwargs - ) -> SMPLXOutput: - ''' - Forward pass for the SMPLX model - Parameters - ---------- - global_orient: torch.tensor, optional, shape Bx3x3 - If given, ignore the member variable and use it as the global - rotation of the body. Useful if someone wishes to predicts this - with an external model. It is expected to be in rotation matrix - format. (default=None) - betas: torch.tensor, optional, shape BxN_b - If given, ignore the member variable `betas` and use it - instead. For example, it can used if shape parameters - `betas` are predicted from some external model. - (default=None) - expression: torch.tensor, optional, shape BxN_e - Expression coefficients. - For example, it can used if expression parameters - `expression` are predicted from some external model. - body_pose: torch.tensor, optional, shape BxJx3x3 - If given, ignore the member variable `body_pose` and use it - instead. For example, it can used if someone predicts the - pose of the body joints are predicted from some external model. - It should be a tensor that contains joint rotations in - rotation matrix format. (default=None) - left_hand_pose: torch.tensor, optional, shape Bx15x3x3 - If given, contains the pose of the left hand. - It should be a tensor that contains joint rotations in - rotation matrix format. (default=None) - right_hand_pose: torch.tensor, optional, shape Bx15x3x3 - If given, contains the pose of the right hand. - It should be a tensor that contains joint rotations in - rotation matrix format. (default=None) - jaw_pose: torch.tensor, optional, shape Bx3x3 - Jaw pose. It should either joint rotations in - rotation matrix format. - transl: torch.tensor, optional, shape Bx3 - Translation vector of the body. - For example, it can used if the translation - `transl` is predicted from some external model. - (default=None) - return_verts: bool, optional - Return the vertices. (default=True) - return_full_pose: bool, optional - Returns the full pose vector (default=False) - Returns - ------- - output: ModelOutput - A data class that contains the posed vertices and joints - ''' - device, dtype = self.shapedirs.device, self.shapedirs.dtype - - model_vars = [betas, global_orient, body_pose, transl, - expression, left_hand_pose, right_hand_pose, jaw_pose] - batch_size = 1 - for var in model_vars: - if var is None: - continue - batch_size = max(batch_size, len(var)) - - if global_orient is None: - global_orient = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, -1, -1, -1).contiguous() - if body_pose is None: - body_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand( - batch_size, self.NUM_BODY_JOINTS, -1, -1).contiguous() - if left_hand_pose is None: - left_hand_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, 15, -1, -1).contiguous() - if right_hand_pose is None: - right_hand_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, 15, -1, -1).contiguous() - if jaw_pose is None: - jaw_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, -1, -1, -1).contiguous() - if leye_pose is None: - leye_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, -1, -1, -1).contiguous() - if reye_pose is None: - reye_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, -1, -1, -1).contiguous() - if expression is None: - expression = torch.zeros([batch_size, self.num_expression_coeffs], - dtype=dtype, device=device) - if betas is None: - betas = torch.zeros([batch_size, self.num_betas], - dtype=dtype, device=device) - if transl is None: - transl = torch.zeros([batch_size, 3], dtype=dtype, device=device) - - # Concatenate all pose vectors - full_pose = torch.cat( - [global_orient.reshape(-1, 1, 3, 3), - body_pose.reshape(-1, self.NUM_BODY_JOINTS, 3, 3), - jaw_pose.reshape(-1, 1, 3, 3), - leye_pose.reshape(-1, 1, 3, 3), - reye_pose.reshape(-1, 1, 3, 3), - left_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3), - right_hand_pose.reshape(-1, self.NUM_HAND_JOINTS, 3, 3)], - dim=1) - shape_components = torch.cat([betas, expression], dim=-1) - - shapedirs = torch.cat([self.shapedirs, self.expr_dirs], dim=-1) - - vertices, joints = lbs(shape_components, full_pose, self.v_template, - shapedirs, self.posedirs, - self.J_regressor, self.parents, - self.lbs_weights, - pose2rot=False, - ) - - lmk_faces_idx = self.lmk_faces_idx.unsqueeze( - dim=0).expand(batch_size, -1).contiguous() - lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat( - batch_size, 1, 1) - if self.use_face_contour: - lmk_idx_and_bcoords = find_dynamic_lmk_idx_and_bcoords( - vertices, full_pose, - self.dynamic_lmk_faces_idx, - self.dynamic_lmk_bary_coords, - self.neck_kin_chain, - pose2rot=False, - ) - dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords - - lmk_faces_idx = torch.cat([lmk_faces_idx, dyn_lmk_faces_idx], 1) - lmk_bary_coords = torch.cat( - [lmk_bary_coords.expand(batch_size, -1, -1), - dyn_lmk_bary_coords], 1) - - landmarks = vertices2landmarks(vertices, self.faces_tensor, - lmk_faces_idx, - lmk_bary_coords) - - # Add any extra joints that might be needed - joints = self.vertex_joint_selector(vertices, joints) - # Add the landmarks to the joints - joints = torch.cat([joints, landmarks], dim=1) - # Map the joints to the current dataset - - if self.joint_mapper is not None: - joints = self.joint_mapper(joints=joints, vertices=vertices) - - if transl is not None: - joints += transl.unsqueeze(dim=1) - vertices += transl.unsqueeze(dim=1) - - output = SMPLXOutput(vertices=vertices if return_verts else None, - joints=joints, - betas=betas, - expression=expression, - global_orient=global_orient, - body_pose=body_pose, - left_hand_pose=left_hand_pose, - right_hand_pose=right_hand_pose, - jaw_pose=jaw_pose, - transl=transl, - full_pose=full_pose if return_full_pose else None) - return output - - -class MANO(SMPL): - # The hand joints are replaced by MANO - NUM_BODY_JOINTS = 1 - NUM_HAND_JOINTS = 15 - NUM_JOINTS = NUM_BODY_JOINTS + NUM_HAND_JOINTS - - def __init__( - self, - model_path: str, - is_rhand: bool = True, - data_struct: Optional[Struct] = None, - create_hand_pose: bool = True, - hand_pose: Optional[Tensor] = None, - use_pca: bool = True, - num_pca_comps: int = 6, - flat_hand_mean: bool = False, - batch_size: int = 1, - dtype=torch.float32, - vertex_ids=None, - use_compressed: bool = True, - ext: str = 'pkl', - **kwargs - ) -> None: - ''' MANO model constructor - Parameters - ---------- - model_path: str - The path to the folder or to the file where the model - parameters are stored - data_struct: Strct - A struct object. If given, then the parameters of the model are - read from the object. Otherwise, the model tries to read the - parameters from the given `model_path`. (default = None) - create_hand_pose: bool, optional - Flag for creating a member variable for the pose of the right - hand. (default = True) - hand_pose: torch.tensor, optional, BxP - The default value for the right hand pose member variable. - (default = None) - num_pca_comps: int, optional - The number of PCA components to use for each hand. - (default = 6) - flat_hand_mean: bool, optional - If False, then the pose of the hand is initialized to False. - batch_size: int, optional - The batch size used for creating the member variables - dtype: torch.dtype, optional - The data type for the created variables - vertex_ids: dict, optional - A dictionary containing the indices of the extra vertices that - will be selected - ''' - - self.num_pca_comps = num_pca_comps - self.is_rhand = is_rhand - # If no data structure is passed, then load the data from the given - # model folder - if data_struct is None: - # Load the model - if osp.isdir(model_path): - model_fn = 'MANO_{}.{ext}'.format( - 'RIGHT' if is_rhand else 'LEFT', ext=ext) - mano_path = os.path.join(model_path, model_fn) - else: - mano_path = model_path - self.is_rhand = True if 'RIGHT' in os.path.basename( - model_path) else False - assert osp.exists(mano_path), 'Path {} does not exist!'.format( - mano_path) - - if ext == 'pkl': - with open(mano_path, 'rb') as mano_file: - model_data = pickle.load(mano_file, encoding='latin1') - elif ext == 'npz': - model_data = np.load(mano_path, allow_pickle=True) - else: - raise ValueError('Unknown extension: {}'.format(ext)) - data_struct = Struct(**model_data) - - if vertex_ids is None: - vertex_ids = VERTEX_IDS['smplh'] - - super(MANO, self).__init__( - model_path=model_path, data_struct=data_struct, - batch_size=batch_size, vertex_ids=vertex_ids, - use_compressed=use_compressed, dtype=dtype, ext=ext, **kwargs) - - # add only MANO tips to the extra joints - self.vertex_joint_selector.extra_joints_idxs = to_tensor( - list(VERTEX_IDS['mano'].values()), dtype=torch.long) - - self.use_pca = use_pca - self.num_pca_comps = num_pca_comps - if self.num_pca_comps == 45: - self.use_pca = False - self.flat_hand_mean = flat_hand_mean - - hand_components = data_struct.hands_components[:num_pca_comps] - - self.np_hand_components = hand_components - - if self.use_pca: - self.register_buffer( - 'hand_components', - torch.tensor(hand_components, dtype=dtype)) - - if self.flat_hand_mean: - hand_mean = np.zeros_like(data_struct.hands_mean) - else: - hand_mean = data_struct.hands_mean - - self.register_buffer('hand_mean', - to_tensor(hand_mean, dtype=self.dtype)) - - # Create the buffers for the pose of the left hand - hand_pose_dim = num_pca_comps if use_pca else 3 * self.NUM_HAND_JOINTS - if create_hand_pose: - if hand_pose is None: - default_hand_pose = torch.zeros([batch_size, hand_pose_dim], - dtype=dtype) - else: - default_hand_pose = torch.tensor(hand_pose, dtype=dtype) - - hand_pose_param = nn.Parameter(default_hand_pose, - requires_grad=True) - self.register_parameter('hand_pose', - hand_pose_param) - - # Create the buffer for the mean pose. - pose_mean = self.create_mean_pose( - data_struct, flat_hand_mean=flat_hand_mean) - pose_mean_tensor = pose_mean.clone().to(dtype) - # pose_mean_tensor = torch.tensor(pose_mean, dtype=dtype) - self.register_buffer('pose_mean', pose_mean_tensor) - - def name(self) -> str: - return 'MANO' - - def create_mean_pose(self, data_struct, flat_hand_mean=False): - # Create the array for the mean pose. If flat_hand is false, then use - # the mean that is given by the data, rather than the flat open hand - global_orient_mean = torch.zeros([3], dtype=self.dtype) - pose_mean = torch.cat([global_orient_mean, self.hand_mean], dim=0) - return pose_mean - - def extra_repr(self): - msg = [super(MANO, self).extra_repr()] - if self.use_pca: - msg.append(f'Number of PCA components: {self.num_pca_comps}') - msg.append(f'Flat hand mean: {self.flat_hand_mean}') - return '\n'.join(msg) - - def forward( - self, - betas: Optional[Tensor] = None, - global_orient: Optional[Tensor] = None, - hand_pose: Optional[Tensor] = None, - transl: Optional[Tensor] = None, - return_verts: bool = True, - return_full_pose: bool = False, - **kwargs - ) -> MANOOutput: - ''' Forward pass for the MANO model - ''' - # If no shape and pose parameters are passed along, then use the - # ones from the module - global_orient = (global_orient if global_orient is not None else - self.global_orient) - betas = betas if betas is not None else self.betas - hand_pose = (hand_pose if hand_pose is not None else - self.hand_pose) - - apply_trans = transl is not None or hasattr(self, 'transl') - if transl is None: - if hasattr(self, 'transl'): - transl = self.transl - - if self.use_pca: - hand_pose = torch.einsum( - 'bi,ij->bj', [hand_pose, self.hand_components]) - - full_pose = torch.cat([global_orient, hand_pose], dim=1) - full_pose += self.pose_mean - - vertices, joints = lbs(betas, full_pose, self.v_template, - self.shapedirs, self.posedirs, - self.J_regressor, self.parents, - self.lbs_weights, pose2rot=True, - ) - - # # Add pre-selected extra joints that might be needed - # joints = self.vertex_joint_selector(vertices, joints) - - if self.joint_mapper is not None: - joints = self.joint_mapper(joints) - - if apply_trans: - joints = joints + transl.unsqueeze(dim=1) - vertices = vertices + transl.unsqueeze(dim=1) - - output = MANOOutput(vertices=vertices if return_verts else None, - joints=joints if return_verts else None, - betas=betas, - global_orient=global_orient, - hand_pose=hand_pose, - full_pose=full_pose if return_full_pose else None) - - return output - - -class MANOLayer(MANO): - def __init__(self, *args, **kwargs) -> None: - ''' MANO as a layer model constructor - ''' - super(MANOLayer, self).__init__( - create_global_orient=False, - create_hand_pose=False, - create_betas=False, - create_transl=False, - *args, **kwargs) - - def name(self) -> str: - return 'MANO' - - def forward( - self, - betas: Optional[Tensor] = None, - global_orient: Optional[Tensor] = None, - hand_pose: Optional[Tensor] = None, - transl: Optional[Tensor] = None, - return_verts: bool = True, - return_full_pose: bool = False, - **kwargs - ) -> MANOOutput: - ''' Forward pass for the MANO model - ''' - device, dtype = self.shapedirs.device, self.shapedirs.dtype - if global_orient is None: - batch_size = 1 - global_orient = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, -1, -1, -1).contiguous() - else: - batch_size = global_orient.shape[0] - if hand_pose is None: - hand_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, 15, -1, -1).contiguous() - if betas is None: - betas = torch.zeros( - [batch_size, self.num_betas], dtype=dtype, device=device) - if transl is None: - transl = torch.zeros([batch_size, 3], dtype=dtype, device=device) - - full_pose = torch.cat([global_orient, hand_pose], dim=1) - vertices, joints = lbs(betas, full_pose, self.v_template, - self.shapedirs, self.posedirs, - self.J_regressor, self.parents, - self.lbs_weights, pose2rot=False) - - if self.joint_mapper is not None: - joints = self.joint_mapper(joints) - - if transl is not None: - joints = joints + transl.unsqueeze(dim=1) - vertices = vertices + transl.unsqueeze(dim=1) - - output = MANOOutput( - vertices=vertices if return_verts else None, - joints=joints if return_verts else None, - betas=betas, - global_orient=global_orient, - hand_pose=hand_pose, - full_pose=full_pose if return_full_pose else None) - - return output - - -class FLAME(SMPL): - NUM_JOINTS = 5 - SHAPE_SPACE_DIM = 300 - EXPRESSION_SPACE_DIM = 100 - NECK_IDX = 0 - - def __init__( - self, - model_path: str, - data_struct=None, - num_expression_coeffs=10, - create_expression: bool = True, - expression: Optional[Tensor] = None, - create_neck_pose: bool = True, - neck_pose: Optional[Tensor] = None, - create_jaw_pose: bool = True, - jaw_pose: Optional[Tensor] = None, - create_leye_pose: bool = True, - leye_pose: Optional[Tensor] = None, - create_reye_pose=True, - reye_pose: Optional[Tensor] = None, - use_face_contour=False, - batch_size: int = 1, - gender: str = 'neutral', - dtype: torch.dtype = torch.float32, - ext='pkl', - **kwargs - ) -> None: - ''' FLAME model constructor - Parameters - ---------- - model_path: str - The path to the folder or to the file where the model - parameters are stored - num_expression_coeffs: int, optional - Number of expression components to use - (default = 10). - create_expression: bool, optional - Flag for creating a member variable for the expression space - (default = True). - expression: torch.tensor, optional, Bx10 - The default value for the expression member variable. - (default = None) - create_neck_pose: bool, optional - Flag for creating a member variable for the neck pose. - (default = False) - neck_pose: torch.tensor, optional, Bx3 - The default value for the neck pose variable. - (default = None) - create_jaw_pose: bool, optional - Flag for creating a member variable for the jaw pose. - (default = False) - jaw_pose: torch.tensor, optional, Bx3 - The default value for the jaw pose variable. - (default = None) - create_leye_pose: bool, optional - Flag for creating a member variable for the left eye pose. - (default = False) - leye_pose: torch.tensor, optional, Bx10 - The default value for the left eye pose variable. - (default = None) - create_reye_pose: bool, optional - Flag for creating a member variable for the right eye pose. - (default = False) - reye_pose: torch.tensor, optional, Bx10 - The default value for the right eye pose variable. - (default = None) - use_face_contour: bool, optional - Whether to compute the keypoints that form the facial contour - batch_size: int, optional - The batch size used for creating the member variables - gender: str, optional - Which gender to load - dtype: torch.dtype - The data type for the created variables - ''' - model_fn = f'FLAME_{gender.upper()}.{ext}' - flame_path = os.path.join(model_path, model_fn) - assert osp.exists(flame_path), 'Path {} does not exist!'.format( - flame_path) - if ext == 'npz': - file_data = np.load(flame_path, allow_pickle=True) - elif ext == 'pkl': - with open(flame_path, 'rb') as smpl_file: - file_data = pickle.load(smpl_file, encoding='latin1') - else: - raise ValueError('Unknown extension: {}'.format(ext)) - data_struct = Struct(**file_data) - - super(FLAME, self).__init__( - model_path=model_path, - data_struct=data_struct, - dtype=dtype, - batch_size=batch_size, - gender=gender, - ext=ext, - **kwargs) - - self.use_face_contour = use_face_contour - - self.vertex_joint_selector.extra_joints_idxs = to_tensor( - [], dtype=torch.long) - - if create_neck_pose: - if neck_pose is None: - default_neck_pose = torch.zeros([batch_size, 3], dtype=dtype) - else: - default_neck_pose = torch.tensor(neck_pose, dtype=dtype) - neck_pose_param = nn.Parameter( - default_neck_pose, requires_grad=True) - self.register_parameter('neck_pose', neck_pose_param) - - if create_jaw_pose: - if jaw_pose is None: - default_jaw_pose = torch.zeros([batch_size, 3], dtype=dtype) - else: - default_jaw_pose = torch.tensor(jaw_pose, dtype=dtype) - jaw_pose_param = nn.Parameter(default_jaw_pose, - requires_grad=True) - self.register_parameter('jaw_pose', jaw_pose_param) - - if create_leye_pose: - if leye_pose is None: - default_leye_pose = torch.zeros([batch_size, 3], dtype=dtype) - else: - default_leye_pose = torch.tensor(leye_pose, dtype=dtype) - leye_pose_param = nn.Parameter(default_leye_pose, - requires_grad=True) - self.register_parameter('leye_pose', leye_pose_param) - - if create_reye_pose: - if reye_pose is None: - default_reye_pose = torch.zeros([batch_size, 3], dtype=dtype) - else: - default_reye_pose = torch.tensor(reye_pose, dtype=dtype) - reye_pose_param = nn.Parameter(default_reye_pose, - requires_grad=True) - self.register_parameter('reye_pose', reye_pose_param) - - shapedirs = data_struct.shapedirs - if len(shapedirs.shape) < 3: - shapedirs = shapedirs[:, :, None] - if (shapedirs.shape[-1] < self.SHAPE_SPACE_DIM + - self.EXPRESSION_SPACE_DIM): - # print(f'WARNING: You are using a {self.name()} model, with only' - # ' 10 shape and 10 expression coefficients.') - expr_start_idx = 10 - expr_end_idx = 20 - num_expression_coeffs = min(num_expression_coeffs, 10) - else: - expr_start_idx = self.SHAPE_SPACE_DIM - expr_end_idx = self.SHAPE_SPACE_DIM + num_expression_coeffs - num_expression_coeffs = min( - num_expression_coeffs, self.EXPRESSION_SPACE_DIM) - - self._num_expression_coeffs = num_expression_coeffs - - expr_dirs = shapedirs[:, :, expr_start_idx:expr_end_idx] - self.register_buffer( - 'expr_dirs', to_tensor(to_np(expr_dirs), dtype=dtype)) - - if create_expression: - if expression is None: - default_expression = torch.zeros( - [batch_size, self.num_expression_coeffs], dtype=dtype) - else: - default_expression = torch.tensor(expression, dtype=dtype) - expression_param = nn.Parameter(default_expression, - requires_grad=True) - self.register_parameter('expression', expression_param) - - # The pickle file that contains the barycentric coordinates for - # regressing the landmarks - landmark_bcoord_filename = osp.join( - model_path, 'flame_static_embedding.pkl') - - with open(landmark_bcoord_filename, 'rb') as fp: - landmarks_data = pickle.load(fp, encoding='latin1') - - lmk_faces_idx = landmarks_data['lmk_face_idx'].astype(np.int64) - self.register_buffer('lmk_faces_idx', - torch.tensor(lmk_faces_idx, dtype=torch.long)) - lmk_bary_coords = landmarks_data['lmk_b_coords'] - self.register_buffer('lmk_bary_coords', - torch.tensor(lmk_bary_coords, dtype=dtype)) - if self.use_face_contour: - face_contour_path = os.path.join( - model_path, 'flame_dynamic_embedding.npy') - contour_embeddings = np.load(face_contour_path, - allow_pickle=True, - encoding='latin1')[()] - - dynamic_lmk_faces_idx = np.array( - contour_embeddings['lmk_face_idx'], dtype=np.int64) - dynamic_lmk_faces_idx = torch.tensor( - dynamic_lmk_faces_idx, - dtype=torch.long) - self.register_buffer('dynamic_lmk_faces_idx', - dynamic_lmk_faces_idx) - - dynamic_lmk_b_coords = torch.tensor( - contour_embeddings['lmk_b_coords'], dtype=dtype) - self.register_buffer( - 'dynamic_lmk_bary_coords', dynamic_lmk_b_coords) - - neck_kin_chain = find_joint_kin_chain(self.NECK_IDX, self.parents) - self.register_buffer( - 'neck_kin_chain', - torch.tensor(neck_kin_chain, dtype=torch.long)) - - @property - def num_expression_coeffs(self): - return self._num_expression_coeffs - - def name(self) -> str: - return 'FLAME' - - def extra_repr(self): - msg = [ - super(FLAME, self).extra_repr(), - f'Number of Expression Coefficients: {self.num_expression_coeffs}', - f'Use face contour: {self.use_face_contour}', - ] - return '\n'.join(msg) - - def forward( - self, - betas: Optional[Tensor] = None, - global_orient: Optional[Tensor] = None, - neck_pose: Optional[Tensor] = None, - transl: Optional[Tensor] = None, - expression: Optional[Tensor] = None, - jaw_pose: Optional[Tensor] = None, - leye_pose: Optional[Tensor] = None, - reye_pose: Optional[Tensor] = None, - return_verts: bool = True, - return_full_pose: bool = False, - pose2rot: bool = True, - **kwargs - ) -> FLAMEOutput: - ''' - Forward pass for the SMPLX model - Parameters - ---------- - global_orient: torch.tensor, optional, shape Bx3 - If given, ignore the member variable and use it as the global - rotation of the body. Useful if someone wishes to predicts this - with an external model. (default=None) - betas: torch.tensor, optional, shape Bx10 - If given, ignore the member variable `betas` and use it - instead. For example, it can used if shape parameters - `betas` are predicted from some external model. - (default=None) - expression: torch.tensor, optional, shape Bx10 - If given, ignore the member variable `expression` and use it - instead. For example, it can used if expression parameters - `expression` are predicted from some external model. - jaw_pose: torch.tensor, optional, shape Bx3 - If given, ignore the member variable `jaw_pose` and - use this instead. It should either joint rotations in - axis-angle format. - jaw_pose: torch.tensor, optional, shape Bx3 - If given, ignore the member variable `jaw_pose` and - use this instead. It should either joint rotations in - axis-angle format. - transl: torch.tensor, optional, shape Bx3 - If given, ignore the member variable `transl` and use it - instead. For example, it can used if the translation - `transl` is predicted from some external model. - (default=None) - return_verts: bool, optional - Return the vertices. (default=True) - return_full_pose: bool, optional - Returns the full axis-angle pose vector (default=False) - Returns - ------- - output: ModelOutput - A named tuple of type `ModelOutput` - ''' - - # If no shape and pose parameters are passed along, then use the - # ones from the module - global_orient = (global_orient if global_orient is not None else - self.global_orient) - jaw_pose = jaw_pose if jaw_pose is not None else self.jaw_pose - neck_pose = neck_pose if neck_pose is not None else self.neck_pose - - leye_pose = leye_pose if leye_pose is not None else self.leye_pose - reye_pose = reye_pose if reye_pose is not None else self.reye_pose - - betas = betas if betas is not None else self.betas - expression = expression if expression is not None else self.expression - - apply_trans = transl is not None or hasattr(self, 'transl') - if transl is None: - if hasattr(self, 'transl'): - transl = self.transl - - full_pose = torch.cat( - [global_orient, neck_pose, jaw_pose, leye_pose, reye_pose], dim=1) - - batch_size = max(betas.shape[0], global_orient.shape[0], - jaw_pose.shape[0]) - # Concatenate the shape and expression coefficients - scale = int(batch_size / betas.shape[0]) - if scale > 1: - betas = betas.expand(scale, -1) - shape_components = torch.cat([betas, expression], dim=-1) - shapedirs = torch.cat([self.shapedirs, self.expr_dirs], dim=-1) - - vertices, joints = lbs(shape_components, full_pose, self.v_template, - shapedirs, self.posedirs, - self.J_regressor, self.parents, - self.lbs_weights, pose2rot=pose2rot, - ) - - lmk_faces_idx = self.lmk_faces_idx.unsqueeze( - dim=0).expand(batch_size, -1).contiguous() - lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat( - self.batch_size, 1, 1) - if self.use_face_contour: - lmk_idx_and_bcoords = find_dynamic_lmk_idx_and_bcoords( - vertices, full_pose, self.dynamic_lmk_faces_idx, - self.dynamic_lmk_bary_coords, - self.neck_kin_chain, - pose2rot=True, - ) - dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords - lmk_faces_idx = torch.cat([lmk_faces_idx, - dyn_lmk_faces_idx], 1) - lmk_bary_coords = torch.cat( - [lmk_bary_coords.expand(batch_size, -1, -1), - dyn_lmk_bary_coords], 1) - - landmarks = vertices2landmarks(vertices, self.faces_tensor, - lmk_faces_idx, - lmk_bary_coords) - - # Add any extra joints that might be needed - joints = self.vertex_joint_selector(vertices, joints) - # Add the landmarks to the joints - joints = torch.cat([joints, landmarks], dim=1) - - # Map the joints to the current dataset - if self.joint_mapper is not None: - joints = self.joint_mapper(joints=joints, vertices=vertices) - - if apply_trans: - joints += transl.unsqueeze(dim=1) - vertices += transl.unsqueeze(dim=1) - - output = FLAMEOutput(vertices=vertices if return_verts else None, - joints=joints, - betas=betas, - expression=expression, - global_orient=global_orient, - neck_pose=neck_pose, - jaw_pose=jaw_pose, - full_pose=full_pose if return_full_pose else None) - return output - - -class FLAMELayer(FLAME): - def __init__(self, *args, **kwargs) -> None: - ''' FLAME as a layer model constructor ''' - super(FLAMELayer, self).__init__( - create_betas=False, - create_expression=False, - create_global_orient=False, - create_neck_pose=False, - create_jaw_pose=False, - create_leye_pose=False, - create_reye_pose=False, - *args, - **kwargs) - - def forward( - self, - betas: Optional[Tensor] = None, - global_orient: Optional[Tensor] = None, - neck_pose: Optional[Tensor] = None, - transl: Optional[Tensor] = None, - expression: Optional[Tensor] = None, - jaw_pose: Optional[Tensor] = None, - leye_pose: Optional[Tensor] = None, - reye_pose: Optional[Tensor] = None, - return_verts: bool = True, - return_full_pose: bool = False, - pose2rot: bool = True, - **kwargs - ) -> FLAMEOutput: - ''' - Forward pass for the SMPLX model - Parameters - ---------- - global_orient: torch.tensor, optional, shape Bx3x3 - Global rotation of the body. Useful if someone wishes to - predicts this with an external model. It is expected to be in - rotation matrix format. (default=None) - betas: torch.tensor, optional, shape BxN_b - Shape parameters. For example, it can used if shape parameters - `betas` are predicted from some external model. - (default=None) - expression: torch.tensor, optional, shape BxN_e - If given, ignore the member variable `expression` and use it - instead. For example, it can used if expression parameters - `expression` are predicted from some external model. - jaw_pose: torch.tensor, optional, shape Bx3x3 - Jaw pose. It should either joint rotations in - rotation matrix format. - transl: torch.tensor, optional, shape Bx3 - Translation vector of the body. - For example, it can used if the translation - `transl` is predicted from some external model. - (default=None) - return_verts: bool, optional - Return the vertices. (default=True) - return_full_pose: bool, optional - Returns the full axis-angle pose vector (default=False) - Returns - ------- - output: ModelOutput - A named tuple of type `ModelOutput` - ''' - device, dtype = self.shapedirs.device, self.shapedirs.dtype - if global_orient is None: - batch_size = 1 - global_orient = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, -1, -1, -1).contiguous() - else: - batch_size = global_orient.shape[0] - if neck_pose is None: - neck_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, 1, -1, -1).contiguous() - if jaw_pose is None: - jaw_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, -1, -1, -1).contiguous() - if leye_pose is None: - leye_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, -1, -1, -1).contiguous() - if reye_pose is None: - reye_pose = torch.eye(3, device=device, dtype=dtype).view( - 1, 1, 3, 3).expand(batch_size, -1, -1, -1).contiguous() - if betas is None: - betas = torch.zeros([batch_size, self.num_betas], - dtype=dtype, device=device) - if expression is None: - expression = torch.zeros([batch_size, self.num_expression_coeffs], - dtype=dtype, device=device) - if transl is None: - transl = torch.zeros([batch_size, 3], dtype=dtype, device=device) - - full_pose = torch.cat( - [global_orient, neck_pose, jaw_pose, leye_pose, reye_pose], dim=1) - - shape_components = torch.cat([betas, expression], dim=-1) - shapedirs = torch.cat([self.shapedirs, self.expr_dirs], dim=-1) - - vertices, joints = lbs(shape_components, full_pose, self.v_template, - shapedirs, self.posedirs, - self.J_regressor, self.parents, - self.lbs_weights, pose2rot=False, - ) - - lmk_faces_idx = self.lmk_faces_idx.unsqueeze( - dim=0).expand(batch_size, -1).contiguous() - lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).repeat( - self.batch_size, 1, 1) - if self.use_face_contour: - lmk_idx_and_bcoords = find_dynamic_lmk_idx_and_bcoords( - vertices, full_pose, self.dynamic_lmk_faces_idx, - self.dynamic_lmk_bary_coords, - self.neck_kin_chain, - pose2rot=False, - ) - dyn_lmk_faces_idx, dyn_lmk_bary_coords = lmk_idx_and_bcoords - lmk_faces_idx = torch.cat([lmk_faces_idx, - dyn_lmk_faces_idx], 1) - lmk_bary_coords = torch.cat( - [lmk_bary_coords.expand(batch_size, -1, -1), - dyn_lmk_bary_coords], 1) - - landmarks = vertices2landmarks(vertices, self.faces_tensor, - lmk_faces_idx, - lmk_bary_coords) - - # Add any extra joints that might be needed - joints = self.vertex_joint_selector(vertices, joints) - # Add the landmarks to the joints - joints = torch.cat([joints, landmarks], dim=1) - - # Map the joints to the current dataset - if self.joint_mapper is not None: - joints = self.joint_mapper(joints=joints, vertices=vertices) - - joints += transl.unsqueeze(dim=1) - vertices += transl.unsqueeze(dim=1) - - output = FLAMEOutput(vertices=vertices if return_verts else None, - joints=joints, - betas=betas, - expression=expression, - global_orient=global_orient, - neck_pose=neck_pose, - jaw_pose=jaw_pose, - full_pose=full_pose if return_full_pose else None) - return output - - -def build_layer( - model_path: str, - model_type: str = 'smpl', - **kwargs -) -> Union[SMPLLayer, SMPLHLayer, SMPLXLayer, MANOLayer, FLAMELayer]: - ''' Method for creating a model from a path and a model type - Parameters - ---------- - model_path: str - Either the path to the model you wish to load or a folder, - where each subfolder contains the differents types, i.e.: - model_path: - | - |-- smpl - |-- SMPL_FEMALE - |-- SMPL_NEUTRAL - |-- SMPL_MALE - |-- smplh - |-- SMPLH_FEMALE - |-- SMPLH_MALE - |-- smplx - |-- SMPLX_FEMALE - |-- SMPLX_NEUTRAL - |-- SMPLX_MALE - |-- mano - |-- MANO RIGHT - |-- MANO LEFT - |-- flame - |-- FLAME_FEMALE - |-- FLAME_MALE - |-- FLAME_NEUTRAL - model_type: str, optional - When model_path is a folder, then this parameter specifies the - type of model to be loaded - **kwargs: dict - Keyword arguments - Returns - ------- - body_model: nn.Module - The PyTorch module that implements the corresponding body model - Raises - ------ - ValueError: In case the model type is not one of SMPL, SMPLH, - SMPLX, MANO or FLAME - ''' - - if osp.isdir(model_path): - model_path = os.path.join(model_path, model_type) - else: - model_type = osp.basename(model_path).split('_')[0].lower() - - if model_type.lower() == 'smpl': - return SMPLLayer(model_path, **kwargs) - elif model_type.lower() == 'smplh': - return SMPLHLayer(model_path, **kwargs) - elif model_type.lower() == 'smplx': - return SMPLXLayer(model_path, **kwargs) - elif 'mano' in model_type.lower(): - return MANOLayer(model_path, **kwargs) - elif 'flame' in model_type.lower(): - return FLAMELayer(model_path, **kwargs) - else: - raise ValueError(f'Unknown model type {model_type}, exiting!') - - -def create( - model_path: str, - model_type: str = 'smpl', - **kwargs -) -> Union[SMPL, SMPLH, SMPLX, MANO, FLAME]: - ''' Method for creating a model from a path and a model type - Parameters - ---------- - model_path: str - Either the path to the model you wish to load or a folder, - where each subfolder contains the differents types, i.e.: - model_path: - | - |-- smpl - |-- SMPL_FEMALE - |-- SMPL_NEUTRAL - |-- SMPL_MALE - |-- smplh - |-- SMPLH_FEMALE - |-- SMPLH_MALE - |-- smplx - |-- SMPLX_FEMALE - |-- SMPLX_NEUTRAL - |-- SMPLX_MALE - |-- mano - |-- MANO RIGHT - |-- MANO LEFT - model_type: str, optional - When model_path is a folder, then this parameter specifies the - type of model to be loaded - **kwargs: dict - Keyword arguments - Returns - ------- - body_model: nn.Module - The PyTorch module that implements the corresponding body model - Raises - ------ - ValueError: In case the model type is not one of SMPL, SMPLH, - SMPLX, MANO or FLAME - ''' - - model_path = os.path.join(model_path, model_type) - - if model_type.lower() == 'smpl': - return SMPL(model_path, **kwargs) - elif model_type.lower() == 'smplh': - return SMPLH(model_path, **kwargs) - elif model_type.lower() == 'smplx': - return SMPLX(model_path, **kwargs) - elif 'mano' in model_type.lower(): - return MANO(model_path, **kwargs) - elif 'flame' in model_type.lower(): - return FLAME(model_path, **kwargs) - else: - raise ValueError(f'Unknown model type {model_type}, exiting!') diff --git a/lib /smplx /joint_names.py b/lib /smplx /joint_names.py deleted file mode 100644 index 0e1ee746bb1bb28e9b966eaf5c5f7c87ba26ac42..0000000000000000000000000000000000000000 --- a/lib /smplx /joint_names.py +++ /dev/null @@ -1,163 +0,0 @@ -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -JOINT_NAMES = [ - 'pelvis', - 'left_hip', - 'right_hip', - 'spine1', - 'left_knee', - 'right_knee', - 'spine2', - 'left_ankle', - 'right_ankle', - 'spine3', - 'left_foot', - 'right_foot', - 'neck', - 'left_collar', - 'right_collar', - 'head', - 'left_shoulder', - 'right_shoulder', - 'left_elbow', - 'right_elbow', - 'left_wrist', - 'right_wrist', - 'jaw', - 'left_eye_smplhf', - 'right_eye_smplhf', - 'left_index1', - 'left_index2', - 'left_index3', - 'left_middle1', - 'left_middle2', - 'left_middle3', - 'left_pinky1', - 'left_pinky2', - 'left_pinky3', - 'left_ring1', - 'left_ring2', - 'left_ring3', - 'left_thumb1', - 'left_thumb2', - 'left_thumb3', - 'right_index1', - 'right_index2', - 'right_index3', - 'right_middle1', - 'right_middle2', - 'right_middle3', - 'right_pinky1', - 'right_pinky2', - 'right_pinky3', - 'right_ring1', - 'right_ring2', - 'right_ring3', - 'right_thumb1', - 'right_thumb2', - 'right_thumb3', - 'nose', - 'right_eye', - 'left_eye', - 'right_ear', - 'left_ear', - 'left_big_toe', - 'left_small_toe', - 'left_heel', - 'right_big_toe', - 'right_small_toe', - 'right_heel', - 'left_thumb', - 'left_index', - 'left_middle', - 'left_ring', - 'left_pinky', - 'right_thumb', - 'right_index', - 'right_middle', - 'right_ring', - 'right_pinky', - 'right_eye_brow1', - 'right_eye_brow2', - 'right_eye_brow3', - 'right_eye_brow4', - 'right_eye_brow5', - 'left_eye_brow5', - 'left_eye_brow4', - 'left_eye_brow3', - 'left_eye_brow2', - 'left_eye_brow1', - 'nose1', - 'nose2', - 'nose3', - 'nose4', - 'right_nose_2', - 'right_nose_1', - 'nose_middle', - 'left_nose_1', - 'left_nose_2', - 'right_eye1', - 'right_eye2', - 'right_eye3', - 'right_eye4', - 'right_eye5', - 'right_eye6', - 'left_eye4', - 'left_eye3', - 'left_eye2', - 'left_eye1', - 'left_eye6', - 'left_eye5', - 'right_mouth_1', - 'right_mouth_2', - 'right_mouth_3', - 'mouth_top', - 'left_mouth_3', - 'left_mouth_2', - 'left_mouth_1', - 'left_mouth_5', # 59 in OpenPose output - 'left_mouth_4', # 58 in OpenPose output - 'mouth_bottom', - 'right_mouth_4', - 'right_mouth_5', - 'right_lip_1', - 'right_lip_2', - 'lip_top', - 'left_lip_2', - 'left_lip_1', - 'left_lip_3', - 'lip_bottom', - 'right_lip_3', - # Face contour - 'right_contour_1', - 'right_contour_2', - 'right_contour_3', - 'right_contour_4', - 'right_contour_5', - 'right_contour_6', - 'right_contour_7', - 'right_contour_8', - 'contour_middle', - 'left_contour_8', - 'left_contour_7', - 'left_contour_6', - 'left_contour_5', - 'left_contour_4', - 'left_contour_3', - 'left_contour_2', - 'left_contour_1', -] \ No newline at end of file diff --git a/lib /smplx /lbs.py b/lib /smplx /lbs.py deleted file mode 100644 index bab88d266c7078d7a341b7dcd9afec15b89fce77..0000000000000000000000000000000000000000 --- a/lib /smplx /lbs.py +++ /dev/null @@ -1,389 +0,0 @@ -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from __future__ import absolute_import -from __future__ import print_function -from __future__ import division - -from typing import Tuple, List, Optional -import numpy as np - -import torch -import torch.nn.functional as F - -from .utils import rot_mat_to_euler, Tensor - - -def find_dynamic_lmk_idx_and_bcoords( - vertices: Tensor, - pose: Tensor, - dynamic_lmk_faces_idx: Tensor, - dynamic_lmk_b_coords: Tensor, - neck_kin_chain: List[int], - pose2rot: bool = True, -) -> Tuple[Tensor, Tensor]: - ''' Compute the faces, barycentric coordinates for the dynamic landmarks - To do so, we first compute the rotation of the neck around the y-axis - and then use a pre-computed look-up table to find the faces and the - barycentric coordinates that will be used. - Special thanks to Soubhik Sanyal (soubhik.sanyal@tuebingen.mpg.de) - for providing the original TensorFlow implementation and for the LUT. - Parameters - ---------- - vertices: torch.tensor BxVx3, dtype = torch.float32 - The tensor of input vertices - pose: torch.tensor Bx(Jx3), dtype = torch.float32 - The current pose of the body model - dynamic_lmk_faces_idx: torch.tensor L, dtype = torch.long - The look-up table from neck rotation to faces - dynamic_lmk_b_coords: torch.tensor Lx3, dtype = torch.float32 - The look-up table from neck rotation to barycentric coordinates - neck_kin_chain: list - A python list that contains the indices of the joints that form the - kinematic chain of the neck. - dtype: torch.dtype, optional - Returns - ------- - dyn_lmk_faces_idx: torch.tensor, dtype = torch.long - A tensor of size BxL that contains the indices of the faces that - will be used to compute the current dynamic landmarks. - dyn_lmk_b_coords: torch.tensor, dtype = torch.float32 - A tensor of size BxL that contains the indices of the faces that - will be used to compute the current dynamic landmarks. - ''' - - dtype = vertices.dtype - batch_size = vertices.shape[0] - - if pose2rot: - aa_pose = torch.index_select(pose.view(batch_size, -1, 3), 1, - neck_kin_chain) - rot_mats = batch_rodrigues( - aa_pose.view(-1, 3)).view(batch_size, -1, 3, 3) - else: - rot_mats = torch.index_select( - pose.view(batch_size, -1, 3, 3), 1, neck_kin_chain) - - rel_rot_mat = torch.eye( - 3, device=vertices.device, dtype=dtype).unsqueeze_(dim=0).repeat( - batch_size, 1, 1) - for idx in range(len(neck_kin_chain)): - rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat) - - y_rot_angle = torch.round( - torch.clamp(-rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi, - max=39)).to(dtype=torch.long) - neg_mask = y_rot_angle.lt(0).to(dtype=torch.long) - mask = y_rot_angle.lt(-39).to(dtype=torch.long) - neg_vals = mask * 78 + (1 - mask) * (39 - y_rot_angle) - y_rot_angle = (neg_mask * neg_vals + - (1 - neg_mask) * y_rot_angle) - - dyn_lmk_faces_idx = torch.index_select(dynamic_lmk_faces_idx, - 0, y_rot_angle) - dyn_lmk_b_coords = torch.index_select(dynamic_lmk_b_coords, - 0, y_rot_angle) - - return dyn_lmk_faces_idx, dyn_lmk_b_coords - - -def vertices2landmarks( - vertices: Tensor, - faces: Tensor, - lmk_faces_idx: Tensor, - lmk_bary_coords: Tensor -) -> Tensor: - ''' Calculates landmarks by barycentric interpolation - Parameters - ---------- - vertices: torch.tensor BxVx3, dtype = torch.float32 - The tensor of input vertices - faces: torch.tensor Fx3, dtype = torch.long - The faces of the mesh - lmk_faces_idx: torch.tensor L, dtype = torch.long - The tensor with the indices of the faces used to calculate the - landmarks. - lmk_bary_coords: torch.tensor Lx3, dtype = torch.float32 - The tensor of barycentric coordinates that are used to interpolate - the landmarks - Returns - ------- - landmarks: torch.tensor BxLx3, dtype = torch.float32 - The coordinates of the landmarks for each mesh in the batch - ''' - # Extract the indices of the vertices for each face - # BxLx3 - batch_size, num_verts = vertices.shape[:2] - device = vertices.device - - lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).view( - batch_size, -1, 3) - - lmk_faces += torch.arange( - batch_size, dtype=torch.long, device=device).view(-1, 1, 1) * num_verts - - lmk_vertices = vertices.view(-1, 3)[lmk_faces].view( - batch_size, -1, 3, 3) - - landmarks = torch.einsum('blfi,blf->bli', [lmk_vertices, lmk_bary_coords]) - return landmarks - - -def lbs( - betas: Tensor, - pose: Tensor, - v_template: Tensor, - shapedirs: Tensor, - posedirs: Tensor, - J_regressor: Tensor, - parents: Tensor, - lbs_weights: Tensor, - pose2rot: bool = True, - return_transformation: bool = False, -) -> Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]: - ''' Performs Linear Blend Skinning with the given shape and pose parameters - Parameters - ---------- - betas : torch.tensor BxNB - The tensor of shape parameters - pose : torch.tensor Bx(J + 1) * 3 - The pose parameters in axis-angle format - v_template torch.tensor BxVx3 - The template mesh that will be deformed - shapedirs : torch.tensor 1xNB - The tensor of PCA shape displacements - posedirs : torch.tensor Px(V * 3) - The pose PCA coefficients - J_regressor : torch.tensor JxV - The regressor array that is used to calculate the joints from - the position of the vertices - parents: torch.tensor J - The array that describes the kinematic tree for the model - lbs_weights: torch.tensor N x V x (J + 1) - The linear blend skinning weights that represent how much the - rotation matrix of each part affects each vertex - pose2rot: bool, optional - Flag on whether to convert the input pose tensor to rotation - matrices. The default value is True. If False, then the pose tensor - should already contain rotation matrices and have a size of - Bx(J + 1)x9 - dtype: torch.dtype, optional - Returns - ------- - verts: torch.tensor BxVx3 - The vertices of the mesh after applying the shape and pose - displacements. - joints: torch.tensor BxJx3 - The joints of the model - ''' - - batch_size = max(betas.shape[0], pose.shape[0]) - device, dtype = betas.device, betas.dtype - - # Add shape contribution - v_shaped = v_template + blend_shapes(betas, shapedirs) - - # Get the joints - # NxJx3 array - J = vertices2joints(J_regressor, v_shaped) - - # 3. Add pose blend shapes - # N x J x 3 x 3 - ident = torch.eye(3, dtype=dtype, device=device) - if pose2rot: - rot_mats = batch_rodrigues(pose.view(-1, 3)).view( - [batch_size, -1, 3, 3]) - - pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1]) - # (N x P) x (P, V * 3) -> N x V x 3 - pose_offsets = torch.matmul( - pose_feature, posedirs).view(batch_size, -1, 3) - else: - pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident - rot_mats = pose.view(batch_size, -1, 3, 3) - - pose_offsets = torch.matmul(pose_feature.view(batch_size, -1), - posedirs).view(batch_size, -1, 3) - - v_posed = pose_offsets + v_shaped - # 4. Get the global joint location - J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype) - - # 5. Do skinning: - # W is N x V x (J + 1) - W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1]) - # (N x V x (J + 1)) x (N x (J + 1) x 16) - num_joints = J_regressor.shape[0] - T = torch.matmul(W, A.view(batch_size, num_joints, 16)) \ - .view(batch_size, -1, 4, 4) - - homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1], - dtype=dtype, device=device) - v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2) - v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1)) - - verts = v_homo[:, :, :3, 0] - - if return_transformation: - return verts, J_transformed, A, T - - return verts, J_transformed - - -def vertices2joints(J_regressor: Tensor, vertices: Tensor) -> Tensor: - ''' Calculates the 3D joint locations from the vertices - Parameters - ---------- - J_regressor : torch.tensor JxV - The regressor array that is used to calculate the joints from the - position of the vertices - vertices : torch.tensor BxVx3 - The tensor of mesh vertices - Returns - ------- - torch.tensor BxJx3 - The location of the joints - ''' - - return torch.einsum('bik,ji->bjk', [vertices, J_regressor]) - - -def blend_shapes(betas: Tensor, shape_disps: Tensor) -> Tensor: - ''' Calculates the per vertex displacement due to the blend shapes - Parameters - ---------- - betas : torch.tensor Bx(num_betas) - Blend shape coefficients - shape_disps: torch.tensor Vx3x(num_betas) - Blend shapes - Returns - ------- - torch.tensor BxVx3 - The per-vertex displacement due to shape deformation - ''' - - # Displacement[b, m, k] = sum_{l} betas[b, l] * shape_disps[m, k, l] - # i.e. Multiply each shape displacement by its corresponding beta and - # then sum them. - blend_shape = torch.einsum('bl,mkl->bmk', [betas, shape_disps]) - return blend_shape - - -def batch_rodrigues( - rot_vecs: Tensor, - epsilon: float = 1e-8, -) -> Tensor: - ''' Calculates the rotation matrices for a batch of rotation vectors - Parameters - ---------- - rot_vecs: torch.tensor Nx3 - array of N axis-angle vectors - Returns - ------- - R: torch.tensor Nx3x3 - The rotation matrices for the given axis-angle parameters - ''' - - batch_size = rot_vecs.shape[0] - device, dtype = rot_vecs.device, rot_vecs.dtype - - angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True) - rot_dir = rot_vecs / angle - - cos = torch.unsqueeze(torch.cos(angle), dim=1) - sin = torch.unsqueeze(torch.sin(angle), dim=1) - - # Bx1 arrays - rx, ry, rz = torch.split(rot_dir, 1, dim=1) - K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device) - - zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device) - K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \ - .view((batch_size, 3, 3)) - - ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0) - rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K) - return rot_mat - - -def transform_mat(R: Tensor, t: Tensor) -> Tensor: - ''' Creates a batch of transformation matrices - Args: - - R: Bx3x3 array of a batch of rotation matrices - - t: Bx3x1 array of a batch of translation vectors - Returns: - - T: Bx4x4 Transformation matrix - ''' - # No padding left or right, only add an extra row - return torch.cat([F.pad(R, [0, 0, 0, 1]), - F.pad(t, [0, 0, 0, 1], value=1)], dim=2) - - -def batch_rigid_transform( - rot_mats: Tensor, - joints: Tensor, - parents: Tensor, - dtype=torch.float32 -) -> Tensor: - """ - Applies a batch of rigid transformations to the joints - Parameters - ---------- - rot_mats : torch.tensor BxNx3x3 - Tensor of rotation matrices - joints : torch.tensor BxNx3 - Locations of joints - parents : torch.tensor BxN - The kinematic tree of each object - dtype : torch.dtype, optional: - The data type of the created tensors, the default is torch.float32 - Returns - ------- - posed_joints : torch.tensor BxNx3 - The locations of the joints after applying the pose rotations - rel_transforms : torch.tensor BxNx4x4 - The relative (with respect to the root joint) rigid transformations - for all the joints - """ - - joints = torch.unsqueeze(joints, dim=-1) - - rel_joints = joints.clone() - rel_joints[:, 1:] -= joints[:, parents[1:]] - - transforms_mat = transform_mat( - rot_mats.reshape(-1, 3, 3), - rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4) - - transform_chain = [transforms_mat[:, 0]] - for i in range(1, parents.shape[0]): - # Subtract the joint location at the rest pose - # No need for rotation, since it's identity when at rest - curr_res = torch.matmul(transform_chain[parents[i]], - transforms_mat[:, i]) - transform_chain.append(curr_res) - - transforms = torch.stack(transform_chain, dim=1) - - # The last column of the transformations contains the posed joints - posed_joints = transforms[:, :, :3, 3] - - joints_homogen = F.pad(joints, [0, 0, 0, 1]) - - rel_transforms = transforms - F.pad( - torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0]) - - return posed_joints, rel_transforms \ No newline at end of file diff --git a/lib /smplx /utils.py b/lib /smplx /utils.py deleted file mode 100644 index bd0c7cdec8b4eadace59a3644767c5e27750de48..0000000000000000000000000000000000000000 --- a/lib /smplx /utils.py +++ /dev/null @@ -1,127 +0,0 @@ -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from typing import NewType, Union, Optional -from dataclasses import dataclass, asdict, fields -import numpy as np -import torch - -Tensor = NewType('Tensor', torch.Tensor) -Array = NewType('Array', np.ndarray) - - -@dataclass -class ModelOutput: - vertices: Optional[Tensor] = None - joints: Optional[Tensor] = None - full_pose: Optional[Tensor] = None - global_orient: Optional[Tensor] = None - transl: Optional[Tensor] = None - - def __getitem__(self, key): - return getattr(self, key) - - def get(self, key, default=None): - return getattr(self, key, default) - - def __iter__(self): - return self.keys() - - def keys(self): - keys = [t.name for t in fields(self)] - return iter(keys) - - def values(self): - values = [getattr(self, t.name) for t in fields(self)] - return iter(values) - - def items(self): - data = [(t.name, getattr(self, t.name)) for t in fields(self)] - return iter(data) - - -@dataclass -class SMPLOutput(ModelOutput): - betas: Optional[Tensor] = None - body_pose: Optional[Tensor] = None - - -@dataclass -class SMPLHOutput(SMPLOutput): - left_hand_pose: Optional[Tensor] = None - right_hand_pose: Optional[Tensor] = None - transl: Optional[Tensor] = None - - -@dataclass -class SMPLXOutput(SMPLHOutput): - expression: Optional[Tensor] = None - jaw_pose: Optional[Tensor] = None - joint_transformation: Optional[Tensor] = None - vertex_transformation: Optional[Tensor] = None - - -@dataclass -class MANOOutput(ModelOutput): - betas: Optional[Tensor] = None - hand_pose: Optional[Tensor] = None - - -@dataclass -class FLAMEOutput(ModelOutput): - betas: Optional[Tensor] = None - expression: Optional[Tensor] = None - jaw_pose: Optional[Tensor] = None - neck_pose: Optional[Tensor] = None - - -def find_joint_kin_chain(joint_id, kinematic_tree): - kin_chain = [] - curr_idx = joint_id - while curr_idx != -1: - kin_chain.append(curr_idx) - curr_idx = kinematic_tree[curr_idx] - return kin_chain - - -def to_tensor( - array: Union[Array, Tensor], dtype=torch.float32 -) -> Tensor: - if torch.is_tensor(array): - return array - else: - return torch.tensor(array, dtype=dtype) - - -class Struct(object): - def __init__(self, **kwargs): - for key, val in kwargs.items(): - setattr(self, key, val) - - -def to_np(array, dtype=np.float32): - if 'scipy.sparse' in str(type(array)): - array = array.todense() - return np.array(array, dtype=dtype) - - -def rot_mat_to_euler(rot_mats): - # Calculates rotation matrix to euler angles - # Careful for extreme cases of eular angles like [0.0, pi, 0.0] - - sy = torch.sqrt(rot_mats[:, 0, 0] * rot_mats[:, 0, 0] + - rot_mats[:, 1, 0] * rot_mats[:, 1, 0]) - return torch.atan2(-rot_mats[:, 2, 0], sy) \ No newline at end of file diff --git a/lib /smplx /vertex_joint_selector.py b/lib /smplx /vertex_joint_selector.py deleted file mode 100644 index 4b8298bd5e087731f86c1c699703b5219e046c5c..0000000000000000000000000000000000000000 --- a/lib /smplx /vertex_joint_selector.py +++ /dev/null @@ -1,77 +0,0 @@ -# -*- coding: utf-8 -*- - -# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is -# holder of all proprietary rights on this computer program. -# You can only use this computer program if you have closed -# a license agreement with MPG or you get the right to use the computer -# program from someone who is authorized to grant you that right. -# Any use of the computer program without a valid license is prohibited and -# liable to prosecution. -# -# Copyright©2019 Max-Planck-Gesellschaft zur Förderung -# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute -# for Intelligent Systems. All rights reserved. -# -# Contact: ps-license@tuebingen.mpg.de - -from __future__ import absolute_import -from __future__ import print_function -from __future__ import division - -import numpy as np - -import torch -import torch.nn as nn - -from .utils import to_tensor - - -class VertexJointSelector(nn.Module): - - def __init__(self, vertex_ids=None, - use_hands=True, - use_feet_keypoints=True, **kwargs): - super(VertexJointSelector, self).__init__() - - extra_joints_idxs = [] - - face_keyp_idxs = np.array([ - vertex_ids['nose'], - vertex_ids['reye'], - vertex_ids['leye'], - vertex_ids['rear'], - vertex_ids['lear']], dtype=np.int64) - - extra_joints_idxs = np.concatenate([extra_joints_idxs, - face_keyp_idxs]) - - if use_feet_keypoints: - feet_keyp_idxs = np.array([vertex_ids['LBigToe'], - vertex_ids['LSmallToe'], - vertex_ids['LHeel'], - vertex_ids['RBigToe'], - vertex_ids['RSmallToe'], - vertex_ids['RHeel']], dtype=np.int32) - - extra_joints_idxs = np.concatenate( - [extra_joints_idxs, feet_keyp_idxs]) - - if use_hands: - self.tip_names = ['thumb', 'index', 'middle', 'ring', 'pinky'] - - tips_idxs = [] - for hand_id in ['l', 'r']: - for tip_name in self.tip_names: - tips_idxs.append(vertex_ids[hand_id + tip_name]) - - extra_joints_idxs = np.concatenate( - [extra_joints_idxs, tips_idxs]) - - self.register_buffer('extra_joints_idxs', - to_tensor(extra_joints_idxs, dtype=torch.long)) - - def forward(self, vertices, joints): - extra_joints = torch.index_select(vertices, 1, self.extra_joints_idxs) - joints = torch.cat([joints, extra_joints], dim=1) - - return joints