electron-rare/SMPLer-X
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Init commit

Browse Source
This commit is contained in:
caizhongang
2023-06-15 00:22:11 +08:00
parent c710d2e93d
commit 857a3ecbae
350 changed files with 65950 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+17
View File
@@ -0,0 +1,17 @@
__pycache__/
common/utils/human_model_files
pretrained_models/
dataset
output
.vscode
*.npy
*.tar
main/vis
main/vis*
tool/AGORA/vis*
tool/SHAPY/vis*
demo/
*egg-info/
.idea/
data/SHAPY/mesh-mesh-intersection/build/
data/SHAPY/mesh-mesh-intersection/mesh_mesh_intersect_cuda.cpython-38-x86_64-linux-gnu.so
LICENSE
+9
View File
@@ -0,0 +1,9 @@
S-Lab License 1.0
Copyright 2022 S-Lab
Redistribution and use for non-commercial purpose in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
2. 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.
3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
4. In the event that redistribution and/or use for commercial purpose in source or binary forms, with or without modification is required, please contact the contributor(s) of the work.
README.md
+92
View File
@@ -1 +1,93 @@
# SMPLer-X
## Install
```bash
conda create -n smplerx python=3.8 -y
conda activate smplerx
conda install pytorch==1.12.0 torchvision==0.13.0 torchaudio==0.12.0 cudatoolkit=11.3 -c pytorch -y
wget http://download.openmmlab.sensetime.com/mmcv/dist/cu113/torch1.12.0/mmcv_full-1.7.1-cp38-cp38-manylinux1_x86_64.whl
pip install mmcv_full-1.7.1-cp38-cp38-manylinux1_x86_64.whl
rm mmcv_full-1.7.1-cp38-cp38-manylinux1_x86_64.whl
pip install -r requirements.txt
# install mmpose
cd main/transformer_utils
pip install -v -e .
cd ../..
```
## Preparation
- process all datasets into [HumanData](https://github.com/open-mmlab/mmhuman3d/blob/main/docs/human_data.md) format, except the following:
- AGORA, MSCOCO, MPII, Human3.6M, UBody
- follow [OSX](https://github.com/IDEA-Research/OSX) in preparing pretrained ViT-Pose models.
- download [SMPL-X](https://smpl-x.is.tue.mpg.de/) body models
The file structure should be like:
```
SMPLer-X/
├── common/
│ └── utils/
│ └── human_model_files/ # body model
├── data/
├── main/
├── pretrained_models/ # pretrained ViT-Pose models
└── dataset/
├── AGORA/
├── ARCTIC/
├── BEDLAM/
├── Behave/
├── cache/
├── CHI3D/
├── CrowdPose/
├── dataset.py/
├── EgoBody/
├── EHF/
├── FIT3D/
├── GTA_Human2/
├── Human36M/
├── HumanSC3D/
├── InstaVariety/
├── LSPET/
├── MPII/
├── MPI_INF_3DHP/
├── MPI_INF_3DHP_folder/
├── MSCOCO/
├── MTP/
├── MuCo/
├── OCHuman/
├── PoseTrack/
├── PROX/
├── PW3D/
├── RenBody/
├── RICH/
├── SHAPY/
├── SPEC/
├── SSP3D/
├── SynBody/
├── Talkshow/
├── UBody/
├── UP3D/
└── preprocessed_datasets/ # HumanData files
```
## Training
```bash
cd main
sh slurm_train.sh {JOB_NAME} {NUM_GPU} {CONFIG_FILE}
# logs and ckpts will be saved to ../output/train_{JOB_NAME}_{DATE_TIME}
# config file is the file name under ./config, e.g. ./config/config_base.py
# a copy of current config file wil be saved to ../output/train_{JOB_NAME}_{DATE_TIME}/code/config_base.py
```
## Testing
```bash
cd main
sh slurm_test.sh {JOB_NAME} {NUM_GPU} {TRAIN_OUTPUT_DIR} {CKPT_ID}
# this will eval the model ../output/train_{JOB_NAME}_{DATE_TIME}/model_dump/snapshot_{CKPT_ID}.pth.tar with confing ../output/train_{JOB_NAME}_{DATE_TIME}/code/config_base.py
# logs and results will be saved to ../output/test_{JOB_NAME}_{DATE_TIME}
```
## References
- [Hand4Whole](https://github.com/mks0601/Hand4Whole_RELEASE)
- [OSX](https://github.com/IDEA-Research/OSX)
- [MMHuman3D](https://github.com/open-mmlab/mmhuman3d)
common/base.py
+340
View File
@@ -0,0 +1,340 @@
import os.path as osp
import math
import abc
from torch.utils.data import DataLoader
import torch.optim
import torchvision.transforms as transforms
from timer import Timer
from logger import colorlogger
from torch.nn.parallel.data_parallel import DataParallel
from config import cfg
from SMPLer_X import get_model
from dataset import MultipleDatasets
# ddp
import torch.distributed as dist
from torch.utils.data import DistributedSampler
import torch.utils.data.distributed
from utils.distribute_utils import (
get_rank, is_main_process, time_synchronized, get_group_idx, get_process_groups
)
from mmcv.runner import get_dist_info
# dynamic dataset import
for i in range(len(cfg.trainset_3d)):
exec('from ' + cfg.trainset_3d[i] + ' import ' + cfg.trainset_3d[i])
for i in range(len(cfg.trainset_2d)):
exec('from ' + cfg.trainset_2d[i] + ' import ' + cfg.trainset_2d[i])
for i in range(len(cfg.trainset_humandata)):
exec('from ' + cfg.trainset_humandata[i] + ' import ' + cfg.trainset_humandata[i])
exec('from ' + cfg.testset + ' import ' + cfg.testset)
class Base(object):
__metaclass__ = abc.ABCMeta
def __init__(self, log_name='logs.txt'):
self.cur_epoch = 0
# timer
self.tot_timer = Timer()
self.gpu_timer = Timer()
self.read_timer = Timer()
# logger
self.logger = colorlogger(cfg.log_dir, log_name=log_name)
@abc.abstractmethod
def _make_batch_generator(self):
return
@abc.abstractmethod
def _make_model(self):
return
class Trainer(Base):
def __init__(self, distributed=False, gpu_idx=None):
super(Trainer, self).__init__(log_name='train_logs.txt')
self.distributed = distributed
self.gpu_idx = gpu_idx
def get_optimizer(self, model):
normal_param = []
special_param = []
for module in model.module.special_trainable_modules:
special_param += list(module.parameters())
# print(module)
for module in model.module.trainable_modules:
normal_param += list(module.parameters())
# self.logger.info(f"N-{self.gpu_idx}, {normal_param}")
# self.logger.info("S", special_param)
optim_params = [
{ # add normal params first
'params': normal_param,
'lr': cfg.lr
},
{
'params': special_param,
'lr': cfg.lr * cfg.lr_mult
},
]
optimizer = torch.optim.Adam(optim_params, lr=cfg.lr)
return optimizer
def save_model(self, state, epoch):
file_path = osp.join(cfg.model_dir, 'snapshot_{}.pth.tar'.format(str(epoch)))
# do not save smplx layer weights
dump_key = []
for k in state['network'].keys():
if 'smplx_layer' in k:
dump_key.append(k)
for k in dump_key:
state['network'].pop(k, None)
torch.save(state, file_path)
self.logger.info("Write snapshot into {}".format(file_path))
def load_model(self, model, optimizer):
if cfg.pretrained_model_path is not None:
ckpt_path = cfg.pretrained_model_path
ckpt = torch.load(ckpt_path, map_location=torch.device('cpu')) # solve CUDA OOM error in DDP
model.load_state_dict(ckpt['network'], strict=False)
self.logger.info('Load checkpoint from {}'.format(ckpt_path))
if not hasattr(cfg, 'start_over') or cfg.start_over:
start_epoch = 0
else:
optimizer.load_state_dict(ckpt['optimizer'])
start_epoch = ckpt['epoch'] + 1
self.logger.info(f'Load optimizer, start from{start_epoch}')
else:
start_epoch = 0
return start_epoch, model, optimizer
def get_lr(self):
for g in self.optimizer.param_groups:
cur_lr = g['lr']
return cur_lr
def _make_batch_generator(self):
# data load and construct batch generator
self.logger_info("Creating dataset...")
trainset3d_loader = []
for i in range(len(cfg.trainset_3d)):
trainset3d_loader.append(eval(cfg.trainset_3d[i])(transforms.ToTensor(), "train"))
trainset2d_loader = []
for i in range(len(cfg.trainset_2d)):
trainset2d_loader.append(eval(cfg.trainset_2d[i])(transforms.ToTensor(), "train"))
trainset_humandata_loader = []
for i in range(len(cfg.trainset_humandata)):
trainset_humandata_loader.append(eval(cfg.trainset_humandata[i])(transforms.ToTensor(), "train"))
data_strategy = getattr(cfg, 'data_strategy', None)
if data_strategy == 'concat':
print("Using [concat] strategy...")
trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader,
make_same_len=False, verbose=True)
elif data_strategy == 'balance':
total_len = getattr(cfg, 'total_data_len', 'auto')
print(f"Using [balance] strategy with total_data_len : {total_len}...")
trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader,
make_same_len=True, total_len=total_len, verbose=True)
else:
# original strategy implementation
valid_loader_num = 0
if len(trainset3d_loader) > 0:
trainset3d_loader = [MultipleDatasets(trainset3d_loader, make_same_len=False)]
valid_loader_num += 1
else:
trainset3d_loader = []
if len(trainset2d_loader) > 0:
trainset2d_loader = [MultipleDatasets(trainset2d_loader, make_same_len=False)]
valid_loader_num += 1
else:
trainset2d_loader = []
if len(trainset_humandata_loader) > 0:
trainset_humandata_loader = [MultipleDatasets(trainset_humandata_loader, make_same_len=False)]
valid_loader_num += 1
if valid_loader_num > 1:
trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader, make_same_len=True)
else:
trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader, make_same_len=False)
self.itr_per_epoch = math.ceil(len(trainset_loader) / cfg.num_gpus / cfg.train_batch_size)
if self.distributed:
self.logger_info(f"Total data length {len(trainset_loader)}.")
rank, world_size = get_dist_info()
self.logger_info("Using distributed data sampler.")
sampler_train = DistributedSampler(trainset_loader, world_size, rank, shuffle=True)
self.batch_generator = DataLoader(dataset=trainset_loader, batch_size=cfg.train_batch_size,
shuffle=False, num_workers=cfg.num_thread, sampler=sampler_train,
pin_memory=True, persistent_workers=True if cfg.num_thread > 0 else False, drop_last=True)
else:
self.batch_generator = DataLoader(dataset=trainset_loader, batch_size=cfg.num_gpus * cfg.train_batch_size,
shuffle=True, num_workers=cfg.num_thread,
pin_memory=True, drop_last=True)
def _make_model(self):
# prepare network
self.logger_info("Creating graph and optimizer...")
model = get_model('train')
if getattr(cfg, 'fine_tune', None) == 'backbone':
print("Fine-tuning [backbone]...")
for module in model.head:
for param in module.parameters():
param.requires_grad = False
for module in model.neck:
for param in module.parameters():
param.requires_grad = False
elif getattr(cfg, 'fine_tune', None) == 'neck_and_head':
print("Fine-tuning [neck and head]...")
for param in model.encoder.parameters():
param.requires_grad = False
elif getattr(cfg, 'fine_tune', None) == 'head':
print("Fine-tuning [head]...")
for param in model.encoder.parameters():
param.requires_grad = False
for module in model.neck:
for param in module.parameters():
param.requires_grad = False
# ddp
if self.distributed:
self.logger_info("Using distributed data parallel.")
model.cuda()
if hasattr(cfg, 'syncbn') and cfg.syncbn:
self.logger_info("Using sync batch norm layers.")
process_groups = get_process_groups()
process_group = process_groups[get_group_idx()]
syncbn_model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model, process_group)
model = torch.nn.parallel.DistributedDataParallel(
syncbn_model, device_ids=[self.gpu_idx],
find_unused_parameters=True)
else:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[self.gpu_idx],
find_unused_parameters=True)
else:
# dp
model = DataParallel(model).cuda()
optimizer = self.get_optimizer(model)
if hasattr(cfg, "scheduler"):
if cfg.scheduler == 'cos':
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, cfg.end_epoch * self.itr_per_epoch,
eta_min=1e-6)
elif cfg.scheduler == 'step':
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, cfg.step_size, gamma=cfg.gamma,
last_epoch=- 1, verbose=False)
else:
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, cfg.end_epoch * self.itr_per_epoch,
eta_min=getattr(cfg,'min_lr',1e-6))
if cfg.continue_train:
if self.distributed:
start_epoch, model, optimizer = self.load_model(model, optimizer)
else:
start_epoch, model, optimizer = self.load_model(model, optimizer)
else:
start_epoch = 0
model.train()
self.scheduler = scheduler
self.start_epoch = start_epoch
self.model = model
self.optimizer = optimizer
def logger_info(self, info):
if self.distributed:
if is_main_process():
self.logger.info(info)
else:
self.logger.info(info)
class Tester(Base):
def __init__(self, test_epoch=None):
if test_epoch is not None:
self.test_epoch = int(test_epoch)
super(Tester, self).__init__(log_name='test_logs.txt')
def _make_batch_generator(self):
# data load and construct batch generator
self.logger.info("Creating dataset...")
testset_loader = eval(cfg.testset)(transforms.ToTensor(), "test")
batch_generator = DataLoader(dataset=testset_loader, batch_size=cfg.num_gpus * cfg.test_batch_size,
shuffle=False, num_workers=cfg.num_thread, pin_memory=True)
self.testset = testset_loader
self.batch_generator = batch_generator
def _make_model(self):
self.logger.info('Load checkpoint from {}'.format(cfg.pretrained_model_path))
# prepare network
self.logger.info("Creating graph...")
model = get_model('test')
model = DataParallel(model).cuda()
if not getattr(cfg, 'random_init', False):
ckpt = torch.load(cfg.pretrained_model_path, map_location=torch.device('cpu'))
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in ckpt['network'].items():
if 'module' not in k:
k = 'module.' + k
k = k.replace('backbone', 'encoder').replace('body_rotation_net', 'body_regressor').replace(
'hand_rotation_net', 'hand_regressor')
new_state_dict[k] = v
self.logger.warning("Attention: Strict=False is set for checkpoint loading. Please check manually.")
model.load_state_dict(new_state_dict, strict=False)
model.eval()
else:
print('Random init!!!!!!!')
self.model = model
def _evaluate(self, outs, cur_sample_idx):
eval_result = self.testset.evaluate(outs, cur_sample_idx)
return eval_result
def _print_eval_result(self, eval_result):
self.testset.print_eval_result(eval_result)
class Demoer(Base):
def __init__(self, test_epoch=None):
if test_epoch is not None:
self.test_epoch = int(test_epoch)
super(Demoer, self).__init__(log_name='test_logs.txt')
def _make_model(self):
self.logger.info('Load checkpoint from {}'.format(cfg.pretrained_model_path))
# prepare network
self.logger.info("Creating graph...")
model = get_model('test')
model = DataParallel(model).cuda()
ckpt = torch.load(cfg.pretrained_model_path, map_location=torch.device('cpu'))
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in ckpt['network'].items():
if 'module' not in k:
k = 'module.' + k
k = k.replace('module.backbone', 'module.encoder').replace('body_rotation_net', 'body_regressor').replace(
'hand_rotation_net', 'hand_regressor')
new_state_dict[k] = v
model.load_state_dict(new_state_dict, strict=False)
model.eval()
self.model = model
common/logger.py
+50
View File
@@ -0,0 +1,50 @@
import logging
import os
OK = '\033[92m'
WARNING = '\033[93m'
FAIL = '\033[91m'
END = '\033[0m'
PINK = '\033[95m'
BLUE = '\033[94m'
GREEN = OK
RED = FAIL
WHITE = END
YELLOW = WARNING
class colorlogger():
def __init__(self, log_dir, log_name='train_logs.txt'):
# set log
self._logger = logging.getLogger(log_name)
self._logger.setLevel(logging.INFO)
log_file = os.path.join(log_dir, log_name)
if not os.path.exists(log_dir):
os.makedirs(log_dir)
file_log = logging.FileHandler(log_file, mode='a')
file_log.setLevel(logging.INFO)
console_log = logging.StreamHandler()
console_log.setLevel(logging.INFO)
formatter = logging.Formatter(
"{}%(asctime)s{} %(message)s".format(GREEN, END),
"%m-%d %H:%M:%S")
file_log.setFormatter(formatter)
console_log.setFormatter(formatter)
self._logger.addHandler(file_log)
self._logger.addHandler(console_log)
def debug(self, msg):
self._logger.debug(str(msg))
def info(self, msg):
self._logger.info(str(msg))
def warning(self, msg):
self._logger.warning(WARNING + 'WRN: ' + str(msg) + END)
def critical(self, msg):
self._logger.critical(RED + 'CRI: ' + str(msg) + END)
def error(self, msg):
self._logger.error(RED + 'ERR: ' + str(msg) + END)
common/nets/layer.py
+53
View File
@@ -0,0 +1,53 @@
import torch.nn as nn
def make_linear_layers(feat_dims, relu_final=True, use_bn=False):
layers = []
for i in range(len(feat_dims)-1):
layers.append(nn.Linear(feat_dims[i], feat_dims[i+1]))
# Do not use ReLU for final estimation
if i < len(feat_dims)-2 or (i == len(feat_dims)-2 and relu_final):
if use_bn:
layers.append(nn.BatchNorm1d(feat_dims[i+1]))
layers.append(nn.ReLU(inplace=True))
return nn.Sequential(*layers)
def make_conv_layers(feat_dims, kernel=3, stride=1, padding=1, bnrelu_final=True):
layers = []
for i in range(len(feat_dims)-1):
layers.append(
nn.Conv2d(
in_channels=feat_dims[i],
out_channels=feat_dims[i+1],
kernel_size=kernel,
stride=stride,
padding=padding
))
# Do not use BN and ReLU for final estimation
if i < len(feat_dims)-2 or (i == len(feat_dims)-2 and bnrelu_final):
layers.append(nn.BatchNorm2d(feat_dims[i+1]))
layers.append(nn.ReLU(inplace=True))
return nn.Sequential(*layers)
def make_deconv_layers(feat_dims, bnrelu_final=True):
layers = []
for i in range(len(feat_dims)-1):
layers.append(
nn.ConvTranspose2d(
in_channels=feat_dims[i],
out_channels=feat_dims[i+1],
kernel_size=4,
stride=2,
padding=1,
output_padding=0,
bias=False))
# Do not use BN and ReLU for final estimation
if i < len(feat_dims)-2 or (i == len(feat_dims)-2 and bnrelu_final):
layers.append(nn.BatchNorm2d(feat_dims[i+1]))
layers.append(nn.ReLU(inplace=True))
return nn.Sequential(*layers)
common/nets/loss.py
+30
View File
@@ -0,0 +1,30 @@
import torch
import torch.nn as nn
class CoordLoss(nn.Module):
def __init__(self):
super(CoordLoss, self).__init__()
def forward(self, coord_out, coord_gt, valid, is_3D=None):
loss = torch.abs(coord_out - coord_gt) * valid
if is_3D is not None:
loss_z = loss[:,:,2:] * is_3D[:,None,None].float()
loss = torch.cat((loss[:,:,:2], loss_z),2)
return loss
class ParamLoss(nn.Module):
def __init__(self):
super(ParamLoss, self).__init__()
def forward(self, param_out, param_gt, valid):
loss = torch.abs(param_out - param_gt) * valid
return loss
class CELoss(nn.Module):
def __init__(self):
super(CELoss, self).__init__()
self.ce_loss = nn.CrossEntropyLoss(reduction='none')
def forward(self, out, gt_index):
loss = self.ce_loss(out, gt_index)
return loss
common/nets/smpler_x.py
+172
View File
@@ -0,0 +1,172 @@
import torch
import torch.nn as nn
from torch.nn import functional as F
from nets.layer import make_conv_layers, make_linear_layers, make_deconv_layers
from utils.transforms import sample_joint_features, soft_argmax_2d, soft_argmax_3d
from utils.human_models import smpl_x
from config import cfg
from mmcv.ops.roi_align import roi_align
class PositionNet(nn.Module):
def __init__(self, part, feat_dim=768):
super(PositionNet, self).__init__()
if part == 'body':
self.joint_num = len(smpl_x.pos_joint_part['body'])
self.hm_shape = cfg.output_hm_shape
elif part == 'hand':
self.joint_num = len(smpl_x.pos_joint_part['rhand'])
self.hm_shape = cfg.output_hand_hm_shape
self.conv = make_conv_layers([feat_dim, self.joint_num * self.hm_shape[0]], kernel=1, stride=1, padding=0, bnrelu_final=False)
def forward(self, img_feat):
joint_hm = self.conv(img_feat).view(-1, self.joint_num, self.hm_shape[0], self.hm_shape[1], self.hm_shape[2])
joint_coord = soft_argmax_3d(joint_hm)
joint_hm = F.softmax(joint_hm.view(-1, self.joint_num, self.hm_shape[0] * self.hm_shape[1] * self.hm_shape[2]), 2)
joint_hm = joint_hm.view(-1, self.joint_num, self.hm_shape[0], self.hm_shape[1], self.hm_shape[2])
return joint_hm, joint_coord
class HandRotationNet(nn.Module):
def __init__(self, part, feat_dim = 768):
super(HandRotationNet, self).__init__()
self.part = part
self.joint_num = len(smpl_x.pos_joint_part['rhand'])
self.hand_conv = make_conv_layers([feat_dim, 512], kernel=1, stride=1, padding=0)
self.hand_pose_out = make_linear_layers([self.joint_num * 515, len(smpl_x.orig_joint_part['rhand']) * 6], relu_final=False)
self.feat_dim = feat_dim
def forward(self, img_feat, joint_coord_img):
batch_size = img_feat.shape[0]
img_feat = self.hand_conv(img_feat)
img_feat_joints = sample_joint_features(img_feat, joint_coord_img[:, :, :2])
feat = torch.cat((img_feat_joints, joint_coord_img), 2) # batch_size, joint_num, 512+3
hand_pose = self.hand_pose_out(feat.view(batch_size, -1))
return hand_pose
class BodyRotationNet(nn.Module):
def __init__(self, feat_dim = 768):
super(BodyRotationNet, self).__init__()
self.joint_num = len(smpl_x.pos_joint_part['body'])
self.body_conv = make_linear_layers([feat_dim, 512], relu_final=False)
self.root_pose_out = make_linear_layers([self.joint_num * (512+3), 6], relu_final=False)
self.body_pose_out = make_linear_layers(
[self.joint_num * (512+3), (len(smpl_x.orig_joint_part['body']) - 1) * 6], relu_final=False) # without root
self.shape_out = make_linear_layers([feat_dim, smpl_x.shape_param_dim], relu_final=False)
self.cam_out = make_linear_layers([feat_dim, 3], relu_final=False)
self.feat_dim = feat_dim
def forward(self, body_pose_token, shape_token, cam_token, body_joint_img):
batch_size = body_pose_token.shape[0]
# shape parameter
shape_param = self.shape_out(shape_token)
# camera parameter
cam_param = self.cam_out(cam_token)
# body pose parameter
body_pose_token = self.body_conv(body_pose_token)
body_pose_token = torch.cat((body_pose_token, body_joint_img), 2)
root_pose = self.root_pose_out(body_pose_token.view(batch_size, -1))
body_pose = self.body_pose_out(body_pose_token.view(batch_size, -1))
return root_pose, body_pose, shape_param, cam_param
class FaceRegressor(nn.Module):
def __init__(self, feat_dim=768):
super(FaceRegressor, self).__init__()
self.expr_out = make_linear_layers([feat_dim, smpl_x.expr_code_dim], relu_final=False)
self.jaw_pose_out = make_linear_layers([feat_dim, 6], relu_final=False)
def forward(self, expr_token, jaw_pose_token):
expr_param = self.expr_out(expr_token) # expression parameter
jaw_pose = self.jaw_pose_out(jaw_pose_token) # jaw pose parameter
return expr_param, jaw_pose
class BoxNet(nn.Module):
def __init__(self, feat_dim=768):
super(BoxNet, self).__init__()
self.joint_num = len(smpl_x.pos_joint_part['body'])
self.deconv = make_deconv_layers([feat_dim + self.joint_num * cfg.output_hm_shape[0], 256, 256, 256])
self.bbox_center = make_conv_layers([256, 3], kernel=1, stride=1, padding=0, bnrelu_final=False)
self.lhand_size = make_linear_layers([256, 256, 2], relu_final=False)
self.rhand_size = make_linear_layers([256, 256, 2], relu_final=False)
self.face_size = make_linear_layers([256, 256, 2], relu_final=False)
def forward(self, img_feat, joint_hm):
joint_hm = joint_hm.view(joint_hm.shape[0], joint_hm.shape[1] * cfg.output_hm_shape[0], cfg.output_hm_shape[1], cfg.output_hm_shape[2])
img_feat = torch.cat((img_feat, joint_hm), 1)
img_feat = self.deconv(img_feat)
# bbox center
bbox_center_hm = self.bbox_center(img_feat)
bbox_center = soft_argmax_2d(bbox_center_hm)
lhand_center, rhand_center, face_center = bbox_center[:, 0, :], bbox_center[:, 1, :], bbox_center[:, 2, :]
# bbox size
lhand_feat = sample_joint_features(img_feat, lhand_center[:, None, :].detach())[:, 0, :]
lhand_size = self.lhand_size(lhand_feat)
rhand_feat = sample_joint_features(img_feat, rhand_center[:, None, :].detach())[:, 0, :]
rhand_size = self.rhand_size(rhand_feat)
face_feat = sample_joint_features(img_feat, face_center[:, None, :].detach())[:, 0, :]
face_size = self.face_size(face_feat)
lhand_center = lhand_center / 8
rhand_center = rhand_center / 8
face_center = face_center / 8
return lhand_center, lhand_size, rhand_center, rhand_size, face_center, face_size
class BoxSizeNet(nn.Module):
def __init__(self):
super(BoxSizeNet, self).__init__()
self.lhand_size = make_linear_layers([256, 256, 2], relu_final=False)
self.rhand_size = make_linear_layers([256, 256, 2], relu_final=False)
self.face_size = make_linear_layers([256, 256, 2], relu_final=False)
def forward(self, box_fea):
# box_fea: [bs, 3, C]
lhand_size = self.lhand_size(box_fea[:, 0])
rhand_size = self.rhand_size(box_fea[:, 1])
face_size = self.face_size(box_fea[:, 2])
return lhand_size, rhand_size, face_size
class HandRoI(nn.Module):
def __init__(self, feat_dim=768, upscale=4):
super(HandRoI, self).__init__()
self.upscale = upscale
if upscale==1:
self.deconv = make_conv_layers([feat_dim, feat_dim], kernel=1, stride=1, padding=0, bnrelu_final=False)
self.conv = make_conv_layers([feat_dim, feat_dim], kernel=1, stride=1, padding=0, bnrelu_final=False)
elif upscale==2:
self.deconv = make_deconv_layers([feat_dim, feat_dim//2])
self.conv = make_conv_layers([feat_dim//2, feat_dim], kernel=1, stride=1, padding=0, bnrelu_final=False)
elif upscale==4:
self.deconv = make_deconv_layers([feat_dim, feat_dim//2, feat_dim//4])
self.conv = make_conv_layers([feat_dim//4, feat_dim], kernel=1, stride=1, padding=0, bnrelu_final=False)
elif upscale==8:
self.deconv = make_deconv_layers([feat_dim, feat_dim//2, feat_dim//4, feat_dim//8])
self.conv = make_conv_layers([feat_dim//8, feat_dim], kernel=1, stride=1, padding=0, bnrelu_final=False)
def forward(self, img_feat, lhand_bbox, rhand_bbox):
lhand_bbox = torch.cat((torch.arange(lhand_bbox.shape[0]).float().cuda()[:, None], lhand_bbox),
1) # batch_idx, xmin, ymin, xmax, ymax
rhand_bbox = torch.cat((torch.arange(rhand_bbox.shape[0]).float().cuda()[:, None], rhand_bbox),
1) # batch_idx, xmin, ymin, xmax, ymax
img_feat = self.deconv(img_feat)
lhand_bbox_roi = lhand_bbox.clone()
lhand_bbox_roi[:, 1] = lhand_bbox_roi[:, 1] / cfg.input_body_shape[1] * cfg.output_hm_shape[2] * self.upscale
lhand_bbox_roi[:, 2] = lhand_bbox_roi[:, 2] / cfg.input_body_shape[0] * cfg.output_hm_shape[1] * self.upscale
lhand_bbox_roi[:, 3] = lhand_bbox_roi[:, 3] / cfg.input_body_shape[1] * cfg.output_hm_shape[2] * self.upscale
lhand_bbox_roi[:, 4] = lhand_bbox_roi[:, 4] / cfg.input_body_shape[0] * cfg.output_hm_shape[1] * self.upscale
assert (cfg.output_hm_shape[1]*self.upscale, cfg.output_hm_shape[2]*self.upscale) == (img_feat.shape[2], img_feat.shape[3])
lhand_img_feat = roi_align(img_feat, lhand_bbox_roi, (cfg.output_hand_hm_shape[1], cfg.output_hand_hm_shape[2]), 1.0, 0, 'avg', False)
lhand_img_feat = torch.flip(lhand_img_feat, [3]) # flip to the right hand
rhand_bbox_roi = rhand_bbox.clone()
rhand_bbox_roi[:, 1] = rhand_bbox_roi[:, 1] / cfg.input_body_shape[1] * cfg.output_hm_shape[2] * self.upscale
rhand_bbox_roi[:, 2] = rhand_bbox_roi[:, 2] / cfg.input_body_shape[0] * cfg.output_hm_shape[1] * self.upscale
rhand_bbox_roi[:, 3] = rhand_bbox_roi[:, 3] / cfg.input_body_shape[1] * cfg.output_hm_shape[2] * self.upscale
rhand_bbox_roi[:, 4] = rhand_bbox_roi[:, 4] / cfg.input_body_shape[0] * cfg.output_hm_shape[1] * self.upscale
rhand_img_feat = roi_align(img_feat, rhand_bbox_roi, (cfg.output_hand_hm_shape[1], cfg.output_hand_hm_shape[2]), 1.0, 0, 'avg', False)
hand_img_feat = torch.cat((lhand_img_feat, rhand_img_feat)) # [bs, c, cfg.output_hand_hm_shape[2]*scale, cfg.output_hand_hm_shape[1]*scale]
hand_img_feat = self.conv(hand_img_feat)
return hand_img_feat
common/timer.py
+38
View File
@@ -0,0 +1,38 @@
# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------
import time
class Timer(object):
"""A simple timer."""
def __init__(self):
self.total_time = 0.
self.calls = 0
self.start_time = 0.
self.diff = 0.
self.average_time = 0.
self.warm_up = 0
def tic(self):
# using time.time instead of time.clock because time time.clock
# does not normalize for multithreading
self.start_time = time.time()
def toc(self, average=True):
self.diff = time.time() - self.start_time
if self.warm_up < 10:
self.warm_up += 1
return self.diff
else:
self.total_time += self.diff
self.calls += 1
self.average_time = self.total_time / self.calls
if average:
return self.average_time
else:
return self.diff
common/utils/__init__.py
View File
common/utils/dir.py
+10
View File
@@ -0,0 +1,10 @@
import os
import sys
def make_folder(folder_name):
os.makedirs(folder_name, exist_ok=True)
def add_pypath(path):
if path not in sys.path:
sys.path.insert(0, path)
common/utils/distribute_utils.py
+217
View File
@@ -0,0 +1,217 @@
import mmcv
import os
import os.path as osp
import pickle
import shutil
import tempfile
import time
import torch
import torch.distributed as dist
from mmcv.runner import get_dist_info
import random
import numpy as np
import subprocess
def set_seed(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
# torch.set_deterministic(True)
def time_synchronized():
torch.cuda.synchronize() if torch.cuda.is_available() else None
return time.time()
def setup_for_distributed(is_master):
"""This function disables printing when not in master process."""
import builtins as __builtin__
builtin_print = __builtin__.print
def print(*args, **kwargs):
force = kwargs.pop('force', False)
if is_master or force:
builtin_print(*args, **kwargs)
__builtin__.print = print
def init_distributed_mode(port = None, master_port=29500):
"""Initialize slurm distributed training environment.
If argument ``port`` is not specified, then the master port will be system
environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
environment variable, then a default port ``29500`` will be used.
Args:
backend (str): Backend of torch.distributed.
port (int, optional): Master port. Defaults to None.
"""
dist_backend = 'nccl'
proc_id = int(os.environ['SLURM_PROCID'])
ntasks = int(os.environ['SLURM_NTASKS'])
node_list = os.environ['SLURM_NODELIST']
num_gpus = torch.cuda.device_count()
torch.cuda.set_device(proc_id % num_gpus)
addr = subprocess.getoutput(
f'scontrol show hostname {node_list} | head -n1')
# specify master port
if port is not None:
os.environ['MASTER_PORT'] = str(port)
elif 'MASTER_PORT' in os.environ:
pass # use MASTER_PORT in the environment variable
else:
# 29500 is torch.distributed default port
os.environ['MASTER_PORT'] = str(master_port)
# use MASTER_ADDR in the environment variable if it already exists
if 'MASTER_ADDR' not in os.environ:
os.environ['MASTER_ADDR'] = addr
os.environ['WORLD_SIZE'] = str(ntasks)
os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
os.environ['RANK'] = str(proc_id)
dist.init_process_group(backend=dist_backend)
distributed = True
gpu_idx = proc_id % num_gpus
return distributed, gpu_idx
def is_dist_avail_and_initialized():
if not dist.is_available():
return False
if not dist.is_initialized():
return False
return True
def get_world_size():
if not is_dist_avail_and_initialized():
return 1
return dist.get_world_size()
def get_rank():
if not is_dist_avail_and_initialized():
return 0
return dist.get_rank()
def get_process_groups():
world_size = int(os.environ['WORLD_SIZE'])
ranks = list(range(world_size))
num_gpus = torch.cuda.device_count()
num_nodes = world_size // num_gpus
if world_size % num_gpus != 0:
raise NotImplementedError('Not implemented for node not fully used.')
groups = []
for node_idx in range(num_nodes):
groups.append(ranks[node_idx*num_gpus : (node_idx+1)*num_gpus])
process_groups = [torch.distributed.new_group(group) for group in groups]
return process_groups
def get_group_idx():
num_gpus = torch.cuda.device_count()
proc_id = get_rank()
group_idx = proc_id // num_gpus
return group_idx
def is_main_process():
return get_rank() == 0
def cleanup():
dist.destroy_process_group()
def collect_results(result_part, size, tmpdir=None):
rank, world_size = get_dist_info()
# create a tmp dir if it is not specified
if tmpdir is None:
MAX_LEN = 512
# 32 is whitespace
dir_tensor = torch.full((MAX_LEN, ),
32,
dtype=torch.uint8,
device='cuda')
if rank == 0:
tmpdir = tempfile.mkdtemp()
tmpdir = torch.tensor(
bytearray(tmpdir.encode()), dtype=torch.uint8, device='cuda')
dir_tensor[:len(tmpdir)] = tmpdir
dist.broadcast(dir_tensor, 0)
tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip()
else:
mmcv.mkdir_or_exist(tmpdir)
# dump the part result to the dir
mmcv.dump(result_part, osp.join(tmpdir, f'part_{rank}.pkl'))
dist.barrier()
# collect all parts
if rank != 0:
return None
else:
# load results of all parts from tmp dir
part_list = []
for i in range(world_size):
part_file = osp.join(tmpdir, f'part_{i}.pkl')
part_list.append(mmcv.load(part_file))
# sort the results
ordered_results = []
for res in zip(*part_list):
ordered_results.extend(list(res))
# the dataloader may pad some samples
ordered_results = ordered_results[:size]
# remove tmp dir
shutil.rmtree(tmpdir)
return ordered_results
def all_gather(data):
"""
Run all_gather on arbitrary picklable data (not necessarily tensors)
Args:
data:
Any picklable object
Returns:
data_list(list):
List of data gathered from each rank
"""
world_size = get_world_size()
if world_size == 1:
return [data]
# serialized to a Tensor
buffer = pickle.dumps(data)
storage = torch.ByteStorage.from_buffer(buffer)
tensor = torch.ByteTensor(storage).to('cuda')
# obtain Tensor size of each rank
local_size = torch.tensor([tensor.numel()], device='cuda')
size_list = [torch.tensor([0], device='cuda') for _ in range(world_size)]
dist.all_gather(size_list, local_size)
size_list = [int(size.item()) for size in size_list]
max_size = max(size_list)
# receiving Tensor from all ranks
# we pad the tensor because torch all_gather does not support
# gathering tensors of different shapes
tensor_list = []
for _ in size_list:
tensor_list.append(
torch.empty((max_size, ), dtype=torch.uint8, device='cuda'))
if local_size != max_size:
padding = torch.empty(
size=(max_size - local_size, ), dtype=torch.uint8, device='cuda')
tensor = torch.cat((tensor, padding), dim=0)
dist.all_gather(tensor_list, tensor)
data_list = []
for size, tensor in zip(size_list, tensor_list):
buffer = tensor.cpu().numpy().tobytes()[:size]
data_list.append(pickle.loads(buffer))
return data_list
common/utils/human_models.py
+176
View File
@@ -0,0 +1,176 @@
import numpy as np
import torch
import os.path as osp
from config import cfg
from utils.smplx import smplx
import pickle
class SMPLX(object):
def __init__(self):
self.layer_arg = {'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}
self.layer = {'neutral': smplx.create(cfg.human_model_path, 'smplx', gender='NEUTRAL', use_pca=False, use_face_contour=True, **self.layer_arg),
'male': smplx.create(cfg.human_model_path, 'smplx', gender='MALE', use_pca=False, use_face_contour=True, **self.layer_arg),
'female': smplx.create(cfg.human_model_path, 'smplx', gender='FEMALE', use_pca=False, use_face_contour=True, **self.layer_arg)
}
self.vertex_num = 10475
self.face = self.layer['neutral'].faces
self.shape_param_dim = 10
self.expr_code_dim = 10
with open(osp.join(cfg.human_model_path, 'smplx', 'SMPLX_to_J14.pkl'), 'rb') as f:
self.j14_regressor = pickle.load(f, encoding='latin1')
with open(osp.join(cfg.human_model_path, 'smplx', 'MANO_SMPLX_vertex_ids.pkl'), 'rb') as f:
self.hand_vertex_idx = pickle.load(f, encoding='latin1')
self.face_vertex_idx = np.load(osp.join(cfg.human_model_path, 'smplx', 'SMPL-X__FLAME_vertex_ids.npy'))
self.J_regressor = self.layer['neutral'].J_regressor.numpy()
self.J_regressor_idx = {'pelvis': 0, 'lwrist': 20, 'rwrist': 21, 'neck': 12}
self.orig_hand_regressor = self.make_hand_regressor()
#self.orig_hand_regressor = {'left': self.layer.J_regressor.numpy()[[20,37,38,39,25,26,27,28,29,30,34,35,36,31,32,33],:], 'right': self.layer.J_regressor.numpy()[[21,52,53,54,40,41,42,43,44,45,49,50,51,46,47,48],:]}
# original SMPLX joint set
self.orig_joint_num = 53 # 22 (body joints) + 30 (hand joints) + 1 (face jaw joint)
self.orig_joints_name = \
('Pelvis', 'L_Hip', 'R_Hip', 'Spine_1', 'L_Knee', 'R_Knee', 'Spine_2', 'L_Ankle', 'R_Ankle', 'Spine_3', 'L_Foot', 'R_Foot', 'Neck', 'L_Collar', 'R_Collar', 'Head', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', # body joints
'L_Index_1', 'L_Index_2', 'L_Index_3', 'L_Middle_1', 'L_Middle_2', 'L_Middle_3', 'L_Pinky_1', 'L_Pinky_2', 'L_Pinky_3', 'L_Ring_1', 'L_Ring_2', 'L_Ring_3', 'L_Thumb_1', 'L_Thumb_2', 'L_Thumb_3', # left hand joints
'R_Index_1', 'R_Index_2', 'R_Index_3', 'R_Middle_1', 'R_Middle_2', 'R_Middle_3', 'R_Pinky_1', 'R_Pinky_2', 'R_Pinky_3', 'R_Ring_1', 'R_Ring_2', 'R_Ring_3', 'R_Thumb_1', 'R_Thumb_2', 'R_Thumb_3', # right hand joints
'Jaw' # face jaw joint
)
self.orig_flip_pairs = \
( (1,2), (4,5), (7,8), (10,11), (13,14), (16,17), (18,19), (20,21), # body joints
(22,37), (23,38), (24,39), (25,40), (26,41), (27,42), (28,43), (29,44), (30,45), (31,46), (32,47), (33,48), (34,49), (35,50), (36,51) # hand joints
)
self.orig_root_joint_idx = self.orig_joints_name.index('Pelvis')
self.orig_joint_part = \
{'body': range(self.orig_joints_name.index('Pelvis'), self.orig_joints_name.index('R_Wrist')+1),
'lhand': range(self.orig_joints_name.index('L_Index_1'), self.orig_joints_name.index('L_Thumb_3')+1),
'rhand': range(self.orig_joints_name.index('R_Index_1'), self.orig_joints_name.index('R_Thumb_3')+1),
'face': range(self.orig_joints_name.index('Jaw'), self.orig_joints_name.index('Jaw')+1)}
# changed SMPLX joint set for the supervision
self.joint_num = 137 # 25 (body joints) + 40 (hand joints) + 72 (face keypoints)
self.joints_name = \
('Pelvis', 'L_Hip', 'R_Hip', 'L_Knee', 'R_Knee', 'L_Ankle', 'R_Ankle', 'Neck', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', 'L_Big_toe', 'L_Small_toe', 'L_Heel', 'R_Big_toe', 'R_Small_toe', 'R_Heel', 'L_Ear', 'R_Ear', 'L_Eye', 'R_Eye', 'Nose',# body joints
'L_Thumb_1', 'L_Thumb_2', 'L_Thumb_3', 'L_Thumb_4', 'L_Index_1', 'L_Index_2', 'L_Index_3', 'L_Index_4', 'L_Middle_1', 'L_Middle_2', 'L_Middle_3', 'L_Middle_4', 'L_Ring_1', 'L_Ring_2', 'L_Ring_3', 'L_Ring_4', 'L_Pinky_1', 'L_Pinky_2', 'L_Pinky_3', 'L_Pinky_4', # left hand joints
'R_Thumb_1', 'R_Thumb_2', 'R_Thumb_3', 'R_Thumb_4', 'R_Index_1', 'R_Index_2', 'R_Index_3', 'R_Index_4', 'R_Middle_1', 'R_Middle_2', 'R_Middle_3', 'R_Middle_4', 'R_Ring_1', 'R_Ring_2', 'R_Ring_3', 'R_Ring_4', 'R_Pinky_1', 'R_Pinky_2', 'R_Pinky_3', 'R_Pinky_4', # right hand joints
*['Face_' + str(i) for i in range(1,73)] # face keypoints (too many keypoints... omit real names. have same name of keypoints defined in FLAME class)
)
self.root_joint_idx = self.joints_name.index('Pelvis')
self.lwrist_idx = self.joints_name.index('L_Wrist')
self.rwrist_idx = self.joints_name.index('R_Wrist')
self.neck_idx = self.joints_name.index('Neck')
self.flip_pairs = \
( (1,2), (3,4), (5,6), (8,9), (10,11), (12,13), (14,17), (15,18), (16,19), (20,21), (22,23), # body joints
(25,45), (26,46), (27,47), (28,48), (29,49), (30,50), (31,51), (32,52), (33,53), (34,54), (35,55), (36,56), (37,57), (38,58), (39,59), (40,60), (41,61), (42,62), (43,63), (44,64), # hand joints
(67,68), # face eyeballs
(69,78), (70,77), (71,76), (72,75), (73,74), # face eyebrow
(83,87), (84,86), # face below nose
(88,97), (89,96), (90,95), (91,94), (92,99), (93,98), # face eyes
(100,106), (101,105), (102,104), (107,111), (108,110), # face mouth
(112,116), (113,115), (117,119), # face lip
(120,136), (121,135), (122,134), (123,133), (124,132), (125,131), (126,130), (127,129) # face contours
)
self.joint_idx = \
(0,1,2,4,5,7,8,12,16,17,18,19,20,21,60,61,62,63,64,65,59,58,57,56,55, # body joints
37,38,39,66,25,26,27,67,28,29,30,68,34,35,36,69,31,32,33,70, # left hand joints
52,53,54,71,40,41,42,72,43,44,45,73,49,50,51,74,46,47,48,75, # right hand joints
22,15, # jaw, head
57,56, # eyeballs
76,77,78,79,80,81,82,83,84,85, # eyebrow
86,87,88,89, # nose
90,91,92,93,94, # below nose
95,96,97,98,99,100,101,102,103,104,105,106, # eyes
107, # right mouth
108,109,110,111,112, # upper mouth
113, # left mouth
114,115,116,117,118, # lower mouth
119, # right lip
120,121,122, # upper lip
123, # left lip
124,125,126, # lower lip
127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143 # face contour
)
self.joint_part = \
{'body': range(self.joints_name.index('Pelvis'), self.joints_name.index('Nose')+1),
'lhand': range(self.joints_name.index('L_Thumb_1'), self.joints_name.index('L_Pinky_4')+1),
'rhand': range(self.joints_name.index('R_Thumb_1'), self.joints_name.index('R_Pinky_4')+1),
'hand': range(self.joints_name.index('L_Thumb_1'), self.joints_name.index('R_Pinky_4')+1),
'face': range(self.joints_name.index('Face_1'), self.joints_name.index('Face_72')+1)}
# changed SMPLX joint set for PositionNet prediction
self.pos_joint_num = 65 # 25 (body joints) + 40 (hand joints)
self.pos_joints_name = \
('Pelvis', 'L_Hip', 'R_Hip', 'L_Knee', 'R_Knee', 'L_Ankle', 'R_Ankle', 'Neck', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', 'L_Big_toe', 'L_Small_toe', 'L_Heel', 'R_Big_toe', 'R_Small_toe', 'R_Heel', 'L_Ear', 'R_Ear', 'L_Eye', 'R_Eye', 'Nose', # body joints
'L_Thumb_1', 'L_Thumb_2', 'L_Thumb_3', 'L_Thumb_4', 'L_Index_1', 'L_Index_2', 'L_Index_3', 'L_Index_4', 'L_Middle_1', 'L_Middle_2', 'L_Middle_3', 'L_Middle_4', 'L_Ring_1', 'L_Ring_2', 'L_Ring_3', 'L_Ring_4', 'L_Pinky_1', 'L_Pinky_2', 'L_Pinky_3', 'L_Pinky_4', # left hand joints
'R_Thumb_1', 'R_Thumb_2', 'R_Thumb_3', 'R_Thumb_4', 'R_Index_1', 'R_Index_2', 'R_Index_3', 'R_Index_4', 'R_Middle_1', 'R_Middle_2', 'R_Middle_3', 'R_Middle_4', 'R_Ring_1', 'R_Ring_2', 'R_Ring_3', 'R_Ring_4', 'R_Pinky_1', 'R_Pinky_2', 'R_Pinky_3', 'R_Pinky_4', # right hand joints
)
self.pos_joint_part = \
{'body': range(self.pos_joints_name.index('Pelvis'), self.pos_joints_name.index('Nose')+1),
'lhand': range(self.pos_joints_name.index('L_Thumb_1'), self.pos_joints_name.index('L_Pinky_4')+1),
'rhand': range(self.pos_joints_name.index('R_Thumb_1'), self.pos_joints_name.index('R_Pinky_4')+1),
'hand': range(self.pos_joints_name.index('L_Thumb_1'), self.pos_joints_name.index('R_Pinky_4')+1)}
self.pos_joint_part['L_MCP'] = [self.pos_joints_name.index('L_Index_1') - len(self.pos_joint_part['body']),
self.pos_joints_name.index('L_Middle_1') - len(self.pos_joint_part['body']),
self.pos_joints_name.index('L_Ring_1') - len(self.pos_joint_part['body']),
self.pos_joints_name.index('L_Pinky_1') - len(self.pos_joint_part['body'])]
self.pos_joint_part['R_MCP'] = [self.pos_joints_name.index('R_Index_1') - len(self.pos_joint_part['body']) - len(self.pos_joint_part['lhand']),
self.pos_joints_name.index('R_Middle_1') - len(self.pos_joint_part['body']) - len(self.pos_joint_part['lhand']),
self.pos_joints_name.index('R_Ring_1') - len(self.pos_joint_part['body']) - len(self.pos_joint_part['lhand']),
self.pos_joints_name.index('R_Pinky_1') - len(self.pos_joint_part['body']) - len(self.pos_joint_part['lhand'])]
def make_hand_regressor(self):
regressor = self.layer['neutral'].J_regressor.numpy()
lhand_regressor = np.concatenate((regressor[[20,37,38,39],:],
np.eye(self.vertex_num)[5361,None],
regressor[[25,26,27],:],
np.eye(self.vertex_num)[4933,None],
regressor[[28,29,30],:],
np.eye(self.vertex_num)[5058,None],
regressor[[34,35,36],:],
np.eye(self.vertex_num)[5169,None],
regressor[[31,32,33],:],
np.eye(self.vertex_num)[5286,None]))
rhand_regressor = np.concatenate((regressor[[21,52,53,54],:],
np.eye(self.vertex_num)[8079,None],
regressor[[40,41,42],:],
np.eye(self.vertex_num)[7669,None],
regressor[[43,44,45],:],
np.eye(self.vertex_num)[7794,None],
regressor[[49,50,51],:],
np.eye(self.vertex_num)[7905,None],
regressor[[46,47,48],:],
np.eye(self.vertex_num)[8022,None]))
hand_regressor = {'left': lhand_regressor, 'right': rhand_regressor}
return hand_regressor
def reduce_joint_set(self, joint):
new_joint = []
for name in self.pos_joints_name:
idx = self.joints_name.index(name)
new_joint.append(joint[:,idx,:])
new_joint = torch.stack(new_joint,1)
return new_joint
class SMPL(object):
def __init__(self):
self.layer_arg = {'create_body_pose': False, 'create_betas': False, 'create_global_orient': False, 'create_transl': False}
self.layer = {'neutral': smplx.create(cfg.human_model_path, 'smpl', gender='NEUTRAL', **self.layer_arg), 'male': smplx.create(cfg.human_model_path, 'smpl', gender='MALE', **self.layer_arg), 'female': smplx.create(cfg.human_model_path, 'smpl', gender='FEMALE', **self.layer_arg)}
self.vertex_num = 6890
self.face = self.layer['neutral'].faces
self.shape_param_dim = 10
self.vposer_code_dim = 32
# original SMPL joint set
self.orig_joint_num = 24
self.orig_joints_name = ('Pelvis', 'L_Hip', 'R_Hip', 'Spine_1', 'L_Knee', 'R_Knee', 'Spine_2', 'L_Ankle', 'R_Ankle', 'Spine_3', 'L_Foot', 'R_Foot', 'Neck', 'L_Collar', 'R_Collar', 'Head', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', 'L_Hand', 'R_Hand')
self.orig_flip_pairs = ( (1,2), (4,5), (7,8), (10,11), (13,14), (16,17), (18,19), (20,21), (22,23) )
self.orig_root_joint_idx = self.orig_joints_name.index('Pelvis')
self.orig_joint_regressor = self.layer['neutral'].J_regressor.numpy().astype(np.float32)
self.joint_num = self.orig_joint_num
self.joints_name = self.orig_joints_name
self.flip_pairs = self.orig_flip_pairs
self.root_joint_idx = self.orig_root_joint_idx
self.joint_regressor = self.orig_joint_regressor
smpl_x = SMPLX()
smpl = SMPL()
common/utils/preprocessing.py
+566
View File
@@ -0,0 +1,566 @@
import numpy as np
import cv2
import random
from config import cfg
import math
from utils.human_models import smpl_x, smpl
from utils.transforms import cam2pixel, transform_joint_to_other_db
from plyfile import PlyData, PlyElement
import torch
def load_img(path, order='RGB'):
img = cv2.imread(path, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION)
if not isinstance(img, np.ndarray):
raise IOError("Fail to read %s" % path)
if order == 'RGB':
img = img[:, :, ::-1].copy()
img = img.astype(np.float32)
return img
def get_bbox(joint_img, joint_valid, extend_ratio=1.2):
x_img, y_img = joint_img[:, 0], joint_img[:, 1]
x_img = x_img[joint_valid == 1];
y_img = y_img[joint_valid == 1];
xmin = min(x_img);
ymin = min(y_img);
xmax = max(x_img);
ymax = max(y_img);
x_center = (xmin + xmax) / 2.;
width = xmax - xmin;
xmin = x_center - 0.5 * width * extend_ratio
xmax = x_center + 0.5 * width * extend_ratio
y_center = (ymin + ymax) / 2.;
height = ymax - ymin;
ymin = y_center - 0.5 * height * extend_ratio
ymax = y_center + 0.5 * height * extend_ratio
bbox = np.array([xmin, ymin, xmax - xmin, ymax - ymin]).astype(np.float32)
return bbox
def sanitize_bbox(bbox, img_width, img_height):
x, y, w, h = bbox
x1 = np.max((0, x))
y1 = np.max((0, y))
x2 = np.min((img_width - 1, x1 + np.max((0, w - 1))))
y2 = np.min((img_height - 1, y1 + np.max((0, h - 1))))
if w * h > 0 and x2 > x1 and y2 > y1:
bbox = np.array([x1, y1, x2 - x1, y2 - y1])
else:
bbox = None
return bbox
def process_bbox(bbox, img_width, img_height, ratio=1.25):
bbox = sanitize_bbox(bbox, img_width, img_height)
if bbox is None:
return bbox
# aspect ratio preserving bbox
w = bbox[2]
h = bbox[3]
c_x = bbox[0] + w / 2.
c_y = bbox[1] + h / 2.
aspect_ratio = cfg.input_img_shape[1] / cfg.input_img_shape[0]
if w > aspect_ratio * h:
h = w / aspect_ratio
elif w < aspect_ratio * h:
w = h * aspect_ratio
bbox[2] = w * ratio
bbox[3] = h * ratio
bbox[0] = c_x - bbox[2] / 2.
bbox[1] = c_y - bbox[3] / 2.
bbox = bbox.astype(np.float32)
return bbox
def get_aug_config():
scale_factor = 0.25
rot_factor = 30
color_factor = 0.2
scale = np.clip(np.random.randn(), -1.0, 1.0) * scale_factor + 1.0
rot = np.clip(np.random.randn(), -2.0,
2.0) * rot_factor if random.random() <= 0.6 else 0
c_up = 1.0 + color_factor
c_low = 1.0 - color_factor
color_scale = np.array([random.uniform(c_low, c_up), random.uniform(c_low, c_up), random.uniform(c_low, c_up)])
do_flip = random.random() <= 0.5
return scale, rot, color_scale, do_flip
def augmentation(img, bbox, data_split):
if getattr(cfg, 'no_aug', False):
scale, rot, color_scale, do_flip = 1.0, 0.0, np.array([1, 1, 1]), False
elif data_split == 'train':
scale, rot, color_scale, do_flip = get_aug_config()
else:
scale, rot, color_scale, do_flip = 1.0, 0.0, np.array([1, 1, 1]), False
img, trans, inv_trans = generate_patch_image(img, bbox, scale, rot, do_flip, cfg.input_img_shape)
img = np.clip(img * color_scale[None, None, :], 0, 255)
return img, trans, inv_trans, rot, do_flip
def generate_patch_image(cvimg, bbox, scale, rot, do_flip, out_shape):
img = cvimg.copy()
img_height, img_width, img_channels = img.shape
bb_c_x = float(bbox[0] + 0.5 * bbox[2])
bb_c_y = float(bbox[1] + 0.5 * bbox[3])
bb_width = float(bbox[2])
bb_height = float(bbox[3])
if do_flip:
img = img[:, ::-1, :]
bb_c_x = img_width - bb_c_x - 1
trans = gen_trans_from_patch_cv(bb_c_x, bb_c_y, bb_width, bb_height, out_shape[1], out_shape[0], scale, rot)
img_patch = cv2.warpAffine(img, trans, (int(out_shape[1]), int(out_shape[0])), flags=cv2.INTER_LINEAR)
img_patch = img_patch.astype(np.float32)
inv_trans = gen_trans_from_patch_cv(bb_c_x, bb_c_y, bb_width, bb_height, out_shape[1], out_shape[0], scale, rot,
inv=True)
return img_patch, trans, inv_trans
def rotate_2d(pt_2d, rot_rad):
x = pt_2d[0]
y = pt_2d[1]
sn, cs = np.sin(rot_rad), np.cos(rot_rad)
xx = x * cs - y * sn
yy = x * sn + y * cs
return np.array([xx, yy], dtype=np.float32)
def gen_trans_from_patch_cv(c_x, c_y, src_width, src_height, dst_width, dst_height, scale, rot, inv=False):
# augment size with scale
src_w = src_width * scale
src_h = src_height * scale
src_center = np.array([c_x, c_y], dtype=np.float32)
# augment rotation
rot_rad = np.pi * rot / 180
src_downdir = rotate_2d(np.array([0, src_h * 0.5], dtype=np.float32), rot_rad)
src_rightdir = rotate_2d(np.array([src_w * 0.5, 0], dtype=np.float32), rot_rad)
dst_w = dst_width
dst_h = dst_height
dst_center = np.array([dst_w * 0.5, dst_h * 0.5], dtype=np.float32)
dst_downdir = np.array([0, dst_h * 0.5], dtype=np.float32)
dst_rightdir = np.array([dst_w * 0.5, 0], dtype=np.float32)
src = np.zeros((3, 2), dtype=np.float32)
src[0, :] = src_center
src[1, :] = src_center + src_downdir
src[2, :] = src_center + src_rightdir
dst = np.zeros((3, 2), dtype=np.float32)
dst[0, :] = dst_center
dst[1, :] = dst_center + dst_downdir
dst[2, :] = dst_center + dst_rightdir
if inv:
trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
else:
trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
trans = trans.astype(np.float32)
return trans
def process_db_coord(joint_img, joint_cam, joint_valid, do_flip, img_shape, flip_pairs, img2bb_trans, rot,
src_joints_name, target_joints_name):
joint_img_original = joint_img.copy()
joint_img, joint_cam, joint_valid = joint_img.copy(), joint_cam.copy(), joint_valid.copy()
# flip augmentation
if do_flip:
joint_cam[:, 0] = -joint_cam[:, 0]
joint_img[:, 0] = img_shape[1] - 1 - joint_img[:, 0]
for pair in flip_pairs:
joint_img[pair[0], :], joint_img[pair[1], :] = joint_img[pair[1], :].copy(), joint_img[pair[0], :].copy()
joint_cam[pair[0], :], joint_cam[pair[1], :] = joint_cam[pair[1], :].copy(), joint_cam[pair[0], :].copy()
joint_valid[pair[0], :], joint_valid[pair[1], :] = joint_valid[pair[1], :].copy(), joint_valid[pair[0],
:].copy()
# 3D data rotation augmentation
rot_aug_mat = 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]], dtype=np.float32)
joint_cam = np.dot(rot_aug_mat, joint_cam.transpose(1, 0)).transpose(1, 0)
# affine transformation
joint_img_xy1 = np.concatenate((joint_img[:, :2], np.ones_like(joint_img[:, :1])), 1)
joint_img[:, :2] = np.dot(img2bb_trans, joint_img_xy1.transpose(1, 0)).transpose(1, 0)
joint_img[:, 0] = joint_img[:, 0] / cfg.input_img_shape[1] * cfg.output_hm_shape[2]
joint_img[:, 1] = joint_img[:, 1] / cfg.input_img_shape[0] * cfg.output_hm_shape[1]
# check truncation
# TODO
joint_trunc = joint_valid * ((joint_img_original[:, 0] > 0) * (joint_img[:, 0] >= 0) * (joint_img[:, 0] < cfg.output_hm_shape[2]) * \
(joint_img_original[:, 1] > 0) *(joint_img[:, 1] >= 0) * (joint_img[:, 1] < cfg.output_hm_shape[1]) * \
(joint_img_original[:, 2] > 0) *(joint_img[:, 2] >= 0) * (joint_img[:, 2] < cfg.output_hm_shape[0])).reshape(-1,
1).astype(
np.float32)
# transform joints to target db joints
joint_img = transform_joint_to_other_db(joint_img, src_joints_name, target_joints_name)
joint_cam_wo_ra = transform_joint_to_other_db(joint_cam, src_joints_name, target_joints_name)
joint_valid = transform_joint_to_other_db(joint_valid, src_joints_name, target_joints_name)
joint_trunc = transform_joint_to_other_db(joint_trunc, src_joints_name, target_joints_name)
# root-alignment, for joint_cam input wo ra
joint_cam_ra = joint_cam_wo_ra.copy()
joint_cam_ra = joint_cam_ra - joint_cam_ra[smpl_x.root_joint_idx, None, :] # root-relative
joint_cam_ra[smpl_x.joint_part['lhand'], :] = joint_cam_ra[smpl_x.joint_part['lhand'], :] - joint_cam_ra[
smpl_x.lwrist_idx, None,
:] # left hand root-relative
joint_cam_ra[smpl_x.joint_part['rhand'], :] = joint_cam_ra[smpl_x.joint_part['rhand'], :] - joint_cam_ra[
smpl_x.rwrist_idx, None,
:] # right hand root-relative
joint_cam_ra[smpl_x.joint_part['face'], :] = joint_cam_ra[smpl_x.joint_part['face'], :] - joint_cam_ra[smpl_x.neck_idx,
None,
:] # face root-relative
return joint_img, joint_cam_wo_ra, joint_cam_ra, joint_valid, joint_trunc
def process_human_model_output(human_model_param, cam_param, do_flip, img_shape, img2bb_trans, rot, human_model_type, joint_img=None):
if human_model_type == 'smplx':
human_model = smpl_x
rotation_valid = np.ones((smpl_x.orig_joint_num), dtype=np.float32)
coord_valid = np.ones((smpl_x.joint_num), dtype=np.float32)
root_pose, body_pose, shape, trans = human_model_param['root_pose'], human_model_param['body_pose'], \
human_model_param['shape'], human_model_param['trans']
if 'lhand_pose' in human_model_param and human_model_param['lhand_valid']:
lhand_pose = human_model_param['lhand_pose']
else:
lhand_pose = np.zeros((3 * len(smpl_x.orig_joint_part['lhand'])), dtype=np.float32)
rotation_valid[smpl_x.orig_joint_part['lhand']] = 0
coord_valid[smpl_x.joint_part['lhand']] = 0
if 'rhand_pose' in human_model_param and human_model_param['rhand_valid']:
rhand_pose = human_model_param['rhand_pose']
else:
rhand_pose = np.zeros((3 * len(smpl_x.orig_joint_part['rhand'])), dtype=np.float32)
rotation_valid[smpl_x.orig_joint_part['rhand']] = 0
coord_valid[smpl_x.joint_part['rhand']] = 0
if 'jaw_pose' in human_model_param and 'expr' in human_model_param and human_model_param['face_valid']:
jaw_pose = human_model_param['jaw_pose']
expr = human_model_param['expr']
expr_valid = True
else:
jaw_pose = np.zeros((3), dtype=np.float32)
expr = np.zeros((smpl_x.expr_code_dim), dtype=np.float32)
rotation_valid[smpl_x.orig_joint_part['face']] = 0
coord_valid[smpl_x.joint_part['face']] = 0
expr_valid = False
if 'gender' in human_model_param:
gender = human_model_param['gender']
else:
gender = 'neutral'
root_pose = torch.FloatTensor(root_pose).view(1, 3) # (1,3)
body_pose = torch.FloatTensor(body_pose).view(-1, 3) # (21,3)
lhand_pose = torch.FloatTensor(lhand_pose).view(-1, 3) # (15,3)
rhand_pose = torch.FloatTensor(rhand_pose).view(-1, 3) # (15,3)
jaw_pose = torch.FloatTensor(jaw_pose).view(-1, 3) # (1,3)
shape = torch.FloatTensor(shape).view(1, -1) # SMPLX shape parameter
expr = torch.FloatTensor(expr).view(1, -1) # SMPLX expression parameter
trans = torch.FloatTensor(trans).view(1, -1) # translation vector
# apply camera extrinsic (rotation)
# merge root pose and camera rotation
if 'R' in cam_param:
R = np.array(cam_param['R'], dtype=np.float32).reshape(3, 3)
root_pose = root_pose.numpy()
root_pose, _ = cv2.Rodrigues(root_pose)
root_pose, _ = cv2.Rodrigues(np.dot(R, root_pose))
root_pose = torch.from_numpy(root_pose).view(1, 3)
# get mesh and joint coordinates
zero_pose = torch.zeros((1, 3)).float() # eye poses
with torch.no_grad():
output = smpl_x.layer[gender](betas=shape, body_pose=body_pose.view(1, -1), global_orient=root_pose,
transl=trans, left_hand_pose=lhand_pose.view(1, -1),
right_hand_pose=rhand_pose.view(1, -1), jaw_pose=jaw_pose.view(1, -1),
leye_pose=zero_pose, reye_pose=zero_pose, expression=expr)
mesh_cam = output.vertices[0].numpy()
joint_cam = output.joints[0].numpy()[smpl_x.joint_idx, :]
### HARDCODE
# joint_cam_orig_ = joint_cam.copy()
# apply camera exrinsic (translation)
# compenstate rotation (translation from origin to root joint was not cancled)
if 'R' in cam_param and 't' in cam_param:
R, t = np.array(cam_param['R'], dtype=np.float32).reshape(3, 3), np.array(cam_param['t'],
dtype=np.float32).reshape(1, 3)
root_cam = joint_cam[smpl_x.root_joint_idx, None, :]
joint_cam = joint_cam - root_cam + np.dot(R, root_cam.transpose(1, 0)).transpose(1, 0) + t
mesh_cam = mesh_cam - root_cam + np.dot(R, root_cam.transpose(1, 0)).transpose(1, 0) + t
# concat root, body, two hands, and jaw pose
pose = torch.cat((root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose))
# joint coordinates
if 'focal' not in cam_param or 'princpt' not in cam_param:
assert joint_img is not None
else:
joint_img = cam2pixel(joint_cam, cam_param['focal'], cam_param['princpt'])
joint_img_original = joint_img.copy()
joint_cam = joint_cam - joint_cam[smpl_x.root_joint_idx, None, :] # root-relative
joint_cam[smpl_x.joint_part['lhand'], :] = joint_cam[smpl_x.joint_part['lhand'], :] - joint_cam[
smpl_x.lwrist_idx, None,
:] # left hand root-relative
joint_cam[smpl_x.joint_part['rhand'], :] = joint_cam[smpl_x.joint_part['rhand'], :] - joint_cam[
smpl_x.rwrist_idx, None,
:] # right hand root-relative
joint_cam[smpl_x.joint_part['face'], :] = joint_cam[smpl_x.joint_part['face'], :] - joint_cam[smpl_x.neck_idx,
None,
:] # face root-relative
joint_img[smpl_x.joint_part['body'], 2] = (joint_cam[smpl_x.joint_part['body'], 2].copy() / (
cfg.body_3d_size / 2) + 1) / 2. * cfg.output_hm_shape[0] # body depth discretize
joint_img[smpl_x.joint_part['lhand'], 2] = (joint_cam[smpl_x.joint_part['lhand'], 2].copy() / (
cfg.hand_3d_size / 2) + 1) / 2. * cfg.output_hm_shape[0] # left hand depth discretize
joint_img[smpl_x.joint_part['rhand'], 2] = (joint_cam[smpl_x.joint_part['rhand'], 2].copy() / (
cfg.hand_3d_size / 2) + 1) / 2. * cfg.output_hm_shape[0] # right hand depth discretize
joint_img[smpl_x.joint_part['face'], 2] = (joint_cam[smpl_x.joint_part['face'], 2].copy() / (
cfg.face_3d_size / 2) + 1) / 2. * cfg.output_hm_shape[0] # face depth discretize
elif human_model_type == 'smpl':
human_model = smpl
pose, shape, trans = human_model_param['pose'], human_model_param['shape'], human_model_param['trans']
if 'gender' in human_model_param:
gender = human_model_param['gender']
else:
gender = 'neutral'
pose = torch.FloatTensor(pose).view(-1, 3)
shape = torch.FloatTensor(shape).view(1, -1);
trans = torch.FloatTensor(trans).view(1, -1) # translation vector
# apply camera extrinsic (rotation)
# merge root pose and camera rotation
if 'R' in cam_param:
R = np.array(cam_param['R'], dtype=np.float32).reshape(3, 3)
root_pose = pose[smpl.orig_root_joint_idx, :].numpy()
root_pose, _ = cv2.Rodrigues(root_pose)
root_pose, _ = cv2.Rodrigues(np.dot(R, root_pose))
pose[smpl.orig_root_joint_idx] = torch.from_numpy(root_pose).view(3)
# get mesh and joint coordinates
root_pose = pose[smpl.orig_root_joint_idx].view(1, 3)
body_pose = torch.cat((pose[:smpl.orig_root_joint_idx, :], pose[smpl.orig_root_joint_idx + 1:, :])).view(1, -1)
with torch.no_grad():
output = smpl.layer[gender](betas=shape, body_pose=body_pose, global_orient=root_pose, transl=trans)
mesh_cam = output.vertices[0].numpy()
joint_cam = np.dot(smpl.joint_regressor, mesh_cam)
# apply camera exrinsic (translation)
# compenstate rotation (translation from origin to root joint was not cancled)
if 'R' in cam_param and 't' in cam_param:
R, t = np.array(cam_param['R'], dtype=np.float32).reshape(3, 3), np.array(cam_param['t'],
dtype=np.float32).reshape(1, 3)
root_cam = joint_cam[smpl.root_joint_idx, None, :]
joint_cam = joint_cam - root_cam + np.dot(R, root_cam.transpose(1, 0)).transpose(1, 0) + t
mesh_cam = mesh_cam - root_cam + np.dot(R, root_cam.transpose(1, 0)).transpose(1, 0) + t
# joint coordinates
if 'focal' not in cam_param or 'princpt' not in cam_param:
assert joint_img is not None
else:
joint_img = cam2pixel(joint_cam, cam_param['focal'], cam_param['princpt'])
joint_img_original = joint_img.copy()
joint_cam = joint_cam - joint_cam[smpl.root_joint_idx, None, :] # body root-relative
joint_img[:, 2] = (joint_cam[:, 2].copy() / (cfg.body_3d_size / 2) + 1) / 2. * cfg.output_hm_shape[
0] # body depth discretize
elif human_model_type == 'mano':
human_model = mano
pose, shape, trans = human_model_param['pose'], human_model_param['shape'], human_model_param['trans']
hand_type = human_model_param['hand_type']
pose = torch.FloatTensor(pose).view(-1, 3)
shape = torch.FloatTensor(shape).view(1, -1);
trans = torch.FloatTensor(trans).view(1, -1) # translation vector
# apply camera extrinsic (rotation)
# merge root pose and camera rotation
if 'R' in cam_param:
R = np.array(cam_param['R'], dtype=np.float32).reshape(3, 3)
root_pose = pose[mano.orig_root_joint_idx, :].numpy()
root_pose, _ = cv2.Rodrigues(root_pose)
root_pose, _ = cv2.Rodrigues(np.dot(R, root_pose))
pose[mano.orig_root_joint_idx] = torch.from_numpy(root_pose).view(3)
# get mesh and joint coordinates
root_pose = pose[mano.orig_root_joint_idx].view(1, 3)
hand_pose = torch.cat((pose[:mano.orig_root_joint_idx, :], pose[mano.orig_root_joint_idx + 1:, :])).view(1, -1)
with torch.no_grad():
output = mano.layer[hand_type](betas=shape, hand_pose=hand_pose, global_orient=root_pose, transl=trans)
mesh_cam = output.vertices[0].numpy()
joint_cam = np.dot(mano.joint_regressor, mesh_cam)
# apply camera exrinsic (translation)
# compenstate rotation (translation from origin to root joint was not cancled)
if 'R' in cam_param and 't' in cam_param:
R, t = np.array(cam_param['R'], dtype=np.float32).reshape(3, 3), np.array(cam_param['t'],
dtype=np.float32).reshape(1, 3)
root_cam = joint_cam[mano.root_joint_idx, None, :]
joint_cam = joint_cam - root_cam + np.dot(R, root_cam.transpose(1, 0)).transpose(1, 0) + t
mesh_cam = mesh_cam - root_cam + np.dot(R, root_cam.transpose(1, 0)).transpose(1, 0) + t
# joint coordinates
if 'focal' not in cam_param or 'princpt' not in cam_param:
assert joint_img is not None
else:
joint_img = cam2pixel(joint_cam, cam_param['focal'], cam_param['princpt'])
joint_cam = joint_cam - joint_cam[mano.root_joint_idx, None, :] # hand root-relative
joint_img[:, 2] = (joint_cam[:, 2].copy() / (cfg.hand_3d_size / 2) + 1) / 2. * cfg.output_hm_shape[
0] # hand depth discretize
mesh_cam_orig = mesh_cam.copy() # back-up the original one
## so far, data augmentations are not applied yet
## now, apply data augmentations
# image projection
if do_flip:
joint_cam[:, 0] = -joint_cam[:, 0]
joint_img[:, 0] = img_shape[1] - 1 - joint_img[:, 0]
for pair in human_model.flip_pairs:
joint_cam[pair[0], :], joint_cam[pair[1], :] = joint_cam[pair[1], :].copy(), joint_cam[pair[0], :].copy()
joint_img[pair[0], :], joint_img[pair[1], :] = joint_img[pair[1], :].copy(), joint_img[pair[0], :].copy()
if human_model_type == 'smplx':
coord_valid[pair[0]], coord_valid[pair[1]] = coord_valid[pair[1]].copy(), coord_valid[pair[0]].copy()
# x,y affine transform, root-relative depth
joint_img_xy1 = np.concatenate((joint_img[:, :2], np.ones_like(joint_img[:, 0:1])), 1)
joint_img[:, :2] = np.dot(img2bb_trans, joint_img_xy1.transpose(1, 0)).transpose(1, 0)[:, :2]
joint_img[:, 0] = joint_img[:, 0] / cfg.input_img_shape[1] * cfg.output_hm_shape[2]
joint_img[:, 1] = joint_img[:, 1] / cfg.input_img_shape[0] * cfg.output_hm_shape[1]
# check truncation
# TODO
joint_trunc = ((joint_img_original[:, 0] > 0) * (joint_img[:, 0] >= 0) * (joint_img[:, 0] < cfg.output_hm_shape[2]) * \
(joint_img_original[:, 1] > 0) * (joint_img[:, 1] >= 0) * (joint_img[:, 1] < cfg.output_hm_shape[1]) * \
(joint_img_original[:, 2] > 0) * (joint_img[:, 2] >= 0) * (joint_img[:, 2] < cfg.output_hm_shape[0])).reshape(-1, 1).astype(
np.float32)
# 3D data rotation augmentation
rot_aug_mat = 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]], dtype=np.float32)
# coordinate
joint_cam = np.dot(rot_aug_mat, joint_cam.transpose(1, 0)).transpose(1, 0)
# parameters
# flip pose parameter (axis-angle)
if do_flip:
for pair in human_model.orig_flip_pairs:
pose[pair[0], :], pose[pair[1], :] = pose[pair[1], :].clone(), pose[pair[0], :].clone()
if human_model_type == 'smplx':
rotation_valid[pair[0]], rotation_valid[pair[1]] = rotation_valid[pair[1]].copy(), rotation_valid[
pair[0]].copy()
pose[:, 1:3] *= -1 # multiply -1 to y and z axis of axis-angle
# rotate root pose
pose = pose.numpy()
root_pose = pose[human_model.orig_root_joint_idx, :]
root_pose, _ = cv2.Rodrigues(root_pose)
root_pose, _ = cv2.Rodrigues(np.dot(rot_aug_mat, root_pose))
pose[human_model.orig_root_joint_idx] = root_pose.reshape(3)
# change to mean shape if beta is too far from it
shape[(shape.abs() > 3).any(dim=1)] = 0.
shape = shape.numpy().reshape(-1)
# return results
if human_model_type == 'smplx':
pose = pose.reshape(-1)
expr = expr.numpy().reshape(-1)
### ### HARDCODE temp save vis for debug
# POSE: torch.cat((root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose))
# pose_ = pose.numpy().copy()
# root_pose = torch.FloatTensor(pose[0:3]).view(1, 3) # (1,3)
# body_pose = torch.FloatTensor(pose[3:66]).view(-1, 3) # (21,3)
# lhand_pose = torch.FloatTensor(pose[66:111]).view(-1, 3) # (15,3)
# rhand_pose = torch.FloatTensor(pose[111:156]).view(-1, 3) # (15,3)
# jaw_pose = torch.FloatTensor(pose[156:159]).view(-1, 3) # (1,3)
# shape = torch.FloatTensor(shape).view(1, -1) # SMPLX shape parameter
# expr = torch.FloatTensor(expr).view(1, -1) # SMPLX expression parameter
# trans = torch.FloatTensor(trans).view(1, -1) # translation vector
# with torch.no_grad():
# output = smpl_x.layer[gender](betas=shape, body_pose=body_pose.view(1, -1), global_orient=root_pose,
# transl=zero_pose, left_hand_pose=lhand_pose.view(1, -1),
# right_hand_pose=rhand_pose.view(1, -1), jaw_pose=jaw_pose.view(1, -1),
# leye_pose=zero_pose, reye_pose=zero_pose, expression=expr)
# mesh_rot = output.vertices[0].numpy()
# joint_rot = output.joints[0].numpy()[smpl_x.joint_idx, :]
# return mesh_rot, joint_cam, joint_trunc, pose, shape, expr, rotation_valid, coord_valid, expr_valid, mesh_cam_orig, joint_cam_orig_
return joint_img, joint_cam, joint_trunc, pose, shape, expr, rotation_valid, coord_valid, expr_valid, mesh_cam_orig
elif human_model_type == 'smpl':
pose = pose.reshape(-1)
return joint_img, joint_cam, joint_trunc, pose, shape, mesh_cam_orig
elif human_model_type == 'mano':
pose = pose.reshape(-1)
return joint_img, joint_cam, joint_trunc, pose, shape, mesh_cam_orig
def get_fitting_error_3D(db_joint, db_joint_from_fit, joint_valid):
# mask coordinate
db_joint = db_joint[np.tile(joint_valid, (1, 3)) == 1].reshape(-1, 3)
db_joint_from_fit = db_joint_from_fit[np.tile(joint_valid, (1, 3)) == 1].reshape(-1, 3)
db_joint_from_fit = db_joint_from_fit - np.mean(db_joint_from_fit, 0)[None, :] + np.mean(db_joint, 0)[None,
:] # translation alignment
error = np.sqrt(np.sum((db_joint - db_joint_from_fit) ** 2, 1)).mean()
return error
def load_obj(file_name):
v = []
obj_file = open(file_name)
for line in obj_file:
words = line.split(' ')
if words[0] == 'v':
x, y, z = float(words[1]), float(words[2]), float(words[3])
v.append(np.array([x, y, z]))
return np.stack(v)
def load_ply(file_name):
plydata = PlyData.read(file_name)
x = plydata['vertex']['x']
y = plydata['vertex']['y']
z = plydata['vertex']['z']
v = np.stack((x, y, z), 1)
return v
def resize_bbox(bbox, scale=1.2):
if isinstance(bbox, list):
x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3]
else:
x1, y1, x2, y2 = bbox
x_center = (x1+x2)/2.0
y_center = (y1+y2)/2.0
x_size, y_size = x2-x1, y2-y1
x1_resize = x_center-x_size/2.0*scale
x2_resize = x_center+x_size/2.0*scale
y1_resize = y_center - y_size / 2.0 * scale
y2_resize = y_center + y_size / 2.0 * scale
bbox[0], bbox[1], bbox[2], bbox[3] = x1_resize, y1_resize, x2_resize, y2_resize
return bbox
common/utils/smplx/LICENSE
+58
View File
@@ -0,0 +1,58 @@
License
Software Copyright License for non-commercial scientific research purposes
Please read carefully the following terms and conditions and any accompanying documentation before you download and/or use the SMPL-X/SMPLify-X model, data and software, (the "Model & Software"), including 3D meshes, blend weights, blend shapes, textures, software, scripts, and animations. By downloading and/or using the Model & Software (including downloading, cloning, installing, and any other use of this github repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Model & Software. Any infringement of the terms of this agreement will automatically terminate your rights under this License
Ownership / Licensees
The Software and the associated materials has been developed at the
Max Planck Institute for Intelligent Systems (hereinafter "MPI").
Any copyright or patent right is owned by and proprietary material of the
Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (hereinafter “MPG”; MPI and MPG hereinafter collectively “Max-Planck”)
hereinafter the “Licensor”.
License Grant
Licensor grants you (Licensee) personally a single-user, non-exclusive, non-transferable, free of charge right:
To install the Model & Software on computers owned, leased or otherwise controlled by you and/or your organization;
To use the Model & Software for the sole purpose of performing non-commercial scientific research, non-commercial education, or non-commercial artistic projects;
Any other use, in particular any use for commercial purposes, is prohibited. This includes, without limitation, incorporation in a commercial product, use in a commercial service, or production of other artifacts for commercial purposes. The Model & Software may not be reproduced, modified and/or made available in any form to any third party without Max-Planck’s prior written permission.
The Model & Software may not be used for pornographic purposes or to generate pornographic material whether commercial or not. This license also prohibits the use of the Model & Software to train methods/algorithms/neural networks/etc. for commercial use of any kind. By downloading the Model & Software, you agree not to reverse engineer it.
No Distribution
The Model & Software and the license herein granted shall not be copied, shared, distributed, re-sold, offered for re-sale, transferred or sub-licensed in whole or in part except that you may make one copy for archive purposes only.
Disclaimer of Representations and Warranties
You expressly acknowledge and agree that the Model & Software results from basic research, is provided “AS IS”, may contain errors, and that any use of the Model & Software is at your sole risk. LICENSOR MAKES NO REPRESENTATIONS OR WARRANTIES OF ANY KIND CONCERNING THE MODEL & SOFTWARE, NEITHER EXPRESS NOR IMPLIED, AND THE ABSENCE OF ANY LEGAL OR ACTUAL DEFECTS, WHETHER DISCOVERABLE OR NOT. Specifically, and not to limit the foregoing, licensor makes no representations or warranties (i) regarding the merchantability or fitness for a particular purpose of the Model & Software, (ii) that the use of the Model & Software will not infringe any patents, copyrights or other intellectual property rights of a third party, and (iii) that the use of the Model & Software will not cause any damage of any kind to you or a third party.
Limitation of Liability
Because this Model & Software License Agreement qualifies as a donation, according to Section 521 of the German Civil Code (Bürgerliches Gesetzbuch – BGB) Licensor as a donor is liable for intent and gross negligence only. If the Licensor fraudulently conceals a legal or material defect, they are obliged to compensate the Licensee for the resulting damage.
Licensor shall be liable for loss of data only up to the amount of typical recovery costs which would have arisen had proper and regular data backup measures been taken. For the avoidance of doubt Licensor shall be liable in accordance with the German Product Liability Act in the event of product liability. The foregoing applies also to Licensor’s legal representatives or assistants in performance. Any further liability shall be excluded.
Patent claims generated through the usage of the Model & Software cannot be directed towards the copyright holders.
The Model & Software is provided in the state of development the licensor defines. If modified or extended by Licensee, the Licensor makes no claims about the fitness of the Model & Software and is not responsible for any problems such modifications cause.
No Maintenance Services
You understand and agree that Licensor is under no obligation to provide either maintenance services, update services, notices of latent defects, or corrections of defects with regard to the Model & Software. Licensor nevertheless reserves the right to update, modify, or discontinue the Model & Software at any time.
Defects of the Model & Software must be notified in writing to the Licensor with a comprehensible description of the error symptoms. The notification of the defect should enable the reproduction of the error. The Licensee is encouraged to communicate any use, results, modification or publication.
Publications using the Model & Software
You acknowledge that the Model & Software is a valuable scientific resource and agree to appropriately reference the following paper in any publication making use of the Model & Software.
Citation:
@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}
}
Commercial licensing opportunities
For commercial uses of the Software, please send email to ps-license@tue.mpg.de
This Agreement shall be governed by the laws of the Federal Republic of Germany except for the UN Sales Convention.
common/utils/smplx/README.md
+186
View File
@@ -0,0 +1,186 @@
## SMPL-X: A new joint 3D model of the human body, face and hands together
[[Paper Page](https://smpl-x.is.tue.mpg.de)] [[Paper](https://ps.is.tuebingen.mpg.de/uploads_file/attachment/attachment/497/SMPL-X.pdf)]
[[Supp. Mat.](https://ps.is.tuebingen.mpg.de/uploads_file/attachment/attachment/498/SMPL-X-supp.pdf)]
![SMPL-X Examples](./images/teaser_fig.png)
## Table of Contents
* [License](#license)
* [Description](#description)
* [Installation](#installation)
* [Downloading the model](#downloading-the-model)
* [Loading SMPL-X, SMPL+H and SMPL](#loading-smpl-x-smplh-and-smpl)
* [SMPL and SMPL+H setup](#smpl-and-smplh-setup)
* [Model loading](https://github.com/vchoutas/smplx#model-loading)
* [MANO and FLAME correspondences](#mano-and-flame-correspondences)
* [Example](#example)
* [Citation](#citation)
* [Acknowledgments](#acknowledgments)
* [Contact](#contact)
## License
Software Copyright License for **non-commercial scientific research purposes**.
Please read carefully the [terms and conditions](https://github.com/vchoutas/smplx/blob/master/LICENSE) and any accompanying documentation before you download and/or use the SMPL-X/SMPLify-X model, data and software, (the "Model & Software"), including 3D meshes, blend weights, blend shapes, textures, software, scripts, and animations. By downloading and/or using the Model & Software (including downloading, cloning, installing, and any other use of this github repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Model & Software. Any infringement of the terms of this agreement will automatically terminate your rights under this [License](./LICENSE).
## Disclaimer
The original images used for the figures 1 and 2 of the paper can be found in this link.
The images in the paper are used under license from gettyimages.com.
We have acquired the right to use them in the publication, but redistribution is not allowed.
Please follow the instructions on the given link to acquire right of usage.
Our results are obtained on the 483 × 724 pixels resolution of the original images.
## Description
*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 = 10, 475 vertices and K = 54 joints,
which include joints for the neck, jaw, eyeballs and fingers.
SMPL-X is defined by a function M(θ, β, ψ), where θ is the pose parameters, β the shape parameters and
ψ the facial expression parameters.
## Installation
To install the model please follow the next steps in the specified order:
1. To install from PyPi simply run:
```Shell
pip install smplx[all]
```
2. Clone this repository and install it using the *setup.py* script:
```Shell
git clone https://github.com/vchoutas/smplx
python setup.py install
```
## Downloading the model
To download the *SMPL-X* model go to [this project website](https://smpl-x.is.tue.mpg.de) and register to get access to the downloads section.
To download the *SMPL+H* model go to [this project website](http://mano.is.tue.mpg.de) and register to get access to the downloads section.
To download the *SMPL* model go to [this](http://smpl.is.tue.mpg.de) (male and female models) and [this](http://smplify.is.tue.mpg.de) (gender neutral model) project website and register to get access to the downloads section.
## Loading SMPL-X, SMPL+H and SMPL
### SMPL and SMPL+H setup
The loader gives the option to use any of the SMPL-X, SMPL+H, SMPL, and MANO models. Depending on the model you want to use, please follow the respective download instructions. To switch between MANO, SMPL, SMPL+H and SMPL-X just change the *model_path* or *model_type* parameters. For more details please check the docs of the model classes.
Before using SMPL and SMPL+H you should follow the instructions in [tools/README.md](./tools/README.md) to remove the
Chumpy objects from both model pkls, as well as merge the MANO parameters with SMPL+H.
### Model loading
You can either use the [create](https://github.com/vchoutas/smplx/blob/c63c02b478c5c6f696491ed9167e3af6b08d89b1/smplx/body_models.py#L54)
function from [body_models](./smplx/body_models.py) or directly call the constructor for the
[SMPL](https://github.com/vchoutas/smplx/blob/c63c02b478c5c6f696491ed9167e3af6b08d89b1/smplx/body_models.py#L106),
[SMPL+H](https://github.com/vchoutas/smplx/blob/c63c02b478c5c6f696491ed9167e3af6b08d89b1/smplx/body_models.py#L395) and
[SMPL-X](https://github.com/vchoutas/smplx/blob/c63c02b478c5c6f696491ed9167e3af6b08d89b1/smplx/body_models.py#L628) model. The path to the model can either be the path to the file with the parameters or a directory with the following structure:
```bash
models
├── smpl
│   ├── SMPL_FEMALE.pkl
│   └── SMPL_MALE.pkl
│   └── SMPL_NEUTRAL.pkl
├── smplh
│   ├── SMPLH_FEMALE.pkl
│   └── SMPLH_MALE.pkl
├── mano
| ├── MANO_RIGHT.pkl
| └── MANO_LEFT.pkl
└── smplx
├── SMPLX_FEMALE.npz
├── SMPLX_FEMALE.pkl
├── SMPLX_MALE.npz
├── SMPLX_MALE.pkl
├── SMPLX_NEUTRAL.npz
└── SMPLX_NEUTRAL.pkl
```
## MANO and FLAME correspondences
The vertex correspondences between SMPL-X and MANO, FLAME can be downloaded
from [the project website](https://smpl-x.is.tue.mpg.de). If you have extracted
the correspondence data in the folder *correspondences*, then use the following
scripts to visualize them:
1. To view MANO correspondences run the following command:
```
python examples/vis_mano_vertices.py --model-folder $SMPLX_FOLDER --corr-fname correspondences/MANO_SMPLX_vertex_ids.pkl
```
2. To view FLAME correspondences run the following command:
```
python examples/vis_flame_vertices.py --model-folder $SMPLX_FOLDER --corr-fname correspondences/SMPL-X__FLAME_vertex_ids.npy
```
## Example
After installing the *smplx* package and downloading the model parameters you should be able to run the *demo.py*
script to visualize the results. For this step you have to install the [pyrender](https://pyrender.readthedocs.io/en/latest/index.html) and [trimesh](https://trimsh.org/) packages.
`python examples/demo.py --model-folder $SMPLX_FOLDER --plot-joints=True --gender="neutral"`
![SMPL-X Examples](./images/example.png)
## 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}
}
```
This repository was originally developed for SMPL-X / SMPLify-X (CVPR 2019), you might be interested in having a look: [https://smpl-x.is.tue.mpg.de](https://smpl-x.is.tue.mpg.de).
## Acknowledgments
### Facial Contour
Special thanks to [Soubhik Sanyal](https://github.com/soubhiksanyal) for sharing the Tensorflow code used for the facial
landmarks.
## Contact
The code of this repository was implemented by [Vassilis Choutas](vassilis.choutas@tuebingen.mpg.de).
For questions, please contact [smplx@tue.mpg.de](smplx@tue.mpg.de).
For commercial licensing (and all related questions for business applications), please contact [ps-licensing@tue.mpg.de](ps-licensing@tue.mpg.de).
common/utils/smplx/examples/demo.py
+180
View File
@@ -0,0 +1,180 @@
# -*- coding: utf-8 -*-
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
import os.path as osp
import argparse
import numpy as np
import torch
import smplx
def main(model_folder,
model_type='smplx',
ext='npz',
gender='neutral',
plot_joints=False,
num_betas=10,
sample_shape=True,
sample_expression=True,
num_expression_coeffs=10,
plotting_module='pyrender',
use_face_contour=False):
model = smplx.create(model_folder, model_type=model_type,
gender=gender, use_face_contour=use_face_contour,
num_betas=num_betas,
num_expression_coeffs=num_expression_coeffs,
ext=ext)
print(model)
betas, expression = None, None
if sample_shape:
betas = torch.randn([1, model.num_betas], dtype=torch.float32)
if sample_expression:
expression = torch.randn(
[1, model.num_expression_coeffs], dtype=torch.float32)
output = model(betas=betas, expression=expression,
return_verts=True)
vertices = output.vertices.detach().cpu().numpy().squeeze()
joints = output.joints.detach().cpu().numpy().squeeze()
print('Vertices shape =', vertices.shape)
print('Joints shape =', joints.shape)
if plotting_module == 'pyrender':
import pyrender
import trimesh
vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8]
tri_mesh = trimesh.Trimesh(vertices, model.faces,
vertex_colors=vertex_colors)
mesh = pyrender.Mesh.from_trimesh(tri_mesh)
scene = pyrender.Scene()
scene.add(mesh)
if plot_joints:
sm = trimesh.creation.uv_sphere(radius=0.005)
sm.visual.vertex_colors = [0.9, 0.1, 0.1, 1.0]
tfs = np.tile(np.eye(4), (len(joints), 1, 1))
tfs[:, :3, 3] = joints
joints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
scene.add(joints_pcl)
pyrender.Viewer(scene, use_raymond_lighting=True)
elif plotting_module == 'matplotlib':
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
mesh = Poly3DCollection(vertices[model.faces], alpha=0.1)
face_color = (1.0, 1.0, 0.9)
edge_color = (0, 0, 0)
mesh.set_edgecolor(edge_color)
mesh.set_facecolor(face_color)
ax.add_collection3d(mesh)
ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], color='r')
if plot_joints:
ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], alpha=0.1)
plt.show()
elif plotting_module == 'open3d':
import open3d as o3d
mesh = o3d.geometry.TriangleMesh()
mesh.vertices = o3d.utility.Vector3dVector(
vertices)
mesh.triangles = o3d.utility.Vector3iVector(model.faces)
mesh.compute_vertex_normals()
mesh.paint_uniform_color([0.3, 0.3, 0.3])
geometry = [mesh]
if plot_joints:
joints_pcl = o3d.geometry.PointCloud()
joints_pcl.points = o3d.utility.Vector3dVector(joints)
joints_pcl.paint_uniform_color([0.7, 0.3, 0.3])
geometry.append(joints_pcl)
o3d.visualization.draw_geometries(geometry)
else:
raise ValueError('Unknown plotting_module: {}'.format(plotting_module))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='SMPL-X Demo')
parser.add_argument('--model-folder', required=True, type=str,
help='The path to the model folder')
parser.add_argument('--model-type', default='smplx', type=str,
choices=['smpl', 'smplh', 'smplx', 'mano', 'flame'],
help='The type of model to load')
parser.add_argument('--gender', type=str, default='neutral',
help='The gender of the model')
parser.add_argument('--num-betas', default=10, type=int,
dest='num_betas',
help='Number of shape coefficients.')
parser.add_argument('--num-expression-coeffs', default=10, type=int,
dest='num_expression_coeffs',
help='Number of expression coefficients.')
parser.add_argument('--plotting-module', type=str, default='pyrender',
dest='plotting_module',
choices=['pyrender', 'matplotlib', 'open3d'],
help='The module to use for plotting the result')
parser.add_argument('--ext', type=str, default='npz',
help='Which extension to use for loading')
parser.add_argument('--plot-joints', default=False,
type=lambda arg: arg.lower() in ['true', '1'],
help='The path to the model folder')
parser.add_argument('--sample-shape', default=True,
dest='sample_shape',
type=lambda arg: arg.lower() in ['true', '1'],
help='Sample a random shape')
parser.add_argument('--sample-expression', default=True,
dest='sample_expression',
type=lambda arg: arg.lower() in ['true', '1'],
help='Sample a random expression')
parser.add_argument('--use-face-contour', default=False,
type=lambda arg: arg.lower() in ['true', '1'],
help='Compute the contour of the face')
args = parser.parse_args()
model_folder = osp.expanduser(osp.expandvars(args.model_folder))
model_type = args.model_type
plot_joints = args.plot_joints
use_face_contour = args.use_face_contour
gender = args.gender
ext = args.ext
plotting_module = args.plotting_module
num_betas = args.num_betas
num_expression_coeffs = args.num_expression_coeffs
sample_shape = args.sample_shape
sample_expression = args.sample_expression
main(model_folder, model_type, ext=ext,
gender=gender, plot_joints=plot_joints,
num_betas=num_betas,
num_expression_coeffs=num_expression_coeffs,
sample_shape=sample_shape,
sample_expression=sample_expression,
plotting_module=plotting_module,
use_face_contour=use_face_contour)
common/utils/smplx/examples/demo_layers.py
+181
View File
@@ -0,0 +1,181 @@
# -*- coding: utf-8 -*-
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
import os.path as osp
import argparse
import numpy as np
import torch
import smplx
def main(model_folder,
model_type='smplx',
ext='npz',
gender='neutral',
plot_joints=False,
num_betas=10,
sample_shape=True,
sample_expression=True,
num_expression_coeffs=10,
plotting_module='pyrender',
use_face_contour=False):
model = smplx.build_layer(
model_folder, model_type=model_type,
gender=gender, use_face_contour=use_face_contour,
num_betas=num_betas,
num_expression_coeffs=num_expression_coeffs,
ext=ext)
print(model)
betas, expression = None, None
if sample_shape:
betas = torch.randn([1, model.num_betas], dtype=torch.float32)
if sample_expression:
expression = torch.randn(
[1, model.num_expression_coeffs], dtype=torch.float32)
output = model(betas=betas, expression=expression,
return_verts=True)
vertices = output.vertices.detach().cpu().numpy().squeeze()
joints = output.joints.detach().cpu().numpy().squeeze()
print('Vertices shape =', vertices.shape)
print('Joints shape =', joints.shape)
if plotting_module == 'pyrender':
import pyrender
import trimesh
vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8]
tri_mesh = trimesh.Trimesh(vertices, model.faces,
vertex_colors=vertex_colors)
mesh = pyrender.Mesh.from_trimesh(tri_mesh)
scene = pyrender.Scene()
scene.add(mesh)
if plot_joints:
sm = trimesh.creation.uv_sphere(radius=0.005)
sm.visual.vertex_colors = [0.9, 0.1, 0.1, 1.0]
tfs = np.tile(np.eye(4), (len(joints), 1, 1))
tfs[:, :3, 3] = joints
joints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
scene.add(joints_pcl)
pyrender.Viewer(scene, use_raymond_lighting=True)
elif plotting_module == 'matplotlib':
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
mesh = Poly3DCollection(vertices[model.faces], alpha=0.1)
face_color = (1.0, 1.0, 0.9)
edge_color = (0, 0, 0)
mesh.set_edgecolor(edge_color)
mesh.set_facecolor(face_color)
ax.add_collection3d(mesh)
ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], color='r')
if plot_joints:
ax.scatter(joints[:, 0], joints[:, 1], joints[:, 2], alpha=0.1)
plt.show()
elif plotting_module == 'open3d':
import open3d as o3d
mesh = o3d.geometry.TriangleMesh()
mesh.vertices = o3d.utility.Vector3dVector(
vertices)
mesh.triangles = o3d.utility.Vector3iVector(model.faces)
mesh.compute_vertex_normals()
mesh.paint_uniform_color([0.3, 0.3, 0.3])
geometry = [mesh]
if plot_joints:
joints_pcl = o3d.geometry.PointCloud()
joints_pcl.points = o3d.utility.Vector3dVector(joints)
joints_pcl.paint_uniform_color([0.7, 0.3, 0.3])
geometry.append(joints_pcl)
o3d.visualization.draw_geometries(geometry)
else:
raise ValueError('Unknown plotting_module: {}'.format(plotting_module))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='SMPL-X Demo')
parser.add_argument('--model-folder', required=True, type=str,
help='The path to the model folder')
parser.add_argument('--model-type', default='smplx', type=str,
choices=['smpl', 'smplh', 'smplx', 'mano', 'flame'],
help='The type of model to load')
parser.add_argument('--gender', type=str, default='neutral',
help='The gender of the model')
parser.add_argument('--num-betas', default=10, type=int,
dest='num_betas',
help='Number of shape coefficients.')
parser.add_argument('--num-expression-coeffs', default=10, type=int,
dest='num_expression_coeffs',
help='Number of expression coefficients.')
parser.add_argument('--plotting-module', type=str, default='pyrender',
dest='plotting_module',
choices=['pyrender', 'matplotlib', 'open3d'],
help='The module to use for plotting the result')
parser.add_argument('--ext', type=str, default='npz',
help='Which extension to use for loading')
parser.add_argument('--plot-joints', default=False,
type=lambda arg: arg.lower() in ['true', '1'],
help='The path to the model folder')
parser.add_argument('--sample-shape', default=True,
dest='sample_shape',
type=lambda arg: arg.lower() in ['true', '1'],
help='Sample a random shape')
parser.add_argument('--sample-expression', default=True,
dest='sample_expression',
type=lambda arg: arg.lower() in ['true', '1'],
help='Sample a random expression')
parser.add_argument('--use-face-contour', default=False,
type=lambda arg: arg.lower() in ['true', '1'],
help='Compute the contour of the face')
args = parser.parse_args()
model_folder = osp.expanduser(osp.expandvars(args.model_folder))
model_type = args.model_type
plot_joints = args.plot_joints
use_face_contour = args.use_face_contour
gender = args.gender
ext = args.ext
plotting_module = args.plotting_module
num_betas = args.num_betas
num_expression_coeffs = args.num_expression_coeffs
sample_shape = args.sample_shape
sample_expression = args.sample_expression
main(model_folder, model_type, ext=ext,
gender=gender, plot_joints=plot_joints,
num_betas=num_betas,
num_expression_coeffs=num_expression_coeffs,
sample_shape=sample_shape,
sample_expression=sample_expression,
plotting_module=plotting_module,
use_face_contour=use_face_contour)
common/utils/smplx/examples/vis_flame_vertices.py
+92
View File
@@ -0,0 +1,92 @@
# -*- coding: utf-8 -*-
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
import os.path as osp
import argparse
import pickle
import numpy as np
import torch
import open3d as o3d
import smplx
def main(model_folder, corr_fname, ext='npz',
head_color=(0.3, 0.3, 0.6),
gender='neutral'):
head_idxs = np.load(corr_fname)
model = smplx.create(model_folder, model_type='smplx',
gender=gender,
ext=ext)
betas = torch.zeros([1, 10], dtype=torch.float32)
expression = torch.zeros([1, 10], dtype=torch.float32)
output = model(betas=betas, expression=expression,
return_verts=True)
vertices = output.vertices.detach().cpu().numpy().squeeze()
joints = output.joints.detach().cpu().numpy().squeeze()
print('Vertices shape =', vertices.shape)
print('Joints shape =', joints.shape)
mesh = o3d.geometry.TriangleMesh()
mesh.vertices = o3d.utility.Vector3dVector(vertices)
mesh.triangles = o3d.utility.Vector3iVector(model.faces)
mesh.compute_vertex_normals()
colors = np.ones_like(vertices) * [0.3, 0.3, 0.3]
colors[head_idxs] = head_color
mesh.vertex_colors = o3d.utility.Vector3dVector(colors)
o3d.visualization.draw_geometries([mesh])
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='SMPL-X Demo')
parser.add_argument('--model-folder', required=True, type=str,
help='The path to the model folder')
parser.add_argument('--corr-fname', required=True, type=str,
dest='corr_fname',
help='Filename with the head correspondences')
parser.add_argument('--gender', type=str, default='neutral',
help='The gender of the model')
parser.add_argument('--ext', type=str, default='npz',
help='Which extension to use for loading')
parser.add_argument('--head', default='right',
choices=['right', 'left'],
type=str, help='Which head to plot')
parser.add_argument('--head-color', type=float, nargs=3, dest='head_color',
default=(0.3, 0.3, 0.6),
help='Color for the head vertices')
args = parser.parse_args()
model_folder = osp.expanduser(osp.expandvars(args.model_folder))
corr_fname = args.corr_fname
gender = args.gender
ext = args.ext
head = args.head
head_color = args.head_color
main(model_folder, corr_fname, ext=ext,
head_color=head_color,
gender=gender
)
common/utils/smplx/examples/vis_mano_vertices.py
+99
View File
@@ -0,0 +1,99 @@
# -*- coding: utf-8 -*-
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
import os.path as osp
import argparse
import pickle
import numpy as np
import torch
import open3d as o3d
import smplx
def main(model_folder, corr_fname, ext='npz',
hand_color=(0.3, 0.3, 0.6),
gender='neutral', hand='right'):
with open(corr_fname, 'rb') as f:
idxs_data = pickle.load(f)
if hand == 'both':
hand_idxs = np.concatenate(
[idxs_data['left_hand'], idxs_data['right_hand']]
)
else:
hand_idxs = idxs_data[f'{hand}_hand']
model = smplx.create(model_folder, model_type='smplx',
gender=gender,
ext=ext)
betas = torch.zeros([1, 10], dtype=torch.float32)
expression = torch.zeros([1, 10], dtype=torch.float32)
output = model(betas=betas, expression=expression,
return_verts=True)
vertices = output.vertices.detach().cpu().numpy().squeeze()
joints = output.joints.detach().cpu().numpy().squeeze()
print('Vertices shape =', vertices.shape)
print('Joints shape =', joints.shape)
mesh = o3d.geometry.TriangleMesh()
mesh.vertices = o3d.utility.Vector3dVector(vertices)
mesh.triangles = o3d.utility.Vector3iVector(model.faces)
mesh.compute_vertex_normals()
colors = np.ones_like(vertices) * [0.3, 0.3, 0.3]
colors[hand_idxs] = hand_color
mesh.vertex_colors = o3d.utility.Vector3dVector(colors)
o3d.visualization.draw_geometries([mesh])
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='SMPL-X Demo')
parser.add_argument('--model-folder', required=True, type=str,
help='The path to the model folder')
parser.add_argument('--corr-fname', required=True, type=str,
dest='corr_fname',
help='Filename with the hand correspondences')
parser.add_argument('--gender', type=str, default='neutral',
help='The gender of the model')
parser.add_argument('--ext', type=str, default='npz',
help='Which extension to use for loading')
parser.add_argument('--hand', default='right',
choices=['right', 'left', 'both'],
type=str, help='Which hand to plot')
parser.add_argument('--hand-color', type=float, nargs=3, dest='hand_color',
default=(0.3, 0.3, 0.6),
help='Color for the hand vertices')
args = parser.parse_args()
model_folder = osp.expanduser(osp.expandvars(args.model_folder))
corr_fname = args.corr_fname
gender = args.gender
ext = args.ext
hand = args.hand
hand_color = args.hand_color
main(model_folder, corr_fname, ext=ext,
hand_color=hand_color,
gender=gender, hand=hand
)
common/utils/smplx/setup.py
+79
View File
@@ -0,0 +1,79 @@
# -*- 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 and the Max Planck Institute for Biological
# Cybernetics. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
import io
import os
from setuptools import setup
# Package meta-data.
NAME = 'smplx'
DESCRIPTION = 'PyTorch module for loading the SMPLX body model'
URL = 'http://smpl-x.is.tuebingen.mpg.de'
EMAIL = 'vassilis.choutas@tuebingen.mpg.de'
AUTHOR = 'Vassilis Choutas'
REQUIRES_PYTHON = '>=3.6.0'
VERSION = '0.1.21'
here = os.path.abspath(os.path.dirname(__file__))
try:
FileNotFoundError
except NameError:
FileNotFoundError = IOError
# Import the README and use it as the long-description.
# Note: this will only work if 'README.md' is present in your MANIFEST.in file!
try:
with io.open(os.path.join(here, 'README.md'), encoding='utf-8') as f:
long_description = '\n' + f.read()
except FileNotFoundError:
long_description = DESCRIPTION
# Load the package's __version__.py module as a dictionary.
about = {}
if not VERSION:
with open(os.path.join(here, NAME, '__version__.py')) as f:
exec(f.read(), about)
else:
about['__version__'] = VERSION
pyrender_reqs = ['pyrender>=0.1.23', 'trimesh>=2.37.6', 'shapely']
matplotlib_reqs = ['matplotlib']
open3d_reqs = ['open3d-python']
setup(name=NAME,
version=about['__version__'],
description=DESCRIPTION,
long_description=long_description,
long_description_content_type='text/markdown',
author=AUTHOR,
author_email=EMAIL,
python_requires=REQUIRES_PYTHON,
url=URL,
install_requires=[
'numpy>=1.16.2',
'torch>=1.0.1.post2',
'torchgeometry>=0.1.2'
],
extras_require={
'pyrender': pyrender_reqs,
'open3d': open3d_reqs,
'matplotlib': matplotlib_reqs,
'all': pyrender_reqs + matplotlib_reqs + open3d_reqs
},
packages=['smplx', 'tools'])
common/utils/smplx/smplx/__init__.py
+30
View File
@@ -0,0 +1,30 @@
# -*- 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,
)
common/utils/smplx/smplx/body_models.py
+2331
View File
File diff suppressed because it is too large Load Diff
common/utils/smplx/smplx/joint_names.py
+163
View File
@@ -0,0 +1,163 @@
# -*- 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',
]
common/utils/smplx/smplx/lbs.py
+404
View File
@@ -0,0 +1,404 @@
# -*- 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
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,
) -> Tuple[Tensor, 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]
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]
# 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
common/utils/smplx/smplx/utils.py
+125
View File
@@ -0,0 +1,125 @@
# -*- 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
@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)
common/utils/smplx/smplx/vertex_ids.py
+77
View File
@@ -0,0 +1,77 @@
# -*- 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,
}
}
common/utils/smplx/smplx/vertex_joint_selector.py
+77
View File
@@ -0,0 +1,77 @@
# -*- 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
common/utils/smplx/tools/README.md
+20
View File
@@ -0,0 +1,20 @@
## Removing Chumpy objects
In a Python 2 virtual environment with [Chumpy](https://github.com/mattloper/chumpy) installed run the following to remove any Chumpy objects from the model data:
```bash
python tools/clean_ch.py --input-models path-to-models/*.pkl --output-folder output-folder
```
## Merging SMPL-H and MANO parameters
In order to use the given PyTorch SMPL-H module we first need to merge the SMPL-H and MANO parameters in a single file. After agreeing to the license and downloading the models, run the following command:
```bash
python tools/merge_smplh_mano.py --smplh-fn SMPLH_FOLDER/SMPLH_GENDER.pkl \
--mano-left-fn MANO_FOLDER/MANO_LEFT.pkl \
--mano-right-fn MANO_FOLDER/MANO_RIGHT.pkl \
--output-folder OUTPUT_FOLDER
```
where SMPLH_FOLDER is the folder with the SMPL-H files and MANO_FOLDER the one for the MANO files.
common/utils/smplx/tools/__init__.py
+19
View File
@@ -0,0 +1,19 @@
# -*- 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 and the Max Planck Institute for Biological
# Cybernetics. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
import clean_ch
import merge_smplh_mano
common/utils/smplx/tools/clean_ch.py
+68
View File
@@ -0,0 +1,68 @@
# -*- 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 and the Max Planck Institute for Biological
# Cybernetics. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
from __future__ import print_function
from __future__ import absolute_import
from __future__ import division
import argparse
import os
import os.path as osp
import pickle
from tqdm import tqdm
import numpy as np
def clean_fn(fn, output_folder='output'):
with open(fn, 'rb') as body_file:
body_data = pickle.load(body_file)
output_dict = {}
for key, data in body_data.iteritems():
if 'chumpy' in str(type(data)):
output_dict[key] = np.array(data)
else:
output_dict[key] = data
out_fn = osp.split(fn)[1]
out_path = osp.join(output_folder, out_fn)
with open(out_path, 'wb') as out_file:
pickle.dump(output_dict, out_file)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--input-models', dest='input_models', nargs='+',
required=True, type=str,
help='The path to the model that will be processed')
parser.add_argument('--output-folder', dest='output_folder',
required=True, type=str,
help='The path to the output folder')
args = parser.parse_args()
input_models = args.input_models
output_folder = args.output_folder
if not osp.exists(output_folder):
print('Creating directory: {}'.format(output_folder))
os.makedirs(output_folder)
for input_model in input_models:
clean_fn(input_model, output_folder=output_folder)
common/utils/smplx/tools/merge_smplh_mano.py
+89
View File
@@ -0,0 +1,89 @@
# -*- 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 and the Max Planck Institute for Biological
# Cybernetics. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
from __future__ import print_function
import os
import os.path as osp
import pickle
import argparse
import numpy as np
def merge_models(smplh_fn, mano_left_fn, mano_right_fn,
output_folder='output'):
with open(smplh_fn, 'rb') as body_file:
body_data = pickle.load(body_file)
with open(mano_left_fn, 'rb') as lhand_file:
lhand_data = pickle.load(lhand_file)
with open(mano_right_fn, 'rb') as rhand_file:
rhand_data = pickle.load(rhand_file)
out_fn = osp.split(smplh_fn)[1]
output_data = body_data.copy()
output_data['hands_componentsl'] = lhand_data['hands_components']
output_data['hands_componentsr'] = rhand_data['hands_components']
output_data['hands_coeffsl'] = lhand_data['hands_coeffs']
output_data['hands_coeffsr'] = rhand_data['hands_coeffs']
output_data['hands_meanl'] = lhand_data['hands_mean']
output_data['hands_meanr'] = rhand_data['hands_mean']
for key, data in output_data.iteritems():
if 'chumpy' in str(type(data)):
output_data[key] = np.array(data)
else:
output_data[key] = data
out_path = osp.join(output_folder, out_fn)
print(out_path)
print('Saving to {}'.format(out_path))
with open(out_path, 'wb') as output_file:
pickle.dump(output_data, output_file)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--smplh-fn', dest='smplh_fn', required=True,
type=str, help='The path to the SMPLH model')
parser.add_argument('--mano-left-fn', dest='mano_left_fn', required=True,
type=str, help='The path to the left hand MANO model')
parser.add_argument('--mano-right-fn', dest='mano_right_fn', required=True,
type=str, help='The path to the right hand MANO model')
parser.add_argument('--output-folder', dest='output_folder',
required=True, type=str,
help='The path to the output folder')
args = parser.parse_args()
smplh_fn = args.smplh_fn
mano_left_fn = args.mano_left_fn
mano_right_fn = args.mano_right_fn
output_folder = args.output_folder
if not osp.exists(output_folder):
print('Creating directory: {}'.format(output_folder))
os.makedirs(output_folder)
merge_models(smplh_fn, mano_left_fn, mano_right_fn, output_folder)
common/utils/transforms.py
+172
View File
@@ -0,0 +1,172 @@
import torch
import numpy as np
import scipy
from config import cfg
from torch.nn import functional as F
import torchgeometry as tgm
def cam2pixel(cam_coord, f, c):
x = cam_coord[:, 0] / cam_coord[:, 2] * f[0] + c[0]
y = cam_coord[:, 1] / cam_coord[:, 2] * f[1] + c[1]
z = cam_coord[:, 2]
return np.stack((x, y, z), 1)
def pixel2cam(pixel_coord, f, c):
x = (pixel_coord[:, 0] - c[0]) / f[0] * pixel_coord[:, 2]
y = (pixel_coord[:, 1] - c[1]) / f[1] * pixel_coord[:, 2]
z = pixel_coord[:, 2]
return np.stack((x, y, z), 1)
def world2cam(world_coord, R, t):
cam_coord = np.dot(R, world_coord.transpose(1, 0)).transpose(1, 0) + t.reshape(1, 3)
return cam_coord
def cam2world(cam_coord, R, t):
world_coord = np.dot(np.linalg.inv(R), (cam_coord - t.reshape(1, 3)).transpose(1, 0)).transpose(1, 0)
return world_coord
def rigid_transform_3D(A, B):
n, dim = A.shape
centroid_A = np.mean(A, axis=0)
centroid_B = np.mean(B, axis=0)
H = np.dot(np.transpose(A - centroid_A), B - centroid_B) / n
U, s, V = np.linalg.svd(H)
R = np.dot(np.transpose(V), np.transpose(U))
if np.linalg.det(R) < 0:
s[-1] = -s[-1]
V[2] = -V[2]
R = np.dot(np.transpose(V), np.transpose(U))
varP = np.var(A, axis=0).sum()
c = 1 / varP * np.sum(s)
t = -np.dot(c * R, np.transpose(centroid_A)) + np.transpose(centroid_B)
return c, R, t
def rigid_align(A, B):
c, R, t = rigid_transform_3D(A, B)
A2 = np.transpose(np.dot(c * R, np.transpose(A))) + t
return A2
def transform_joint_to_other_db(src_joint, src_name, dst_name):
src_joint_num = len(src_name)
dst_joint_num = len(dst_name)
new_joint = np.zeros(((dst_joint_num,) + src_joint.shape[1:]), dtype=np.float32)
for src_idx in range(len(src_name)):
name = src_name[src_idx]
if name in dst_name:
dst_idx = dst_name.index(name)
new_joint[dst_idx] = src_joint[src_idx]
return new_joint
def rot6d_to_axis_angle(x):
batch_size = x.shape[0]
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)
rot_mat = torch.stack((b1, b2, b3), dim=-1) # 3x3 rotation matrix
rot_mat = torch.cat([rot_mat, torch.zeros((batch_size, 3, 1)).cuda().float()], 2) # 3x4 rotation matrix
axis_angle = tgm.rotation_matrix_to_angle_axis(rot_mat).reshape(-1, 3) # axis-angle
axis_angle[torch.isnan(axis_angle)] = 0.0
return axis_angle
def sample_joint_features(img_feat, joint_xy):
height, width = img_feat.shape[2:]
x = joint_xy[:, :, 0] / (width - 1) * 2 - 1
y = joint_xy[:, :, 1] / (height - 1) * 2 - 1
grid = torch.stack((x, y), 2)[:, :, None, :]
img_feat = F.grid_sample(img_feat, grid, align_corners=True)[:, :, :, 0] # batch_size, channel_dim, joint_num
img_feat = img_feat.permute(0, 2, 1).contiguous() # batch_size, joint_num, channel_dim
return img_feat
def soft_argmax_2d(heatmap2d):
batch_size = heatmap2d.shape[0]
height, width = heatmap2d.shape[2:]
heatmap2d = heatmap2d.reshape((batch_size, -1, height * width))
heatmap2d = F.softmax(heatmap2d, 2)
heatmap2d = heatmap2d.reshape((batch_size, -1, height, width))
accu_x = heatmap2d.sum(dim=(2))
accu_y = heatmap2d.sum(dim=(3))
accu_x = accu_x * torch.arange(width).float().cuda()[None, None, :]
accu_y = accu_y * torch.arange(height).float().cuda()[None, None, :]
accu_x = accu_x.sum(dim=2, keepdim=True)
accu_y = accu_y.sum(dim=2, keepdim=True)
coord_out = torch.cat((accu_x, accu_y), dim=2)
return coord_out
def soft_argmax_3d(heatmap3d):
batch_size = heatmap3d.shape[0]
depth, height, width = heatmap3d.shape[2:]
heatmap3d = heatmap3d.reshape((batch_size, -1, depth * height * width))
heatmap3d = F.softmax(heatmap3d, 2)
heatmap3d = heatmap3d.reshape((batch_size, -1, depth, height, width))
accu_x = heatmap3d.sum(dim=(2, 3))
accu_y = heatmap3d.sum(dim=(2, 4))
accu_z = heatmap3d.sum(dim=(3, 4))
accu_x = accu_x * torch.arange(width).float().cuda()[None, None, :]
accu_y = accu_y * torch.arange(height).float().cuda()[None, None, :]
accu_z = accu_z * torch.arange(depth).float().cuda()[None, None, :]
accu_x = accu_x.sum(dim=2, keepdim=True)
accu_y = accu_y.sum(dim=2, keepdim=True)
accu_z = accu_z.sum(dim=2, keepdim=True)
coord_out = torch.cat((accu_x, accu_y, accu_z), dim=2)
return coord_out
def restore_bbox(bbox_center, bbox_size, aspect_ratio, extension_ratio):
bbox = bbox_center.view(-1, 1, 2) + torch.cat((-bbox_size.view(-1, 1, 2) / 2., bbox_size.view(-1, 1, 2) / 2.),
1) # xyxy in (cfg.output_hm_shape[2], cfg.output_hm_shape[1]) space
bbox[:, :, 0] = bbox[:, :, 0] / cfg.output_hm_shape[2] * cfg.input_body_shape[1]
bbox[:, :, 1] = bbox[:, :, 1] / cfg.output_hm_shape[1] * cfg.input_body_shape[0]
bbox = bbox.view(-1, 4)
# xyxy -> xywh
bbox[:, 2] = bbox[:, 2] - bbox[:, 0]
bbox[:, 3] = bbox[:, 3] - bbox[:, 1]
# aspect ratio preserving bbox
w = bbox[:, 2]
h = bbox[:, 3]
c_x = bbox[:, 0] + w / 2.
c_y = bbox[:, 1] + h / 2.
mask1 = w > (aspect_ratio * h)
mask2 = w < (aspect_ratio * h)
h[mask1] = w[mask1] / aspect_ratio
w[mask2] = h[mask2] * aspect_ratio
bbox[:, 2] = w * extension_ratio
bbox[:, 3] = h * extension_ratio
bbox[:, 0] = c_x - bbox[:, 2] / 2.
bbox[:, 1] = c_y - bbox[:, 3] / 2.
# xywh -> xyxy
bbox[:, 2] = bbox[:, 2] + bbox[:, 0]
bbox[:, 3] = bbox[:, 3] + bbox[:, 1]
return bbox
common/utils/vis.py
+183
View File
@@ -0,0 +1,183 @@
import os
import cv2
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import matplotlib as mpl
import os
os.environ["PYOPENGL_PLATFORM"] = "egl"
import pyrender
import trimesh
from config import cfg
def vis_keypoints_with_skeleton(img, kps, kps_lines, kp_thresh=0.4, alpha=1):
# Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kps_lines) + 2)]
colors = [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors]
# Perform the drawing on a copy of the image, to allow for blending.
kp_mask = np.copy(img)
# Draw the keypoints.
for l in range(len(kps_lines)):
i1 = kps_lines[l][0]
i2 = kps_lines[l][1]
p1 = kps[0, i1].astype(np.int32), kps[1, i1].astype(np.int32)
p2 = kps[0, i2].astype(np.int32), kps[1, i2].astype(np.int32)
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
cv2.line(
kp_mask, p1, p2,
color=colors[l], thickness=2, lineType=cv2.LINE_AA)
if kps[2, i1] > kp_thresh:
cv2.circle(
kp_mask, p1,
radius=3, color=colors[l], thickness=-1, lineType=cv2.LINE_AA)
if kps[2, i2] > kp_thresh:
cv2.circle(
kp_mask, p2,
radius=3, color=colors[l], thickness=-1, lineType=cv2.LINE_AA)
# Blend the keypoints.
return cv2.addWeighted(img, 1.0 - alpha, kp_mask, alpha, 0)
def vis_keypoints(img, kps, alpha=1, radius=3, color=None):
# Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
cmap = plt.get_cmap('rainbow')
if color is None:
colors = [cmap(i) for i in np.linspace(0, 1, len(kps) + 2)]
colors = [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors]
# Perform the drawing on a copy of the image, to allow for blending.
kp_mask = np.copy(img)
# Draw the keypoints.
for i in range(len(kps)):
p = kps[i][0].astype(np.int32), kps[i][1].astype(np.int32)
if color is None:
cv2.circle(kp_mask, p, radius=radius, color=colors[i], thickness=-1, lineType=cv2.LINE_AA)
else:
cv2.circle(kp_mask, p, radius=radius, color=color, thickness=-1, lineType=cv2.LINE_AA)
# Blend the keypoints.
return cv2.addWeighted(img, 1.0 - alpha, kp_mask, alpha, 0)
def vis_mesh(img, mesh_vertex, alpha=0.5):
# Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(mesh_vertex))]
colors = [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors]
# Perform the drawing on a copy of the image, to allow for blending.
mask = np.copy(img)
# Draw the mesh
for i in range(len(mesh_vertex)):
p = mesh_vertex[i][0].astype(np.int32), mesh_vertex[i][1].astype(np.int32)
cv2.circle(mask, p, radius=1, color=colors[i], thickness=-1, lineType=cv2.LINE_AA)
# Blend the keypoints.
return cv2.addWeighted(img, 1.0 - alpha, mask, alpha, 0)
def vis_3d_skeleton(kpt_3d, kpt_3d_vis, kps_lines, filename=None):
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
# Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kps_lines) + 2)]
colors = [np.array((c[2], c[1], c[0])) for c in colors]
for l in range(len(kps_lines)):
i1 = kps_lines[l][0]
i2 = kps_lines[l][1]
x = np.array([kpt_3d[i1,0], kpt_3d[i2,0]])
y = np.array([kpt_3d[i1,1], kpt_3d[i2,1]])
z = np.array([kpt_3d[i1,2], kpt_3d[i2,2]])
if kpt_3d_vis[i1,0] > 0 and kpt_3d_vis[i2,0] > 0:
ax.plot(x, z, -y, c=colors[l], linewidth=2)
if kpt_3d_vis[i1,0] > 0:
ax.scatter(kpt_3d[i1,0], kpt_3d[i1,2], -kpt_3d[i1,1], c=colors[l], marker='o')
if kpt_3d_vis[i2,0] > 0:
ax.scatter(kpt_3d[i2,0], kpt_3d[i2,2], -kpt_3d[i2,1], c=colors[l], marker='o')
x_r = np.array([0, cfg.input_shape[1]], dtype=np.float32)
y_r = np.array([0, cfg.input_shape[0]], dtype=np.float32)
z_r = np.array([0, 1], dtype=np.float32)
if filename is None:
ax.set_title('3D vis')
else:
ax.set_title(filename)
ax.set_xlabel('X Label')
ax.set_ylabel('Z Label')
ax.set_zlabel('Y Label')
ax.legend()
plt.show()
cv2.waitKey(0)
def save_obj(v, f, file_name='output.obj'):
obj_file = open(file_name, 'w')
for i in range(len(v)):
obj_file.write('v ' + str(v[i][0]) + ' ' + str(v[i][1]) + ' ' + str(v[i][2]) + '\n')
for i in range(len(f)):
obj_file.write('f ' + str(f[i][0]+1) + '/' + str(f[i][0]+1) + ' ' + str(f[i][1]+1) + '/' + str(f[i][1]+1) + ' ' + str(f[i][2]+1) + '/' + str(f[i][2]+1) + '\n')
obj_file.close()
def perspective_projection(vertices, cam_param):
# vertices: [N, 3]
# cam_param: [3]
fx, fy= cam_param['focal']
cx, cy = cam_param['princpt']
vertices[:, 0] = vertices[:, 0] * fx / vertices[:, 2] + cx
vertices[:, 1] = vertices[:, 1] * fy / vertices[:, 2] + cy
return vertices
def render_mesh(img, mesh, face, cam_param, mesh_as_vertices=False):
if mesh_as_vertices:
# to run on cluster where headless pyrender is not supported for A100/V100
vertices_2d = perspective_projection(mesh, cam_param)
img = vis_keypoints(img, vertices_2d, alpha=0.8, radius=2, color=(0, 0, 255))
else:
# mesh
mesh = trimesh.Trimesh(mesh, face)
rot = trimesh.transformations.rotation_matrix(
np.radians(180), [1, 0, 0])
mesh.apply_transform(rot)
material = pyrender.MetallicRoughnessMaterial(metallicFactor=0.0, alphaMode='OPAQUE', baseColorFactor=(1.0, 1.0, 0.9, 1.0))
mesh = pyrender.Mesh.from_trimesh(mesh, material=material, smooth=False)
scene = pyrender.Scene(ambient_light=(0.3, 0.3, 0.3))
scene.add(mesh, 'mesh')
focal, princpt = cam_param['focal'], cam_param['princpt']
camera = pyrender.IntrinsicsCamera(fx=focal[0], fy=focal[1], cx=princpt[0], cy=princpt[1])
scene.add(camera)
# renderer
renderer = pyrender.OffscreenRenderer(viewport_width=img.shape[1], viewport_height=img.shape[0], point_size=1.0)
# light
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=0.8)
light_pose = np.eye(4)
light_pose[:3, 3] = np.array([0, -1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([0, 1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([1, 1, 2])
scene.add(light, pose=light_pose)
# render
rgb, depth = renderer.render(scene, flags=pyrender.RenderFlags.RGBA)
rgb = rgb[:,:,:3].astype(np.float32)
valid_mask = (depth > 0)[:,:,None]
# save to image
img = rgb * valid_mask + img * (1-valid_mask)
return img
data/.DS_Store
Vendored
BIN
View File
Binary file not shown.
data/AGORA/AGORA.py
+853
View File
@@ -0,0 +1,853 @@
import os
import os.path as osp
from glob import glob
import numpy as np
from config import cfg
import copy
import json
import pickle
import cv2
import torch
from pycocotools.coco import COCO
from utils.human_models import smpl_x
from utils.preprocessing import load_img, sanitize_bbox, process_bbox, augmentation, process_db_coord, \
process_human_model_output, load_ply, load_obj
from utils.transforms import rigid_align
import tqdm
import random
from humandata import Cache
class AGORA(torch.utils.data.Dataset):
def __init__(self, transform, data_split):
self.transform = transform
self.data_split = data_split
if getattr(cfg, 'eval_on_train', False):
self.data_split = 'eval_train'
print("Evaluate on train set.")
self.data_path = osp.join(cfg.data_dir, 'AGORA', 'data')
self.save_idx = 0
self.resolution = (2160, 3840) # height, width. one of (720, 1280) and (2160, 3840)
if cfg.agora_benchmark == 'agora_model_test' or cfg.agora_benchmark == 'test_only':
self.test_set = 'test'
else:
self.test_set = 'val' # val, test
# AGORA joint set
self.joint_set = {
'joint_num': 127,
'joints_name': \
('Pelvis', 'L_Hip', 'R_Hip', 'Spine_1', 'L_Knee', 'R_Knee', 'Spine_2', 'L_Ankle', 'R_Ankle', 'Spine_3',
'L_Foot', 'R_Foot', 'Neck', 'L_Collar', 'R_Collar', 'Head', 'L_Shoulder', 'R_Shoulder', 'L_Elbow',
'R_Elbow', 'L_Wrist', 'R_Wrist', # body
'Jaw', 'L_Eye_SMPLH', 'R_Eye_SMPLH', # SMPLH
'L_Index_1', 'L_Index_2', 'L_Index_3', 'L_Middle_1', 'L_Middle_2', 'L_Middle_3', 'L_Pinky_1',
'L_Pinky_2', 'L_Pinky_3', 'L_Ring_1', 'L_Ring_2', 'L_Ring_3', 'L_Thumb_1', 'L_Thumb_2', 'L_Thumb_3',
# fingers
'R_Index_1', 'R_Index_2', 'R_Index_3', 'R_Middle_1', 'R_Middle_2', 'R_Middle_3', 'R_Pinky_1',
'R_Pinky_2', 'R_Pinky_3', 'R_Ring_1', 'R_Ring_2', 'R_Ring_3', 'R_Thumb_1', 'R_Thumb_2', 'R_Thumb_3',
# fingers
'Nose', 'R_Eye', 'L_Eye', 'R_Ear', 'L_Ear', # face in body
'L_Big_toe', 'L_Small_toe', 'L_Heel', 'R_Big_toe', 'R_Small_toe', 'R_Heel', # feet
'L_Thumb_4', 'L_Index_4', 'L_Middle_4', 'L_Ring_4', 'L_Pinky_4', # finger tips
'R_Thumb_4', 'R_Index_4', 'R_Middle_4', 'R_Ring_4', 'R_Pinky_4', # finger tips
*['Face_' + str(i) for i in range(5, 56)] # face
),
'flip_pairs': \
((1, 2), (4, 5), (7, 8), (10, 11), (13, 14), (16, 17), (18, 19), (20, 21), # body
(23, 24), # SMPLH
(25, 40), (26, 41), (27, 42), (28, 43), (29, 44), (30, 45), (31, 46), (32, 47), (33, 48), (34, 49),
(35, 50), (36, 51), (37, 52), (38, 53), (39, 54), # fingers
(56, 57), (58, 59), # face in body
(60, 63), (61, 64), (62, 65), # feet
(66, 71), (67, 72), (68, 73), (69, 74), (70, 75), # fingertips
(76, 85), (77, 84), (78, 83), (79, 82), (80, 81), # face eyebrow
(90, 94), (91, 93), # face below nose
(95, 104), (96, 103), (97, 102), (98, 101), (99, 106), (100, 105), # face eyes
(107, 113), (108, 112), (109, 111), (114, 118), (115, 117), # face mouth
(119, 123), (120, 122), (124, 126) # face lip
)
}
self.joint_set['joint_part'] = {
'body': list(range(self.joint_set['joints_name'].index('Pelvis'),
self.joint_set['joints_name'].index('R_Eye_SMPLH') + 1)) + list(
range(self.joint_set['joints_name'].index('Nose'), self.joint_set['joints_name'].index('R_Heel') + 1)),
'lhand': list(range(self.joint_set['joints_name'].index('L_Index_1'),
self.joint_set['joints_name'].index('L_Thumb_3') + 1)) + list(
range(self.joint_set['joints_name'].index('L_Thumb_4'),
self.joint_set['joints_name'].index('L_Pinky_4') + 1)),
'rhand': list(range(self.joint_set['joints_name'].index('R_Index_1'),
self.joint_set['joints_name'].index('R_Thumb_3') + 1)) + list(
range(self.joint_set['joints_name'].index('R_Thumb_4'),
self.joint_set['joints_name'].index('R_Pinky_4') + 1)),
'face': list(range(self.joint_set['joints_name'].index('Face_5'),
self.joint_set['joints_name'].index('Face_55') + 1))}
self.joint_set['root_joint_idx'] = self.joint_set['joints_name'].index('Pelvis')
self.joint_set['lwrist_idx'] = self.joint_set['joints_name'].index('L_Wrist')
self.joint_set['rwrist_idx'] = self.joint_set['joints_name'].index('R_Wrist')
self.joint_set['neck_idx'] = self.joint_set['joints_name'].index('Neck')
# self.datalist = self.load_data()
# load data or cache
self.use_cache = getattr(cfg, 'use_cache', False)
if 'train'in self.data_split or (self.data_split == 'test' and self.test_set == 'val'):
if 'train' in self.data_split:
if getattr(cfg, 'agora_fix_betas', False):
assert getattr(cfg, 'agora_fix_global_orient_transl')
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', f'AGORA_{self.data_split}_fix_betas.npz')
elif getattr(cfg, 'agora_fix_global_orient_transl', False):
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', f'AGORA_{self.data_split}_fix_global_orient_transl.npz')
else:
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', f'AGORA_{self.data_split}.npz')
else:
if getattr(cfg, 'agora_fix_betas', False):
assert getattr(cfg, 'agora_fix_global_orient_transl')
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'AGORA_validation_fix_betas.npz')
elif getattr(cfg, 'agora_fix_global_orient_transl', False):
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'AGORA_validation_fix_global_orient_transl.npz')
else:
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'AGORA_validation.npz')
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
datalist = Cache(self.annot_path_cache)
assert datalist.data_strategy == getattr(cfg, 'data_strategy', None), \
f'Cache data strategy {datalist.data_strategy} does not match current data strategy ' \
f'{getattr(cfg, "data_strategy", None)}'
self.datalist = datalist
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data()
if self.use_cache:
print(f'[{self.__class__.__name__}] Caching datalist to {self.annot_path_cache}...')
Cache.save(
self.annot_path_cache,
self.datalist,
data_strategy=getattr(cfg, 'data_strategy', None)
)
else: # test
self.datalist = self.load_data()
def load_data(self):
datalist = []
if 'train' in self.data_split or (self.data_split == 'test' and self.test_set == 'val'):
print('dataset settings:')
print('agora_fix_betas', getattr(cfg, 'agora_fix_betas', False))
print('agora_fix_global_orient_transl', getattr(cfg, 'agora_fix_global_orient_transl', False))
print('agora_valid_root_pose', getattr(cfg, 'agora_valid_root_pose', False))
if 'train' in self.data_split:
if getattr(cfg, 'agora_fix_betas', False):
assert getattr(cfg, 'agora_fix_global_orient_transl')
db = COCO(osp.join(self.data_path, 'AGORA_train_fix_betas.json'))
elif getattr(cfg, 'agora_fix_global_orient_transl', False):
db = COCO(osp.join(self.data_path, 'AGORA_train_fix_global_orient_transl.json'))
else:
db = COCO(osp.join(self.data_path, 'AGORA_train.json'))
else:
if getattr(cfg, 'agora_fix_betas', False):
assert getattr(cfg, 'agora_fix_global_orient_transl')
db = COCO(osp.join(self.data_path, 'AGORA_validation_fix_betas.json'))
elif getattr(cfg, 'agora_fix_global_orient_transl', False):
db = COCO(osp.join(self.data_path, 'AGORA_validation_fix_global_orient_transl.json'))
else:
db = COCO(osp.join(self.data_path, 'AGORA_validation.json'))
### HARDCODE vis for debug
# count = 0
i = 0
for aid in tqdm.tqdm(list(db.anns.keys())):
# if count > 50:
# continue
# count += 1
i += 1
if self.data_split == 'train' and i % getattr(cfg, 'AGORA_train_sample_interval', 1) != 0:
continue
ann = db.anns[aid]
image_id = ann['image_id']
img = db.loadImgs(image_id)[0]
if not ann['is_valid']:
continue
joints_2d_path = osp.join(self.data_path, ann['smplx_joints_2d_path'])
joints_3d_path = osp.join(self.data_path, ann['smplx_joints_3d_path'])
verts_path = osp.join(self.data_path, ann['smplx_verts_path'])
smplx_param_path = osp.join(self.data_path, ann['smplx_param_path'])
kid = ann['kid']
gender = ann['gender']
if not osp.exists(smplx_param_path): print(smplx_param_path)
if self.resolution == (720, 1280):
img_shape = self.resolution
img_path = osp.join(self.data_path, img['file_name_1280x720'])
# convert to current resolution
bbox = np.array(ann['bbox']).reshape(2, 2)
bbox[:, 0] = bbox[:, 0] / 3840 * 1280
bbox[:, 1] = bbox[:, 1] / 2160 * 720
bbox = bbox.reshape(4)
if hasattr(cfg, 'bbox_ratio'):
bbox_ratio = cfg.bbox_ratio * 0.833 # agora preprocess is giving 1.2 box padding
else:
bbox_ratio = 1.25
bbox = process_bbox(bbox, img_shape[1], img_shape[0], ratio=bbox_ratio)
if bbox is None:
continue
lhand_bbox = np.array(ann['lhand_bbox']).reshape(2, 2)
lhand_bbox[:, 0] = lhand_bbox[:, 0] / 3840 * 1280
lhand_bbox[:, 1] = lhand_bbox[:, 1] / 2160 * 720
lhand_bbox = lhand_bbox.reshape(4)
lhand_bbox = sanitize_bbox(lhand_bbox, img_shape[1], img_shape[0])
if lhand_bbox is not None:
lhand_bbox[2:] += lhand_bbox[:2] # xywh -> xyxy
rhand_bbox = np.array(ann['rhand_bbox']).reshape(2, 2)
rhand_bbox[:, 0] = rhand_bbox[:, 0] / 3840 * 1280
rhand_bbox[:, 1] = rhand_bbox[:, 1] / 2160 * 720
rhand_bbox = rhand_bbox.reshape(4)
rhand_bbox = sanitize_bbox(rhand_bbox, img_shape[1], img_shape[0])
if rhand_bbox is not None:
rhand_bbox[2:] += rhand_bbox[:2] # xywh -> xyxy
face_bbox = np.array(ann['face_bbox']).reshape(2, 2)
face_bbox[:, 0] = face_bbox[:, 0] / 3840 * 1280
face_bbox[:, 1] = face_bbox[:, 1] / 2160 * 720
face_bbox = face_bbox.reshape(4)
face_bbox = sanitize_bbox(face_bbox, img_shape[1], img_shape[0])
if face_bbox is not None:
face_bbox[2:] += face_bbox[:2] # xywh -> xyxy
data_dict = {'img_path': img_path, 'img_shape': img_shape, 'bbox': bbox, 'lhand_bbox': lhand_bbox,
'rhand_bbox': rhand_bbox, 'face_bbox': face_bbox, 'joints_2d_path': joints_2d_path,
'joints_3d_path': joints_3d_path, 'verts_path': verts_path,
'smplx_param_path': smplx_param_path, 'ann_id': str(aid), 'kid': kid, 'gender': gender}
datalist.append(data_dict)
elif self.resolution == (2160,
3840): # use cropped and resized images. loading 4K images in pytorch dataloader takes too much time...
img_path = osp.join(self.data_path, '3840x2160',
img['file_name_3840x2160'].split('/')[-2] + '_crop',
img['file_name_3840x2160'].split('/')[-1][:-4] + '_ann_id_' + str(aid) + '.png')
json_path = osp.join(self.data_path, '3840x2160',
img['file_name_3840x2160'].split('/')[-2] + '_crop',
img['file_name_3840x2160'].split('/')[-1][:-4] + '_ann_id_' + str(
aid) + '.json')
if not osp.isfile(json_path):
continue
with open(json_path) as f:
crop_resize_info = json.load(f)
img2bb_trans_from_orig = np.array(crop_resize_info['img2bb_trans'], dtype=np.float32)
resized_height, resized_width = crop_resize_info['resized_height'], crop_resize_info[
'resized_width']
img_shape = (resized_height, resized_width)
bbox = np.array([0, 0, resized_width, resized_height], dtype=np.float32)
# transform from original image to crop_and_resize image
lhand_bbox = np.array(ann['lhand_bbox']).reshape(2, 2)
lhand_bbox[1] += lhand_bbox[0] # xywh -> xyxy
lhand_bbox = np.dot(img2bb_trans_from_orig,
np.concatenate((lhand_bbox, np.ones_like(lhand_bbox[:, :1])), 1).transpose(1,
0)).transpose(
1, 0)
lhand_bbox[1] -= lhand_bbox[0] # xyxy -> xywh
lhand_bbox = lhand_bbox.reshape(4)
lhand_bbox = sanitize_bbox(lhand_bbox, self.resolution[1], self.resolution[0])
if lhand_bbox is not None:
lhand_bbox[2:] += lhand_bbox[:2] # xywh -> xyxy
# transform from original image to crop_and_resize image
rhand_bbox = np.array(ann['rhand_bbox']).reshape(2, 2)
rhand_bbox[1] += rhand_bbox[0] # xywh -> xyxy
rhand_bbox = np.dot(img2bb_trans_from_orig,
np.concatenate((rhand_bbox, np.ones_like(rhand_bbox[:, :1])), 1).transpose(1,
0)).transpose(
1, 0)
rhand_bbox[1] -= rhand_bbox[0] # xyxy -> xywh
rhand_bbox = rhand_bbox.reshape(4)
rhand_bbox = sanitize_bbox(rhand_bbox, self.resolution[1], self.resolution[0])
if rhand_bbox is not None:
rhand_bbox[2:] += rhand_bbox[:2] # xywh -> xyxy
# transform from original image to crop_and_resize image
face_bbox = np.array(ann['face_bbox']).reshape(2, 2)
face_bbox[1] += face_bbox[0] # xywh -> xyxy
face_bbox = np.dot(img2bb_trans_from_orig,
np.concatenate((face_bbox, np.ones_like(face_bbox[:, :1])), 1).transpose(1,
0)).transpose(
1, 0)
face_bbox[1] -= face_bbox[0] # xyxy -> xywh
face_bbox = face_bbox.reshape(4)
face_bbox = sanitize_bbox(face_bbox, self.resolution[1], self.resolution[0])
if face_bbox is not None:
face_bbox[2:] += face_bbox[:2] # xywh -> xyxy
data_dict = {'img_path': img_path, 'img_shape': img_shape, 'bbox': bbox, 'lhand_bbox': lhand_bbox,
'rhand_bbox': rhand_bbox, 'face_bbox': face_bbox,
'img2bb_trans_from_orig': img2bb_trans_from_orig, 'joints_2d_path': joints_2d_path,
'joints_3d_path': joints_3d_path, 'verts_path': verts_path,
'smplx_param_path': smplx_param_path, 'ann_id': str(aid), 'kid': kid, 'gender': gender}
datalist.append(data_dict)
print('[AGORA train] original size:', len(db.anns.keys()),
'. Sample interval:', getattr(cfg, 'AGORA_train_sample_interval', 1),
'. Sampled size:', len(datalist))
elif self.data_split == 'test' and self.test_set == 'test':
with open(osp.join(self.data_path, 'AGORA_test_bbox.json')) as f:
bboxs = json.load(f)
for filename in tqdm.tqdm(bboxs.keys()):
if self.resolution == (720, 1280):
img_path = osp.join(self.data_path, 'test', filename)
img_shape = self.resolution
person_num = len(bboxs[filename])
for pid in range(person_num):
# change bbox from (2160,3840) to target resoution
bbox = np.array(bboxs[filename][pid]['bbox']).reshape(2, 2)
bbox[:, 0] = bbox[:, 0] / 3840 * 1280
bbox[:, 1] = bbox[:, 1] / 2160 * 720
bbox = bbox.reshape(4)
if hasattr(cfg, 'bbox_ratio'):
bbox_ratio = cfg.bbox_ratio * 0.833 # agora preprocess is giving 1.2 box padding
else:
bbox_ratio = 1.25
bbox = process_bbox(bbox, img_shape[1], img_shape[0], ratio=bbox_ratio)
if bbox is None:
continue
datalist.append({'img_path': img_path, 'img_shape': img_shape, 'bbox': bbox, 'person_idx': pid})
elif self.resolution == (2160,
3840): # use cropped and resized images. loading 4K images in pytorch dataloader takes too much time...
person_num = len(bboxs[filename])
for pid in range(person_num):
img_path = osp.join(self.data_path, '3840x2160', 'test_crop',
filename[:-4] + '_pid_' + str(pid) + '.png')
json_path = osp.join(self.data_path, '3840x2160', 'test_crop',
filename[:-4] + '_pid_' + str(pid) + '.json')
if not osp.isfile(json_path):
continue
with open(json_path) as f:
crop_resize_info = json.load(f)
img2bb_trans_from_orig = np.array(crop_resize_info['img2bb_trans'], dtype=np.float32)
resized_height, resized_width = crop_resize_info['resized_height'], crop_resize_info[
'resized_width']
img_shape = (resized_height, resized_width)
bbox = np.array([0, 0, resized_width, resized_height], dtype=np.float32)
datalist.append({'img_path': img_path, 'img_shape': img_shape,
'img2bb_trans_from_orig': img2bb_trans_from_orig, 'bbox': bbox,
'person_idx': pid})
if (getattr(cfg, 'data_strategy', None) == 'balance' and self.data_split == 'train') or \
self.data_split == 'eval_train':
print(f"[Agora] Using [balance] strategy with datalist shuffled...")
random.seed(2023)
random.shuffle(datalist)
if self.data_split == "eval_train":
return datalist[:10000]
return datalist
def process_hand_face_bbox(self, bbox, do_flip, img_shape, img2bb_trans):
if bbox is None:
bbox = np.array([0, 0, 1, 1], dtype=np.float32).reshape(2, 2) # dummy value
bbox_valid = float(False) # dummy value
else:
# reshape to top-left (x,y) and bottom-right (x,y)
bbox = bbox.reshape(2, 2)
# flip augmentation
if do_flip:
bbox[:, 0] = img_shape[1] - bbox[:, 0] - 1
bbox[0, 0], bbox[1, 0] = bbox[1, 0].copy(), bbox[0, 0].copy() # xmin <-> xmax swap
# make four points of the bbox
bbox = bbox.reshape(4).tolist()
xmin, ymin, xmax, ymax = bbox
bbox = np.array([[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]], dtype=np.float32).reshape(4, 2)
# affine transformation (crop, rotation, scale)
bbox_xy1 = np.concatenate((bbox, np.ones_like(bbox[:, :1])), 1)
bbox = np.dot(img2bb_trans, bbox_xy1.transpose(1, 0)).transpose(1, 0)[:, :2]
bbox[:, 0] = bbox[:, 0] / cfg.input_img_shape[1] * cfg.output_hm_shape[2]
bbox[:, 1] = bbox[:, 1] / cfg.input_img_shape[0] * cfg.output_hm_shape[1]
# make box a rectangle without rotation
xmin = np.min(bbox[:, 0]);
xmax = np.max(bbox[:, 0]);
ymin = np.min(bbox[:, 1]);
ymax = np.max(bbox[:, 1]);
bbox = np.array([xmin, ymin, xmax, ymax], dtype=np.float32)
bbox_valid = float(True)
bbox = bbox.reshape(2, 2)
return bbox, bbox_valid
def __len__(self):
return len(self.datalist)
def __getitem__(self, idx):
data = copy.deepcopy(self.datalist[idx])
img_path, img_shape, bbox = data['img_path'], data['img_shape'], data['bbox']
# image load
img = load_img(img_path)
# affine transform
img, img2bb_trans, bb2img_trans, rot, do_flip = augmentation(img, bbox, self.data_split)
img = self.transform(img.astype(np.float32)) / 255.
if self.data_split == 'train':
# gt load
with open(data['joints_2d_path']) as f:
joint_img = np.array(json.load(f)).reshape(-1, 2)
if self.resolution == (2160, 3840):
joint_img[:, :2] = np.dot(data['img2bb_trans_from_orig'],
np.concatenate((joint_img, np.ones_like(joint_img[:, :1])), 1).transpose(
1, 0)).transpose(1,
0) # transform from original image to crop_and_resize image
joint_img[:, 0] = joint_img[:, 0] / 3840 * self.resolution[1]
joint_img[:, 1] = joint_img[:, 1] / 2160 * self.resolution[0]
with open(data['joints_3d_path']) as f:
joint_cam = np.array(json.load(f)).reshape(-1, 3)
### HARDCODE vis for debug
# joint_cam_orig = joint_cam.copy()
with open(data['smplx_param_path'], 'rb') as f:
smplx_param = pickle.load(f, encoding='latin1')
# hand and face bbox transform
lhand_bbox, rhand_bbox, face_bbox = data['lhand_bbox'], data['rhand_bbox'], data['face_bbox']
lhand_bbox, lhand_bbox_valid = self.process_hand_face_bbox(lhand_bbox, do_flip, img_shape, img2bb_trans)
rhand_bbox, rhand_bbox_valid = self.process_hand_face_bbox(rhand_bbox, do_flip, img_shape, img2bb_trans)
face_bbox, face_bbox_valid = self.process_hand_face_bbox(face_bbox, do_flip, img_shape, img2bb_trans)
if do_flip:
lhand_bbox, rhand_bbox = rhand_bbox, lhand_bbox
lhand_bbox_valid, rhand_bbox_valid = rhand_bbox_valid, lhand_bbox_valid
lhand_bbox_center = (lhand_bbox[0] + lhand_bbox[1]) / 2.;
rhand_bbox_center = (rhand_bbox[0] + rhand_bbox[1]) / 2.;
face_bbox_center = (face_bbox[0] + face_bbox[1]) / 2.
lhand_bbox_size = lhand_bbox[1] - lhand_bbox[0];
rhand_bbox_size = rhand_bbox[1] - rhand_bbox[0];
face_bbox_size = face_bbox[1] - face_bbox[0];
"""
# for debug
_img = img.numpy().transpose(1,2,0)[:,:,::-1].copy() * 255
if lhand_bbox_valid:
_tmp = lhand_bbox.copy().reshape(2,2)
_tmp[:,0] = _tmp[:,0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1]
_tmp[:,1] = _tmp[:,1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0]
cv2.rectangle(_img, (int(_tmp[0,0]), int(_tmp[0,1])), (int(_tmp[1,0]), int(_tmp[1,1])), (255,0,0), 3)
cv2.imwrite('agora_' + str(idx) + '_lhand.jpg', _img)
if rhand_bbox_valid:
_tmp = rhand_bbox.copy().reshape(2,2)
_tmp[:,0] = _tmp[:,0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1]
_tmp[:,1] = _tmp[:,1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0]
cv2.rectangle(_img, (int(_tmp[0,0]), int(_tmp[0,1])), (int(_tmp[1,0]), int(_tmp[1,1])), (255,0,0), 3)
cv2.imwrite('agora_' + str(idx) + '_rhand.jpg', _img)
if face_bbox_valid:
_tmp = face_bbox.copy().reshape(2,2)
_tmp[:,0] = _tmp[:,0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1]
_tmp[:,1] = _tmp[:,1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0]
cv2.rectangle(_img, (int(_tmp[0,0]), int(_tmp[0,1])), (int(_tmp[1,0]), int(_tmp[1,1])), (255,0,0), 3)
cv2.imwrite('agora_' + str(idx) + '_face.jpg', _img)
#cv2.imwrite('agora_' + str(idx) + '.jpg', _img)
"""
# coordinates
joint_cam = joint_cam - joint_cam[self.joint_set['root_joint_idx'], None, :] # root-relative
joint_cam[self.joint_set['joint_part']['lhand'], :] = joint_cam[self.joint_set['joint_part']['lhand'],
:] - joint_cam[self.joint_set['lwrist_idx'], None,
:] # left hand root-relative
joint_cam[self.joint_set['joint_part']['rhand'], :] = joint_cam[self.joint_set['joint_part']['rhand'],
:] - joint_cam[self.joint_set['rwrist_idx'], None,
:] # right hand root-relative
joint_cam[self.joint_set['joint_part']['face'], :] = joint_cam[self.joint_set['joint_part']['face'],
:] - joint_cam[self.joint_set['neck_idx'], None,
:] # face root-relative
joint_img = np.concatenate((joint_img[:, :2], joint_cam[:, 2:]), 1) # x, y, depth
joint_img[self.joint_set['joint_part']['body'], 2] = (joint_cam[self.joint_set['joint_part'][
'body'], 2].copy() / (
cfg.body_3d_size / 2) + 1) / 2. * \
cfg.output_hm_shape[0] # body depth discretize
joint_img[self.joint_set['joint_part']['lhand'], 2] = (joint_cam[self.joint_set['joint_part'][
'lhand'], 2].copy() / (
cfg.hand_3d_size / 2) + 1) / 2. * \
cfg.output_hm_shape[0] # left hand depth discretize
joint_img[self.joint_set['joint_part']['rhand'], 2] = (joint_cam[self.joint_set['joint_part'][
'rhand'], 2].copy() / (
cfg.hand_3d_size / 2) + 1) / 2. * \
cfg.output_hm_shape[0] # right hand depth discretize
joint_img[self.joint_set['joint_part']['face'], 2] = (joint_cam[self.joint_set['joint_part'][
'face'], 2].copy() / (
cfg.face_3d_size / 2) + 1) / 2. * \
cfg.output_hm_shape[0] # face depth discretize
joint_valid = np.ones_like(joint_img[:, :1])
# alr ra when passed into this function
joint_img, joint_cam_ra, _, joint_valid, joint_trunc = process_db_coord(joint_img, joint_cam, joint_valid,
do_flip, img_shape,
self.joint_set['flip_pairs'],
img2bb_trans, rot,
self.joint_set['joints_name'],
smpl_x.joints_name)
# reverse ra
joint_cam_wo_ra = joint_cam_ra.copy()
joint_cam_wo_ra[smpl_x.joint_part['lhand'], :] = joint_cam_wo_ra[smpl_x.joint_part['lhand'], :] \
+ joint_cam_wo_ra[smpl_x.lwrist_idx, None, :] # left hand root-relative
joint_cam_wo_ra[smpl_x.joint_part['rhand'], :] = joint_cam_wo_ra[smpl_x.joint_part['rhand'], :] \
+ joint_cam_wo_ra[smpl_x.rwrist_idx, None, :] # right hand root-relative
joint_cam_wo_ra[smpl_x.joint_part['face'], :] = joint_cam_wo_ra[smpl_x.joint_part['face'], :] \
+ joint_cam_wo_ra[smpl_x.neck_idx, None,: ] # face root-relative
"""
# for debug
_tmp = joint_img.copy()
_tmp[:,0] = _tmp[:,0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1]
_tmp[:,1] = _tmp[:,1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0]
_img = img.numpy().transpose(1,2,0)[:,:,::-1] * 255
_img = vis_keypoints(_img.copy(), _tmp)
cv2.imwrite('agora_' + str(idx) + '.jpg', _img)
"""
"""
# for debug
_tmp = joint_cam.copy()[:,:2]
_tmp[:,0] = _tmp[:,0] / (cfg.body_3d_size / 2) * cfg.input_img_shape[1] + cfg.input_img_shape[1]/2
_tmp[:,1] = _tmp[:,1] / (cfg.body_3d_size / 2) * cfg.input_img_shape[0] + cfg.input_img_shape[0]/2
_img = np.zeros((cfg.input_img_shape[0], cfg.input_img_shape[1], 3), dtype=np.float32)
_img = vis_keypoints(_img.copy(), _tmp)
cv2.imwrite('agora_' + str(idx) + '_cam.jpg', _img)
"""
# smplx parameters
root_pose = np.array(smplx_param['global_orient'], dtype=np.float32).reshape(
-1) # rotation to world coordinate
body_pose = np.array(smplx_param['body_pose'], dtype=np.float32).reshape(-1)
# use adapted shape for adults
if getattr(cfg, 'agora_fix_betas', False) and not data['kid']:
shape = np.array(smplx_param['betas_neutral'], dtype=np.float32).reshape(-1)[:10]
else:
shape = np.array(smplx_param['betas'], dtype=np.float32).reshape(-1)[:10] # bug?
lhand_pose = np.array(smplx_param['left_hand_pose'], dtype=np.float32).reshape(-1)
rhand_pose = np.array(smplx_param['right_hand_pose'], dtype=np.float32).reshape(-1)
jaw_pose = np.array(smplx_param['jaw_pose'], dtype=np.float32).reshape(-1)
expr = np.array(smplx_param['expression'], dtype=np.float32).reshape(-1)
trans = np.array(smplx_param['transl'], dtype=np.float32).reshape(-1) # translation to world coordinate
cam_param = {'focal': cfg.focal,
'princpt': cfg.princpt} # put random camera paraemter as we do not use coordinates from smplx parameters
smplx_param = {'root_pose': root_pose, 'body_pose': body_pose, 'shape': shape,
'lhand_pose': lhand_pose, 'lhand_valid': True,
'rhand_pose': rhand_pose, 'rhand_valid': True,
'jaw_pose': jaw_pose, 'expr': expr, 'face_valid': True,
'trans': trans}
_, _, _, smplx_pose, smplx_shape, smplx_expr, smplx_pose_valid, _, smplx_expr_valid, _ = process_human_model_output(
smplx_param, cam_param, do_flip, img_shape, img2bb_trans, rot, 'smplx')
### HARDCODE vis for debug
# mesh_rot_, joint_cam_, _, smplx_pose, smplx_shape, smplx_expr, smplx_pose_valid, _, smplx_expr_valid, mesh_orig, joint_cam_orig_ = process_human_model_output(
# smplx_param, cam_param, do_flip, img_shape, img2bb_trans, rot, 'smplx')
smplx_pose_valid = np.tile(smplx_pose_valid[:, None], (1, 3)).reshape(-1)
if not getattr(cfg, 'agora_valid_root_pose', False):
smplx_pose_valid[:3] = 0 # global orient of the provided parameter is a rotation to world coordinate system. I want camera coordinate system.
smplx_shape_valid = True
inputs = {'img': img}
targets = {'joint_img': joint_img, 'joint_cam': joint_cam_wo_ra, #from annot
'smplx_joint_img': joint_img, 'smplx_joint_cam': joint_cam_ra, #_smplx_joint_cam, # from smplx param w/ ra
'smplx_pose': smplx_pose, 'smplx_shape': smplx_shape, 'smplx_expr': smplx_expr,
'lhand_bbox_center': lhand_bbox_center, 'lhand_bbox_size': lhand_bbox_size,
'rhand_bbox_center': rhand_bbox_center, 'rhand_bbox_size': rhand_bbox_size,
'face_bbox_center': face_bbox_center, 'face_bbox_size': face_bbox_size}
meta_info = {'joint_valid': joint_valid, 'joint_trunc': joint_trunc,
'smplx_joint_valid': joint_valid, 'smplx_joint_trunc': joint_trunc,
'smplx_pose_valid': smplx_pose_valid, 'smplx_shape_valid': float(smplx_shape_valid),
'smplx_expr_valid': float(smplx_expr_valid), 'is_3D': float(True),
'lhand_bbox_valid': lhand_bbox_valid, 'rhand_bbox_valid': rhand_bbox_valid, 'face_bbox_valid': face_bbox_valid}
### HARDCODE vis for debug
# 'gt_3d_path': data['joints_3d_path'], 'smplx_path': data['smplx_param_path'], 'id': idx}
return inputs, targets, meta_info
else:
# load crop and resize information (for the 4K setting)
if self.resolution == (2160, 3840):
img2bb_trans = np.dot(
np.concatenate((img2bb_trans,
np.array([0, 0, 1], dtype=np.float32).reshape(1, 3))),
np.concatenate((data['img2bb_trans_from_orig'],
np.array([0, 0, 1], dtype=np.float32).reshape(1, 3)))
)
bb2img_trans = np.linalg.inv(img2bb_trans)[:2, :]
img2bb_trans = img2bb_trans[:2, :]
if self.test_set == 'val':
# gt load
with open(data['verts_path']) as f:
verts = np.array(json.load(f)).reshape(-1, 3)
with open(data['smplx_param_path'], 'rb') as f:
smplx_param = pickle.load(f, encoding='latin1')
transl = np.array(smplx_param['transl'], dtype=np.float32).reshape(-1)
inputs = {'img': img}
targets = {'smplx_mesh_cam': verts}
meta_info = {'bb2img_trans': bb2img_trans, 'img_path': img_path, 'gt_smplx_transl':transl}
else:
inputs = {'img': img}
targets = {'smplx_mesh_cam': np.zeros((smpl_x.vertex_num, 3), dtype=np.float32)} # dummy vertex
meta_info = {'bb2img_trans': bb2img_trans, 'img_path': img_path}
return inputs, targets, meta_info
def evaluate(self, outs, cur_sample_idx):
annots = self.datalist
sample_num = len(outs)
eval_result = {'pa_mpvpe_all': [], 'pa_mpvpe_l_hand': [], 'pa_mpvpe_r_hand': [], 'pa_mpvpe_hand': [], 'pa_mpvpe_face': [],
'mpvpe_all': [], 'mpvpe_l_hand': [], 'mpvpe_r_hand': [], 'mpvpe_hand': [], 'mpvpe_face': []}
vis = getattr(cfg, 'vis', False)
vis_save_dir = cfg.vis_dir
if getattr(cfg, 'vis', False):
import csv
csv_file = f'{cfg.vis_dir}/agora_smplx_error.csv'
file = open(csv_file, 'a', newline='')
writer = csv.writer(file)
for n in range(sample_num):
annot = annots[cur_sample_idx + n]
out = outs[n]
mesh_gt = out['smplx_mesh_cam_target']
mesh_out = out['smplx_mesh_cam']
# MPVPE from all vertices
mesh_out_align = mesh_out - np.dot(smpl_x.J_regressor, mesh_out)[smpl_x.J_regressor_idx['pelvis'], None,
:] + np.dot(smpl_x.J_regressor, mesh_gt)[smpl_x.J_regressor_idx['pelvis'], None,
:]
mpvpe_all = np.sqrt(np.sum((mesh_out_align - mesh_gt) ** 2, 1)).mean() * 1000
eval_result['mpvpe_all'].append(mpvpe_all)
mesh_out_align = rigid_align(mesh_out, mesh_gt)
pa_mpvpe_all = np.sqrt(np.sum((mesh_out_align - mesh_gt) ** 2, 1)).mean() * 1000
eval_result['pa_mpvpe_all'].append(pa_mpvpe_all)
# MPVPE from hand vertices
mesh_gt_lhand = mesh_gt[smpl_x.hand_vertex_idx['left_hand'], :]
mesh_out_lhand = mesh_out[smpl_x.hand_vertex_idx['left_hand'], :]
mesh_gt_rhand = mesh_gt[smpl_x.hand_vertex_idx['right_hand'], :]
mesh_out_rhand = mesh_out[smpl_x.hand_vertex_idx['right_hand'], :]
mesh_out_lhand_align = mesh_out_lhand - np.dot(smpl_x.J_regressor, mesh_out)[
smpl_x.J_regressor_idx['lwrist'], None, :] + np.dot(
smpl_x.J_regressor, mesh_gt)[smpl_x.J_regressor_idx['lwrist'], None, :]
mesh_out_rhand_align = mesh_out_rhand - np.dot(smpl_x.J_regressor, mesh_out)[
smpl_x.J_regressor_idx['rwrist'], None, :] + np.dot(
smpl_x.J_regressor, mesh_gt)[smpl_x.J_regressor_idx['rwrist'], None, :]
eval_result['mpvpe_l_hand'].append(np.sqrt(
np.sum((mesh_out_lhand_align - mesh_gt_lhand) ** 2, 1)).mean() * 1000)
eval_result['mpvpe_r_hand'].append(np.sqrt(
np.sum((mesh_out_rhand_align - mesh_gt_rhand) ** 2, 1)).mean() * 1000)
eval_result['mpvpe_hand'].append((np.sqrt(
np.sum((mesh_out_lhand_align - mesh_gt_lhand) ** 2, 1)).mean() * 1000 + np.sqrt(
np.sum((mesh_out_rhand_align - mesh_gt_rhand) ** 2, 1)).mean() * 1000) / 2.)
mesh_out_lhand_align = rigid_align(mesh_out_lhand, mesh_gt_lhand)
mesh_out_rhand_align = rigid_align(mesh_out_rhand, mesh_gt_rhand)
eval_result['pa_mpvpe_l_hand'].append(np.sqrt(
np.sum((mesh_out_lhand_align - mesh_gt_lhand) ** 2, 1)).mean() * 1000)
eval_result['pa_mpvpe_r_hand'].append(np.sqrt(
np.sum((mesh_out_rhand_align - mesh_gt_rhand) ** 2, 1)).mean() * 1000)
eval_result['pa_mpvpe_hand'].append((np.sqrt(
np.sum((mesh_out_lhand_align - mesh_gt_lhand) ** 2, 1)).mean() * 1000 + np.sqrt(
np.sum((mesh_out_rhand_align - mesh_gt_rhand) ** 2, 1)).mean() * 1000) / 2.)
# MPVPE from face vertices
mesh_gt_face = mesh_gt[smpl_x.face_vertex_idx, :]
mesh_out_face = mesh_out[smpl_x.face_vertex_idx, :]
mesh_out_face_align = mesh_out_face - np.dot(smpl_x.J_regressor, mesh_out)[smpl_x.J_regressor_idx['neck'],
None, :] + np.dot(smpl_x.J_regressor, mesh_gt)[
smpl_x.J_regressor_idx['neck'], None, :]
eval_result['mpvpe_face'].append(
np.sqrt(np.sum((mesh_out_face_align - mesh_gt_face) ** 2, 1)).mean() * 1000)
mesh_out_face_align = rigid_align(mesh_out_face, mesh_gt_face)
eval_result['pa_mpvpe_face'].append(
np.sqrt(np.sum((mesh_out_face_align - mesh_gt_face) ** 2, 1)).mean() * 1000)
### HARDCODE
if vis:
# from utils.vis import vis_keypoints, vis_mesh, save_obj, render_mesh
# # img = (out['img'].transpose(1,2,0)[:,:,::-1] * 255).copy()
# # joint_img = out['joint_img'].copy()
# # joint_img[:,0] = joint_img[:,0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1]
# # joint_img[:,1] = joint_img[:,1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0]
# # for j in range(len(joint_img)):
# # cv2.circle(img, (int(joint_img[j][0]), int(joint_img[j][1])), 3, (0,0,255), -1)
# # cv2.imwrite(str(cur_sample_idx + n) + '.jpg', img)
# img_path = annot['img_path']
# img_id = img_path.split('/')[-1][:-4]
# ann_id = 0
# # ann_id = annot['ann_id']
# img = load_img(img_path)[:, :, ::-1]
# bbox = annot['bbox']
# focal = list(cfg.focal)
# princpt = list(cfg.princpt)
# focal[0] = focal[0] / cfg.input_body_shape[1] * bbox[2]
# focal[1] = focal[1] / cfg.input_body_shape[0] * bbox[3]
# princpt[0] = princpt[0] / cfg.input_body_shape[1] * bbox[2] + bbox[0]
# princpt[1] = princpt[1] / cfg.input_body_shape[0] * bbox[3] + bbox[1]
# img = render_mesh(img, out['smplx_mesh_cam'], smpl_x.face, {'focal': focal, 'princpt': princpt}, mesh_as_vertices=True)
# # img = cv2.resize(img, (512,512))
# cv2.imwrite(osp.join(vis_save_dir, img_id + '_' + str(ann_id) + '.jpg'), img)
# vis_mesh_out = out['smplx_mesh_cam']
# vis_mesh_out = vis_mesh_out - np.dot(smpl_x.layer['neutral'].J_regressor, vis_mesh_out)[
# smpl_x.J_regressor_idx['pelvis'], None, :]
# # vis_mesh_gt = out['smplx_mesh_cam_target']
# # vis_mesh_gt = vis_mesh_gt - np.dot(smpl_x.layer['neutral'].J_regressor, vis_mesh_gt)[smpl_x.J_regressor_idx['pelvis'],None,:]
# # save_obj(vis_mesh_out, smpl_x.face, osp.join(img_id + '_' + str(ann_id) + '.obj'))
# # save_obj(vis_mesh_gt, smpl_x.face, str(cur_sample_idx + n) + '_gt.obj')
img_path = out['img_path']
rel_img_path = img_path.split('..')[-1]
smplx_pred = {}
smplx_pred['global_orient'] = out['smplx_root_pose'].reshape(-1,3)
smplx_pred['body_pose'] = out['smplx_body_pose'].reshape(-1,3)
smplx_pred['left_hand_pose'] = out['smplx_lhand_pose'].reshape(-1,3)
smplx_pred['right_hand_pose'] = out['smplx_rhand_pose'].reshape(-1,3)
smplx_pred['jaw_pose'] = out['smplx_jaw_pose'].reshape(-1,3)
smplx_pred['leye_pose'] = np.zeros((1, 3))
smplx_pred['reye_pose'] = np.zeros((1, 3))
smplx_pred['betas'] = out['smplx_shape'].reshape(-1,10)
smplx_pred['expression'] = out['smplx_expr'].reshape(-1,10)
smplx_pred['transl'] = out['gt_smplx_transl'].reshape(-1,3)
smplx_pred['img_path'] = rel_img_path
npz_path = os.path.join(cfg.vis_dir, f'{self.save_idx}.npz')
np.savez(npz_path, **smplx_pred)
# save img path and error
new_line = [self.save_idx, rel_img_path, mpvpe_all, pa_mpvpe_all]
# Append the new line to the CSV file
writer.writerow(new_line)
self.save_idx += 1
# save_obj(out['smplx_mesh_cam'], smpl_x.face, str(cur_sample_idx + n) + '.obj')
# save results for the official evaluation codes/server
save_name = annot['img_path'].split('/')[-1][:-4]
if self.data_split == 'test' and self.test_set == 'test':
if self.resolution == (2160, 3840):
save_name = save_name.split('_pid')[0]
elif self.data_split == 'test' and self.test_set == 'val':
if self.resolution == (2160, 3840):
save_name = save_name.split('_ann_id')[0]
else:
save_name = save_name.split('_1280x720')[0]
if 'person_idx' in annot:
person_idx = annot['person_idx']
else:
exist_result_path = glob(osp.join(cfg.result_dir, 'AGORA', save_name + '*'))
if len(exist_result_path) == 0:
person_idx = 0
else:
last_person_idx = max(
[int(name.split('personId_')[1].split('.pkl')[0]) for name in exist_result_path])
person_idx = last_person_idx + 1
save_name += '_personId_' + str(person_idx) + '.pkl'
joint_proj = out['smplx_joint_proj']
joint_proj[:, 0] = joint_proj[:, 0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1]
joint_proj[:, 1] = joint_proj[:, 1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0]
joint_proj = np.concatenate((joint_proj, np.ones_like(joint_proj[:, :1])), 1)
joint_proj = np.dot(out['bb2img_trans'], joint_proj.transpose(1, 0)).transpose(1, 0)
joint_proj[:, 0] = joint_proj[:, 0] / self.resolution[1] * 3840 # restore to original resolution
joint_proj[:, 1] = joint_proj[:, 1] / self.resolution[0] * 2160 # restore to original resolution
save_dict = {'params':
{'transl': out['cam_trans'].reshape(1, -1),
'global_orient': out['smplx_root_pose'].reshape(1, -1),
'body_pose': out['smplx_body_pose'].reshape(1, -1),
'left_hand_pose': out['smplx_lhand_pose'].reshape(1, -1),
'right_hand_pose': out['smplx_rhand_pose'].reshape(1, -1),
'reye_pose': np.zeros((1, 3)),
'leye_pose': np.zeros((1, 3)),
'jaw_pose': out['smplx_jaw_pose'].reshape(1, -1),
'expression': out['smplx_expr'].reshape(1, -1),
'betas': out['smplx_shape'].reshape(1, -1)},
'joints': joint_proj.reshape(1, -1, 2)
}
os.makedirs(osp.join(cfg.result_dir, 'predictions'), exist_ok=True)
with open(osp.join(cfg.result_dir, 'predictions', save_name), 'wb') as f:
pickle.dump(save_dict, f)
"""
# for debug
img_path = annot['img_path']
img_path = osp.join(self.data_path, '3840x2160', 'test', img_path.split('/')[-1].split('_')[0] + '.png')
img = cv2.imread(img_path)
img = vis_keypoints(img.copy(), joint_proj)
cv2.imwrite(img_path.split('/')[-1], img)
"""
if getattr(cfg, 'vis', False):
file.close()
return eval_result
def print_eval_result(self, eval_result):
print('AGORA test results are dumped at: ' + osp.join(cfg.result_dir, 'predictions'))
if self.data_split == 'test' and self.test_set == 'test': # do not print. just submit the results to the official evaluation server
return
print('======AGORA-val======')
print(f'{cfg.vis_dir}')
print('PA MPVPE (All): %.2f mm' % np.mean(eval_result['pa_mpvpe_all']))
print('PA MPVPE (L-Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_l_hand']))
print('PA MPVPE (R-Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_r_hand']))
print('PA MPVPE (Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_hand']))
print('PA MPVPE (Face): %.2f mm' % np.mean(eval_result['pa_mpvpe_face']))
print()
print('MPVPE (All): %.2f mm' % np.mean(eval_result['mpvpe_all']))
print('MPVPE (L-Hands): %.2f mm' % np.mean(eval_result['mpvpe_l_hand']))
print('MPVPE (R-Hands): %.2f mm' % np.mean(eval_result['mpvpe_r_hand']))
print('MPVPE (Hands): %.2f mm' % np.mean(eval_result['mpvpe_hand']))
print('MPVPE (Face): %.2f mm' % np.mean(eval_result['mpvpe_face']))
print()
print(f"{np.mean(eval_result['pa_mpvpe_all'])},{np.mean(eval_result['pa_mpvpe_l_hand'])},{np.mean(eval_result['pa_mpvpe_r_hand'])},{np.mean(eval_result['pa_mpvpe_hand'])},{np.mean(eval_result['pa_mpvpe_face'])},"
f"{np.mean(eval_result['mpvpe_all'])},{np.mean(eval_result['mpvpe_l_hand'])},{np.mean(eval_result['mpvpe_r_hand'])},{np.mean(eval_result['mpvpe_hand'])},{np.mean(eval_result['mpvpe_face'])}")
print()
f = open(os.path.join(cfg.result_dir, 'result.txt'), 'w')
f.write(f'AGORA-val dataset: \n')
f.write('PA MPVPE (All): %.2f mm\n' % np.mean(eval_result['pa_mpvpe_all']))
f.write('PA MPVPE (L-Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_l_hand']))
f.write('PA MPVPE (R-Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_r_hand']))
f.write('PA MPVPE (Hands): %.2f mm\n' % np.mean(eval_result['pa_mpvpe_hand']))
f.write('PA MPVPE (Face): %.2f mm\n' % np.mean(eval_result['pa_mpvpe_face']))
f.write('MPVPE (All): %.2f mm\n' % np.mean(eval_result['mpvpe_all']))
f.write('MPVPE (L-Hands): %.2f mm' % np.mean(eval_result['mpvpe_l_hand']))
f.write('MPVPE (R-Hands): %.2f mm' % np.mean(eval_result['mpvpe_r_hand']))
f.write('MPVPE (Hands): %.2f mm' % np.mean(eval_result['mpvpe_hand']))
f.write('MPVPE (Face): %.2f mm\n' % np.mean(eval_result['mpvpe_face']))
f.write(f"{np.mean(eval_result['pa_mpvpe_all'])},{np.mean(eval_result['pa_mpvpe_l_hand'])},{np.mean(eval_result['pa_mpvpe_r_hand'])},{np.mean(eval_result['pa_mpvpe_hand'])},{np.mean(eval_result['pa_mpvpe_face'])},"
f"{np.mean(eval_result['mpvpe_all'])},{np.mean(eval_result['mpvpe_l_hand'])},{np.mean(eval_result['mpvpe_r_hand'])},{np.mean(eval_result['mpvpe_hand'])},{np.mean(eval_result['mpvpe_face'])}")
if getattr(cfg, 'eval_on_train', False):
import csv
csv_file = f'{cfg.root_dir}/output/agora_eval_on_train.csv'
exp_id = cfg.exp_name.split('_')[1]
new_line = [exp_id,np.mean(eval_result['pa_mpvpe_all']),np.mean(eval_result['pa_mpvpe_l_hand']),np.mean(eval_result['pa_mpvpe_r_hand']),np.mean(eval_result['pa_mpvpe_hand']),np.mean(eval_result['pa_mpvpe_face']),
np.mean(eval_result['mpvpe_all']),np.mean(eval_result['mpvpe_l_hand']),np.mean(eval_result['mpvpe_r_hand']),np.mean(eval_result['mpvpe_hand']),np.mean(eval_result['mpvpe_face'])]
# Append the new line to the CSV file
with open(csv_file, 'a', newline='') as file:
writer = csv.writer(file)
writer.writerow(new_line)
data/ARCTIC/ARCTIC.py
+51
View File
@@ -0,0 +1,51 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class ARCTIC(HumanDataset):
def __init__(self, transform, data_split):
super(ARCTIC, self).__init__(transform, data_split)
if getattr(cfg, 'eval_on_train', False):
self.data_split = 'eval_train'
print("Evaluate on train set.")
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', f'arctic_{self.data_split}.npz')
self.img_shape = None # (h, w)
self.cam_param = {}
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
if 'train' in data_split:
filename = getattr(cfg, 'filename', 'p1_train.npz')
else:
filename = getattr(cfg, 'filename', 'p1_val.npz')
self.img_dir = osp.join(cfg.data_dir, 'ARCTIC')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
# load data
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1),
test_sample_interval=getattr(cfg, f'{self.__class__.__name__}_test_sample_interval', 10))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/BEDLAM/BEDLAM.py
+45
View File
@@ -0,0 +1,45 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class BEDLAM(HumanDataset):
def __init__(self, transform, data_split):
super(BEDLAM, self).__init__(transform, data_split)
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'bedlam_train.npz')
self.img_shape = None # (h, w)
self.cam_param = {}
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
if self.data_split == 'train':
filename = getattr(cfg, 'filename', 'bedlam_train.npz')
else:
raise ValueError('BEDLAM test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'BEDLAM')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
# load data
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/Behave/Behave.py
+50
View File
@@ -0,0 +1,50 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class Behave(HumanDataset):
def __init__(self, transform, data_split):
super(Behave, self).__init__(transform, data_split)
pre_prc_file_train = 'behave_train_230516_231_downsampled.npz'
pre_prc_file_test = 'behave_test_230516_090_downsampled.npz'
if self.data_split == 'train':
filename = getattr(cfg, 'filename', pre_prc_file_train)
else:
filename = getattr(cfg, 'filename', pre_prc_file_test)
self.img_dir = osp.join(cfg.data_dir, 'Behave')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = (1536, 2048) # (h, w)
self.cam_param = {}
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/CHI3D/CHI3D.py
+56
View File
@@ -0,0 +1,56 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class CHI3D(HumanDataset):
def __init__(self, transform, data_split):
super(CHI3D, self).__init__(transform, data_split)
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'CHI3D_train_230511_1492.npz')
self.img_shape = (900, 900) # (h, w)
self.cam_param = {}
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = []
for pre_prc_file in ['CHI3D_train_230511_1492_0.npz','CHI3D_train_230511_1492_1.npz']:
if self.data_split == 'train':
filename = getattr(cfg, 'filename', pre_prc_file)
else:
raise ValueError('CHI3D test set is not support')
self.img_dir = cfg.data_dir # CHI3D included
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data
datalist_slice = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
self.datalist.extend(datalist_slice)
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/CrowdPose/CrowdPose.py
+47
View File
@@ -0,0 +1,47 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class CrowdPose(HumanDataset):
def __init__(self, transform, data_split):
super(CrowdPose, self).__init__(transform, data_split)
self.datalist = []
pre_prc_file = 'crowdpose_neural_annot_train_new.npz'
if self.data_split == 'train':
filename = getattr(cfg, 'filename', pre_prc_file)
else:
raise ValueError('CrowdPose test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'CrowdPose')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = None
self.cam_param = {}
print("Various image shape in CrowdPose dataset.")
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/EHF/EHF.py
+373
View File
@@ -0,0 +1,373 @@
import os
import os.path as osp
from glob import glob
import numpy as np
from config import cfg
import copy
import json
import cv2
import torch
from pycocotools.coco import COCO
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, load_ply
from utils.transforms import rigid_align
class EHF(torch.utils.data.Dataset):
def __init__(self, transform, data_split):
self.transform = transform
self.data_split = data_split
# self.data_path = osp.join('..', 'data', 'EHF', 'data')
self.data_path = osp.join(cfg.data_dir, 'EHF', 'data')
self.datalist = self.load_data()
self.cam_param = {'R': [-2.98747896, 0.01172457, -0.05704687]}
self.cam_param['R'], _ = cv2.Rodrigues(np.array(self.cam_param['R']))
self.save_idx = 0
def load_data(self):
datalist = []
db = COCO(osp.join(self.data_path, 'EHF.json'))
for aid in db.anns.keys():
ann = db.anns[aid]
image_id = ann['image_id']
img = db.loadImgs(image_id)[0]
img_shape = (img['height'], img['width'])
img_path = osp.join(self.data_path, img['file_name'])
bbox = ann['body_bbox']
bbox = process_bbox(bbox, img['width'], img['height'], ratio=getattr(cfg, 'bbox_ratio', 1.25))
if bbox is None:
continue
lhand_bbox = np.array(ann['lefthand_bbox']).reshape(4)
if hasattr(cfg, 'bbox_ratio'):
lhand_bbox = process_bbox(lhand_bbox, img['width'], img['height'], ratio=cfg.bbox_ratio)
lhand_bbox[2:] += lhand_bbox[:2] # xywh -> xyxy
rhand_bbox = np.array(ann['righthand_bbox']).reshape(4)
if hasattr(cfg, 'bbox_ratio'):
rhand_bbox = process_bbox(rhand_bbox, img['width'], img['height'], ratio=cfg.bbox_ratio)
rhand_bbox[2:] += rhand_bbox[:2] # xywh -> xyxy
face_bbox = np.array(ann['face_bbox']).reshape(4)
if hasattr(cfg, 'bbox_ratio'):
face_bbox = process_bbox(face_bbox, img['width'], img['height'], ratio=cfg.bbox_ratio)
face_bbox[2:] += face_bbox[:2] # xywh -> xyxy
mesh_gt_path = osp.join(self.data_path, img['file_name'].split('_')[0] + '_align.ply')
data_dict = {'img_path': img_path, 'img_shape': img_shape, 'bbox': bbox, 'lhand_bbox': lhand_bbox,
'rhand_bbox': rhand_bbox, 'face_bbox': face_bbox, 'mesh_gt_path': mesh_gt_path}
datalist.append(data_dict)
return datalist
def process_hand_face_bbox(self, bbox, do_flip, img_shape, img2bb_trans):
if bbox is None:
bbox = np.zeros((2, 2), dtype=np.float32) # dummy value
bbox_valid = float(False) # dummy value
else:
# reshape to top-left (x,y) and bottom-right (x,y)
bbox = bbox.reshape(2, 2)
# flip augmentation
if do_flip:
bbox[:, 0] = img_shape[1] - bbox[:, 0] - 1
bbox[0, 0], bbox[1, 0] = bbox[1, 0].copy(), bbox[0, 0].copy() # xmin <-> xmax swap
# make four points of the bbox
bbox = bbox.reshape(4).tolist()
xmin, ymin, xmax, ymax = bbox
bbox = np.array([[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]], dtype=np.float32).reshape(4, 2)
# affine transformation (crop, rotation, scale)
bbox_xy1 = np.concatenate((bbox, np.ones_like(bbox[:, :1])), 1)
bbox = np.dot(img2bb_trans, bbox_xy1.transpose(1, 0)).transpose(1, 0)[:, :2]
bbox[:, 0] = bbox[:, 0] / cfg.input_img_shape[1] * cfg.output_hm_shape[1]
bbox[:, 1] = bbox[:, 1] / cfg.input_img_shape[0] * cfg.output_hm_shape[0]
# make box a rectangle without rotation
xmin = np.min(bbox[:, 0]);
xmax = np.max(bbox[:, 0]);
ymin = np.min(bbox[:, 1]);
ymax = np.max(bbox[:, 1]);
bbox = np.array([xmin, ymin, xmax, ymax], dtype=np.float32)
bbox_valid = float(True)
bbox = bbox.reshape(2, 2)
return bbox, bbox_valid
def __len__(self):
return len(self.datalist)
def __getitem__(self, idx):
data = copy.deepcopy(self.datalist[idx])
img_path, img_shape, bbox, mesh_gt_path = data['img_path'], data['img_shape'], data['bbox'], data[
'mesh_gt_path']
# image load
img = load_img(img_path)
# affine transform
img, img2bb_trans, bb2img_trans, rot, do_flip = augmentation(img, bbox, self.data_split)
img = self.transform(img.astype(np.float32)) / 255.
# hand and face bbox transform
lhand_bbox, lhand_bbox_valid = self.process_hand_face_bbox(data['lhand_bbox'], do_flip, img_shape, img2bb_trans)
rhand_bbox, rhand_bbox_valid = self.process_hand_face_bbox(data['rhand_bbox'], do_flip, img_shape, img2bb_trans)
face_bbox, face_bbox_valid = self.process_hand_face_bbox(data['face_bbox'], do_flip, img_shape, img2bb_trans)
if do_flip:
lhand_bbox, rhand_bbox = rhand_bbox, lhand_bbox
lhand_bbox_valid, rhand_bbox_valid = rhand_bbox_valid, lhand_bbox_valid
lhand_bbox_center = (lhand_bbox[0] + lhand_bbox[1]) / 2.;
rhand_bbox_center = (rhand_bbox[0] + rhand_bbox[1]) / 2.;
face_bbox_center = (face_bbox[0] + face_bbox[1]) / 2.
lhand_bbox_size = lhand_bbox[1] - lhand_bbox[0];
rhand_bbox_size = rhand_bbox[1] - rhand_bbox[0];
face_bbox_size = face_bbox[1] - face_bbox[0];
# mesh gt load
mesh_gt = load_ply(mesh_gt_path)
inputs = {'img': img}
targets = {'smplx_mesh_cam': mesh_gt, 'lhand_bbox_center': lhand_bbox_center,
'rhand_bbox_center': rhand_bbox_center, 'face_bbox_center': face_bbox_center,
'lhand_bbox_size': lhand_bbox_size, 'rhand_bbox_size': rhand_bbox_size,
'face_bbox_size': face_bbox_size}
meta_info = {'bb2img_trans': bb2img_trans, 'lhand_bbox_valid': float(True), 'rhand_bbox_valid': float(True),
'face_bbox_valid': float(True),
'img_path': img_path}
return inputs, targets, meta_info
def evaluate(self, outs, cur_sample_idx):
annots = self.datalist
sample_num = len(outs)
eval_result = {'pa_mpvpe_all': [], 'pa_mpvpe_l_hand': [], 'pa_mpvpe_r_hand': [], 'pa_mpvpe_hand': [], 'pa_mpvpe_face': [],
'mpvpe_all': [], 'mpvpe_l_hand': [], 'mpvpe_r_hand': [], 'mpvpe_hand': [], 'mpvpe_face': [],
'pa_mpjpe_body': [], 'pa_mpjpe_l_hand': [], 'pa_mpjpe_r_hand': [], 'pa_mpjpe_hand': []}
if getattr(cfg, 'vis', False):
import csv
csv_file = f'{cfg.vis_dir}/ehf_smplx_error.csv'
file = open(csv_file, 'a', newline='')
writer = csv.writer(file)
for n in range(sample_num):
annot = annots[cur_sample_idx + n]
ann_id = annot['img_path'].split('/')[-1].split('_')[0]
# print(annot['img_path'])
# ann_id = annot['ann_id']
out = outs[n]
# MPVPE from all vertices
mesh_gt = np.dot(self.cam_param['R'], out['smplx_mesh_cam_target'].transpose(1, 0)).transpose(1, 0)
mesh_out = out['smplx_mesh_cam']
# mesh_gt_align = rigid_align(mesh_gt, mesh_out)
# print(mesh_out.shape)
mesh_out_align = rigid_align(mesh_out, mesh_gt)
pa_mpvpe_all = np.sqrt(np.sum((mesh_out_align - mesh_gt) ** 2, 1)).mean() * 1000
eval_result['pa_mpvpe_all'].append(pa_mpvpe_all)
mesh_out_align = mesh_out - np.dot(smpl_x.J_regressor, mesh_out)[smpl_x.J_regressor_idx['pelvis'], None,
:] + np.dot(smpl_x.J_regressor, mesh_gt)[smpl_x.J_regressor_idx['pelvis'], None,
:]
mpvpe_all = np.sqrt(np.sum((mesh_out_align - mesh_gt) ** 2, 1)).mean() * 1000
eval_result['mpvpe_all'].append(mpvpe_all)
# MPVPE from hand vertices
mesh_gt_lhand = mesh_gt[smpl_x.hand_vertex_idx['left_hand'], :]
mesh_out_lhand = mesh_out[smpl_x.hand_vertex_idx['left_hand'], :]
mesh_out_lhand_align = rigid_align(mesh_out_lhand, mesh_gt_lhand)
mesh_gt_rhand = mesh_gt[smpl_x.hand_vertex_idx['right_hand'], :]
mesh_out_rhand = mesh_out[smpl_x.hand_vertex_idx['right_hand'], :]
mesh_out_rhand_align = rigid_align(mesh_out_rhand, mesh_gt_rhand)
eval_result['pa_mpvpe_l_hand'].append(np.sqrt(
np.sum((mesh_out_lhand_align - mesh_gt_lhand) ** 2, 1)).mean() * 1000)
eval_result['pa_mpvpe_r_hand'].append(np.sqrt(
np.sum((mesh_out_rhand_align - mesh_gt_rhand) ** 2, 1)).mean() * 1000)
eval_result['pa_mpvpe_hand'].append((np.sqrt(
np.sum((mesh_out_lhand_align - mesh_gt_lhand) ** 2, 1)).mean() * 1000 + np.sqrt(
np.sum((mesh_out_rhand_align - mesh_gt_rhand) ** 2, 1)).mean() * 1000) / 2.)
mesh_out_lhand_align = mesh_out_lhand - np.dot(smpl_x.J_regressor, mesh_out)[
smpl_x.J_regressor_idx['lwrist'], None, :] + np.dot(
smpl_x.J_regressor, mesh_gt)[smpl_x.J_regressor_idx['lwrist'], None, :]
mesh_out_rhand_align = mesh_out_rhand - np.dot(smpl_x.J_regressor, mesh_out)[
smpl_x.J_regressor_idx['rwrist'], None, :] + np.dot(
smpl_x.J_regressor, mesh_gt)[smpl_x.J_regressor_idx['rwrist'], None, :]
eval_result['mpvpe_l_hand'].append(np.sqrt(
np.sum((mesh_out_lhand_align - mesh_gt_lhand) ** 2, 1)).mean() * 1000)
eval_result['mpvpe_r_hand'].append(np.sqrt(
np.sum((mesh_out_rhand_align - mesh_gt_rhand) ** 2, 1)).mean() * 1000)
eval_result['mpvpe_hand'].append((np.sqrt(
np.sum((mesh_out_lhand_align - mesh_gt_lhand) ** 2, 1)).mean() * 1000 + np.sqrt(
np.sum((mesh_out_rhand_align - mesh_gt_rhand) ** 2, 1)).mean() * 1000) / 2.)
# MPVPE from face vertices
mesh_gt_face = mesh_gt[smpl_x.face_vertex_idx, :]
mesh_out_face = mesh_out[smpl_x.face_vertex_idx, :]
mesh_out_face_align = rigid_align(mesh_out_face, mesh_gt_face)
eval_result['pa_mpvpe_face'].append(
np.sqrt(np.sum((mesh_out_face_align - mesh_gt_face) ** 2, 1)).mean() * 1000)
mesh_out_face_align = mesh_out_face - np.dot(smpl_x.J_regressor, mesh_out)[smpl_x.J_regressor_idx['neck'],
None, :] + np.dot(smpl_x.J_regressor, mesh_gt)[
smpl_x.J_regressor_idx['neck'], None, :]
eval_result['mpvpe_face'].append(
np.sqrt(np.sum((mesh_out_face_align - mesh_gt_face) ** 2, 1)).mean() * 1000)
# MPJPE from body joints
joint_gt_body = np.dot(smpl_x.j14_regressor, mesh_gt)
joint_out_body = np.dot(smpl_x.j14_regressor, mesh_out)
joint_out_body_align = rigid_align(joint_out_body, joint_gt_body)
eval_result['pa_mpjpe_body'].append(
np.sqrt(np.sum((joint_out_body_align - joint_gt_body) ** 2, 1)).mean() * 1000)
# MPJPE from hand joints
joint_gt_lhand = np.dot(smpl_x.orig_hand_regressor['left'], mesh_gt)
joint_out_lhand = np.dot(smpl_x.orig_hand_regressor['left'], mesh_out)
joint_out_lhand_align = rigid_align(joint_out_lhand, joint_gt_lhand)
joint_gt_rhand = np.dot(smpl_x.orig_hand_regressor['right'], mesh_gt)
joint_out_rhand = np.dot(smpl_x.orig_hand_regressor['right'], mesh_out)
joint_out_rhand_align = rigid_align(joint_out_rhand, joint_gt_rhand)
eval_result['pa_mpjpe_l_hand'].append(np.sqrt(
np.sum((joint_out_lhand_align - joint_gt_lhand) ** 2, 1)).mean() * 1000)
eval_result['pa_mpjpe_r_hand'].append(np.sqrt(
np.sum((joint_out_rhand_align - joint_gt_rhand) ** 2, 1)).mean() * 1000)
eval_result['pa_mpjpe_hand'].append((np.sqrt(
np.sum((joint_out_lhand_align - joint_gt_lhand) ** 2, 1)).mean() * 1000 + np.sqrt(
np.sum((joint_out_rhand_align - joint_gt_rhand) ** 2, 1)).mean() * 1000) / 2.)
vis = cfg.vis
if vis:
# save_folder = cfg.vis_dir
# kpt_save_folder = os.path.join(save_folder, 'KPT')
# os.makedirs(kpt_save_folder, exist_ok=True)
# mesh_save_folder = os.path.join(save_folder, 'mesh_origin')
# os.makedirs(mesh_save_folder, exist_ok=True)
# # from utils.vis import vis_keypoints, render_mesh, save_obj
# img = (out['img'].transpose(1, 2, 0)[:, :, ::-1] * 255).copy()
# joint_img = out['joint_img'].copy()
# joint_img[:, 0] = joint_img[:, 0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1]
# joint_img[:, 1] = joint_img[:, 1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0]
# for j in range(len(joint_img)):
# cv2.circle(img, (int(joint_img[j][0]), int(joint_img[j][1])), 3, (0, 0, 255), -1)
# lhand_bbox = out['lhand_bbox'].reshape(2, 2).copy()
# cv2.rectangle(img, (int(lhand_bbox[0][0]), int(lhand_bbox[0][1])),
# (int(lhand_bbox[1][0]), int(lhand_bbox[1][1])), (255, 0, 0), 3)
# rhand_bbox = out['rhand_bbox'].reshape(2, 2).copy()
# cv2.rectangle(img, (int(rhand_bbox[0][0]), int(rhand_bbox[0][1])),
# (int(rhand_bbox[1][0]), int(rhand_bbox[1][1])), (255, 0, 0), 3)
# face_bbox = out['face_bbox'].reshape(2, 2).copy()
# cv2.rectangle(img, (int(face_bbox[0][0]), int(face_bbox[0][1])),
# (int(face_bbox[1][0]), int(face_bbox[1][1])), (255, 0, 0), 3)
# cv2.imwrite(os.path.join(kpt_save_folder, str(cur_sample_idx + n) + '.jpg'), img)
# vis_img = img.copy()
# focal = [cfg.focal[0] / cfg.input_body_shape[1] * cfg.input_img_shape[1],
# cfg.focal[1] / cfg.input_body_shape[0] * cfg.input_img_shape[0]]
# princpt = [cfg.princpt[0] / cfg.input_body_shape[1] * cfg.input_img_shape[1],
# cfg.princpt[1] / cfg.input_body_shape[0] * cfg.input_img_shape[0]]
# rendered_img = render_mesh(vis_img, mesh_out, smpl_x.face, {'focal': focal, 'princpt': princpt})
# vis_img = img.copy()
# # rendered_img_gt = render_mesh(vis_img, mesh_gt_align, smpl_x.face, {'focal': focal, 'princpt': princpt})
# cv2.imwrite(os.path.join(mesh_save_folder, f'{ann_id}_render.jpg'), rendered_img)
# # cv2.imwrite(os.path.join(mesh_save_folder, f'{ann_id}_render_gt.jpg'), rendered_img_gt)
# cv2.imwrite(os.path.join(mesh_save_folder, f'{ann_id}.jpg'), vis_img)
# np.save(os.path.join(mesh_save_folder, f'{ann_id}.npy'), mesh_out)
# save smplx param
smplx_pred = {}
smplx_pred['global_orient'] = out['smplx_root_pose'].reshape(-1,3)
smplx_pred['body_pose'] = out['smplx_body_pose'].reshape(-1,3)
smplx_pred['left_hand_pose'] = out['smplx_lhand_pose'].reshape(-1,3)
smplx_pred['right_hand_pose'] = out['smplx_rhand_pose'].reshape(-1,3)
smplx_pred['jaw_pose'] = out['smplx_jaw_pose'].reshape(-1,3)
smplx_pred['leye_pose'] = np.zeros((1, 3))
smplx_pred['reye_pose'] = np.zeros((1, 3))
smplx_pred['betas'] = out['smplx_shape'].reshape(-1,10)
smplx_pred['expression'] = out['smplx_expr'].reshape(-1,10)
smplx_pred['transl'] = out['cam_trans'].reshape(-1,3)
# import pdb; pdb.set_trace()
np.savez(os.path.join(cfg.vis_dir, f'{self.save_idx}.npz'), **smplx_pred)
# save img path and error
img_path = out['img_path']
rel_img_path = img_path.split('..')[-1]
new_line = [self.save_idx, rel_img_path, mpvpe_all, pa_mpvpe_all]
# Append the new line to the CSV file
writer.writerow(new_line)
self.save_idx += 1
# save_obj(out['smplx_mesh_cam'], smpl_x.face, str(cur_sample_idx + n) + '.obj')
if getattr(cfg, 'vis', False):
file.close()
return eval_result
def print_eval_result(self, eval_result):
print('======EHF======')
print('PA MPVPE (All): %.2f mm' % np.mean(eval_result['pa_mpvpe_all']))
print('PA MPVPE (L-Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_l_hand']))
print('PA MPVPE (R-Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_r_hand']))
print('PA MPVPE (Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_hand']))
print('PA MPVPE (Face): %.2f mm' % np.mean(eval_result['pa_mpvpe_face']))
print()
print('MPVPE (All): %.2f mm' % np.mean(eval_result['mpvpe_all']))
print('MPVPE (L-Hands): %.2f mm' % np.mean(eval_result['mpvpe_l_hand']))
print('MPVPE (R-Hands): %.2f mm' % np.mean(eval_result['mpvpe_r_hand']))
print('MPVPE (Hands): %.2f mm' % np.mean(eval_result['mpvpe_hand']))
print('MPVPE (Face): %.2f mm' % np.mean(eval_result['mpvpe_face']))
print()
print('PA MPJPE (Body): %.2f mm' % np.mean(eval_result['pa_mpjpe_body']))
print('PA MPJPE (L-Hands): %.2f mm' % np.mean(eval_result['pa_mpjpe_l_hand']))
print('PA MPJPE (R-Hands): %.2f mm' % np.mean(eval_result['pa_mpjpe_r_hand']))
print('PA MPJPE (Hands): %.2f mm' % np.mean(eval_result['pa_mpjpe_hand']))
print()
print(f"{np.mean(eval_result['pa_mpvpe_all'])},{np.mean(eval_result['pa_mpvpe_l_hand'])},{np.mean(eval_result['pa_mpvpe_r_hand'])},{np.mean(eval_result['pa_mpvpe_hand'])},{np.mean(eval_result['pa_mpvpe_face'])},"
f"{np.mean(eval_result['mpvpe_all'])},{np.mean(eval_result['mpvpe_l_hand'])},{np.mean(eval_result['mpvpe_r_hand'])},{np.mean(eval_result['mpvpe_hand'])},{np.mean(eval_result['mpvpe_face'])}")
print()
f = open(os.path.join(cfg.result_dir, 'result.txt'), 'w')
f.write(f'EHF dataset: \n')
f.write('PA MPVPE (All): %.2f mm\n' % np.mean(eval_result['pa_mpvpe_all']))
f.write('PA MPVPE (L-Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_l_hand']))
f.write('PA MPVPE (R-Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_r_hand']))
f.write('PA MPVPE (Hands): %.2f mm\n' % np.mean(eval_result['pa_mpvpe_hand']))
f.write('PA MPVPE (Face): %.2f mm\n' % np.mean(eval_result['pa_mpvpe_face']))
f.write('MPVPE (All): %.2f mm\n' % np.mean(eval_result['mpvpe_all']))
f.write('MPVPE (L-Hands): %.2f mm' % np.mean(eval_result['mpvpe_l_hand']))
f.write('MPVPE (R-Hands): %.2f mm' % np.mean(eval_result['mpvpe_r_hand']))
f.write('MPVPE (Hands): %.2f mm' % np.mean(eval_result['mpvpe_hand']))
f.write('MPVPE (Face): %.2f mm\n' % np.mean(eval_result['mpvpe_face']))
f.write('PA MPJPE (Body): %.2f mm\n' % np.mean(eval_result['pa_mpjpe_body']))
f.write('PA MPJPE (L-Hands): %.2f mm' % np.mean(eval_result['pa_mpjpe_l_hand']))
f.write('PA MPJPE (R-Hands): %.2f mm' % np.mean(eval_result['pa_mpjpe_r_hand']))
f.write('PA MPJPE (Hands): %.2f mm\n' % np.mean(eval_result['pa_mpjpe_hand']))
f.write(f"{np.mean(eval_result['pa_mpvpe_all'])},{np.mean(eval_result['pa_mpvpe_l_hand'])},{np.mean(eval_result['pa_mpvpe_r_hand'])},{np.mean(eval_result['pa_mpvpe_hand'])},{np.mean(eval_result['pa_mpvpe_face'])},"
f"{np.mean(eval_result['mpvpe_all'])},{np.mean(eval_result['mpvpe_l_hand'])},{np.mean(eval_result['mpvpe_r_hand'])},{np.mean(eval_result['mpvpe_hand'])},{np.mean(eval_result['mpvpe_face'])}")
# for i in range(len(eval_result['pa_mpvpe_all'])):
# f.write(f'{i+1:02d}.jpg\n')
# f.write('PA MPVPE (All): %.2f mm\n' % eval_result['pa_mpvpe_all'][i])
# f.write('PA MPVPE (Hands): %.2f mm\n' % eval_result['pa_mpvpe_hand'][i])
# f.write('PA MPVPE (Face): %.2f mm\n' % eval_result['pa_mpvpe_face'][i])
# f.write('MPVPE (All): %.2f mm\n' % eval_result['mpvpe_all'][i])
# f.write('MPVPE (Hands): %.2f mm\n' % eval_result['mpvpe_hand'][i])
# f.write('MPVPE (Face): %.2f mm\n' % eval_result['mpvpe_face'][i])
# f.write('PA MPJPE (Body): %.2f mm\n' % eval_result['pa_mpjpe_body'][i])
# f.write('PA MPJPE (Hands): %.2f mm\n' % eval_result['pa_mpjpe_hand'][i])
data/EgoBody_Egocentric/EgoBody_Egocentric.py
+53
View File
@@ -0,0 +1,53 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class EgoBody_Egocentric(HumanDataset):
def __init__(self, transform, data_split):
super(EgoBody_Egocentric, self).__init__(transform, data_split)
if getattr(cfg, 'eval_on_train', False):
self.data_split = 'eval_train'
print("Evaluate on train set.")
if 'train' in self.data_split:
filename = getattr(cfg, 'filename', 'egobody_egocentric_train_230425_065_fix_betas.npz')
else:
filename = getattr(cfg, 'filename', 'egobody_egocentric_test_230425_043_fix_betas.npz')
self.use_betas_neutral = getattr(cfg, 'egobody_fix_betas', False)
self.img_dir = osp.join(cfg.data_dir, 'EgoBody')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = (1080, 1920) # (h, w)
self.cam_param = {}
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/EgoBody_Kinect/EgoBody_Kinect.py
+50
View File
@@ -0,0 +1,50 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class EgoBody_Kinect(HumanDataset):
def __init__(self, transform, data_split):
super(EgoBody_Kinect, self).__init__(transform, data_split)
if self.data_split == 'train':
filename = getattr(cfg, 'filename', 'egobody_kinect_train_230503_065_fix_betas.npz')
else:
filename = getattr(cfg, 'filename', 'egobody_kinect_test_230503_043_fix_betas.npz')
self.use_betas_neutral = getattr(cfg, 'egobody_fix_betas', False)
self.img_dir = osp.join(cfg.data_dir, 'EgoBody')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = (1080, 1920) # (h, w)
self.cam_param = {}
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_{self.data_split}_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/FIT3D/FIT3D.py
+58
View File
@@ -0,0 +1,58 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class FIT3D(HumanDataset):
def __init__(self, transform, data_split):
super(FIT3D, self).__init__(transform, data_split)
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'FIT3D_train_230511_1504.npz')
self.img_shape = (900, 900) # (h, w)
self.cam_param = {}
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = []
for pre_prc_file in ['FIT3D_train_230511_1504_0.npz','FIT3D_train_230511_1504_1.npz',
'FIT3D_train_230511_1504_2.npz', 'FIT3D_train_230511_1504_3.npz',
'FIT3D_train_230511_1504_4.npz', 'FIT3D_train_230511_1504_5.npz']:
if self.data_split == 'train':
filename = getattr(cfg, 'filename', pre_prc_file)
else:
raise ValueError('FIT3D test set is not support')
self.img_dir = cfg.data_dir # FIT3D included
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data
datalist_slice = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
self.datalist.extend(datalist_slice)
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/GTA_Human2/GTA_Human2.py
+47
View File
@@ -0,0 +1,47 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class GTA_Human2(HumanDataset):
def __init__(self, transform, data_split):
super(GTA_Human2, self).__init__(transform, data_split)
filename = 'gta_human2multiple_230406_04000_0.npz'
self.img_dir = osp.join(cfg.data_dir, 'GTA_Human2')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = (1080, 1920) # (h, w)
self.cam_param = {
'focal': (1158.0337, 1158.0337), # (fx, fy)
'princpt': (960, 540) # (cx, cy)
}
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/Human36M/Human36M.py
+286
View File
@@ -0,0 +1,286 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
import random
from humandata import Cache
class Human36M(torch.utils.data.Dataset):
def __init__(self, transform, data_split):
self.transform = transform
self.data_split = data_split
self.img_dir = osp.join(cfg.data_dir, 'Human36M', 'images')
self.annot_path = osp.join(cfg.data_dir, 'Human36M', 'annotations')
self.action_name = ['Directions', 'Discussion', 'Eating', 'Greeting', 'Phoning', 'Posing', 'Purchases', 'Sitting', 'SittingDown', 'Smoking', 'Photo', 'Waiting', 'Walking', 'WalkDog', 'WalkTogether']
# H36M joint set
self.joint_set = {'joint_num': 17,
'joints_name': ('Pelvis', 'R_Hip', 'R_Knee', 'R_Ankle', 'L_Hip', 'L_Knee', 'L_Ankle', 'Torso', 'Neck', 'Head', 'Head_top', 'L_Shoulder', 'L_Elbow', 'L_Wrist', 'R_Shoulder', 'R_Elbow', 'R_Wrist'),
'flip_pairs': ( (1, 4), (2, 5), (3, 6), (14, 11), (15, 12), (16, 13) ),
'eval_joint': (1, 2, 3, 4, 5, 6, 8, 10, 11, 12, 13, 14, 15, 16),
'regressor': np.load(osp.join(cfg.data_dir, 'Human36M', 'J_regressor_h36m_smplx.npy'))
}
self.joint_set['root_joint_idx'] = self.joint_set['joints_name'].index('Pelvis')
# self.datalist = self.load_data()
# load data or cache
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', f'Human36M_{data_split}.npz')
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
datalist = Cache(self.annot_path_cache)
assert datalist.data_strategy == getattr(cfg, 'data_strategy', None), \
f'Cache data strategy {datalist.data_strategy} does not match current data strategy ' \
f'{getattr(cfg, "data_strategy", None)}'
self.datalist = datalist
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data()
if self.use_cache:
print(f'[{self.__class__.__name__}] Caching datalist to {self.annot_path_cache}...')
Cache.save(
self.annot_path_cache,
self.datalist,
data_strategy=getattr(cfg, 'data_strategy', None)
)
def get_subsampling_ratio(self):
if self.data_split == 'train':
return 5
elif self.data_split == 'test':
return 64
else:
assert 0, print('Unknown subset')
def get_subject(self):
if self.data_split == 'train':
subject = [1,5,6,7,8]
elif self.data_split == 'test':
subject = [9,11]
else:
assert 0, print("Unknown subset")
return subject
def load_data(self):
subject_list = self.get_subject()
sampling_ratio = self.get_subsampling_ratio()
# aggregate annotations from each subject
db = COCO()
cameras = {}
joints = {}
smplx_params = {}
for subject in subject_list:
# data load
with open(osp.join(self.annot_path, 'Human36M_subject' + str(subject) + '_data.json'),'r') as f:
annot = json.load(f)
if len(db.dataset) == 0:
for k,v in annot.items():
db.dataset[k] = v
else:
for k,v in annot.items():
db.dataset[k] += v
# camera load
with open(osp.join(self.annot_path, 'Human36M_subject' + str(subject) + '_camera.json'),'r') as f:
cameras[str(subject)] = json.load(f)
# joint coordinate load
with open(osp.join(self.annot_path, 'Human36M_subject' + str(subject) + '_joint_3d.json'),'r') as f:
joints[str(subject)] = json.load(f)
# smplx parameter load
with open(osp.join(self.annot_path, 'Human36M_subject' + str(subject) + '_SMPLX_NeuralAnnot.json'),'r') as f:
smplx_params[str(subject)] = json.load(f)
db.createIndex()
datalist = []
i = 0
for aid in db.anns.keys():
i += 1
if self.data_split == 'train' and i % getattr(cfg, 'Human36M_train_sample_interval', 1) != 0:
continue
ann = db.anns[aid]
image_id = ann['image_id']
img = db.loadImgs(image_id)[0]
img_path = osp.join(self.img_dir, img['file_name'])
img_shape = (img['height'], img['width'])
# check subject and frame_idx
frame_idx = img['frame_idx'];
if frame_idx % sampling_ratio != 0:
continue
# smplx parameter
subject = img['subject']; action_idx = img['action_idx']; subaction_idx = img['subaction_idx']; frame_idx = img['frame_idx']; cam_idx = img['cam_idx'];
smplx_param = smplx_params[str(subject)][str(action_idx)][str(subaction_idx)][str(frame_idx)]
# camera parameter
cam_param = cameras[str(subject)][str(cam_idx)]
R,t,f,c = np.array(cam_param['R'], dtype=np.float32), np.array(cam_param['t'], dtype=np.float32), np.array(cam_param['f'], dtype=np.float32), np.array(cam_param['c'], dtype=np.float32)
cam_param = {'R': R, 't': t, 'focal': f, 'princpt': c}
# only use frontal camera following previous works (HMR and SPIN)
if self.data_split == 'test' and str(cam_idx) != '4':
continue
# project world coordinate to cam, image coordinate space
joint_world = np.array(joints[str(subject)][str(action_idx)][str(subaction_idx)][str(frame_idx)], dtype=np.float32)
joint_cam = world2cam(joint_world, R, t)
joint_img = cam2pixel(joint_cam, f, c)[:,:2]
joint_valid = np.ones((self.joint_set['joint_num'],1))
bbox = process_bbox(np.array(ann['bbox']), img['width'], img['height'], ratio=getattr(cfg, 'bbox_ratio', 1.25))
if bbox is None: continue
datalist.append({
'img_path': img_path,
'img_shape': img_shape,
'bbox': bbox,
'joint_img': joint_img,
'joint_cam': joint_cam,
'joint_valid': joint_valid,
'smplx_param': smplx_param,
'cam_param': cam_param})
if self.data_split == 'train':
print('[Human36M train] original size:', len(db.anns.keys()),
'. Sample interval:', getattr(cfg, 'Human36M_train_sample_interval', 1),
'. Sampled size:', len(datalist))
if getattr(cfg, 'data_strategy', None) == 'balance' and self.data_split == 'train':
print(f'[Human36M] Using [balance] strategy with datalist shuffled...')
random.shuffle(datalist)
return datalist
def __len__(self):
return len(self.datalist)
def __getitem__(self, idx):
data = copy.deepcopy(self.datalist[idx])
img_path, img_shape, bbox, cam_param = data['img_path'], data['img_shape'], data['bbox'], data['cam_param']
# img
img = load_img(img_path)
img, img2bb_trans, bb2img_trans, rot, do_flip = augmentation(img, bbox, self.data_split)
img = self.transform(img.astype(np.float32))/255.
if self.data_split == 'train':
# h36m gt
joint_cam = data['joint_cam']
joint_cam = (joint_cam - joint_cam[self.joint_set['root_joint_idx'],None,:]) / 1000 # root-relative. milimeter to meter.
joint_img = data['joint_img']
joint_img = np.concatenate((joint_img[:,:2], joint_cam[:,2:]),1) # x, y, depth
joint_img[:,2] = (joint_img[:,2] / (cfg.body_3d_size / 2) + 1)/2. * cfg.output_hm_shape[0] # discretize depth
joint_img, joint_cam, joint_cam_ra, joint_valid, joint_trunc = process_db_coord(joint_img, joint_cam, data['joint_valid'], do_flip, img_shape, self.joint_set['flip_pairs'], img2bb_trans, rot, self.joint_set['joints_name'], smpl_x.joints_name)
# smplx coordinates and parameters
smplx_param = data['smplx_param']
cam_param['t'] /= 1000 # milimeter to meter
smplx_joint_img, smplx_joint_cam, smplx_joint_trunc, smplx_pose, smplx_shape, smplx_expr, \
smplx_pose_valid, smplx_joint_valid, smplx_expr_valid, smplx_mesh_cam_orig = \
process_human_model_output(smplx_param, cam_param, do_flip, img_shape, img2bb_trans, rot, 'smplx')
"""
# for debug
_tmp = joint_img.copy()
_tmp[:,0] = _tmp[:,0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1]
_tmp[:,1] = _tmp[:,1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0]
_img = img.numpy().transpose(1,2,0)[:,:,::-1] * 255
_img = vis_keypoints(_img, _tmp)
cv2.imwrite('h36m_' + str(idx) + '.jpg', _img)
"""
# reverse ra
smplx_joint_cam_wo_ra = smplx_joint_cam.copy()
smplx_joint_cam_wo_ra[smpl_x.joint_part['lhand'], :] = smplx_joint_cam_wo_ra[smpl_x.joint_part['lhand'], :] \
+ smplx_joint_cam_wo_ra[smpl_x.lwrist_idx, None, :] # left hand root-relative
smplx_joint_cam_wo_ra[smpl_x.joint_part['rhand'], :] = smplx_joint_cam_wo_ra[smpl_x.joint_part['rhand'], :] \
+ smplx_joint_cam_wo_ra[smpl_x.rwrist_idx, None, :] # right hand root-relative
smplx_joint_cam_wo_ra[smpl_x.joint_part['face'], :] = smplx_joint_cam_wo_ra[smpl_x.joint_part['face'], :] \
+ smplx_joint_cam_wo_ra[smpl_x.neck_idx, None,: ] # face root-relative
# SMPLX pose parameter validity
for name in ('L_Ankle', 'R_Ankle', 'L_Wrist', 'R_Wrist'):
smplx_pose_valid[smpl_x.orig_joints_name.index(name)] = 0
smplx_pose_valid = np.tile(smplx_pose_valid[:,None], (1,3)).reshape(-1)
# SMPLX joint coordinate validity
for name in ('L_Big_toe', 'L_Small_toe', 'L_Heel', 'R_Big_toe', 'R_Small_toe', 'R_Heel'):
smplx_joint_valid[smpl_x.joints_name.index(name)] = 0
smplx_joint_valid = smplx_joint_valid[:,None]
smplx_joint_trunc = smplx_joint_valid * smplx_joint_trunc
smplx_shape_valid = True
# dummy hand/face bbox
dummy_center = np.zeros((2), dtype=np.float32)
dummy_size = np.zeros((2), dtype=np.float32)
inputs = {'img': img}
targets = {'joint_img': smplx_joint_img, 'smplx_joint_img': smplx_joint_img,
'joint_cam': smplx_joint_cam_wo_ra, 'smplx_joint_cam': smplx_joint_cam,
'smplx_pose': smplx_pose, 'smplx_shape': smplx_shape, 'smplx_expr': smplx_expr,
'lhand_bbox_center': dummy_center, 'lhand_bbox_size': dummy_size,
'rhand_bbox_center': dummy_center, 'rhand_bbox_size': dummy_size,
'face_bbox_center': dummy_center, 'face_bbox_size': dummy_size}
meta_info = {'joint_valid': smplx_joint_valid, 'joint_trunc': smplx_joint_trunc,
'smplx_joint_valid': smplx_joint_valid, 'smplx_joint_trunc': smplx_joint_trunc,
'smplx_pose_valid': smplx_pose_valid, 'smplx_shape_valid': float(smplx_shape_valid),
'smplx_expr_valid': float(smplx_expr_valid), 'is_3D': float(True),
'lhand_bbox_valid': float(False), 'rhand_bbox_valid': float(False),
'face_bbox_valid': float(False)}
return inputs, targets, meta_info
else:
inputs = {'img': img}
targets = {}
meta_info = {}
return inputs, targets, meta_info
def evaluate(self, outs, cur_sample_idx):
annots = self.datalist
sample_num = len(outs)
eval_result = {'mpjpe': [], 'pa_mpjpe': []}
for n in range(sample_num):
annot = annots[cur_sample_idx + n]
out = outs[n]
# h36m joint from gt mesh
joint_gt = annot['joint_cam']
joint_gt = joint_gt - joint_gt[self.joint_set['root_joint_idx'],None] # root-relative
joint_gt = joint_gt[self.joint_set['eval_joint'],:]
# h36m joint from param mesh
mesh_out = out['smpl_mesh_cam'] * 1000 # meter to milimeter
joint_out = np.dot(self.joint_set['regressor'], mesh_out) # meter to milimeter
joint_out = joint_out - joint_out[self.joint_set['root_joint_idx'],None] # root-relative
joint_out = joint_out[self.joint_set['eval_joint'],:]
joint_out_aligned = rigid_align(joint_out, joint_gt)
eval_result['mpjpe'].append(np.sqrt(np.sum((joint_out - joint_gt)**2,1)).mean())
eval_result['pa_mpjpe'].append(np.sqrt(np.sum((joint_out_aligned - joint_gt)**2,1)).mean())
vis = False
if vis:
from utils.vis import vis_keypoints, vis_mesh, save_obj
filename = annot['img_path'].split('/')[-1][:-4]
img = load_img(annot['img_path'])[:,:,::-1]
img = vis_mesh(img, mesh_out_img, 0.5)
cv2.imwrite(filename + '.jpg', img)
save_obj(mesh_out, smpl_x.face, filename + '.obj')
return eval_result
def print_eval_result(self, eval_result):
print('MPJPE: %.2f mm' % np.mean(eval_result['mpjpe']))
print('PA MPJPE: %.2f mm' % np.mean(eval_result['pa_mpjpe']))
data/HumanArt/HumanArt.py
View File
data/HumanSC3D/HumanSC3D.py
+57
View File
@@ -0,0 +1,57 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class HumanSC3D(HumanDataset):
def __init__(self, transform, data_split):
super(HumanSC3D, self).__init__(transform, data_split)
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'HumanSC3D_train_230511_2752.npz')
self.img_shape = (900, 900) # (h, w)
self.cam_param = {}
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = []
for pre_prc_file in ['HumanSC3D_train_230511_2752_0.npz','HumanSC3D_train_230511_2752_1.npz',
'HumanSC3D_train_230511_2752_2.npz']:
if self.data_split == 'train':
filename = getattr(cfg, 'filename', pre_prc_file)
else:
raise ValueError('HumanSC3D test set is not support')
self.img_dir = cfg.data_dir # HumanSC3D included
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data
datalist_slice = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
self.datalist.extend(datalist_slice)
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/InstaVariety/InstaVariety.py
+51
View File
@@ -0,0 +1,51 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class InstaVariety(HumanDataset):
def __init__(self, transform, data_split):
super(InstaVariety, self).__init__(transform, data_split)
self.datalist = []
pre_prc_file = 'insta_variety_neural_annot_train.npz'
if self.data_split == 'train':
filename = getattr(cfg, 'filename', pre_prc_file)
else:
raise ValueError('InstaVariety test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'InstaVariety')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = (224,224) # (h, w)
self.cam_param = {}
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/LSPET/LSPET.py
+43
View File
@@ -0,0 +1,43 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class LSPET(HumanDataset):
def __init__(self, transform, data_split):
super(LSPET, self).__init__(transform, data_split)
if self.data_split == 'train':
filename = getattr(cfg, 'filename', 'eft_lspet.npz')
else:
raise ValueError('LSPET test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'LSPET')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = None # (h, w)
self.cam_param = {}
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/MPII/MPII.py
+195
View File
@@ -0,0 +1,195 @@
import os
import os.path as osp
import numpy as np
from config import cfg
import copy
import json
import cv2
import torch
from pycocotools.coco import COCO
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, \
process_human_model_output
import random
from humandata import Cache
# from utils.vis import vis_keypoints, vis_mesh, save_obj
class MPII(torch.utils.data.Dataset):
def __init__(self, transform, data_split):
self.transform = transform
self.data_split = data_split
self.img_path = osp.join(cfg.data_dir, 'MPII', 'data')
self.annot_path = osp.join(cfg.data_dir, 'MPII', 'data', 'annotations')
# mpii skeleton
self.joint_set = {
'joint_num': 16,
'joints_name': ('R_Ankle', 'R_Knee', 'R_Hip', 'L_Hip', 'L_Knee', 'L_Ankle', 'Pelvis', 'Thorax', 'Neck', 'Head_top', 'R_Wrist', 'R_Elbow', 'R_Shoulder', 'L_Shoulder', 'L_Elbow', 'L_Wrist'),
'flip_pairs': ( (0,5), (1,4), (2,3), (10,15), (11,14), (12,13) ),
}
# self.datalist = self.load_data()
# load data or cache
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', f'MPII_{data_split}.npz')
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
datalist = Cache(self.annot_path_cache)
assert datalist.data_strategy == getattr(cfg, 'data_strategy', None), \
f'Cache data strategy {datalist.data_strategy} does not match current data strategy ' \
f'{getattr(cfg, "data_strategy", None)}'
self.datalist = datalist
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data()
if self.use_cache:
print(f'[{self.__class__.__name__}] Caching datalist to {self.annot_path_cache}...')
Cache.save(
self.annot_path_cache,
self.datalist,
data_strategy=getattr(cfg, 'data_strategy', None)
)
def load_data(self):
db = COCO(osp.join(self.annot_path, 'train.json'))
with open(osp.join(self.annot_path, 'MPII_train_SMPLX_NeuralAnnot.json')) as f:
smplx_params = json.load(f)
datalist = []
i = 0
for aid in db.anns.keys():
i += 1
if self.data_split == 'train' and i % getattr(cfg, 'MPII_train_sample_interval', 1) != 0:
continue
ann = db.anns[aid]
img = db.loadImgs(ann['image_id'])[0]
imgname = img['file_name']
img_path = osp.join(self.img_path, imgname)
# bbox
bbox = process_bbox(ann['bbox'], img['width'], img['height'], ratio=getattr(cfg, 'bbox_ratio', 1.25))
if bbox is None: continue
# joint coordinates
joint_img = np.array(ann['keypoints'], dtype=np.float32).reshape(-1,3)
joint_valid = joint_img[:,2:].copy()
joint_img[:,2] = 0
# smplx parameter
if str(aid) in smplx_params:
smplx_param = smplx_params[str(aid)]
else:
smplx_param = None
datalist.append({
'img_path': img_path,
'img_shape': (img['height'], img['width']),
'bbox': bbox,
'joint_img': joint_img,
'joint_valid': joint_valid,
'smplx_param': smplx_param
})
if self.data_split == 'train':
print('[MPII train] original size:', len(db.anns.keys()),
'. Sample interval:', getattr(cfg, 'MPII_train_sample_interval', 1),
'. Sampled size:', len(datalist))
if getattr(cfg, 'data_strategy', None) == 'balance' and self.data_split == 'train':
print(f'[MPII] Using [balance] strategy with datalist shuffled...')
random.shuffle(datalist)
return datalist
def __len__(self):
return len(self.datalist)
def __getitem__(self, idx):
data = copy.deepcopy(self.datalist[idx])
img_path, img_shape, bbox = data['img_path'], data['img_shape'], data['bbox']
# image load and affine transform
img = load_img(img_path)
img, img2bb_trans, bb2img_trans, rot, do_flip = augmentation(img, bbox, self.data_split)
img = self.transform(img.astype(np.float32))/255.
# mpii gt
dummy_coord = np.zeros((self.joint_set['joint_num'],3), dtype=np.float32)
joint_img = data['joint_img']
joint_img = np.concatenate((joint_img[:,:2], np.zeros_like(joint_img[:,:1])),1) # x, y, dummy depth
joint_img, joint_cam, joint_cam_ra, joint_valid, joint_trunc = process_db_coord(joint_img, dummy_coord, data['joint_valid'], do_flip, img_shape, self.joint_set['flip_pairs'], img2bb_trans, rot, self.joint_set['joints_name'], smpl_x.joints_name)
# smplx coordinates and parameters
smplx_param = data['smplx_param']
if smplx_param is not None:
smplx_joint_img, smplx_joint_cam, smplx_joint_trunc, smplx_pose, smplx_shape, smplx_expr, smplx_pose_valid, smplx_joint_valid, smplx_expr_valid, smplx_mesh_cam_orig = process_human_model_output(smplx_param['smplx_param'], smplx_param['cam_param'], do_flip, img_shape, img2bb_trans, rot, 'smplx')
is_valid_fit = True
"""
# for debug
_tmp = joint_img.copy()
_tmp[:,0] = _tmp[:,0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1]
_tmp[:,1] = _tmp[:,1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0]
_img = img.numpy().transpose(1,2,0)[:,:,::-1] * 255
_img = vis_keypoints(_img.copy(), _tmp)
cv2.imwrite('mpii_' + str(idx) + '.jpg', _img)
"""
else:
# dummy values
smplx_joint_img = np.zeros((smpl_x.joint_num,3), dtype=np.float32)
smplx_joint_cam = np.zeros((smpl_x.joint_num,3), dtype=np.float32)
smplx_joint_trunc = np.zeros((smpl_x.joint_num,1), dtype=np.float32)
smplx_joint_valid = np.zeros((smpl_x.joint_num), dtype=np.float32)
smplx_pose = np.zeros((smpl_x.orig_joint_num*3), dtype=np.float32)
smplx_shape = np.zeros((smpl_x.shape_param_dim), dtype=np.float32)
smplx_expr = np.zeros((smpl_x.expr_code_dim), dtype=np.float32)
smplx_pose_valid = np.zeros((smpl_x.orig_joint_num), dtype=np.float32)
smplx_expr_valid = False
is_valid_fit = False
# SMPLX pose parameter validity
for name in ('L_Ankle', 'R_Ankle', 'L_Wrist', 'R_Wrist'):
smplx_pose_valid[smpl_x.orig_joints_name.index(name)] = 0
smplx_pose_valid = np.tile(smplx_pose_valid[:,None], (1,3)).reshape(-1)
# SMPLX joint coordinate validity
for name in ('L_Big_toe', 'L_Small_toe', 'L_Heel', 'R_Big_toe', 'R_Small_toe', 'R_Heel'):
smplx_joint_valid[smpl_x.joints_name.index(name)] = 0
smplx_joint_valid = smplx_joint_valid[:,None]
smplx_joint_trunc = smplx_joint_valid * smplx_joint_trunc
# make zero mask for invalid fit
if not is_valid_fit:
smplx_pose_valid[:] = 0
smplx_joint_valid[:] = 0
smplx_joint_trunc[:] = 0
smplx_shape_valid = False
else:
smplx_shape_valid = True
# dummy hand/face bbox
dummy_center = np.zeros((2), dtype=np.float32)
dummy_size = np.zeros((2), dtype=np.float32)
inputs = {'img': img}
targets = {'joint_img': joint_img, 'smplx_joint_img': smplx_joint_img,
'joint_cam': joint_cam, 'smplx_joint_cam': smplx_joint_cam,
'smplx_pose': smplx_pose, 'smplx_shape': smplx_shape, 'smplx_expr': smplx_expr,
'lhand_bbox_center': dummy_center, 'lhand_bbox_size': dummy_size,
'rhand_bbox_center': dummy_center, 'rhand_bbox_size': dummy_size,
'face_bbox_center': dummy_center, 'face_bbox_size': dummy_size}
meta_info = {'joint_valid': joint_valid, 'joint_trunc': joint_trunc,
'smplx_joint_valid': smplx_joint_valid,
'smplx_joint_trunc': smplx_joint_trunc, 'smplx_pose_valid': smplx_pose_valid,
'smplx_shape_valid': float(smplx_shape_valid),
'smplx_expr_valid': float(smplx_expr_valid), 'is_3D': float(False),
'lhand_bbox_valid': float(False), 'rhand_bbox_valid': float(False),
'face_bbox_valid': float(False)}
return inputs, targets, meta_info
data/MPI_INF_3DHP/MPI_INF_3DHP.py
+56
View File
@@ -0,0 +1,56 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class MPI_INF_3DHP(HumanDataset):
def __init__(self, transform, data_split):
super(MPI_INF_3DHP, self).__init__(transform, data_split)
if data_split != 'train':
raise NotImplementedError('MPI_INF_3DHP test set is not supported')
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'mpi_inf_3dhp_neural_annot.npz')
self.img_shape = None # (h, w)
self.cam_param = {}
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = []
for pre_prc_file in ['mpi_inf_3dhp_neural_annot_part1.npz', 'mpi_inf_3dhp_neural_annot_part2.npz',
'mpi_inf_3dhp_neural_annot_part3.npz']:
if self.data_split == 'train':
filename = getattr(cfg, 'filename', pre_prc_file)
else:
raise ValueError('MPI_INF_3DHP test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'MPI_INF_3DHP')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
# load data
datalist_slice = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
self.datalist.extend(datalist_slice)
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/MSCOCO/MSCOCO.py
+478
View File
@@ -0,0 +1,478 @@
import os
import os.path as osp
import numpy as np
from config import cfg
import copy
import json
import cv2
import torch
from pycocotools.coco import COCO
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output
import random
from humandata import Cache
class MSCOCO(torch.utils.data.Dataset):
def __init__(self, transform, data_split):
self.transform = transform
self.data_split = data_split
if os.path.exists(osp.join(cfg.data_dir, 'MSCOCO', 'images')):
self.img_path = osp.join(cfg.data_dir, 'MSCOCO', 'images')
self.annot_path = osp.join(cfg.data_dir, 'MSCOCO', 'annotations')
else:
self.img_path = osp.join(cfg.data_dir, 'coco')
self.annot_path = osp.join(cfg.data_dir, 'coco', 'annotations')
# mscoco joint set
self.joint_set = {
'joint_num': 134, # body 24 (23 + pelvis), lhand 21, rhand 21, face 68
'joints_name': \
(
'Nose', 'L_Eye', 'R_Eye', 'L_Ear', 'R_Ear', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist',
'R_Wrist', 'L_Hip', 'R_Hip', 'L_Knee', 'R_Knee', 'L_Ankle', 'R_Ankle', 'Pelvis', 'L_Big_toe',
'L_Small_toe', 'L_Heel', 'R_Big_toe', 'R_Small_toe', 'R_Heel', # body part
'L_Wrist_Hand', 'L_Thumb_1', 'L_Thumb_2', 'L_Thumb_3', 'L_Thumb_4', 'L_Index_1', 'L_Index_2',
'L_Index_3', 'L_Index_4', 'L_Middle_1', 'L_Middle_2', 'L_Middle_3', 'L_Middle_4', 'L_Ring_1',
'L_Ring_2', 'L_Ring_3', 'L_Ring_4', 'L_Pinky_1', 'L_Pinky_2', 'L_Pinky_3', 'L_Pinky_4', # left hand
'R_Wrist_Hand', 'R_Thumb_1', 'R_Thumb_2', 'R_Thumb_3', 'R_Thumb_4', 'R_Index_1', 'R_Index_2',
'R_Index_3', 'R_Index_4', 'R_Middle_1', 'R_Middle_2', 'R_Middle_3', 'R_Middle_4', 'R_Ring_1',
'R_Ring_2', 'R_Ring_3', 'R_Ring_4', 'R_Pinky_1', 'R_Pinky_2', 'R_Pinky_3', 'R_Pinky_4', # right hand
*['Face_' + str(i) for i in range(56, 73)], # face contour
*['Face_' + str(i) for i in range(5, 56)] # face
),
'flip_pairs': \
((1, 2), (3, 4), (5, 6), (7, 8), (9, 10), (11, 12), (13, 14), (15, 16), (18, 21), (19, 22), (20, 23),
# body part
(24, 45), (25, 46), (26, 47), (27, 48), (28, 49), (29, 50), (30, 51), (31, 52), (32, 53), (33, 54),
(34, 55), (35, 56), (36, 57), (37, 58), (38, 59), (39, 60), (40, 61), (41, 62), (42, 63), (43, 64),
(44, 65), # hand part
(66, 82), (67, 81), (68, 80), (69, 79), (70, 78), (71, 77), (72, 76), (73, 75), # face contour
(83, 92), (84, 91), (85, 90), (86, 89), (87, 88), # face eyebrow
(97, 101), (98, 100), # face below nose
(102, 111), (103, 110), (104, 109), (105, 108), (106, 113), (107, 112), # face eyes
(114, 120), (115, 119), (116, 118), (121, 125), (122, 124), # face mouth
(126, 130), (127, 129), (131, 133) # face lip
)
}
# self.datalist = self.load_data()
# load data or cache
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', f'MSCOCO_{data_split}.npz')
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
datalist = Cache(self.annot_path_cache)
assert datalist.data_strategy == getattr(cfg, 'data_strategy', None), \
f'Cache data strategy {datalist.data_strategy} does not match current data strategy ' \
f'{getattr(cfg, "data_strategy", None)}'
self.datalist = datalist
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data()
if self.use_cache:
print(f'[{self.__class__.__name__}] Caching datalist to {self.annot_path_cache}...')
Cache.save(
self.annot_path_cache,
self.datalist,
data_strategy=getattr(cfg, 'data_strategy', None)
)
def merge_joint(self, joint_img, feet_img, lhand_img, rhand_img, face_img):
# pelvis
lhip_idx = self.joint_set['joints_name'].index('L_Hip')
rhip_idx = self.joint_set['joints_name'].index('R_Hip')
pelvis = (joint_img[lhip_idx, :] + joint_img[rhip_idx, :]) * 0.5
pelvis[2] = joint_img[lhip_idx, 2] * joint_img[rhip_idx, 2] # joint_valid
pelvis = pelvis.reshape(1, 3)
# feet
lfoot = feet_img[:3, :]
rfoot = feet_img[3:, :]
joint_img = np.concatenate((joint_img, pelvis, lfoot, rfoot, lhand_img, rhand_img, face_img)).astype(
np.float32) # self.joint_set['joint_num'], 3
return joint_img
def load_data(self):
if self.data_split == 'train':
db = COCO(osp.join(self.annot_path, 'coco_wholebody_train_v1.0.json'))
smplx_json_path = osp.join(self.annot_path, 'MSCOCO_train_SMPLX_all_NeuralAnnot.json') # MSCOCO_train_SMPLX.json
with open(smplx_json_path) as f:
print(f'load SMPLX parameters from {smplx_json_path}')
smplx_params = json.load(f)
else:
db = COCO(osp.join(self.annot_path, 'coco_wholebody_val_v1.0.json'))
# train mode
if self.data_split == 'train':
datalist = []
i = 0
for aid in db.anns.keys():
i += 1
if self.data_split == 'train' and i % getattr(cfg, 'MSCOCO_train_sample_interval', 1) != 0:
continue
ann = db.anns[aid]
img = db.loadImgs(ann['image_id'])[0]
imgname = osp.join('train2017', img['file_name'])
img_path = osp.join(self.img_path, imgname)
# exclude the samples that are crowd or have few visible keypoints
if ann['iscrowd'] or (ann['num_keypoints'] == 0): continue
# bbox
bbox = process_bbox(ann['bbox'], img['width'], img['height'], ratio=getattr(cfg, 'bbox_ratio', 1.25))
if bbox is None: continue
# joint coordinates
joint_img = np.array(ann['keypoints'], dtype=np.float32).reshape(-1, 3)
foot_img = np.array(ann['foot_kpts'], dtype=np.float32).reshape(-1, 3)
lhand_img = np.array(ann['lefthand_kpts'], dtype=np.float32).reshape(-1, 3)
rhand_img = np.array(ann['righthand_kpts'], dtype=np.float32).reshape(-1, 3)
face_img = np.array(ann['face_kpts'], dtype=np.float32).reshape(-1, 3)
joint_img = self.merge_joint(joint_img, foot_img, lhand_img, rhand_img, face_img)
joint_valid = (joint_img[:, 2].copy().reshape(-1, 1) > 0).astype(np.float32)
joint_img[:, 2] = 0
# use body annotation to fill hand/face annotation
for body_name, part_name in (
('L_Wrist', 'L_Wrist_Hand'), ('R_Wrist', 'R_Wrist_Hand'), ('Nose', 'Face_18')):
if joint_valid[self.joint_set['joints_name'].index(part_name), 0] == 0:
joint_img[self.joint_set['joints_name'].index(part_name)] = joint_img[
self.joint_set['joints_name'].index(body_name)]
joint_valid[self.joint_set['joints_name'].index(part_name)] = joint_valid[
self.joint_set['joints_name'].index(body_name)]
# hand/face bbox
if ann['lefthand_valid']:
lhand_bbox = np.array(ann['lefthand_box']).reshape(4)
if hasattr(cfg, 'bbox_ratio'):
lhand_bbox = process_bbox(lhand_bbox, img['width'], img['height'], ratio=cfg.bbox_ratio)
if lhand_bbox is not None:
lhand_bbox[2:] += lhand_bbox[:2] # xywh -> xyxy
else:
lhand_bbox = None
if ann['righthand_valid']:
rhand_bbox = np.array(ann['righthand_box']).reshape(4)
if hasattr(cfg, 'bbox_ratio'):
rhand_bbox = process_bbox(rhand_bbox, img['width'], img['height'], ratio=cfg.bbox_ratio)
if rhand_bbox is not None:
rhand_bbox[2:] += rhand_bbox[:2] # xywh -> xyxy
else:
rhand_bbox = None
if ann['face_valid']:
face_bbox = np.array(ann['face_box']).reshape(4)
if hasattr(cfg, 'bbox_ratio'):
face_bbox = process_bbox(face_bbox, img['width'], img['height'], ratio=cfg.bbox_ratio)
if face_bbox is not None:
face_bbox[2:] += face_bbox[:2] # xywh -> xyxy
else:
face_bbox = None
if str(aid) in smplx_params:
smplx_param = smplx_params[str(aid)]
if 'lhand_valid' not in smplx_param['smplx_param']:
smplx_param['smplx_param']['lhand_valid'] = ann['lefthand_valid']
smplx_param['smplx_param']['rhand_valid'] = ann['righthand_valid']
smplx_param['smplx_param']['face_valid'] = ann['face_valid']
else:
smplx_param = None
data_dict = {'img_path': img_path, 'img_shape': (img['height'], img['width']), 'bbox': bbox,
'joint_img': joint_img, 'joint_valid': joint_valid, 'smplx_param': smplx_param,
'lhand_bbox': lhand_bbox, 'rhand_bbox': rhand_bbox, 'face_bbox': face_bbox}
datalist.append(data_dict)
print('[MSCOCO train] original size:', len(db.anns.keys()),
'. Sample interval:', getattr(cfg, 'MSCOCO_train_sample_interval', 1),
'. Sampled size:', len(datalist))
if getattr(cfg, 'data_strategy', None) == 'balance':
print(f"[MSCOCO] Using [balance] strategy with datalist shuffled...")
random.shuffle(datalist)
return datalist
# test mode
else:
datalist = []
for aid in db.anns.keys():
ann = db.anns[aid]
img = db.loadImgs(ann['image_id'])[0]
imgname = osp.join('val2017', img['file_name'])
img_path = osp.join(self.img_path, imgname)
# bbox
bbox = process_bbox(ann['bbox'], img['width'], img['height'], ratio=getattr(cfg, 'bbox_ratio', 1.25))
if bbox is None: continue
# hand/face bbox
if ann['lefthand_valid']:
lhand_bbox = np.array(ann['lefthand_box']).reshape(4)
lhand_bbox[2:] += lhand_bbox[:2] # xywh -> xyxy
else:
lhand_bbox = None
if ann['righthand_valid']:
rhand_bbox = np.array(ann['righthand_box']).reshape(4)
rhand_bbox[2:] += rhand_bbox[:2] # xywh -> xyxy
else:
rhand_bbox = None
if ann['face_valid']:
face_bbox = np.array(ann['face_box']).reshape(4)
face_bbox[2:] += face_bbox[:2] # xywh -> xyxy
else:
face_bbox = None
data_dict = {'img_path': img_path, 'ann_id': aid, 'img_shape': (img['height'], img['width']),
'bbox': bbox, 'lhand_bbox': lhand_bbox, 'rhand_bbox': rhand_bbox, 'face_bbox': face_bbox}
datalist.append(data_dict)
return datalist
def process_hand_face_bbox(self, bbox, do_flip, img_shape, img2bb_trans):
if bbox is None:
bbox = np.array([0, 0, 1, 1], dtype=np.float32).reshape(2, 2) # dummy value
bbox_valid = float(False) # dummy value
else:
# reshape to top-left (x,y) and bottom-right (x,y)
bbox = bbox.reshape(2, 2)
# flip augmentation
if do_flip:
bbox[:, 0] = img_shape[1] - bbox[:, 0] - 1
bbox[0, 0], bbox[1, 0] = bbox[1, 0].copy(), bbox[0, 0].copy() # xmin <-> xmax swap
# make four points of the bbox
bbox = bbox.reshape(4).tolist()
xmin, ymin, xmax, ymax = bbox
bbox = np.array([[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]], dtype=np.float32).reshape(4, 2)
# affine transformation (crop, rotation, scale)
bbox_xy1 = np.concatenate((bbox, np.ones_like(bbox[:, :1])), 1)
bbox = np.dot(img2bb_trans, bbox_xy1.transpose(1, 0)).transpose(1, 0)[:, :2]
bbox[:, 0] = bbox[:, 0] / cfg.input_img_shape[1] * cfg.output_hm_shape[2]
bbox[:, 1] = bbox[:, 1] / cfg.input_img_shape[0] * cfg.output_hm_shape[1]
# make box a rectangle without rotation
xmin = np.min(bbox[:, 0]);
xmax = np.max(bbox[:, 0]);
ymin = np.min(bbox[:, 1]);
ymax = np.max(bbox[:, 1]);
bbox = np.array([xmin, ymin, xmax, ymax], dtype=np.float32)
bbox_valid = float(True)
bbox = bbox.reshape(2, 2)
return bbox, bbox_valid
def __len__(self):
return len(self.datalist)
def __getitem__(self, idx):
data = copy.deepcopy(self.datalist[idx])
# train mode
if self.data_split == 'train':
img_path, img_shape = data['img_path'], data['img_shape']
# image load
img = load_img(img_path)
bbox = data['bbox']
img, img2bb_trans, bb2img_trans, rot, do_flip = augmentation(img, bbox, self.data_split)
img = self.transform(img.astype(np.float32)) / 255.
# hand and face bbox transform
lhand_bbox, lhand_bbox_valid = self.process_hand_face_bbox(data['lhand_bbox'], do_flip, img_shape,
img2bb_trans)
rhand_bbox, rhand_bbox_valid = self.process_hand_face_bbox(data['rhand_bbox'], do_flip, img_shape,
img2bb_trans)
face_bbox, face_bbox_valid = self.process_hand_face_bbox(data['face_bbox'], do_flip, img_shape,
img2bb_trans)
if do_flip:
lhand_bbox, rhand_bbox = rhand_bbox, lhand_bbox
lhand_bbox_valid, rhand_bbox_valid = rhand_bbox_valid, lhand_bbox_valid
lhand_bbox_center = (lhand_bbox[0] + lhand_bbox[1]) / 2.;
rhand_bbox_center = (rhand_bbox[0] + rhand_bbox[1]) / 2.;
face_bbox_center = (face_bbox[0] + face_bbox[1]) / 2.
lhand_bbox_size = lhand_bbox[1] - lhand_bbox[0];
rhand_bbox_size = rhand_bbox[1] - rhand_bbox[0];
face_bbox_size = face_bbox[1] - face_bbox[0];
# coco gt
dummy_coord = np.zeros((self.joint_set['joint_num'], 3), dtype=np.float32)
joint_img = data['joint_img']
joint_img = np.concatenate((joint_img[:, :2], np.zeros_like(joint_img[:, :1])), 1) # x, y, dummy depth
joint_img, joint_cam, joint_cam_ra, joint_valid, joint_trunc = process_db_coord(joint_img, dummy_coord,
data['joint_valid'], do_flip, img_shape,
self.joint_set['flip_pairs'],
img2bb_trans, rot,
self.joint_set['joints_name'],
smpl_x.joints_name)
# smplx coordinates and parameters
smplx_param = data['smplx_param']
if smplx_param is not None:
smplx_joint_img, smplx_joint_cam, smplx_joint_trunc, smplx_pose, smplx_shape, smplx_expr, smplx_pose_valid, smplx_joint_valid, smplx_expr_valid, smplx_mesh_cam_orig \
= process_human_model_output(smplx_param['smplx_param'], smplx_param['cam_param'], do_flip,
img_shape, img2bb_trans, rot, 'smplx')
is_valid_fit = True
"""
# for debug
_tmp = joint_img.copy()
_tmp[:,0] = _tmp[:,0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1]
_tmp[:,1] = _tmp[:,1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0]
_img = img.numpy().transpose(1,2,0)[:,:,::-1] * 255
_img = vis_keypoints(_img, _tmp)
cv2.imwrite('coco_' + str(idx) + '.jpg', _img)
"""
else:
# dummy values
smplx_joint_img = np.zeros((smpl_x.joint_num, 3), dtype=np.float32)
smplx_joint_cam = np.zeros((smpl_x.joint_num, 3), dtype=np.float32)
smplx_joint_trunc = np.zeros((smpl_x.joint_num, 1), dtype=np.float32)
smplx_joint_valid = np.zeros((smpl_x.joint_num), dtype=np.float32)
smplx_pose = np.zeros((smpl_x.orig_joint_num * 3), dtype=np.float32)
smplx_shape = np.zeros((smpl_x.shape_param_dim), dtype=np.float32)
smplx_expr = np.zeros((smpl_x.expr_code_dim), dtype=np.float32)
smplx_pose_valid = np.zeros((smpl_x.orig_joint_num), dtype=np.float32)
smplx_expr_valid = False
is_valid_fit = False
# SMPLX pose parameter validity
smplx_pose_valid = np.tile(smplx_pose_valid[:, None], (1, 3)).reshape(-1)
# SMPLX joint coordinate validity
smplx_joint_valid = smplx_joint_valid[:, None]
smplx_joint_trunc = smplx_joint_valid * smplx_joint_trunc
# make zero mask for invalid fit
if not is_valid_fit:
smplx_pose_valid[:] = 0
smplx_joint_valid[:] = 0
smplx_joint_trunc[:] = 0
smplx_shape_valid = False
else:
smplx_shape_valid = True
inputs = {'img': img}
targets = {'joint_img': joint_img, 'joint_cam': joint_cam, 'smplx_joint_img': smplx_joint_img,
'smplx_joint_cam': smplx_joint_cam,
'smplx_pose': smplx_pose, 'smplx_shape': smplx_shape, 'smplx_expr': smplx_expr,
'lhand_bbox_center': lhand_bbox_center,
'lhand_bbox_size': lhand_bbox_size, 'rhand_bbox_center': rhand_bbox_center,
'rhand_bbox_size': rhand_bbox_size,
'face_bbox_center': face_bbox_center, 'face_bbox_size': face_bbox_size}
meta_info = {'joint_valid': joint_valid, 'joint_trunc': joint_trunc, 'smplx_joint_valid': smplx_joint_valid,
'smplx_joint_trunc': smplx_joint_trunc,
'smplx_pose_valid': smplx_pose_valid, 'smplx_shape_valid': float(smplx_shape_valid),
'smplx_expr_valid': float(smplx_expr_valid), 'is_3D': float(False),
# 'lhand_bbox_valid': float(False), 'rhand_bbox_valid': float(False),
# 'face_bbox_valid': float(False)}
'lhand_bbox_valid': lhand_bbox_valid,
'rhand_bbox_valid': rhand_bbox_valid, 'face_bbox_valid': face_bbox_valid}
return inputs, targets, meta_info
# test mode
else:
img_path, img_shape = data['img_path'], data['img_shape']
# image load
img = load_img(img_path)
bbox = data['bbox']
img, img2bb_trans, bb2img_trans, rot, do_flip = augmentation(img, bbox, self.data_split)
img = self.transform(img.astype(np.float32)) / 255.
inputs = {'img': img}
targets = {}
meta_info = {'bb2img_trans': bb2img_trans}
return inputs, targets, meta_info
def evaluate(self, outs, cur_sample_idx):
annots = self.datalist
sample_num = len(outs)
for n in range(sample_num):
annot = annots[cur_sample_idx + n]
ann_id = annot['ann_id']
out = outs[n]
if annot['lhand_bbox'] is not None:
lhand_bbox = out['lhand_bbox'].copy().reshape(2, 2)
lhand_bbox = np.concatenate((lhand_bbox, np.ones((2, 1))), 1)
lhand_bbox = np.dot(out['bb2img_trans'], lhand_bbox.transpose(1, 0)).transpose(1, 0)[:, :2]
if annot['rhand_bbox'] is not None:
rhand_bbox = out['rhand_bbox'].copy().reshape(2, 2)
rhand_bbox = np.concatenate((rhand_bbox, np.ones((2, 1))), 1)
rhand_bbox = np.dot(out['bb2img_trans'], rhand_bbox.transpose(1, 0)).transpose(1, 0)[:, :2]
if annot['face_bbox'] is not None:
face_bbox = out['face_bbox'].copy().reshape(2, 2)
face_bbox = np.concatenate((face_bbox, np.ones((2, 1))), 1)
face_bbox = np.dot(out['bb2img_trans'], face_bbox.transpose(1, 0)).transpose(1, 0)[:, :2]
vis = False
if vis:
img_path = annot['img_path']
"""
img = (out['img'].transpose(1,2,0)[:,:,::-1] * 255).copy()
joint_img = out['joint_img'].copy()
joint_img[:,0] = joint_img[:,0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1]
joint_img[:,1] = joint_img[:,1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0]
for j in range(len(joint_img)):
if j in smpl_x.pos_joint_part['body']:
cv2.circle(img, (int(joint_img[j][0]), int(joint_img[j][1])), 3, (0,0,255), -1)
lhand_bbox = out['lhand_bbox'].reshape(2,2).copy()
cv2.rectangle(img, (int(lhand_bbox[0][0]), int(lhand_bbox[0][1])), (int(lhand_bbox[1][0]), int(lhand_bbox[1][1])), (255,0,0), 3)
rhand_bbox = out['rhand_bbox'].reshape(2,2).copy()
cv2.rectangle(img, (int(rhand_bbox[0][0]), int(rhand_bbox[0][1])), (int(rhand_bbox[1][0]), int(rhand_bbox[1][1])), (255,0,0), 3)
face_bbox = out['face_bbox'].reshape(2,2).copy()
cv2.rectangle(img, (int(face_bbox[0][0]), int(face_bbox[0][1])), (int(face_bbox[1][0]), int(face_bbox[1][1])), (255,0,0), 3)
cv2.imwrite(str(ann_id) + '.jpg', img)
"""
# save_obj(out['smplx_mesh_cam'], smpl_x.face, img_id + '_' + str(ann_id) + '.obj')
"""
img = load_img(img_path)[:,:,::-1]
bbox = annot['bbox']
focal = list(cfg.focal)
princpt = list(cfg.princpt)
focal[0] = focal[0] / cfg.input_body_shape[1] * bbox[2]
focal[1] = focal[1] / cfg.input_body_shape[0] * bbox[3]
princpt[0] = princpt[0] / cfg.input_body_shape[1] * bbox[2] + bbox[0]
princpt[1] = princpt[1] / cfg.input_body_shape[0] * bbox[3] + bbox[1]
img = render_mesh(img, out['smplx_mesh_cam'], smpl_x.face, {'focal': focal, 'princpt': princpt})
#img = cv2.resize(img, (512,512))
cv2.imwrite(img_id + '_' + str(ann_id) + '.jpg', img)
"""
bbox = annot['bbox']
focal = list(cfg.focal)
princpt = list(cfg.princpt)
focal[0] = focal[0] / cfg.input_body_shape[1] * bbox[2]
focal[1] = focal[1] / cfg.input_body_shape[0] * bbox[3]
princpt[0] = princpt[0] / cfg.input_body_shape[1] * bbox[2] + bbox[0]
princpt[1] = princpt[1] / cfg.input_body_shape[0] * bbox[3] + bbox[1]
param_save = {'smplx_param': {'root_pose': out['smplx_root_pose'].tolist(),
'body_pose': out['smplx_body_pose'].tolist(),
'lhand_pose': out['smplx_lhand_pose'].tolist(),
'rhand_pose': out['smplx_rhand_pose'].tolist(),
'jaw_pose': out['smplx_jaw_pose'].tolist(),
'shape': out['smplx_shape'].tolist(), 'expr': out['smplx_expr'].tolist(),
'trans': out['cam_trans'].tolist()},
'cam_param': {'focal': focal, 'princpt': princpt}
}
with open(str(ann_id) + '.json', 'w') as f:
json.dump(param_save, f)
return {}
def print_eval_result(self, eval_result):
return
data/MTP/MTP.py
+45
View File
@@ -0,0 +1,45 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class MTP(HumanDataset):
def __init__(self, transform, data_split):
super(MTP, self).__init__(transform, data_split)
pre_prc_file = 'mtp_smplx_train.npz'
if self.data_split == 'train':
filename = getattr(cfg, 'filename', pre_prc_file)
else:
raise ValueError('MTP test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'MTP')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = None
self.cam_param = {}
print("Various image shape in MTP dataset.")
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/MuCo/MuCo.py
+48
View File
@@ -0,0 +1,48 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class MuCo(HumanDataset):
def __init__(self, transform, data_split):
super(MuCo, self).__init__(transform, data_split)
if self.data_split == 'train':
filename = getattr(cfg, 'filename', 'muco3dhp_train.npz')
else:
raise ValueError('MoCo test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'MuCo')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = (1024, 1024) # (h, w)
self.cam_param = {}
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/OCHuman/OCHuman.py
+47
View File
@@ -0,0 +1,47 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class OCHuman(HumanDataset):
def __init__(self, transform, data_split):
super(OCHuman, self).__init__(transform, data_split)
self.datalist = []
pre_prc_file = 'eft_ochuman.npz'
if self.data_split == 'train':
filename = getattr(cfg, 'filename', pre_prc_file)
else:
raise ValueError('OCHuman test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'OCHuman')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = None
self.cam_param = {}
print("Various image shape in OCHuman dataset.")
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/PROX/PROX.py
+48
View File
@@ -0,0 +1,48 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class PROX(HumanDataset):
def __init__(self, transform, data_split):
super(PROX, self).__init__(transform, data_split)
if self.data_split == 'train':
filename = getattr(cfg, 'filename', 'prox_train_smplx_new.npz')
else:
raise ValueError('PROX test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'PROX')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = (1080, 1920) # (h, w)
self.cam_param = {}
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/PW3D/PW3D.py
+244
View File
@@ -0,0 +1,244 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x, smpl
from utils.preprocessing import load_img, process_bbox, augmentation, process_human_model_output, process_db_coord
from utils.transforms import rigid_align
import numpy as np
import random
class PW3D(torch.utils.data.Dataset):
def __init__(self, transform, data_split):
self.transform = transform
self.data_split = data_split
self.data_path = osp.join(cfg.data_dir, 'PW3D', 'data')
# 3dpw skeleton
self.joint_set = {
'joint_num': smpl_x.joint_num,
'joints_name': smpl_x.joints_name,
'flip_pairs': smpl_x.flip_pairs}
self.datalist = self.load_data()
def load_data(self):
db = COCO(osp.join(self.data_path, '3DPW_' + self.data_split + '.json'))
datalist = []
i = 0
if getattr(cfg, 'eval_on_train', False):
self.data_split = 'eval_train'
print("Evaluate on train set.")
for aid in db.anns.keys():
i += 1
if 'train' in self.data_split and i % getattr(cfg, 'PW3D_train_sample_interval', 1) != 0:
continue
ann = db.anns[aid]
image_id = ann['image_id']
img = db.loadImgs(image_id)[0]
sequence_name = img['sequence']
img_name = img['file_name']
img_path = osp.join(self.data_path, 'imageFiles', sequence_name, img_name)
cam_param = {k: np.array(v, dtype=np.float32) for k,v in img['cam_param'].items()}
smpl_param = ann['smpl_param']
bbox = process_bbox(np.array(ann['bbox']), img['width'], img['height'], ratio=getattr(cfg, 'bbox_ratio', 1.25))
if bbox is None: continue
data_dict = {'img_path': img_path, 'ann_id': aid, 'img_shape': (img['height'], img['width']),
'bbox': bbox, 'smpl_param': smpl_param, 'cam_param': cam_param}
datalist.append(data_dict)
if self.data_split == 'train':
print('[PW3D train] original size:', len(db.anns.keys()),
'. Sample interval:', getattr(cfg, 'PW3D_train_sample_interval', 1),
'. Sampled size:', len(datalist))
if (getattr(cfg, 'data_strategy', None) == 'balance' and self.data_split == 'train') or \
self.data_split == 'eval_train':
print(f'[PW3D] Using [balance] strategy with datalist shuffled...')
random.seed(2023)
random.shuffle(datalist)
if self.data_split == "eval_train":
return datalist[:10000]
return datalist
def __len__(self):
return len(self.datalist)
def __getitem__(self, idx):
data = copy.deepcopy(self.datalist[idx])
img_path, img_shape = data['img_path'], data['img_shape']
# img
img = load_img(img_path)
bbox, smpl_param, cam_param = data['bbox'], data['smpl_param'], data['cam_param']
img, img2bb_trans, bb2img_trans, rot, do_flip = augmentation(img, bbox, self.data_split)
img = self.transform(img.astype(np.float32))/255.
cam_param = data['cam_param']
if self.data_split == 'train':
smplx_param = {}
smplx_param['root_pose'] = np.array(smpl_param['pose']).reshape(-1,3)[:1, :]
smplx_param['body_pose'] = np.array(smpl_param['pose']).reshape(-1,3)[1:22, :]
smplx_param['trans'] = np.array(smpl_param['trans']).reshape(-1,3)
smplx_param['shape'] = np.zeros(10, dtype=np.float32) # drop smpl betas for smplx
# smpl coordinates
smplx_joint_img, smplx_joint_cam, smplx_joint_trunc, smplx_pose, smplx_shape, smplx_expr, \
smplx_pose_valid, smplx_joint_valid, smplx_expr_valid, smplx_mesh_cam_orig = process_human_model_output(
smplx_param, cam_param, do_flip, img_shape, img2bb_trans, rot, 'smplx',
joint_img=None)
# reverse ra
smplx_joint_cam_wo_ra = smplx_joint_cam.copy()
smplx_joint_cam_wo_ra[smpl_x.joint_part['lhand'], :] = smplx_joint_cam_wo_ra[smpl_x.joint_part['lhand'], :] \
+ smplx_joint_cam_wo_ra[smpl_x.lwrist_idx, None, :] # left hand root-relative
smplx_joint_cam_wo_ra[smpl_x.joint_part['rhand'], :] = smplx_joint_cam_wo_ra[smpl_x.joint_part['rhand'], :] \
+ smplx_joint_cam_wo_ra[smpl_x.rwrist_idx, None, :] # right hand root-relative
smplx_joint_cam_wo_ra[smpl_x.joint_part['face'], :] = smplx_joint_cam_wo_ra[smpl_x.joint_part['face'], :] \
+ smplx_joint_cam_wo_ra[smpl_x.neck_idx, None,: ] # face root-relative
smplx_pose_valid = np.tile(smplx_pose_valid[:, None], (1, 3)).reshape(-1)
smplx_joint_valid = smplx_joint_valid[:, None]
smplx_joint_trunc = smplx_joint_valid * smplx_joint_trunc
# smpl coordinates
smpl_joint_img, _, _, _, _, _ = process_human_model_output(
smpl_param, cam_param, do_flip, img_shape, img2bb_trans, rot, 'smpl',
joint_img=None)
joint_img = np.zeros_like(smplx_joint_img)
joint_img[:22] = smpl_joint_img[:22, :]
# dummy hand/face bbox
dummy_center = np.zeros((2), dtype=np.float32)
dummy_size = np.zeros((2), dtype=np.float32)
inputs = {'img': img}
targets = {'joint_img': joint_img, 'smplx_joint_img': joint_img,
'joint_cam': smplx_joint_cam_wo_ra, 'smplx_joint_cam': smplx_joint_cam,
'smplx_pose': smplx_pose, 'smplx_shape': smplx_shape, 'smplx_expr': smplx_expr,
'lhand_bbox_center': dummy_center, 'lhand_bbox_size': dummy_size,
'rhand_bbox_center': dummy_center, 'rhand_bbox_size': dummy_size,
'face_bbox_center': dummy_center, 'face_bbox_size': dummy_size}
meta_info = {'joint_valid': smplx_joint_valid, 'joint_trunc': smplx_joint_trunc,
'smplx_joint_valid': smplx_joint_valid, 'smplx_joint_trunc': smplx_joint_trunc,
'smplx_pose_valid': smplx_pose_valid, 'smplx_shape_valid': float(False),
'smplx_expr_valid': float(smplx_expr_valid), 'is_3D': float(True),
'lhand_bbox_valid': float(False), 'rhand_bbox_valid': float(False),
'face_bbox_valid': float(False)}
return inputs, targets, meta_info
else:
# smpl coordinates
smpl_joint_img, smpl_joint_cam, smpl_joint_trunc, smpl_pose, smpl_shape, smpl_mesh_cam_orig = process_human_model_output(smpl_param, cam_param, do_flip, img_shape, img2bb_trans, rot, 'smpl')
inputs = {'img': img}
targets = {'smpl_mesh_cam': smpl_mesh_cam_orig}
meta_info = {}
return inputs, targets, meta_info
def evaluate(self, outs, cur_sample_idx):
annots = self.datalist
sample_num = len(outs)
eval_result = {'mpjpe_body': [], 'pa_mpjpe_body': [], }
## smpl/smplx -> lsp
# ['left_hip', 'right_hip', 'left_knee', 'right_knee', 'left_ankle',
# 'right_ankle', 'neck', 'head', 'left_shoulder', 'right_shoulder',
# 'left_elbow', 'right_elbow', 'left_wrist', 'right_wrist']
joint_mapper = [1, 2, 4, 5, 7, 8, 12, 15, 16, 17, 18, 19, 20, 21]
### Save vis for debug
# joint_gt_body_to_save = np.zeros((sample_num, len(joint_mapper), 3))
# joint_out_body_root_align_to_save = np.zeros((sample_num, len(joint_mapper), 3))
# joint_out_body_pa_align_to_save = np.zeros((sample_num, len(joint_mapper), 3))
for n in range(sample_num):
out = outs[n]
# MPVPE from all vertices
mesh_gt = out['smpl_mesh_cam_target']
mesh_out = out['smplx_mesh_cam']
# MPJPE from body joints
mesh_out_align = mesh_out - np.dot(smpl_x.J_regressor, mesh_out)[smpl_x.J_regressor_idx['pelvis'], None, :] \
+ np.dot(smpl.joint_regressor, mesh_gt)[smpl.root_joint_idx, None, :]
# only eval point0-21 since only smpl gt is given
# joint_gt_body = np.dot(smpl.joint_regressor, mesh_gt)[:22, :]
# joint_out_body = np.dot(smpl_x.J_regressor, mesh_out)[:22, :]
# joint_out_body_root_align = np.dot(smpl_x.J_regressor, mesh_out_align)[:22, :]
# only test 14 keypoints
joint_gt_body = np.dot(smpl.joint_regressor, mesh_gt)[joint_mapper, :]
joint_out_body = np.dot(smpl_x.J_regressor, mesh_out)[joint_mapper, :]
joint_out_body_root_align = np.dot(smpl_x.J_regressor, mesh_out_align)[joint_mapper, :]
eval_result['mpjpe_body'].append(
np.sqrt(np.sum((joint_out_body_root_align - joint_gt_body) ** 2, 1)).mean() * 1000)
# PAMPJPE from body joints
joint_out_body_pa_align = rigid_align(joint_out_body, joint_gt_body)
eval_result['pa_mpjpe_body'].append(
np.sqrt(np.sum((joint_out_body_pa_align - joint_gt_body) ** 2, 1)).mean() * 1000)
### Save vis for debug
# joint_gt_body_to_save[n, ...] = joint_gt_body
# joint_out_body_root_align_to_save[n, ...] = joint_out_body_root_align
# joint_out_body_pa_align_to_save[n, ...] = joint_out_body_pa_align
### Save vis for debug
# import numpy as np
# np.save(f'./vis/val_0509_joint_gt_body.npy', joint_gt_body_to_save)
# np.save(f'./vis/val_0509_joint_out_body_root_align.npy', joint_out_body_root_align_to_save)
# np.save(f'./vis/val_0509_joint_out_body_pa_align.npy', joint_out_body_pa_align_to_save)
# import pdb; pdb.set_trace()
return eval_result
def print_eval_result(self, eval_result):
print('======3DPW-test======')
print('MPJPE (Body): %.2f mm' % np.mean(eval_result['mpjpe_body']))
print('PA MPJPE (Body): %.2f mm' % np.mean(eval_result['pa_mpjpe_body']))
print()
print(f"{np.mean(eval_result['mpjpe_body'])},{np.mean(eval_result['pa_mpjpe_body'])}")
print()
f = open(os.path.join(cfg.result_dir, 'result.txt'), 'w')
f.write(f'3DPW-test dataset: \n')
f.write('MPJPE (Body): %.2f mm\n' % np.mean(eval_result['mpjpe_body']))
f.write('PA MPJPE (Body): %.2f mm\n' % np.mean(eval_result['pa_mpjpe_body']))
f.write(f"{np.mean(eval_result['mpjpe_body'])},{np.mean(eval_result['pa_mpjpe_body'])}")
if getattr(cfg, 'eval_on_train', False):
import csv
csv_file = f'{cfg.root_dir}/output/pw3d_eval_on_train.csv'
exp_id = cfg.exp_name.split('_')[1]
new_line = [exp_id,np.mean(eval_result['mpjpe_body']), np.mean(eval_result['pa_mpjpe_body'])]
# Append the new line to the CSV file
with open(csv_file, 'a', newline='') as file:
writer = csv.writer(file)
writer.writerow(new_line)
data/PoseTrack/PoseTrack.py
+47
View File
@@ -0,0 +1,47 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class PoseTrack(HumanDataset):
def __init__(self, transform, data_split):
super(PoseTrack, self).__init__(transform, data_split)
self.datalist = []
pre_prc_file = 'eft_posetrack.npz'
if self.data_split == 'train':
filename = getattr(cfg, 'filename', pre_prc_file)
else:
raise ValueError('PoseTrack test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'PoseTrack/data/images')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = None
self.cam_param = {}
print("Various image shape in PoseTrack dataset.")
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print('loading cache from {}'.format(self.annot_path_cache))
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/RICH/RICH.py
+46
View File
@@ -0,0 +1,46 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class RICH(HumanDataset):
def __init__(self, transform, data_split):
super(RICH, self).__init__(transform, data_split)
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'rich_train_fix_betas.npz')
self.img_shape = None # (h, w)
self.cam_param = {}
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
if self.data_split == 'train':
filename = getattr(cfg, 'filename', 'rich_train_fix_betas.npz')
else:
raise ValueError('RICH test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'RICH')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
# load data
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/RenBody/RenBody.py
+60
View File
@@ -0,0 +1,60 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
# issue: 4 IndexError: index 432000 is out of bounds for axis 0 with size 432000 (bbox = bbox_xywh[i][:4])
class RenBody(HumanDataset):
def __init__(self, transform, data_split):
super(RenBody, self).__init__(transform, data_split)
self.use_cache = getattr(cfg, 'use_cache', False)
if self.data_split == 'train':
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'renbody_train_230525_399_ds10_fix_betas.npz')
else:
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'renbody_test_230525_78_ds10_fix_betas.npz')
self.img_shape = None # (h, w)
self.cam_param = {}
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
# load data or cache
self.datalist = []
for idx in range(10):
if self.data_split == 'train':
pre_prc_file_train = f'renbody_train_230525_399_{idx}.npz'
filename = getattr(cfg, 'filename', pre_prc_file_train)
else:
if idx > 1: continue
pre_prc_file_test = f'renbody_test_230525_78_{idx}.npz'
filename = getattr(cfg, 'filename', pre_prc_file_test)
self.img_dir = osp.join(cfg.data_dir, 'RenBody')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
# load data
datalist_slice = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1),
test_sample_interval=getattr(cfg, f'{self.__class__.__name__}_test_sample_interval', 1))
self.datalist.extend(datalist_slice)
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/RenBody_HiRes/RenBody_HiRes.py
+61
View File
@@ -0,0 +1,61 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class RenBody_HiRes(HumanDataset):
def __init__(self, transform, data_split):
super(RenBody_HiRes, self).__init__(transform, data_split)
self.datalist = []
if getattr(cfg, 'eval_on_train', False):
self.data_split = 'eval_train'
print("Evaluate on train set.")
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', f'renbody_{self.data_split}_highrescam_230517_399_fix_betas.npz')
self.img_shape = None # (h, w)
self.cam_param = {}
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
for idx in range(2):
if 'train' in self.data_split:
pre_prc_file_train = f'renbody_train_highrescam_230517_399_{idx}_fix_betas.npz'
filename = getattr(cfg, 'filename', pre_prc_file_train)
else:
if idx > 0: continue
pre_prc_file_test = f'renbody_test_highrescam_230517_78_{idx}_fix_betas.npz'
filename = getattr(cfg, 'filename', pre_prc_file_test)
self.img_dir = osp.join(cfg.data_dir, 'RenBody')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
data_split = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_{self.data_split}_sample_interval', 1),
test_sample_interval=getattr(cfg, f'{self.__class__.__name__}_{self.data_split}_sample_interval', 10))
self.datalist.extend(data_split)
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/SHAPY/.gitignore
+3
View File
@@ -0,0 +1,3 @@
__pycache__/
build/
*.so
data/SHAPY/SHAPY.py
+218
View File
@@ -0,0 +1,218 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
import pickle
from body_measurements import BodyMeasurements
import smplx
from test_submission_format import test_submission_file_format
def point_error(x, y, align=True):
""" Ref: https://github.com/muelea/shapy/blob/master/regressor/hbw_evaluation/evaluate_hbw.py#LL44C1-L58C31 """
t = 0.0
if align:
t = x.mean(0, keepdims=True) - y.mean(0, keepdims=True)
x_hat = x - t
error = np.sqrt(np.power(x_hat - y, 2).sum(axis=-1))
return error.mean().item()
class SHAPY(HumanDataset):
def __init__(self, transform, data_split):
super(SHAPY, self).__init__(transform, data_split)
self.eval_split = getattr(cfg, 'shapy_eval_split')
if self.data_split == 'train':
raise NotImplementedError('Shapy train not implemented yet. Need to consider invalid parameters')
if self.data_split == 'test' and self.eval_split == 'test':
filename = getattr(cfg, 'filename', 'shapy_test_230512_1631.npz')
elif self.data_split == 'test' and self.eval_split == 'val':
filename = getattr(cfg, 'filename', 'shapy_val_230512_705.npz')
else:
raise ValueError(f'Undefined. data split: {self.data_split}; eval_split: {self.test_set}')
self.img_dir = osp.join(cfg.data_dir, 'SHAPY')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.v_shape_load_dir = osp.join(cfg.data_dir, 'SHAPY', 'HBW', 'smplx', 'val')
self.img_shape = None # variable img_shape
self.cam_param = {}
# load data
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
### SHAPY utils
### ref: https://github.com/muelea/shapy/blob/master/regressor/hbw_evaluation/evaluate_hbw.py#L28
# load body model
# ref: common/utils/human_models.py
self.layer_arg = {'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}
self.smplx_layer = smplx.create(cfg.human_model_path,
'smplx',
gender='NEUTRAL',
use_pca=False,
use_face_contour=True,
flat_hand_mean=True, # critical!
**self.layer_arg).cuda()
# self.smplx_layer = copy.deepcopy(smpl_x.layer['neutral']).cuda()
self.faces_tensor_smplx = self.smplx_layer.faces_tensor.detach().cpu().numpy()
# load files to compute P2P-20K Error
point_reg = osp.join(cfg.data_dir, 'SHAPY', 'utility_files', 'evaluation', 'eval_point_set', 'HD_SMPLX_from_SMPL.pkl')
with open(point_reg, 'rb') as f:
self.point_regressor = pickle.load(f)
# load files to compute Measurements Error
body_measurement_folder = osp.join(cfg.data_dir, 'SHAPY', 'utility_files', 'measurements')
meas_def_path = osp.join(body_measurement_folder, 'measurement_defitions.yaml')
meas_verts_path_gt = osp.join(body_measurement_folder, 'smplx_measurements.yaml')
self.body_measurements = BodyMeasurements(
{'meas_definition_path': meas_def_path,
'meas_vertices_path': meas_verts_path_gt},
).to('cuda')
self.v_shaped_gt = {}
# to save preditions
self.images_names = []
self.v_shaped = []
def evaluate(self, outs, cur_sample_idx):
annots = self.datalist
sample_num = len(outs)
eval_result = {'v2v_t_errors': [], 'point_t_errors': [], 'height': [], 'chest': [], 'waist': [], 'hips': [], 'mass': []}
# sample_num = sample_num // 10 # TODO: debug only
for n in range(sample_num):
annot = annots[cur_sample_idx + n]
out = outs[n]
betas_fit = out['smplx_shape']
img_path = out['img_path']
# compute v_shaped
betas_fit = torch.tensor(betas_fit.reshape(-1, 10)).cuda()
# betas_fit = torch.tensor(out['smplx_shape_target'].reshape(-1, 10)).cuda() # TODO: debug only
# betas_fit = torch.zeros((1, 10)).cuda() # TODO: debug only
output = self.smplx_layer(
betas=betas_fit,
body_pose=torch.zeros((1, 63)).to(betas_fit.device),
global_orient=torch.zeros((1, 3)).to(betas_fit.device),
right_hand_pose=torch.zeros((1, 45)).to(betas_fit.device),
left_hand_pose=torch.zeros((1, 45)).to(betas_fit.device),
jaw_pose=torch.zeros((1, 3)).to(betas_fit.device),
leye_pose=torch.zeros((1, 3)).to(betas_fit.device),
reye_pose=torch.zeros((1, 3)).to(betas_fit.device),
expression=torch.zeros((1, 10)).to(betas_fit.device),
return_verts=True
)
v_shaped_fit = output.vertices.detach().cpu().numpy().squeeze()
# v_shaped_gt = v_shaped_fit # TODO: debug only
# v_shaped_fit = self.smplx_layer.forward_shape(betas=betas_fit)
image_name = '/'.join(img_path.split('/')[-4:])
self.images_names.append(image_name)
self.v_shaped.append(v_shaped_fit)
if self.eval_split == 'val':
# load gt vertices
subject = img_path.split('/')[-3]
subject_id_npy = subject.split('_')[0] + '.npy'
v_shaped_gt_path = osp.join(self.v_shape_load_dir, subject_id_npy)
if v_shaped_gt_path not in self.v_shaped_gt:
v_shaped_gt = np.load(v_shaped_gt_path)
self.v_shaped_gt[v_shaped_gt_path] = v_shaped_gt
else:
v_shaped_gt = self.v_shaped_gt[v_shaped_gt_path]
# compute vertex-to-vertex error (SMPL-X only)
# ref: https://github.com/muelea/shapy/blob/master/regressor/hbw_evaluation/evaluate_hbw.py#LL142C1-L171C48
v2v_error = point_error(v_shaped_fit, v_shaped_gt, align=True)
eval_result['v2v_t_errors'].append(v2v_error)
# compute P2P-20k error
points_gt = self.point_regressor.dot(v_shaped_gt)
points_fit = self.point_regressor.dot(v_shaped_fit)
p2p_error = point_error(points_gt, points_fit, align=True)
eval_result['point_t_errors'].append(p2p_error)
# compute height/chest/waist/hip error
shaped_triangles_gt = v_shaped_gt[self.faces_tensor_smplx]
shaped_triangles_gt = torch.from_numpy(shaped_triangles_gt).unsqueeze(0).to('cuda')
measurements_gt = self.body_measurements(shaped_triangles_gt)['measurements']
shaped_triangles_fit = v_shaped_fit[self.faces_tensor_smplx]
shaped_triangles_fit = torch.from_numpy(shaped_triangles_fit).unsqueeze(0).to('cuda')
measurements_fit = self.body_measurements(shaped_triangles_fit)['measurements']
for k in ['height', 'chest', 'waist', 'hips', 'mass']:
error = abs(measurements_gt[k]['tensor'].item() - measurements_fit[k]['tensor'].item())
eval_result[k].append(error)
return eval_result
def print_eval_result(self, eval_result):
# print('SHAPY results are dumped at: ' + osp.join(cfg.result_dir, 'predictions'))
if self.data_split == 'test' and self.eval_split == 'test': # do not print. just submit the results to the official evaluation server
# save predictions in the format of HBW challenge
# ref: https://github.com/muelea/shapy/blob/master/regressor/hbw_evaluation/README_HBW_EVAL.md#hbw-challenge
save_dir = osp.join(cfg.result_dir, 'predictions')
os.makedirs(save_dir, exist_ok=True)
save_name = osp.join(save_dir, 'hbw_prediction')
images_names = np.array(self.images_names).reshape(1631, )
v_shaped = np.array(self.v_shaped).reshape(1631, 10475, 3)
np.savez(save_name,
image_name=images_names,
v_shaped=v_shaped)
print('predictions saved at: ' + save_name + '.npz')
# run format test
test_submission_file_format(save_name + '.npz')
return
v2v_t_errors = np.mean(eval_result['v2v_t_errors']) * 1000
point_t_errors = np.mean(eval_result['point_t_errors']) * 1000
chest = np.mean(eval_result['chest']) * 1000
waist = np.mean(eval_result['waist']) * 1000
hips = np.mean(eval_result['hips']) * 1000
height = np.mean(eval_result['height']) * 1000
mass = np.mean(eval_result['mass'])
print('======SHAPY-val======')
print('Height Error: %.2f mm' % height)
print('Chest Error: %.2f mm' % chest)
print('Waist Error: %.2f mm' % waist)
print('Hips Error: %.2f mm' % hips)
print('P2P-20k Error: %.2f mm' % point_t_errors)
print('V2V Error: %.2f mm' % v2v_t_errors)
print('Mass Error: %.2f kg' % mass)
f = open(os.path.join(cfg.result_dir, 'result.txt'), 'w')
f.write(f'SHAPY-val dataset: \n')
f.write('Height Error: %.2f mm\n' % height)
f.write('Chest Error: %.2f mm' % chest)
f.write('Waist Error: %.2f mm\n' % waist)
f.write('Hips Error: %.2f mm\n' % hips)
f.write('P2P-20k Error: %.2f mm' % point_t_errors)
f.write('V2V Error: %.2f mm\n' % v2v_t_errors)
f.write('Mass Error: %.2f kg\n' % mass)
f.close()
data/SHAPY/mesh-mesh-intersection/LICENSE
+58
View File
@@ -0,0 +1,58 @@
License
Software Copyright License for non-commercial scientific research purposes
Please read carefully the following terms and conditions and any accompanying documentation before you download and/or use the SMPL-X/SMPLify-X model, data and software, (the "Model & Software"), including 3D meshes, blend weights, blend shapes, textures, software, scripts, and animations. By downloading and/or using the Model & Software (including downloading, cloning, installing, and any other use of this github repository), you acknowledge that you have read these terms and conditions, understand them, and agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use the Model & Software. Any infringement of the terms of this agreement will automatically terminate your rights under this License
Ownership / Licensees
The Software and the associated materials has been developed at the
Max Planck Institute for Intelligent Systems (hereinafter "MPI").
Any copyright or patent right is owned by and proprietary material of the
Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (hereinafter “MPG”; MPI and MPG hereinafter collectively “Max-Planck”)
hereinafter the “Licensor”.
License Grant
Licensor grants you (Licensee) personally a single-user, non-exclusive, non-transferable, free of charge right:
To install the Model & Software on computers owned, leased or otherwise controlled by you and/or your organization;
To use the Model & Software for the sole purpose of performing non-commercial scientific research, non-commercial education, or non-commercial artistic projects;
Any other use, in particular any use for commercial purposes, is prohibited. This includes, without limitation, incorporation in a commercial product, use in a commercial service, or production of other artifacts for commercial purposes. The Model & Software may not be reproduced, modified and/or made available in any form to any third party without Max-Planck’s prior written permission.
The Model & Software may not be used for pornographic purposes or to generate pornographic material whether commercial or not. This license also prohibits the use of the Model & Software to train methods/algorithms/neural networks/etc. for commercial use of any kind. By downloading the Model & Software, you agree not to reverse engineer it.
No Distribution
The Model & Software and the license herein granted shall not be copied, shared, distributed, re-sold, offered for re-sale, transferred or sub-licensed in whole or in part except that you may make one copy for archive purposes only.
Disclaimer of Representations and Warranties
You expressly acknowledge and agree that the Model & Software results from basic research, is provided “AS IS”, may contain errors, and that any use of the Model & Software is at your sole risk. LICENSOR MAKES NO REPRESENTATIONS OR WARRANTIES OF ANY KIND CONCERNING THE MODEL & SOFTWARE, NEITHER EXPRESS NOR IMPLIED, AND THE ABSENCE OF ANY LEGAL OR ACTUAL DEFECTS, WHETHER DISCOVERABLE OR NOT. Specifically, and not to limit the foregoing, licensor makes no representations or warranties (i) regarding the merchantability or fitness for a particular purpose of the Model & Software, (ii) that the use of the Model & Software will not infringe any patents, copyrights or other intellectual property rights of a third party, and (iii) that the use of the Model & Software will not cause any damage of any kind to you or a third party.
Limitation of Liability
Because this Model & Software License Agreement qualifies as a donation, according to Section 521 of the German Civil Code (Bürgerliches Gesetzbuch – BGB) Licensor as a donor is liable for intent and gross negligence only. If the Licensor fraudulently conceals a legal or material defect, they are obliged to compensate the Licensee for the resulting damage.
Licensor shall be liable for loss of data only up to the amount of typical recovery costs which would have arisen had proper and regular data backup measures been taken. For the avoidance of doubt Licensor shall be liable in accordance with the German Product Liability Act in the event of product liability. The foregoing applies also to Licensor’s legal representatives or assistants in performance. Any further liability shall be excluded.
Patent claims generated through the usage of the Model & Software cannot be directed towards the copyright holders.
The Model & Software is provided in the state of development the licensor defines. If modified or extended by Licensee, the Licensor makes no claims about the fitness of the Model & Software and is not responsible for any problems such modifications cause.
No Maintenance Services
You understand and agree that Licensor is under no obligation to provide either maintenance services, update services, notices of latent defects, or corrections of defects with regard to the Model & Software. Licensor nevertheless reserves the right to update, modify, or discontinue the Model & Software at any time.
Defects of the Model & Software must be notified in writing to the Licensor with a comprehensible description of the error symptoms. The notification of the defect should enable the reproduction of the error. The Licensee is encouraged to communicate any use, results, modification or publication.
Publications using the Model & Software
You acknowledge that the Model & Software is a valuable scientific resource and agree to appropriately reference the following paper in any publication making use of the Model & Software.
Citation:
@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}
}
Commercial licensing opportunities
For commercial uses of the Software, please send email to ps-license@tue.mpg.de
This Agreement shall be governed by the laws of the Federal Republic of Germany except for the UN Sales Convention.
data/SHAPY/mesh-mesh-intersection/README.md
+82
View File
@@ -0,0 +1,82 @@
# Computing mesh-mesh intersection
This package provides a PyTorch module that can efficiently compute mesh-mesh
intersections using a BVH.
## Table of Contents
* [Description](#description)
* [Installation](#installation)
* [Examples](#examples)
* [Citation](#citation)
* [License](#license)
* [Contact](#contact)
## Description
This repository provides a PyTorch wrapper around a CUDA kernel that implements
the method described in [Maximizing parallelism in the construction of BVHs,
octrees, and k-d trees](https://dl.acm.org/citation.cfm?id=2383801). More
specifically, given an input mesh it builds a
BVH tree for each one and queries it for self-intersections.
## Installation
See the instructions [here](docs/install.md) on how to install the package.
## Examples
### Fitting to measurements
To fit a 3D human body model to height, weight and circumenference measurements
use the following command:
```python
python examples/fit_measurements.py --model-folder PATH_TO_BODY_MODELS \
--model-type [smpl/smplh/star/smplx] --gender neutral/female/male --num-betas 30 \
--meas-vertices-path data/smpl_measurement_vertices.yaml
```
If you are using SMPL-X then set `--meas-vertices-path data/smplx_measurements.yaml`.
## Citation
If you find this code useful in your research please cite the relevant work(s) of the following list, for detecting and penalizing mesh intersections accordingly:
```
@inproceedings{Karras:2012:MPC:2383795.2383801,
author = {Karras, Tero},
title = {Maximizing Parallelism in the Construction of BVHs, Octrees, and K-d Trees},
booktitle = {Proceedings of the Fourth ACM SIGGRAPH / Eurographics Conference on High-Performance Graphics},
year = {2012},
pages = {33--37},
numpages = {5},
url = {https://doi.org/10.2312/EGGH/HPG12/033-037},
doi = {10.2312/EGGH/HPG12/033-037},
publisher = {Eurographics Association}
}
```
## License
Software Copyright License for **non-commercial scientific research purposes**.
Please read carefully the [terms and
conditions](https://github.com/vchoutas/mesh-mesh-intersection/blob/master/LICENSE) and any
accompanying documentation before you download and/or use the SMPL-X/SMPLify-X
model, data and software, (the "Model & Software"), including 3D meshes, blend
weights, blend shapes, textures, software, scripts, and animations. By
downloading and/or using the Model & Software (including downloading, cloning,
installing, and any other use of this github repository), you acknowledge that
you have read these terms and conditions, understand them, and agree to be bound
by them. If you do not agree with these terms and conditions, you must not
download and/or use the Model & Software. Any infringement of the terms of this
agreement will automatically terminate your rights under this
[License](./LICENSE).
## Contact
The code of this repository was implemented by [Vassilis Choutas](vassilis.choutas@tuebingen.mpg.de).
For questions, please contact [smplx@tue.mpg.de](smplx@tue.mpg.de).
For commercial licensing (and all related questions for business applications), please contact [ps-licensing@tue.mpg.de](ps-licensing@tue.mpg.de). Please note that the method for this component has been [patented by NVidia](https://patents.google.com/patent/US9396512B2/en) and a license needs to be obtained also by them.
data/SHAPY/mesh-mesh-intersection/body_measurements/__init__.py
+2
View File
@@ -0,0 +1,2 @@
from .body_measurements import BodyMeasurements
from .cwh_measurements import ChestWaistHipsMeasurements
data/SHAPY/mesh-mesh-intersection/body_measurements/body_measurements.py
+246
View File
@@ -0,0 +1,246 @@
from typing import NewType, Dict, Tuple
import os.path as osp
import yaml
import numpy as np
from mesh_mesh_intersection import MeshMeshIntersection
import torch
import torch.autograd as autograd
import torch.nn as nn
import torch.nn.functional as F
from scipy.spatial import ConvexHull
# from loguru import logger
Tensor = NewType('Tensor', torch.Tensor)
class BodyMeasurements(nn.Module):
# The density of the human body is 985 kg / m^3
DENSITY = 985
def __init__(self, cfg, **kwargs):
''' Loss that penalizes deviations in weight and height
'''
super(BodyMeasurements, self).__init__()
meas_definition_path = cfg.get('meas_definition_path', '')
meas_definition_path = osp.expanduser(
osp.expandvars(meas_definition_path))
meas_vertices_path = cfg.get('meas_vertices_path', '')
meas_vertices_path = osp.expanduser(
osp.expandvars(meas_vertices_path))
with open(meas_definition_path, 'r') as f:
measurements_definitions = yaml.safe_load(f, )
with open(meas_vertices_path, 'r') as f:
meas_vertices = yaml.safe_load(f)
head_top = meas_vertices['HeadTop']
left_heel = meas_vertices['HeelLeft']
left_heel_bc = left_heel['bc']
self.left_heel_face_idx = left_heel['face_idx']
left_heel_bc = torch.tensor(left_heel['bc'], dtype=torch.float32)
self.register_buffer('left_heel_bc', left_heel_bc)
head_top_bc = torch.tensor(head_top['bc'], dtype=torch.float32)
self.register_buffer('head_top_bc', head_top_bc)
self.head_top_face_idx = head_top['face_idx']
action = measurements_definitions['CW_p']
chest_periphery_data = meas_vertices[action[0]]
self.chest_face_index = chest_periphery_data['face_idx']
chest_bcs = torch.tensor(
chest_periphery_data['bc'], dtype=torch.float32)
self.register_buffer('chest_bcs', chest_bcs)
action = measurements_definitions['BW_p']
belly_periphery_data = meas_vertices[action[0]]
self.belly_face_index = belly_periphery_data['face_idx']
belly_bcs = torch.tensor(
belly_periphery_data['bc'], dtype=torch.float32)
self.register_buffer('belly_bcs', belly_bcs)
action = measurements_definitions['IW_p']
hips_periphery_data = meas_vertices[action[0]]
self.hips_face_index = hips_periphery_data['face_idx']
hips_bcs = torch.tensor(
hips_periphery_data['bc'], dtype=torch.float32)
self.register_buffer('hips_bcs', hips_bcs)
max_collisions = cfg.get('max_collisions', 256)
self.isect_module = MeshMeshIntersection(max_collisions=max_collisions)
def extra_repr(self) -> str:
msg = []
msg.append(f'Human Body Density: {self.DENSITY}')
return '\n'.join(msg)
def _get_plane_at_heights(self, height: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
device = height.device
batch_size = height.shape[0]
verts = torch.tensor(
[[-1., 0, -1], [1, 0, -1], [1, 0, 1], [-1, 0, 1]],
device=device).unsqueeze(dim=0).expand(batch_size, -1, -1).clone()
verts[:, :, 1] = height.reshape(batch_size, -1)
faces = torch.tensor([[0, 1, 2], [0, 2, 3]], device=device,
dtype=torch.long)
return verts, faces, verts[:, faces]
def compute_peripheries(
self,
triangles: Tensor,
compute_chest: bool = True,
compute_waist: bool = True,
compute_hips: bool = True,
) -> Dict[str, Tensor]:
'''
Parameters
----------
triangles: BxFx3x3 torch.Tensor
Contains the triangle coordinates for a batch of meshes with
the same topology
'''
batch_size, num_triangles = triangles.shape[:2]
device = triangles.device
batch_indices = torch.arange(
batch_size, dtype=torch.long,
device=device).reshape(-1, 1) * num_triangles
meas_data = {}
if compute_chest:
meas_data['chest'] = (self.chest_face_index, self.chest_bcs)
if compute_waist:
meas_data['waist'] = (self.belly_face_index, self.belly_bcs)
if compute_hips:
meas_data['hips'] = (self.hips_face_index, self.hips_bcs)
output = {}
for name, (face_index, bcs) in meas_data.items():
vertex = (
triangles[:, face_index] * bcs.reshape(1, 3, 1)).sum(axis=1)
_, _, plane_tris = self._get_plane_at_heights(vertex[:, 1])
with torch.no_grad():
collision_faces, collision_bcs = self.isect_module(
plane_tris, triangles)
selected_triangles = triangles.view(-1, 3, 3)[
(collision_faces + batch_indices).view(-1)].reshape(
batch_size, -1, 3, 3)
points = (
selected_triangles[:, :, None] *
collision_bcs[:, :, :, :, None]).sum(
axis=-2).reshape(batch_size, -1, 2, 3)
np_points = points.detach().cpu().numpy()
collision_faces = collision_faces.detach().cpu().numpy()
collision_bcs = collision_bcs.detach().cpu().numpy()
output[name] = {
'points': [],
'valid_points': [],
'value': [],
'plane_height': vertex[:, 1],
}
for ii in range(batch_size):
valid_face_idxs = np.where(collision_faces[ii] > 0)[0]
points_in_plane = np_points[
ii, valid_face_idxs, :, ][:, :, [0, 2]].reshape(
-1, 2)
hull = ConvexHull(points_in_plane)
point_indices = hull.simplices.reshape(-1)
hull_points = points[ii][valid_face_idxs].view(
-1, 3)[point_indices].reshape(-1, 2, 3)
meas_value = (
hull_points[:, 1] - hull_points[:, 0]).pow(2).sum(
dim=-1).sqrt().sum()
output[name]['valid_points'].append(
np_points[ii, valid_face_idxs])
output[name]['points'].append(hull_points)
output[name]['value'].append(meas_value)
output[name]['tensor'] = torch.stack(output[name]['value'])
return output
def compute_height(self, shaped_triangles: Tensor) -> Tuple[Tensor, Tensor]:
''' Compute the height using the heel and the top of the head
'''
head_top_tri = shaped_triangles[:, self.head_top_face_idx]
head_top = (head_top_tri[:, 0, :] * self.head_top_bc[0] +
head_top_tri[:, 1, :] * self.head_top_bc[1] +
head_top_tri[:, 2, :] * self.head_top_bc[2])
head_top = (
head_top_tri * self.head_top_bc.reshape(1, 3, 1)
).sum(dim=1)
left_heel_tri = shaped_triangles[:, self.left_heel_face_idx]
left_heel = (
left_heel_tri * self.left_heel_bc.reshape(1, 3, 1)
).sum(dim=1)
return (torch.abs(head_top[:, 1] - left_heel[:, 1]),
torch.stack([head_top, left_heel], axis=0)
)
def compute_mass(self, tris: Tensor) -> Tensor:
''' Computes the mass from volume and average body density
'''
x = tris[:, :, :, 0]
y = tris[:, :, :, 1]
z = tris[:, :, :, 2]
volume = (
-x[:, :, 2] * y[:, :, 1] * z[:, :, 0] +
x[:, :, 1] * y[:, :, 2] * z[:, :, 0] +
x[:, :, 2] * y[:, :, 0] * z[:, :, 1] -
x[:, :, 0] * y[:, :, 2] * z[:, :, 1] -
x[:, :, 1] * y[:, :, 0] * z[:, :, 2] +
x[:, :, 0] * y[:, :, 1] * z[:, :, 2]
).sum(dim=1).abs() / 6.0
return volume * self.DENSITY
def forward(
self,
triangles: Tensor,
compute_mass: bool = True,
compute_height: bool = True,
compute_chest: bool = True,
compute_waist: bool = True,
compute_hips: bool = True,
**kwargs
):
measurements = {}
if compute_mass:
measurements['mass'] = {}
mesh_mass = self.compute_mass(triangles)
measurements['mass']['tensor'] = mesh_mass
if compute_height:
measurements['height'] = {}
mesh_height, points = self.compute_height(triangles)
measurements['height']['tensor'] = mesh_height
measurements['height']['points'] = points
output = self.compute_peripheries(triangles,
compute_chest=compute_chest,
compute_waist=compute_waist,
compute_hips=compute_hips,
)
measurements.update(output)
return {'measurements': measurements}
data/SHAPY/mesh-mesh-intersection/body_measurements/cwh_measurements.py
+182
View File
@@ -0,0 +1,182 @@
# -*- 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 sys
import os
import os.path as osp
from typing import NewType, Dict
import time
import yaml
import numpy as np
import torch
import torch.nn as nn
import torch.autograd as autograd
# from loguru import logger
from mesh_mesh_intersection import MeshMeshIntersection
from scipy.spatial import ConvexHull
Tensor = NewType('Tensor', torch.Tensor)
class ChestWaistHipsMeasurements(nn.Module):
def __init__(
self, meas_definition_path: str, meas_vertices_path: str,
max_collisions=256,
*args, **kwargs
) -> None:
super(ChestWaistHipsMeasurements, self).__init__()
meas_definition_path = osp.expanduser(
osp.expandvars(meas_definition_path))
meas_vertices_path = osp.expanduser(
osp.expandvars(meas_vertices_path))
assert osp.exists(meas_definition_path), (
'Measurement definition path does not exist:'
f' {meas_definition_path}'
)
assert osp.exists(meas_definition_path), (
'Measurement vertex path does not exist:'
f' {meas_vertices_path}'
)
with open(meas_definition_path, 'r') as f:
measurements_definitions = yaml.load(f)
with open(meas_vertices_path, 'r') as f:
meas_vertices = yaml.load(f)
action = measurements_definitions['CW_p']
chest_periphery_data = meas_vertices[action[0]]
self.chest_face_index = chest_periphery_data['face_idx']
chest_bcs = torch.tensor(
chest_periphery_data['bc'], dtype=torch.float32)
self.register_buffer('chest_bcs', chest_bcs)
action = measurements_definitions['BW_p']
belly_periphery_data = meas_vertices[action[0]]
self.belly_face_index = belly_periphery_data['face_idx']
belly_bcs = torch.tensor(
belly_periphery_data['bc'], dtype=torch.float32)
self.register_buffer('belly_bcs', belly_bcs)
action = measurements_definitions['IW_p']
hips_periphery_data = meas_vertices[action[0]]
self.hips_face_index = hips_periphery_data['face_idx']
hips_bcs = torch.tensor(
hips_periphery_data['bc'], dtype=torch.float32)
self.register_buffer('hips_bcs', hips_bcs)
self.isect_module = MeshMeshIntersection(max_collisions=max_collisions)
def _get_plane_at_heights(self, height: Tensor):
device = height.device
batch_size = height.shape[0]
verts = torch.tensor(
[[-1., 0, -1], [1, 0, -1], [1, 0, 1], [-1, 0, 1]],
device=device).unsqueeze(dim=0).expand(batch_size, -1, -1).clone()
verts[:, :, 1] = height.reshape(batch_size, -1)
faces = torch.tensor([[0, 1, 2], [0, 2, 3]], device=device,
dtype=torch.long)
return verts, faces, verts[:, faces]
def forward(
self,
triangles: Tensor
) -> Dict[str, Tensor]:
'''
Parameters
----------
triangles: BxFx3x3 torch.Tensor
Contains the triangle coordinates for a batch of meshes with
the same topology
'''
batch_size, num_triangles = triangles.shape[:2]
device = triangles.device
batch_indices = torch.arange(
batch_size, dtype=torch.long,
device=device).reshape(-1, 1) * num_triangles
meas_data = {
'chest': (self.chest_face_index, self.chest_bcs),
'belly': (self.belly_face_index, self.belly_bcs),
'hips': (self.hips_face_index, self.hips_bcs),
}
output = {}
for name, (face_index, bcs) in meas_data.items():
vertex = (
triangles[:, face_index] * bcs.reshape(1, 3, 1)).sum(axis=1)
_, _, plane_tris = self._get_plane_at_heights(vertex[:, 1])
with torch.no_grad():
collision_faces, collision_bcs = self.isect_module(
plane_tris, triangles)
selected_triangles = triangles.view(-1, 3, 3)[
(collision_faces + batch_indices).view(-1)].reshape(
batch_size, -1, 3, 3)
points = (
selected_triangles[:, :, None] *
collision_bcs[:, :, :, :, None]).sum(
axis=-2).reshape(batch_size, -1, 2, 3)
np_points = points.detach().cpu().numpy()
collision_faces = collision_faces.detach().cpu().numpy()
collision_bcs = collision_bcs.detach().cpu().numpy()
output[name] = {
'points': [],
'valid_points': [],
'value': [],
'plane_height': vertex[:, 1],
}
for ii in range(batch_size):
valid_face_idxs = np.where(collision_faces[ii] > 0)[0]
points_in_plane = np_points[
ii, valid_face_idxs, :, ][:, :, [0, 2]].reshape(
-1, 2)
hull = ConvexHull(points_in_plane)
point_indices = hull.simplices.reshape(-1)
hull_points = points[ii][valid_face_idxs].view(
-1, 3)[point_indices]
meas_value = (
hull_points[1::2] - hull_points[:-1:2]).pow(2).sum(
dim=-1).sqrt().sum()
# logger.info(f'{ii}: {name}, {meas_value}')
output[name]['valid_points'].append(
np_points[ii, valid_face_idxs])
output[name]['points'].append(hull_points)
output[name]['value'].append(meas_value)
# values.append(
# )
return output
data/SHAPY/mesh-mesh-intersection/data/measurement_defitions.yaml
+111
View File
@@ -0,0 +1,111 @@
A:
- FingerTipRight
- FingerTipLeft
- 0
- - 0
- 0
- 0.2
BW:
- BellyButton
- 0
- - 0.0
- 0
- 0.15
BW_p:
- BellyButton
- 0
- - 0.0
- 0
- 0
CW:
- NippleRight
- 0
- - 0.0
- 0
- 0.2
CW_p:
- NippleRight
- 0
- - 0.0
- 0
- 0
DBB:
- BackBellyButton
- BellyButton
- 2
- - 0.15
- 0
- 0
H:
- HeelLeft
- HeadTop
- 1
- - -0.15
- 0
- 0.0
HB:
- HeelLeft
- NippleRight
- NippleLeft
- 1
- - -0.1
- 0
- 0.0
HBB:
- HeelLeft
- BellyButton
- 1
- - -0.05
- 0
- 0.0
HI:
- HeelLeft
- Crotch
- 1
- - 0
- 0
- 0
IW:
- Crotch
- 0
- - 0.0
- 0
- 0.15
IW_p:
- Crotch
- 0
- - 0.0
- 0
- 0
SW:
- ShoulderApose
- 0
- - 0
- 0
- 0.1
V: []
W2E:
- ElbowRight
- WristRight
- ElbowLeft
- WristLeft
- 0
- - 0
- 0.0
- 0.1
W2S:
- ShoulderRight
- WristRight
- ShoulderLeft
- WristLeft
- 0
- - 0
- 0.0
- 0.1
W2W:
- WristRight
- WristLeft
- 0
- - 0
- -0.0
- 0.15
data/SHAPY/mesh-mesh-intersection/data/smpl_measurement_vertices.yaml
+112
View File
@@ -0,0 +1,112 @@
BackBellyButton:
bc:
- 1.0
- 0.0
- 0.0
face_idx: 4971
vertex_id: 3022
BellyButton:
bc:
- 0.0
- 0.0
- 1.0
face_idx: 6833
vertex_id: 3501
Crotch:
bc:
- 0.0
- 1.0
- 0.0
face_idx: 1341
vertex_id: 1210
ElbowLeft:
bc:
- 1.0
- 0.0
- 0.0
face_idx: 1867
vertex_id: 1658
ElbowRight:
bc:
- 0.0
- 0.0
- 1.0
face_idx: 8756
vertex_id: 5129
FingerTipLeft:
bc:
- 0.0
- 0.0
- 1.0
face_idx: 3259
vertex_id: 2445
FingerTipRight:
bc:
- 0.0
- 0.0
- 1.0
face_idx: 10147
vertex_id: 5905
HeadTop:
bc:
- 0.0
- 1.0
- 0.0
face_idx: 435
vertex_id: 411
HeelLeft:
bc:
- 0.0
- 0.0
- 1.0
face_idx: 5975
vertex_id: 3466
NippleLeft:
bc:
- 0.0
- 0.0
- 1.0
face_idx: 4997
vertex_id: 3042
NippleRight:
bc:
- 0.0
- 0.0
- 1.0
face_idx: 11885
vertex_id: 6489
ShoulderApose:
bc:
- 1.0
- 0.0
- 0.0
face_idx: 11937
vertex_id: 6496
ShoulderLeft:
bc:
- 0.0
- 0.0
- 1.0
face_idx: 4572
vertex_id: 2893
ShoulderRight:
bc:
- 1.0
- 0.0
- 0.0
face_idx: 9117
vertex_id: 5291
WristLeft:
bc:
- 0.0
- 0.0
- 1.0
face_idx: 2603
vertex_id: 2099
WristRight:
bc:
- 0.0
- 1.0
- 0.0
face_idx: 9491
vertex_id: 5559
data/SHAPY/mesh-mesh-intersection/data/smplx_measurements.yaml
+160
View File
@@ -0,0 +1,160 @@
BackBellyButton:
bc:
- 0.0
- 0.0
- 1.0
closest_points:
- 0.0
- -0.261382
- -0.102003
face_idx: 7861
BellyButton:
bc:
- 0.0
- 1.0
- 0.0
closest_points:
- -0.0
- -0.271119
- 0.144329
face_idx: 19229
Crotch:
bc:
- 0.0
- 0.0
- 1.0
closest_points:
- 0.0
- -0.53204
- 0.036471
face_idx: 6194
ElbowLeft:
bc:
- 0.0
- 1.0
- 0.0
closest_points:
- 0.447234
- 0.075995
- -0.099974
face_idx: 3959
ElbowRight:
bc:
- 0.0
- 1.0
- 0.0
closest_points:
- -0.447233
- 0.075995
- -0.099973
face_idx: 7846
FingerTipLeft:
bc:
- 0.0
- 0.0
- 1.0
closest_points:
- 0.918755
- 0.085413
- -0.084745
face_idx: 3469
FingerTipRight:
bc:
- 1.3918288479186636e-05
- 0.9998084353210817
- 0.00017764639043920215
closest_points:
- -0.9187545563631815
- 0.08541207136535614
- -0.08474471930198169
face_idx: 17602
HeadTop:
bc:
- 0.8277337276382795
- 0.1422200962169292
- 0.030046176144791284
closest_points:
- -0.0017716260945737938
- 0.4363661265424736
- -0.015488867245138805
face_idx: 2581
HeelLeft:
bc:
- 0.0
- 1.0
- 0.0
closest_points:
- 0.103005
- -1.346656
- -0.082664
face_idx: 15605
NippleLeft:
bc:
- 1.0
- 0.0
- 0.0
closest_points:
- 0.09747
- -0.032798
- 0.103687
face_idx: 16306
NippleRight:
bc:
- 0.0
- 0.0
- 1.0
closest_points:
- -0.09747
- -0.032798
- 0.103687
face_idx: 18402
ShoulderApose:
bc:
- 1.0
- 0.0
- 0.0
closest_points:
- -0.013092
- 0.117681
- 0.03777
face_idx: 18412
ShoulderLeft:
bc:
- 0.0
- 0.0
- 1.0
closest_points:
- 0.183378
- 0.139588
- -0.066495
face_idx: 3865
ShoulderRight:
bc:
- 0.0
- 0.0
- 1.0
closest_points:
- -0.196498
- 0.140399
- -0.063292
face_idx: 18186
WristLeft:
bc:
- 1.0
- 0.0
- 0.0
closest_points:
- 0.72226
- 0.065775
- -0.070361
face_idx: 3363
WristRight:
bc:
- 0.0
- 0.0
- 1.0
closest_points:
- -0.72226
- 0.065775
- -0.070361
face_idx: 6722
data/SHAPY/mesh-mesh-intersection/docs/install.md
+31
View File
@@ -0,0 +1,31 @@
# Installation
Before installing anything please make sure to set the environment variable
*$CUDA_SAMPLES_INC* to the path that contains the header `helper_math.h`, which
can be found in the repo [CUDA Samples repository](https://github.com/NVIDIA/cuda-samples).
To install the module run the following commands:
**1. Clone this repository**
```Shell
git clone https://github.com/vchoutas/torch-mesh-isect
cd torch-mesh-isect
```
**2. Install the dependencies**
```Shell
pip install -r requirements.txt
```
**3. Run the *setup.py* script**
```Shell
python setup.py install
```
## Dependencies
1. [PyTorch](https://pytorch.org)
### Optional Dependencies
1. [Trimesh](https://trimsh.org) for loading triangular meshes
2. [open3d](http://www.open3d.org/) for visualization
The code has been tested with Python 3.6, CUDA 10.0, CuDNN 7.3 and PyTorch 1.0.
data/SHAPY/mesh-mesh-intersection/examples/fit_measurements.py
+304
View File
@@ -0,0 +1,304 @@
# -*- 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 sys
import os.path as osp
import argparse
import numpy as np
import torch
import torch.nn.functional as F
import torch.optim as optim
import smplx
import open3d as o3d
import time
import cv2
from tqdm import tqdm
import trimesh
from loguru import logger
from star.pytorch.star import STAR
from star.config import cfg as star_cfg
from body_measurements import BodyMeasurements
from torchtrustncg import TrustRegion
def get_plane_at_height(h):
verts = np.array([[-1., h, -1], [1, h, -1], [1, h, 1], [-1, h, 1]])
faces = np.array([[0, 1, 2], [0, 2, 3]])
normal = np.array([0.0, 1.0, 0.0])
return verts, faces, (verts[0], normal)
def main(
model_folder,
height: float = 1.76,
mass: float = -1,
chest: float = 1.12,
waist: float = 0.93,
hips: float = 1.14,
model_type='smplx',
ext='npz',
gender='neutral',
num_betas=10,
meas_definition_path: str = 'data/measurement_defitions.yaml',
meas_vertices_path: str = 'data/smpl_measurement_vertices.yaml',
summary_steps: int = 50,
num_iterations: int = 500,
betas_weight: float = 0.0,
):
device = torch.device('cuda')
dtype = torch.float32
cfg = {
'meas_definition_path': meas_definition_path,
'meas_vertices_path': meas_vertices_path,
}
meas_module = BodyMeasurements(cfg)
meas_module = meas_module.to(device=device)
num_samples = 1
trans, pose = None, None
logger.info(f'Model type: {model_type}')
if 'star' in model_type:
star_cfg.path_male_star = osp.expandvars(
osp.join(model_folder, 'star', 'STAR_MALE.npz'))
star_cfg.path_female_star = osp.expandvars(
osp.join(model_folder, 'star', 'STAR_FEMALE.npz'))
model = STAR(gender=gender, num_betas=num_betas)
trans = torch.zeros([num_samples, 3], dtype=dtype, device=device)
pose = torch.zeros([num_samples, 72], dtype=dtype, device=device)
else:
model = smplx.build_layer(
model_folder, model_type=model_type,
gender=gender,
num_betas=num_betas,
ext=ext)
logger.info(model)
model = model.to(device=device)
betas = torch.zeros(
[num_samples, model.num_betas],
requires_grad=True, dtype=torch.float32, device=device)
dtype = torch.float32
gt = {
'height': torch.tensor(height, dtype=dtype, device=device),
'mass': torch.tensor(mass, dtype=dtype, device=device),
'chest': torch.tensor(chest, dtype=dtype, device=device),
'waist': torch.tensor(waist, dtype=dtype, device=device),
'hips': torch.tensor(hips, dtype=dtype, device=device),
}
weights = {
'height': 100.0 if height > 0 else 0.0,
'mass': 1.0 if mass > 0 else 0.0,
'chest': 2000.0 if chest > 0 else 0.0,
'waist': 1000.0 if waist > 0 else 0.0,
'hips': 1000.0 if hips > 0 else 0.0,
}
optimizer = TrustRegion([betas])
def compute_loss(gt, output, weights):
losses = {}
for key, gt_val in gt.items():
if weights[key] <= 1e-3 or gt_val.item() < 0:
continue
est_val = output[key]['tensor']
if isinstance(est_val, (tuple, list)):
est_val = torch.stack(output[key]['value'])
curr_loss = (gt_val - est_val).pow(2).sum() * weights[key]
losses[key] = curr_loss
losses['betas'] = betas_weight * betas.pow(2).sum()
return losses
def closure(backward=True):
if backward:
optimizer.zero_grad()
if model_type == 'star':
vertices = model(pose=pose, trans=trans, betas=betas)
model_tris = vertices[:, model.faces]
else:
output = model(betas=betas, return_verts=True)
model_tris = output.vertices[:, model.faces_tensor]
output = meas_module(model_tris)['measurements']
losses = compute_loss(gt, output, weights)
loss = sum(losses.values())
if backward:
loss.backward(create_graph=True)
return loss
Y_OFFSET = -1.10
for n in tqdm(range(num_iterations)):
loss = optimizer.step(closure)
if n % summary_steps == 0:
if model_type == 'star':
vertices = model(pose=pose, trans=trans, betas=betas)
model_tris = vertices[:, model.faces]
vertices = vertices.detach().cpu().numpy().squeeze()
faces = model.faces.detach().cpu().numpy()
else:
output = model(betas=betas, return_verts=True)
vertices = output.vertices.detach().cpu().numpy().squeeze()
faces = model.faces
model_tris = output.vertices[:, model.faces_tensor]
y_offset = - vertices[:, 1].min() + Y_OFFSET
vertices[:, 1] = vertices[:, 1] + y_offset
# for key, val in losses.items():
mesh = o3d.geometry.TriangleMesh()
mesh.vertices = o3d.utility.Vector3dVector(vertices)
mesh.triangles = o3d.utility.Vector3iVector(faces)
mesh.compute_vertex_normals()
colors = np.ones_like(vertices) * [0.3, 0.3, 0.3]
mesh.vertex_colors = o3d.utility.Vector3dVector(colors)
geometry = []
geometry.append(mesh)
output = meas_module(model_tris)['measurements']
for key, val in gt.items():
est_val = output[key]["tensor"][0].item()
logger.info(
f'[{n:04d}]: {key}: est = {est_val}, gt = {val}')
losses = compute_loss(gt, output, weights)
for key, val in losses.items():
logger.info(f'[{n:04d}]: {key} loss = {val:.3f}')
for meas_name in output:
pcl = o3d.geometry.PointCloud()
if 'points' not in output[meas_name]:
continue
points = output[meas_name]['points']
if isinstance(points, (tuple, list)):
points = torch.stack(points)
if torch.is_tensor(points):
points = points.detach().cpu().numpy()
points = points.reshape(-1, 3)
points[:, 1] = points[:, 1] + y_offset
pcl.points = o3d.utility.Vector3dVector(points)
pcl.paint_uniform_color([1.0, 0.0, 0.0])
geometry.append(pcl)
lineset = o3d.geometry.LineSet()
line_ids = np.arange(len(points)).reshape(-1, 2)
lineset.points = o3d.utility.Vector3dVector(points)
lineset.lines = o3d.utility.Vector2iVector(line_ids)
lineset.paint_uniform_color([0.0, 0.0, 0.0])
geometry.append(lineset)
o3d.visualization.draw_geometries(
geometry,
lookat=np.array([0.0, 0.0, 0.0]).reshape(3, 1),
up=np.array([0.0, 1.0, 0.0]).reshape(3, 1),
front=np.array([0.0, 0.0, 1.0]).reshape(3, 1),
zoom=1.0,
)
if __name__ == '__main__':
logger.remove()
logger.add(lambda x: tqdm.write(x, end=''), colorize=True)
parser = argparse.ArgumentParser(description='SMPL-X Demo')
parser.add_argument('--model-folder', required=True, type=str,
help='The path to the model folder')
parser.add_argument('--model-type', default='smpl', type=str,
choices=['smpl', 'smplh', 'smplx', 'mano', 'flame',
'star', ],
help='The type of model to load')
parser.add_argument('--gender', type=str, default='neutral',
help='The gender of the model')
parser.add_argument('--num-betas', default=10, type=int,
dest='num_betas',
help='Number of shape coefficients.')
parser.add_argument('--ext', type=str, default='npz',
help='Which extension to use for loading')
parser.add_argument('--height', type=float, default=1.80,
help='Height of the subject in meters')
parser.add_argument('--mass', type=float, default=-1,
help='Mass of the subject in kilograms')
parser.add_argument('--chest', type=float, default=-1,
help='Chest circumference in meters')
parser.add_argument('--waist', type=float, default=-1,
help='Waist circumference in meters')
parser.add_argument('--hips', type=float, default=-1,
help='Hips circumference in meters')
parser.add_argument('--meas-definition-path',
dest='meas_definition_path',
default='data/measurement_defitions.yaml',
type=str,
help='The definitions of the measurements')
parser.add_argument('--meas-vertices-path', dest='meas_vertices_path',
type=str,
default='data/smpl_measurement_vertices.yaml',
help='The indices of the vertices used for the'
' the measurements')
parser.add_argument('--betas-weight', dest='betas_weight', default=0.0,
type=float,
help='The weight of the shape prior term.')
args = parser.parse_args()
model_folder = osp.expanduser(osp.expandvars(args.model_folder))
model_type = args.model_type
gender = args.gender
ext = args.ext
num_betas = args.num_betas
height = args.height
mass = args.mass
chest = args.chest
waist = args.waist
hips = args.hips
meas_definition_path = args.meas_definition_path
meas_vertices_path = args.meas_vertices_path
betas_weight = args.betas_weight
main(model_folder,
height=height,
mass=mass,
chest=chest,
waist=waist,
hips=hips,
model_type=model_type,
ext=ext,
gender=gender,
num_betas=num_betas,
meas_definition_path=meas_definition_path,
meas_vertices_path=meas_vertices_path,
betas_weight=betas_weight,
)
data/SHAPY/mesh-mesh-intersection/examples/vis_measurements.py
+245
View File
@@ -0,0 +1,245 @@
# -*- 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 sys
import os.path as osp
import argparse
import numpy as np
import torch
import smplx
import open3d as o3d
import time
import cv2
from scipy.spatial import ConvexHull
import trimesh
# from meas_definitions import measurements_definitions, measures_vertex
from loguru import logger
from star.pytorch.star import STAR
from star.config import cfg as star_cfg
from mesh_mesh_intersection import MeshMeshIntersection
from body_measurements import ChestWaistHipsMeasurements
def get_plane_at_height(h):
verts = np.array([[-1., h, -1], [1, h, -1], [1, h, 1], [-1, h, 1]])
faces = np.array([[0, 1, 2], [0, 2, 3]])
normal = np.array([0.0, 1.0, 0.0])
return verts, faces, (verts[0], normal)
def main(
model_folder,
model_type='smplx',
ext='npz',
gender='neutral',
plot_joints=False,
num_betas=10,
sample_shape=False,
num_expression_coeffs=10,
plotting_module='pyrender',
num_samples=1,
use_face_contour=False,
meas_definition_path: str = 'data/measurement_defitions.yaml',
meas_vertices_path: str = 'data/smpl_measurement_vertices.yaml',
):
device = torch.device('cuda')
meas_module = ChestWaistHipsMeasurements(
meas_definition_path=meas_definition_path,
meas_vertices_path=meas_vertices_path,
# '$HOME/workspace/caesar_betas_exps/measurement_defitions.yaml',
# '$HOME/workspace/caesar_betas_exps/smpl_measurement_vertices.yaml',
)
# meas_module = ChestWaistHipsMeasurements(
# '$HOME/workspace/caesar_betas_exps/measurement_defitions.yaml',
# '$HOME/workspace/caesar_betas_exps/smplx_measurements.yaml',
# )
meas_module = meas_module.to(device=device)
dtype = torch.float32
trans, pose = None, None
if model_type == 'star':
star_cfg.path_male_star = osp.expandvars(
'$HOME/workspace/body_models/star/STAR_MALE.npz')
star_cfg.path_female_star = osp.expandvars(
'$HOME/workspace/body_models/star/STAR_FEMALE.npz')
model = STAR(gender=gender, num_betas=num_betas)
trans = torch.zeros([num_samples, 3], dtype=dtype, device=device)
pose = torch.zeros([num_samples, 72], dtype=dtype, device=device)
else:
model = smplx.build_layer(
model_folder, model_type=model_type,
gender=gender, use_face_contour=use_face_contour,
num_betas=num_betas,
num_expression_coeffs=num_expression_coeffs,
ext=ext)
model = model.to(device=device)
# meas_to_vis = ['CW_p', 'BW_p', 'IW_p']
# meas_to_vis = ['CW_p']
if sample_shape:
# betas = shape_dist.sample().reshape(1, -1)
betas = torch.randn(
[num_samples, model.num_betas],
requires_grad=True,
dtype=torch.float32, device=device)
else:
betas = torch.zeros(
[num_samples, model.num_betas],
requires_grad=True, dtype=torch.float32, device=device)
model.zero_grad()
if model_type == 'star':
vertices = model(pose=pose, trans=trans, betas=betas)
model_tris = vertices[:, model.faces]
vertices = vertices.detach().cpu().numpy()
faces = model.faces.detach().cpu().numpy()
else:
output = model(betas=betas, return_verts=True)
vertices = output.vertices.detach().cpu().numpy()
model_tris = output.vertices[:, model.faces_tensor]
faces = model.faces
output = meas_module(model_tris)
# loss = sum(v.pow(2) for v in output['chest']['value'])
for n in range(num_samples):
mesh = o3d.geometry.TriangleMesh()
mesh.vertices = o3d.utility.Vector3dVector(vertices[n])
mesh.triangles = o3d.utility.Vector3iVector(faces)
mesh.compute_vertex_normals()
colors = np.ones_like(vertices[n]) * [0.3, 0.3, 0.3]
mesh.vertex_colors = o3d.utility.Vector3dVector(colors)
geometry = []
geometry.append(mesh)
for meas_name in output:
pcl = o3d.geometry.PointCloud()
if 'points' not in output[meas_name]:
continue
points = output[meas_name]['points']
if isinstance(points, (tuple, list)):
points = torch.stack(points)
if torch.is_tensor(points):
points = points.detach().cpu().numpy()
points = points.reshape(-1, 3)
pcl.points = o3d.utility.Vector3dVector(points)
pcl.paint_uniform_color([1.0, 0.0, 0.0])
geometry.append(pcl)
lineset = o3d.geometry.LineSet()
line_ids = np.arange(len(points)).reshape(-1, 2)
lineset.points = o3d.utility.Vector3dVector(points)
lineset.lines = o3d.utility.Vector2iVector(line_ids)
lineset.paint_uniform_color([0.0, 0.0, 0.0])
geometry.append(lineset)
o3d.visualization.draw_geometries(geometry)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='SMPL-X Demo')
parser.add_argument('--model-folder', required=True, type=str,
help='The path to the model folder')
parser.add_argument('--model-type', default='smplx', type=str,
choices=['smpl', 'smplh', 'smplx', 'mano', 'flame',
'star', ],
help='The type of model to load')
parser.add_argument('--gender', type=str, default='neutral',
help='The gender of the model')
parser.add_argument('--num-betas', default=10, type=int,
dest='num_betas',
help='Number of shape coefficients.')
parser.add_argument('--num-expression-coeffs', default=10, type=int,
dest='num_expression_coeffs',
help='Number of expression coefficients.')
parser.add_argument('--plotting-module', type=str, default='pyrender',
dest='plotting_module',
choices=['pyrender', 'matplotlib', 'open3d'],
help='The module to use for plotting the result')
parser.add_argument('--ext', type=str, default='npz',
help='Which extension to use for loading')
parser.add_argument('--plot-joints', default=False,
type=lambda arg: arg.lower() in ['true', '1'],
help='The path to the model folder')
parser.add_argument('--sample-shape', default=False,
dest='sample_shape',
type=lambda arg: arg.lower() in ['true', '1'],
help='Sample a random shape')
parser.add_argument('--sample-expression', default=True,
dest='sample_expression',
type=lambda arg: arg.lower() in ['true', '1'],
help='Sample a random expression')
parser.add_argument('--use-face-contour', default=False,
type=lambda arg: arg.lower() in ['true', '1'],
help='Compute the contour of the face')
parser.add_argument('--num-samples', default=1, type=int,
dest='num_samples',
help='Number of samples to draw.')
parser.add_argument('--meas-definition-path',
dest='meas_definition_path',
default='data/measurement_defitions.yaml',
type=str,
help='The definitions of the measurements')
parser.add_argument('--meas-vertices-path', dest='meas_vertices_path',
type=str,
default='data/smpl_measurement_vertices.yaml',
help='The indices of the vertices used for the'
' the measurements')
args = parser.parse_args()
model_folder = osp.expanduser(osp.expandvars(args.model_folder))
model_type = args.model_type
plot_joints = args.plot_joints
use_face_contour = args.use_face_contour
gender = args.gender
ext = args.ext
plotting_module = args.plotting_module
num_betas = args.num_betas
num_expression_coeffs = args.num_expression_coeffs
sample_shape = args.sample_shape
sample_expression = args.sample_expression
num_samples = args.num_samples
meas_definition_path = args.meas_definition_path
meas_vertices_path = args.meas_vertices_path
main(model_folder, model_type, ext=ext,
gender=gender, plot_joints=plot_joints,
num_betas=num_betas,
num_samples=num_samples,
num_expression_coeffs=num_expression_coeffs,
sample_shape=sample_shape,
plotting_module=plotting_module,
meas_definition_path=meas_definition_path,
meas_vertices_path=meas_vertices_path,
use_face_contour=use_face_contour,
)
data/SHAPY/mesh-mesh-intersection/include/aabb.hpp
+117
View File
@@ -0,0 +1,117 @@
/*
* 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.
*
* @author Vasileios Choutas
* Contact: vassilis.choutas@tuebingen.mpg.de
* Contact: ps-license@tuebingen.mpg.de
*
*/
#ifndef AABB_H
#define AABB_H
#include <cuda.h>
#include "device_launch_parameters.h"
#include <cuda_runtime.h>
#include "defs.hpp"
#include "math_utils.hpp"
#include "double_vec_ops.hpp"
#include "helper_math.h"
template <typename T>
__align__(32)
struct AABB {
public:
__host__ __device__ AABB() {
min_t.x = std::is_same<T, float>::value ? FLT_MAX : DBL_MAX;
min_t.y = std::is_same<T, float>::value ? FLT_MAX : DBL_MAX;
min_t.z = std::is_same<T, float>::value ? FLT_MAX : DBL_MAX;
max_t.x = std::is_same<T, float>::value ? -FLT_MAX : -DBL_MAX;
max_t.y = std::is_same<T, float>::value ? -FLT_MAX : -DBL_MAX;
max_t.z = std::is_same<T, float>::value ? -FLT_MAX : -DBL_MAX;
};
__host__ __device__ AABB(const vec3<T> &min_t, const vec3<T> &max_t)
: min_t(min_t), max_t(max_t){};
__host__ __device__ ~AABB(){};
__host__ __device__ AABB(T min_t_x, T min_t_y, T min_t_z, T max_t_x,
T max_t_y, T max_t_z) {
min_t.x = min_t_x;
min_t.y = min_t_y;
min_t.z = min_t_z;
max_t.x = max_t_x;
max_t.y = max_t_y;
max_t.z = max_t_z;
}
__host__ __device__ AABB<T> operator+(const AABB<T> &bbox2) const {
return AABB<T>(
min(this->min_t.x, bbox2.min_t.x), min(this->min_t.y, bbox2.min_t.y),
min(this->min_t.z, bbox2.min_t.z), max(this->max_t.x, bbox2.max_t.x),
max(this->max_t.y, bbox2.max_t.y), max(this->max_t.z, bbox2.max_t.z));
};
__host__ __device__ T distance(const vec3<T> point) const {
};
__host__ __device__ T operator*(const AABB<T> &bbox2) const {
return (min(this->max_t.x, bbox2.max_t.x) -
max(this->min_t.x, bbox2.min_t.x)) *
(min(this->max_t.y, bbox2.max_t.y) -
max(this->min_t.y, bbox2.min_t.y)) *
(min(this->max_t.z, bbox2.max_t.z) -
max(this->min_t.z, bbox2.min_t.z));
};
vec3<T> min_t;
vec3<T> max_t;
};
template <typename T>
std::ostream &operator<<(std::ostream &os, const AABB<T> &x) {
os << x.min_t << std::endl;
os << x.max_t << std::endl;
return os;
}
template <typename T> struct MergeAABB {
public:
__host__ __device__ MergeAABB(){};
// Create an operator Struct that will be used by thrust::reduce
// to calculate the bounding box of the scene.
__host__ __device__ AABB<T> operator()(const AABB<T> &bbox1,
const AABB<T> &bbox2) {
return bbox1 + bbox2;
};
};
template <typename T>
__forceinline__
__host__ __device__ T pointToAABBDistance(vec3<T> point, const AABB<T>& bbox ) {
T diff_x = point.x - clamp<T>(point.x, bbox.min_t.x, bbox.max_t.x);
T diff_y = point.y - clamp<T>(point.y, bbox.min_t.y, bbox.max_t.y);
T diff_z = point.z - clamp<T>(point.z, bbox.min_t.z, bbox.max_t.z);
return diff_x * diff_x + diff_y * diff_y + diff_z * diff_z;
}
#endif // ifndef AABB_H
data/SHAPY/mesh-mesh-intersection/include/defs.hpp
+68
View File
@@ -0,0 +1,68 @@
/*
* 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.
*
* @author Vasileios Choutas
* Contact: vassilis.choutas@tuebingen.mpg.de
* Contact: ps-license@tuebingen.mpg.de
*
*/
#ifndef DEFINITIONS_H
#define DEFINITIONS_H
#include <cuda.h>
#include "device_launch_parameters.h"
#include <cuda_runtime.h>
template <typename T>
using vec3 = typename std::conditional<std::is_same<T, float>::value, float3,
double3>::type;
template <typename T>
using vec2 = typename std::conditional<std::is_same<T, float>::value, float2,
double2>::type;
float3 make_float3(double3 vec) {
return make_float3(vec.x, vec.y, vec.z);
}
float3 make_float3(double x, double y, double z) {
return make_float3(x, y, z);
}
double3 make_double3(float3 vec) {
return make_double3(vec.x, vec.y, vec.z);
}
double3 make_double3(float x, float y, float z) {
return make_double3(x, y, z);
}
template <typename T>
__host__ __device__
vec3<T> make_vec3(T x, T y, T z) {
}
template <>
__host__ __device__
vec3<float> make_vec3(float x, float y, float z) {
return make_float3(static_cast<float>(x), static_cast<float>(y), static_cast<float>(z));
}
template <>
__host__ __device__
vec3<double> make_vec3(double x, double y, double z) {
return make_double3(static_cast<double>(x), static_cast<double>(y), static_cast<double>(z));
}
#endif // ifndef DEFINITIONS_H
data/SHAPY/mesh-mesh-intersection/include/double_vec_ops.h
+117
View File
@@ -0,0 +1,117 @@
/*
* 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.
*
* @author Vasileios Choutas
* Contact: vassilis.choutas@tuebingen.mpg.de
* Contact: ps-license@tuebingen.mpg.de
*
*/
#ifndef DOUBLE_VEC_OPS_H
#define DOUBLE_VEC_OPS_H
#include "cuda_runtime.h"
inline __host__ __device__ double2 operator+(double2 a, double2 b) {
return make_double2(a.x + b.x, a.y + b.y);
}
inline __host__ __device__ double3 operator+(double3 a, double3 b) {
return make_double3(a.x + b.x, a.y + b.y, a.z + b.z);
}
inline __host__ __device__ void operator/=(double2 &a, double2 b) {
a.x /= b.x;
a.y /= b.y;
}
inline __host__ __device__ double2 operator/(double2 a, double b) {
return make_double2(a.x / b, a.y / b);
}
inline __host__ __device__ double3 operator/(double3 a, double3 b) {
return make_double3(a.x / b.x, a.y / b.y, a.z / b.z);
}
inline __host__ __device__ double3 operator*(double a, double3 b) {
return make_double3(a * b.x, a * b.y, a * b.z);
}
inline __host__ __device__ double3 operator*(double3 a, double3 b) {
return make_double3(a.x * b.x, a.y * b.y, a.z * b.z);
}
inline __host__ __device__ void operator/=(double3 &a, double3 b) {
a.x /= b.x;
a.y /= b.y;
a.z /= b.z;
}
inline __host__ __device__ double3 operator/(double3 a, double b) {
return make_double3(a.x / b, a.y / b, a.z / b);
}
inline __host__ __device__ double dot(double2 a, double2 b) {
return a.x * b.x + a.y * b.y;
}
inline __host__ __device__ double dot(double3 a, double3 b) {
return a.x * b.x + a.y * b.y + a.z * b.z;
}
inline __host__ __device__ double3 cross(double3 a, double3 b)
{
return make_double3(a.y*b.z - a.z*b.y, a.z*b.x - a.x*b.z, a.x*b.y - a.y*b.x);
}
inline __host__ __device__ double2 operator-(double2 a, double2 b)
{
return make_double2(a.x - b.x, a.y - b.y);
}
inline __host__ __device__ void operator-=(double2 &a, double2 b)
{
a.x -= b.x;
a.y -= b.y;
}
inline __host__ __device__ double2 operator-(double2 a, double b)
{
return make_double2(a.x - b, a.y - b);
}
inline __host__ __device__ double2 operator-(double b, double2 a)
{
return make_double2(b - a.x, b - a.y);
}
inline __host__ __device__ double3 operator-(double3 a, double3 b)
{
return make_double3(a.x - b.x, a.y - b.y, a.z - b.z);
}
inline __host__ __device__ void operator-=(double3 &a, double3 b)
{
a.x -= b.x;
a.y -= b.y;
a.z -= b.z;
}
inline __host__ __device__ double3 operator-(double3 a, double b)
{
return make_double3(a.x - b, a.y - b, a.z - b);
}
inline __host__ __device__ double3 operator-(double b, double3 a)
{
return make_double3(b - a.x, b - a.y, b - a.z);
}
#endif // ifndef DOUBLE_VEC_OPS_H
data/SHAPY/mesh-mesh-intersection/include/double_vec_ops.hpp
+117
View File
@@ -0,0 +1,117 @@
/*
* 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.
*
* @author Vasileios Choutas
* Contact: vassilis.choutas@tuebingen.mpg.de
* Contact: ps-license@tuebingen.mpg.de
*
*/
#ifndef DOUBLE_VEC_OPS_H
#define DOUBLE_VEC_OPS_H
#include "cuda_runtime.h"
inline __host__ __device__ double2 operator+(double2 a, double2 b) {
return make_double2(a.x + b.x, a.y + b.y);
}
inline __host__ __device__ double3 operator+(double3 a, double3 b) {
return make_double3(a.x + b.x, a.y + b.y, a.z + b.z);
}
inline __host__ __device__ void operator/=(double2 &a, double2 b) {
a.x /= b.x;
a.y /= b.y;
}
inline __host__ __device__ double2 operator/(double2 a, double b) {
return make_double2(a.x / b, a.y / b);
}
inline __host__ __device__ double3 operator/(double3 a, double3 b) {
return make_double3(a.x / b.x, a.y / b.y, a.z / b.z);
}
inline __host__ __device__ double3 operator*(double a, double3 b) {
return make_double3(a * b.x, a * b.y, a * b.z);
}
inline __host__ __device__ double3 operator*(double3 a, double3 b) {
return make_double3(a.x * b.x, a.y * b.y, a.z * b.z);
}
inline __host__ __device__ void operator/=(double3 &a, double3 b) {
a.x /= b.x;
a.y /= b.y;
a.z /= b.z;
}
inline __host__ __device__ double3 operator/(double3 a, double b) {
return make_double3(a.x / b, a.y / b, a.z / b);
}
inline __host__ __device__ double dot(double2 a, double2 b) {
return a.x * b.x + a.y * b.y;
}
inline __host__ __device__ double dot(double3 a, double3 b) {
return a.x * b.x + a.y * b.y + a.z * b.z;
}
inline __host__ __device__ double3 cross(double3 a, double3 b)
{
return make_double3(a.y*b.z - a.z*b.y, a.z*b.x - a.x*b.z, a.x*b.y - a.y*b.x);
}
inline __host__ __device__ double2 operator-(double2 a, double2 b)
{
return make_double2(a.x - b.x, a.y - b.y);
}
inline __host__ __device__ void operator-=(double2 &a, double2 b)
{
a.x -= b.x;
a.y -= b.y;
}
inline __host__ __device__ double2 operator-(double2 a, double b)
{
return make_double2(a.x - b, a.y - b);
}
inline __host__ __device__ double2 operator-(double b, double2 a)
{
return make_double2(b - a.x, b - a.y);
}
inline __host__ __device__ double3 operator-(double3 a, double3 b)
{
return make_double3(a.x - b.x, a.y - b.y, a.z - b.z);
}
inline __host__ __device__ void operator-=(double3 &a, double3 b)
{
a.x -= b.x;
a.y -= b.y;
a.z -= b.z;
}
inline __host__ __device__ double3 operator-(double3 a, double b)
{
return make_double3(a.x - b, a.y - b, a.z - b);
}
inline __host__ __device__ double3 operator-(double b, double3 a)
{
return make_double3(b - a.x, b - a.y, b - a.z);
}
#endif // ifndef DOUBLE_VEC_OPS_H
data/SHAPY/mesh-mesh-intersection/include/helper_math.h
+1453
View File
File diff suppressed because it is too large Load Diff
data/SHAPY/mesh-mesh-intersection/include/math_utils.hpp
+58
View File
@@ -0,0 +1,58 @@
/*
* 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.
*
* @author Vasileios Choutas
* Contact: vassilis.choutas@tuebingen.mpg.de
* Contact: ps-license@tuebingen.mpg.de
*
*/
#ifndef MATH_UTILS_H
#define MATH_UTILS_H
#include <cuda.h>
#include "device_launch_parameters.h"
#include <cuda_runtime.h>
#include "defs.hpp"
#include "double_vec_ops.hpp"
#include "helper_math.h"
template <typename T>
__host__ __device__ T sign(T x) {
return x > 0 ? 1 : (x < 0 ? -1 : 0);
}
template <typename T>
__host__ __device__ __forceinline__ float vec_abs_diff(const vec3<T> &vec1,
const vec3<T> &vec2) {
return fabs(vec1.x - vec2.x) + fabs(vec1.y - vec2.y) + fabs(vec1.z - vec2.z);
}
template <typename T>
__host__ __device__ __forceinline__ float vec_sq_diff(const vec3<T> &vec1,
const vec3<T> &vec2) {
return dot(vec1 - vec2, vec1 - vec2);
}
template <typename T>
__host__ __device__ __forceinline__ T clamp(T value, T min_value, T max_value) {
return min(max(value, min_value), max_value);
}
template <typename T> __host__ __device__ T dot2(vec3<T> v) {
return dot(v, v);
}
#endif // ifndef MATH_UTILS_H
data/SHAPY/mesh-mesh-intersection/include/moller.h
+704
View File
@@ -0,0 +1,704 @@
/* Triangle/triangle intersection test routine,
* by Tomas Moller, 1997.
* See article "A Fast Triangle-Triangle Intersection Test",
* Journal of Graphics Tools, 2(2), 1997
* updated: 2001-06-20 (added line of intersection)
*
* int tri_tri_intersect(float V0[3],float V1[3],float V2[3],
* float U0[3],float U1[3],float U2[3])
*
* parameters: vertices of triangle 1: V0,V1,V2
* vertices of triangle 2: U0,U1,U2
* result : returns 1 if the triangles intersect, otherwise 0
*
* Here is a version withouts divisions (a little faster)
* int NoDivTriTriIsect(float V0[3],float V1[3],float V2[3],
* float U0[3],float U1[3],float U2[3]);
*
* This version computes the line of intersection as well (if they are not
*coplanar): int tri_tri_intersect_with_isectline(float V0[3],float V1[3],float
*V2[3], float U0[3],float U1[3],float U2[3],int *coplanar, float
*isect_point1[3],float isect_point2[3]); coplanar returns whether the tris are
*coplanar isect_point1, isect_point2 are the endpoints of the line of
*intersection
*/
#include <math.h>
#include "defs.h"
#include "double_vec_ops.hpp"
#include "helper_math.h"
#define FABS(x) ((float)fabs(x)) /* implement as is fastest on your machine */
/* if USE_EPSILON_TEST is true then we do a check:
if |dv|<EPSILON then dv=0.0;
else no check is done (which is less robust)
*/
#define USE_EPSILON_TEST TRUE
#define EPSILON 0.000001
template <typename T> __host__ __device__ inline void sort(T *a, T *b) {
if (a > b) {
T c;
c = *a;
*a = *b;
b = *c;
}
return;
}
template <typename T> __host__ __device__ inline int sort(T *a, T *b) {
if (a > b) {
T c;
c = *a;
*a = *b;
b = *c;
return 1;
} else
return 0;
}
#define ISECT(VV0, VV1, VV2, D0, D1, D2, isect0, isect1) \
isect0 = VV0 + (VV1 - VV0) * D0 / (D0 - D1); \
isect1 = VV0 + (VV2 - VV0) * D0 / (D0 - D2);
#define COMPUTE_INTERVALS(VV0, VV1, VV2, D0, D1, D2, D0D1, D0D2, isect0, \
isect1) \
if (D0D1 > 0.0f) { \
/* here we know that D0D2<=0.0 */ \
/* that is D0, D1 are on the same side, D2 on the other or on the plane */ \
ISECT(VV2, VV0, VV1, D2, D0, D1, isect0, isect1); \
} else if (D0D2 > 0.0f) { \
/* here we know that d0d1<=0.0 */ \
ISECT(VV1, VV0, VV2, D1, D0, D2, isect0, isect1); \
} else if (D1 * D2 > 0.0f || D0 != 0.0f) { \
/* here we know that d0d1<=0.0 or that D0!=0.0 */ \
ISECT(VV0, VV1, VV2, D0, D1, D2, isect0, isect1); \
} else if (D1 != 0.0f) { \
ISECT(VV1, VV0, VV2, D1, D0, D2, isect0, isect1); \
} else if (D2 != 0.0f) { \
ISECT(VV2, VV0, VV1, D2, D0, D1, isect0, isect1); \
} else { \
/* triangles are coplanar */ \
return coplanar_tri_tri(N1, V0, V1, V2, U0, U1, U2); \
}
template <typename T>
inline void compute_intervals(
)
/* this edge to edge test is based on Franlin Antonio's gem:
"Faster Line Segment Intersection", in Graphics Gems III,
pp. 199-202 */
#define EDGE_EDGE_TEST(V0, U0, U1) \
Bx = U0[i0] - U1[i0]; \
By = U0[i1] - U1[i1]; \
Cx = V0[i0] - U0[i0]; \
Cy = V0[i1] - U0[i1]; \
f = Ay * Bx - Ax * By; \
d = By * Cx - Bx * Cy; \
if ((f > 0 && d >= 0 && d <= f) || (f < 0 && d <= 0 && d >= f)) { \
e = Ax * Cy - Ay * Cx; \
if (f > 0) { \
if (e >= 0 && e <= f) \
return 1; \
} else { \
if (e <= 0 && e >= f) \
return 1; \
} \
}
#define EDGE_AGAINST_TRI_EDGES(V0, V1, U0, U1, U2) \
{ \
float Ax, Ay, Bx, By, Cx, Cy, e, d, f; \
Ax = V1[i0] - V0[i0]; \
Ay = V1[i1] - V0[i1]; \
/* test edge U0,U1 against V0,V1 */ \
EDGE_EDGE_TEST(V0, U0, U1); \
/* test edge U1,U2 against V0,V1 */ \
EDGE_EDGE_TEST(V0, U1, U2); \
/* test edge U2,U1 against V0,V1 */ \
EDGE_EDGE_TEST(V0, U2, U0); \
}
#define POINT_IN_TRI(V0, U0, U1, U2) \
{ \
float a, b, c, d0, d1, d2; \
/* is T1 completly inside T2? */ \
/* check if V0 is inside tri(U0,U1,U2) */ \
a = U1[i1] - U0[i1]; \
b = -(U1[i0] - U0[i0]); \
c = -a * U0[i0] - b * U0[i1]; \
d0 = a * V0[i0] + b * V0[i1] + c; \
\
a = U2[i1] - U1[i1]; \
b = -(U2[i0] - U1[i0]); \
c = -a * U1[i0] - b * U1[i1]; \
d1 = a * V0[i0] + b * V0[i1] + c; \
\
a = U0[i1] - U2[i1]; \
b = -(U0[i0] - U2[i0]); \
c = -a * U2[i0] - b * U2[i1]; \
d2 = a * V0[i0] + b * V0[i1] + c; \
if (d0 * d1 > 0.0) { \
if (d0 * d2 > 0.0) \
return 1; \
} \
}
template <typename T>
__host__ __device__
bool point_in_tri(vec3<T> V0, vec3<T>U0, vec3<T>U1, vec3<T>U2)
{
T a, b, c, d0, d1, d2;
/* is T1 completly inside T2? */
/* check if V0 is inside tri(U0,U1,U2) */
a = U1[i1] - U0[i1];
b = -(U1[i0] - U0[i0]);
c = -a * U0[i0] - b * U0[i1];
d0 = a * V0[i0] + b * V0[i1] + c;
a = U2[i1] - U1[i1];
b = -(U2[i0] - U1[i0]);
c = -a * U1[i0] - b * U1[i1];
d1 = a * V0[i0] + b * V0[i1] + c;
a = U0[i1] - U2[i1];
b = -(U0[i0] - U2[i0]);
c = -a * U2[i0] - b * U2[i1];
d2 = a * V0[i0] + b * V0[i1] + c;
if (d0 * d1 > 0.0) {
if (d0 * d2 > 0.0)
return true;
}
return false
}
template <typename T>
__host__ __device__
bool coplanar_tri_tri(
vec3<T> N, vec3<T> V0, vec3<T> V1, vec3<T> V2,
vec3<T> U0, vec3<T> U1, vec3<T> U2) {
vec3<T> A;
short i0, i1;
/* first project onto an axis-aligned plane, that maximizes the area */
/* of the triangles, compute indices: i0,i1. */
A.x = fabs(N.x);
A.y = fabs(N.y);
A.z = fabs(N.z);
if (A.x > A.y) {
if (A.x > A.z) {
i0 = 1; /* A[0] is greatest */
i1 = 2;
} else {
i0 = 0; /* A[2] is greatest */
i1 = 1;
}
} else /* A[0]<=A[1] */
{
if (A.z > A.y) {
i0 = 0; /* A[2] is greatest */
i1 = 1;
} else {
i0 = 0; /* A[1] is greatest */
i1 = 2;
}
}
/* test all edges of triangle 1 against the edges of triangle 2 */
EDGE_AGAINST_TRI_EDGES(V0, V1, U0, U1, U2);
EDGE_AGAINST_TRI_EDGES(V1, V2, U0, U1, U2);
EDGE_AGAINST_TRI_EDGES(V2, V0, U0, U1, U2);
/* finally, test if tri1 is totally contained in tri2 or vice versa */
POINT_IN_TRI(V0, U0, U1, U2);
POINT_IN_TRI(U0, V0, V1, V2);
return false
}
int tri_tri_intersect(float V0[3], float V1[3], float V2[3], float U0[3],
float U1[3], float U2[3]) {
float E1[3], E2[3];
float N1[3], N2[3], d1, d2;
float du0, du1, du2, dv0, dv1, dv2;
float D[3];
float isect1[2], isect2[2];
float du0du1, du0du2, dv0dv1, dv0dv2;
short index;
float vp0, vp1, vp2;
float up0, up1, up2;
float b, c, max;
/* compute plane equation of triangle(V0,V1,V2) */
SUB(E1, V1, V0);
SUB(E2, V2, V0);
CROSS(N1, E1, E2);
d1 = -DOT(N1, V0);
/* plane equation 1: N1.X+d1=0 */
/* put U0,U1,U2 into plane equation 1 to compute signed distances to the
* plane*/
du0 = DOT(N1, U0) + d1;
du1 = DOT(N1, U1) + d1;
du2 = DOT(N1, U2) + d1;
/* coplanarity robustness check */
#if USE_EPSILON_TEST == TRUE
if (fabs(du0) < EPSILON)
du0 = 0.0;
if (fabs(du1) < EPSILON)
du1 = 0.0;
if (fabs(du2) < EPSILON)
du2 = 0.0;
#endif
du0du1 = du0 * du1;
du0du2 = du0 * du2;
if (du0du1 > 0.0f &&
du0du2 > 0.0f) /* same sign on all of them + not equal 0 ? */
return 0; /* no intersection occurs */
/* compute plane of triangle (U0,U1,U2) */
SUB(E1, U1, U0);
SUB(E2, U2, U0);
CROSS(N2, E1, E2);
d2 = -DOT(N2, U0);
/* plane equation 2: N2.X+d2=0 */
/* put V0,V1,V2 into plane equation 2 */
dv0 = DOT(N2, V0) + d2;
dv1 = DOT(N2, V1) + d2;
dv2 = DOT(N2, V2) + d2;
#if USE_EPSILON_TEST == TRUE
if (fabs(dv0) < EPSILON)
dv0 = 0.0;
if (fabs(dv1) < EPSILON)
dv1 = 0.0;
if (fabs(dv2) < EPSILON)
dv2 = 0.0;
#endif
dv0dv1 = dv0 * dv1;
dv0dv2 = dv0 * dv2;
if (dv0dv1 > 0.0f &&
dv0dv2 > 0.0f) /* same sign on all of them + not equal 0 ? */
return 0; /* no intersection occurs */
/* compute direction of intersection line */
CROSS(D, N1, N2);
/* compute and index to the largest component of D */
max = fabs(D[0]);
index = 0;
b = fabs(D[1]);
c = fabs(D[2]);
if (b > max)
max = b, index = 1;
if (c > max)
max = c, index = 2;
/* this is the simplified projection onto L*/
vp0 = V0[index];
vp1 = V1[index];
vp2 = V2[index];
up0 = U0[index];
up1 = U1[index];
up2 = U2[index];
/* compute interval for triangle 1 */
COMPUTE_INTERVALS(vp0, vp1, vp2, dv0, dv1, dv2, dv0dv1, dv0dv2, isect1[0],
isect1[1]);
/* compute interval for triangle 2 */
COMPUTE_INTERVALS(up0, up1, up2, du0, du1, du2, du0du1, du0du2, isect2[0],
isect2[1]);
sort(&isect1[0], &isect1[1]);
sort(&isect2[0], &isect2[1]);
if (isect1[1] < isect2[0] || isect2[1] < isect1[0])
return 0;
return 1;
}
#define NEWCOMPUTE_INTERVALS(VV0, VV1, VV2, D0, D1, D2, D0D1, D0D2, A, B, C, \
X0, X1) \
{ \
if (D0D1 > 0.0f) { \
/* here we know that D0D2<=0.0 */ \
/* that is D0, D1 are on the same side, D2 on the other or on the plane \
*/ \
A = VV2; \
B = (VV0 - VV2) * D2; \
C = (VV1 - VV2) * D2; \
X0 = D2 - D0; \
X1 = D2 - D1; \
} else if (D0D2 > 0.0f) { \
/* here we know that d0d1<=0.0 */ \
A = VV1; \
B = (VV0 - VV1) * D1; \
C = (VV2 - VV1) * D1; \
X0 = D1 - D0; \
X1 = D1 - D2; \
} else if (D1 * D2 > 0.0f || D0 != 0.0f) { \
/* here we know that d0d1<=0.0 or that D0!=0.0 */ \
A = VV0; \
B = (VV1 - VV0) * D0; \
C = (VV2 - VV0) * D0; \
X0 = D0 - D1; \
X1 = D0 - D2; \
} else if (D1 != 0.0f) { \
A = VV1; \
B = (VV0 - VV1) * D1; \
C = (VV2 - VV1) * D1; \
X0 = D1 - D0; \
X1 = D1 - D2; \
} else if (D2 != 0.0f) { \
A = VV2; \
B = (VV0 - VV2) * D2; \
C = (VV1 - VV2) * D2; \
X0 = D2 - D0; \
X1 = D2 - D1; \
} else { \
/* triangles are coplanar */ \
return coplanar_tri_tri(N1, V0, V1, V2, U0, U1, U2); \
} \
}
int NoDivTriTriIsect(float V0[3], float V1[3], float V2[3], float U0[3],
float U1[3], float U2[3]) {
float E1[3], E2[3];
float N1[3], N2[3], d1, d2;
float du0, du1, du2, dv0, dv1, dv2;
float D[3];
float isect1[2], isect2[2];
float du0du1, du0du2, dv0dv1, dv0dv2;
short index;
float vp0, vp1, vp2;
float up0, up1, up2;
float bb, cc, max;
float a, b, c, x0, x1;
float d, e, f, y0, y1;
float xx, yy, xxyy, tmp;
/* compute plane equation of triangle(V0,V1,V2) */
SUB(E1, V1, V0);
SUB(E2, V2, V0);
CROSS(N1, E1, E2);
d1 = -DOT(N1, V0);
/* plane equation 1: N1.X+d1=0 */
/* put U0,U1,U2 into plane equation 1 to compute signed distances to the
* plane*/
du0 = DOT(N1, U0) + d1;
du1 = DOT(N1, U1) + d1;
du2 = DOT(N1, U2) + d1;
/* coplanarity robustness check */
#if USE_EPSILON_TEST == TRUE
if (FABS(du0) < EPSILON)
du0 = 0.0;
if (FABS(du1) < EPSILON)
du1 = 0.0;
if (FABS(du2) < EPSILON)
du2 = 0.0;
#endif
du0du1 = du0 * du1;
du0du2 = du0 * du2;
if (du0du1 > 0.0f &&
du0du2 > 0.0f) /* same sign on all of them + not equal 0 ? */
return 0; /* no intersection occurs */
/* compute plane of triangle (U0,U1,U2) */
SUB(E1, U1, U0);
SUB(E2, U2, U0);
CROSS(N2, E1, E2);
d2 = -DOT(N2, U0);
/* plane equation 2: N2.X+d2=0 */
/* put V0,V1,V2 into plane equation 2 */
dv0 = DOT(N2, V0) + d2;
dv1 = DOT(N2, V1) + d2;
dv2 = DOT(N2, V2) + d2;
#if USE_EPSILON_TEST == TRUE
if (FABS(dv0) < EPSILON)
dv0 = 0.0;
if (FABS(dv1) < EPSILON)
dv1 = 0.0;
if (FABS(dv2) < EPSILON)
dv2 = 0.0;
#endif
dv0dv1 = dv0 * dv1;
dv0dv2 = dv0 * dv2;
if (dv0dv1 > 0.0f &&
dv0dv2 > 0.0f) /* same sign on all of them + not equal 0 ? */
return 0; /* no intersection occurs */
/* compute direction of intersection line */
CROSS(D, N1, N2);
/* compute and index to the largest component of D */
max = (float)FABS(D[0]);
index = 0;
bb = (float)FABS(D[1]);
cc = (float)FABS(D[2]);
if (bb > max)
max = bb, index = 1;
if (cc > max)
max = cc, index = 2;
/* this is the simplified projection onto L*/
vp0 = V0[index];
vp1 = V1[index];
vp2 = V2[index];
up0 = U0[index];
up1 = U1[index];
up2 = U2[index];
/* compute interval for triangle 1 */
NEWCOMPUTE_INTERVALS(vp0, vp1, vp2, dv0, dv1, dv2, dv0dv1, dv0dv2, a, b, c,
x0, x1);
/* compute interval for triangle 2 */
NEWCOMPUTE_INTERVALS(up0, up1, up2, du0, du1, du2, du0du1, du0du2, d, e, f,
y0, y1);
xx = x0 * x1;
yy = y0 * y1;
xxyy = xx * yy;
tmp = a * xxyy;
isect1[0] = tmp + b * x1 * yy;
isect1[1] = tmp + c * x0 * yy;
tmp = d * xxyy;
isect2[0] = tmp + e * xx * y1;
isect2[1] = tmp + f * xx * y0;
SORT(isect1[0], isect1[1]);
SORT(isect2[0], isect2[1]);
if (isect1[1] < isect2[0] || isect2[1] < isect1[0])
return 0;
return 1;
}
template <typename T>
inline void isect2(vec3<T> VTX0, vec3<T> VTX1, vec3<T> VTX2, T VV0, T VV1,
T VV2, T D0, T D1, T D2, T *isect0, T *isect1,
vec3<T> isectpoint0, vec3<T> isectpoint1) {
T tmp = D0 / (D0 - D1);
T diff[3];
*isect0 = VV0 + (VV1 - VV0) * tmp;
SUB(diff, VTX1, VTX0);
MULT(diff, diff, tmp);
ADD(isectpoint0, diff, VTX0);
tmp = D0 / (D0 - D2);
*isect1 = VV0 + (VV2 - VV0) * tmp;
SUB(diff, VTX2, VTX0);
MULT(diff, diff, tmp);
ADD(isectpoint1, VTX0, diff);
}
template <typename T>
__host__ __device__ inline bool
compute_intervals_isectline(vec3<T> VERT0, vec3<T> VERT1, vec3<T> VERT2, T VV0,
T VV1, T VV2, T D0, T D1, T D2, T D0D1, T D0D2,
T *isect0, T *isect1, vec3<T> isectpoint0,
vec3<T> isectpoint1) {
if (D0D1 > 0.0f) {
/* here we know that D0D2<=0.0 */
/* that is D0, D1 are on the same side, D2 on the other or on the plane */
isect2(VERT2, VERT0, VERT1, VV2, VV0, VV1, D2, D0, D1, isect0, isect1,
isectpoint0, isectpoint1);
} else if (D0D2 > 0.0f) {
/* here we know that d0d1<=0.0 */
isect2(VERT1, VERT0, VERT2, VV1, VV0, VV2, D1, D0, D2, isect0, isect1,
isectpoint0, isectpoint1);
} else if (D1 * D2 > 0.0f || D0 != 0.0f) {
/* here we know that d0d1<=0.0 or that D0!=0.0 */
isect2(VERT0, VERT1, VERT2, VV0, VV1, VV2, D0, D1, D2, isect0, isect1,
isectpoint0, isectpoint1);
} else if (D1 != 0.0f) {
isect2(VERT1, VERT0, VERT2, VV1, VV0, VV2, D1, D0, D2, isect0, isect1,
isectpoint0, isectpoint1);
} else if (D2 != 0.0f) {
isect2(VERT2, VERT0, VERT1, VV2, VV0, VV1, D2, D0, D1, isect0, isect1,
isectpoint0, isectpoint1);
} else {
/* triangles are coplanar */
return 1;
}
return 0;
}
template <typename T>
__host__ __device__ inline bool tri_tri_intersect_with_isectline(
vec3<T> V0, vec3<T> V1, vec3<T> V2, vec3<T> U0, vec3<T> U1, vec3<T> U2,
bool *coplanar, vec3<T> isect_point1[3], vec3<T> isect_point2[3]) {
vec3<T> E1, E2;
vec3<T> N1, N2, d1, d2;
vec3<T> D;
T du0, du1, du2, dv0, dv1, dv2;
vec2<T> isect1, isect2;
vec3<T> isectpointA1, isectpointA2;
vec3<T> isectpointB1, isectpointB2;
T du0du1, du0du2, dv0dv1, dv0dv2;
short index;
T vp0, vp1, vp2;
T up0, up1, up2;
T b, c, max;
T tmp, diff[3];
int smallest1, smallest2;
/* compute plane equation of triangle(V0,V1,V2) */
E1 = V1 - V0;
E2 = V2 - V0;
N1 = cross(E1, E2) d1 = -dot(N1, V0);
/* plane equation 1: N1.X+d1=0 */
/* put U0,U1,U2 into plane equation 1 to compute signed distances to the
* plane*/
du0 = dot(N1, U0) + d1;
du1 = dot(N1, U1) + d1;
du2 = dot(N1, U2) + d1;
/* coplanarity robustness check */
#if USE_EPSILON_TEST == TRUE
if (fabs(du0) < EPSILON)
du0 = 0.0;
if (fabs(du1) < EPSILON)
du1 = 0.0;
if (fabs(du2) < EPSILON)
du2 = 0.0;
#endif
du0du1 = du0 * du1;
du0du2 = du0 * du2;
if (du0du1 > 0.0f &&
du0du2 > 0.0f) /* same sign on all of them + not equal 0 ? */
return 0; /* no intersection occurs */
/* compute plane of triangle (U0,U1,U2) */
E1 = U1 - U0;
E2 = U2 - U0;
N2 = cross(E1, E2);
d2 = -dot(N2, U0);
/* plane equation 2: N2.X+d2=0 */
/* put V0,V1,V2 into plane equation 2 */
dv0 = dot(N2, V0) + d2;
dv1 = dot(N2, V1) + d2;
dv2 = dot(N2, V2) + d2;
#if USE_EPSILON_TEST == TRUE
if (fabs(dv0) < EPSILON)
dv0 = 0.0;
if (fabs(dv1) < EPSILON)
dv1 = 0.0;
if (fabs(dv2) < EPSILON)
dv2 = 0.0;
#endif
dv0dv1 = dv0 * dv1;
dv0dv2 = dv0 * dv2;
/* same sign on all of them + not equal 0 ? */
if (dv0dv1 > 0.0f && dv0dv2 > 0.0f)
return 0; /* no intersection occurs */
/* compute direction of intersection line */
D = cross(N1, N2);
/* compute and index to the largest component of D */
max = fabs(D.x);
index = 0;
b = fabs(D.y);
c = fabs(D.z);
if (b > max)
max = b, index = 1;
if (c > max)
max = c, index = 2;
/* this is the simplified projection onto L*/
vp0 = V0[index];
vp1 = V1[index];
vp2 = V2[index];
up0 = U0[index];
up1 = U1[index];
up2 = U2[index];
/* compute interval for triangle 1 */
*coplanar = compute_intervals_isectline(
V0, V1, V2, vp0, vp1, vp2, dv0, dv1, dv2, dv0dv1, dv0dv2, &isect1[0],
&isect1[1], isectpointA1, isectpointA2);
if (*coplanar)
return coplanar_tri_tri(N1, V0, V1, V2, U0, U1, U2);
/* compute interval for triangle 2 */
compute_intervals_isectline(U0, U1, U2, up0, up1, up2, du0, du1, du2, du0du1,
du0du2, &isect2[0], &isect2[1], isectpointB1,
isectpointB2);
smallest1 = sort(&isect1[0], &isect1[1]);
smallest2 = sort(&isect2[0], &isect2[1]);
if (isect1[1] < isect2[0] || isect2[1] < isect1[0])
return 0;
/* at this point, we know that the triangles intersect */
if (isect2[0] < isect1[0]) {
if (smallest1 == 0) {
SET(isect_point1, isectpointA1);
} else {
SET(isect_point1, isectpointA2);
}
if (isect2[1] < isect1[1]) {
if (smallest2 == 0) {
SET(isect_point2, isectpointB2);
} else {
SET(isect_point2, isectpointB1);
}
} else {
if (smallest1 == 0) {
SET(isect_point2, isectpointA2);
} else {
SET(isect_point2, isectpointA1);
}
}
} else {
if (smallest2 == 0) {
SET(isect_point1, isectpointB1);
} else {
SET(isect_point1, isectpointB2);
}
if (isect2[1] > isect1[1]) {
if (smallest1 == 0) {
SET(isect_point2, isectpointA2);
} else {
SET(isect_point2, isectpointA1);
}
} else {
if (smallest2 == 0) {
SET(isect_point2, isectpointB2);
} else {
SET(isect_point2, isectpointB1);
}
}
}
return 1;
}
data/SHAPY/mesh-mesh-intersection/include/priority_queue.hpp
+169
View File
@@ -0,0 +1,169 @@
/*
* 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.
*
* @author Vasileios Choutas
* Contact: vassilis.choutas@tuebingen.mpg.de
* Contact: ps-license@tuebingen.mpg.de
*
*/
#ifndef PRIORITY_QUEUE_H
#define PRIORITY_QUEUE_H
#include <float.h>
#include <stdio.h>
#include <utility>
#include <cuda.h>
#include "device_launch_parameters.h"
#include <cuda_runtime.h>
template<typename T>
__host__ __device__
void swap_array_els(T* array, int i, int j) {
T tmp = array[i];
array[i] = array[j];
array[j] = tmp;
}
template <typename T, typename Obj, int QueueSize = 128, bool recursive = false>
class PriorityQueue {
public:
__host__ __device__
PriorityQueue() : heap_size(0) {}
inline
__host__ __device__
int get_size() {
return heap_size;
}
inline
__host__ __device__
int parent(int i) {
return (i - 1) / 2;
}
inline
__host__ __device__
int left_child(int i) {
return 2 * i + 1;
}
inline
__host__ __device__
int right_child(int i) {
return 2 * i + 2;
}
__host__ __device__
std::pair<T, Obj> get_min() {
if (heap_size > 0) {
return std::pair<T, Obj>(priority_heap[0], obj_heap[0]);
}
else {
return std::pair<T, Obj>(
std::is_same<T, float>::value ? FLT_MAX : DBL_MAX, nullptr);
}
}
__host__ __device__
void min_heapify(int index) {
if (recursive) {
int left = left_child(index);
int right = right_child(index);
int smallest = index;
if (left < heap_size && priority_heap[left] < priority_heap[index])
smallest = left;
if (right < heap_size && priority_heap[right] < priority_heap[index])
smallest = right;
if (smallest != index) {
swap_array_els(priority_heap, index, smallest);
swap_array_els(obj_heap, index, smallest);
min_heapify(smallest);
}
} else {
int ii = index;
int smallest;
while (true) {
int left = left_child(ii);
int right = right_child(ii);
smallest = ii;
if (left < heap_size && priority_heap[left] < priority_heap[ii])
smallest = left;
if (right < heap_size && priority_heap[right] < priority_heap[ii])
smallest = right;
if (smallest != ii) {
swap_array_els(priority_heap, ii, smallest);
swap_array_els(obj_heap, ii, smallest);
ii = smallest;
}
else
break;
}
}
}
__host__ __device__
void insert_key(T key, Obj obj) {
if (heap_size == QueueSize) {
printf("The queue has exceed its maximum size\n");
return;
}
heap_size++;
int ii = heap_size - 1;
priority_heap[ii] = key;
obj_heap[ii] = obj;
// Fix the min heap property if it is violated
min_heapify(0);
// while (ii != 0 && priority_heap[parent(ii)] > priority_heap[ii]) {
// swap_array_els(priority_heap, ii, parent(ii));
// swap_array_els(obj_heap, ii, parent(ii));
// ii = parent(ii);
// }
}
// void print() {
// for (int i = 0; i < heap_size; i++) {
// std::cout << i << ": " << heap[i] << std::endl;
// }
// }
__host__ __device__
std::pair<T, Obj> extract() {
if (heap_size <= 0)
return std::pair<T, Obj>(
std::is_same<T, float>::value ? FLT_MAX : DBL_MAX, nullptr);
T root_prio = priority_heap[0];
Obj root_obj = obj_heap[0];
// Replace the root with the last element
priority_heap[0] = priority_heap[heap_size - 1];
obj_heap[0] = obj_heap[heap_size - 1];
// Decrease the size of the heap
heap_size--;
min_heapify(0);
return std::pair<T, Obj>(root_prio, root_obj);
}
private:
T priority_heap[QueueSize];
Obj obj_heap[QueueSize];
int heap_size;
};
#endif // #ifndef PRIORITY_QUEUE_H
data/SHAPY/mesh-mesh-intersection/include/triangle.hpp
+66
View File
@@ -0,0 +1,66 @@
/*
* 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.
*
* @author Vasileios Choutas
* Contact: vassilis.choutas@tuebingen.mpg.de
* Contact: ps-license@tuebingen.mpg.de
*
*/
#ifndef TRIANGLE_H
#define TRIANGLE_H
#include "defs.hpp"
#include "double_vec_ops.hpp"
#include "helper_math.h"
#include "math_utils.hpp"
#include <cuda.h>
#include <cuda_profiler_api.h>
#include <cuda_runtime.h>
template <typename T>
__align__(48)
struct Triangle {
public:
vec3<T> v0;
vec3<T> v1;
vec3<T> v2;
__host__ __device__ Triangle() {}
__host__ __device__ Triangle(vec3<T> vertex0, vec3<T> vertex1,
vec3<T> vertex2)
: v0(vertex0), v1(vertex1), v2(vertex2){};
__host__ __device__ Triangle(const vec3<T> &vertex0, const vec3<T> &vertex1,
const vec3<T> &vertex2)
: v0(vertex0), v1(vertex1), v2(vertex2){};
__host__ __device__ AABB<T> bbox() const {
return AABB<T>(min(v0.x, min(v1.x, v2.x)), min(v0.y, min(v1.y, v2.y)),
min(v0.z, min(v1.z, v2.z)), max(v0.x, max(v1.x, v2.x)),
max(v0.y, max(v1.y, v2.y)), max(v0.z, max(v1.z, v2.z)));
}
};
template <typename T> using TrianglePtr = Triangle<T> *;
template <typename T>
std::ostream &operator<<(std::ostream &os, const Triangle<T> &x) {
os << x.v0 << std::endl;
os << x.v1 << std::endl;
os << x.v2 << std::endl;
return os;
}
#endif // TRIANGLE_H
data/SHAPY/mesh-mesh-intersection/mesh_mesh_intersection/__init__.py
+18
View File
@@ -0,0 +1,18 @@
# -*- 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©2020 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 .mesh_mesh_intersection import MeshMeshIntersection
data/SHAPY/mesh-mesh-intersection/mesh_mesh_intersection/loss.py
+317
View File
@@ -0,0 +1,317 @@
# -*- 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 and the Max Planck Institute for Biological
# Cybernetics. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
from __future__ import absolute_import
from __future__ import print_function
from __future__ import division
import sys
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
def calc_circumcircle(triangles, edge_cross_prod, idx=None):
''' Calculate the circumscribed circle for the given triangles
Args:
- triangles (torch.tensor BxTx3x3): The tensor that contains the
coordinates of the triangle vertices
- edge_cross_prod (torch.tensor BxCx3): Contains the unnormalized
perpendicular vector to the surface of the triangle.
Returns:
- circumradius (torch.tensor BxTx1): The radius of the
circumscribed circle
- circumcenter (torch.tensor BxTx3): The center of the
circumscribed circel
'''
alpha = triangles[:, :, 0] - triangles[:, :, 2]
beta = triangles[:, :, 1] - triangles[:, :, 2]
# Calculate the radius of the circumscribed circle
# Should be BxF
circumradius = (torch.norm(alpha - beta, dim=2, keepdim=True) /
(2 * torch.norm(edge_cross_prod, dim=2, keepdim=True)) *
torch.norm(alpha, dim=2, keepdim=True) *
torch.norm(beta, dim=2, keepdim=True))
# Calculate the coordinates of the circumcenter of each triangle
# Should BxFx3
circumcenter = torch.cross(
torch.sum(alpha ** 2, dim=2, keepdim=True) * beta -
torch.sum(beta ** 2, dim=2, keepdim=True) * alpha,
torch.cross(alpha, beta, dim=-1), dim=2)
circumcenter /= (2 * torch.sum(edge_cross_prod ** 2, dim=2, keepdim=True))
return circumradius, circumcenter + triangles[:, :, 2]
def repulsion_intensity(x, sigma=0.5, penalize_outside=True, linear_max=1000):
''' Penalizer function '''
quad_penalty = (-(1.0 - 2.0 * sigma) / (4.0 * sigma ** 2) *
x ** 2 - 1 / (2.0 * sigma) * x +
0.25 * (3 - 2 * sigma))
linear_region_mask = (x.le(-sigma) * x.gt(-linear_max)).to(dtype=x.dtype)
if penalize_outside:
quad_region_mask = (x.gt(-sigma) * x.lt(sigma)).to(dtype=x.dtype)
else:
quad_region_mask = (x.gt(-sigma) * x.lt(0)).to(dtype=x.dtype)
return (linear_region_mask * (-x + 1 - sigma) +
quad_region_mask * quad_penalty)
def dist_to_cone_axis(points_rel, dot_prod, cone_axis, cone_radius,
sigma=0.5, epsilon=1e-6, vectorized=True):
''' Computes the distance of each point to the axis
This function projects the points on the plane of the base of the cone
and computes the distance to the axis. This is subsequently normalized
by the radius of the cone at the height level of the point, so that
points with distance < 1 are in the code, distance == 1 means that the
point is on the surface and distance > 1 means that the point is
outside the cone.
Args:
- points_rel (torch.tensor BxCxNx3): The coordinates of the points
relative to the center of the cone
- dot_prod (torch.tensor BxCxN): The dot product of the points (in
relative coordinates with respect to the cone center) with the
axis of the cone
- cone_axis (torch.tensor BxCx3): The axis of the cone
- cone_radius (torch.tensor BxCx1): The radius of the cone
Keyword args:
- sigma (float = 0.5): The height of the cone
- epsilon (float = 1e-6): Numerical stability constant for the
float division
- vectorized (bool = True): Whether to use an iterative or a
vectorized version of the function
'''
if vectorized:
batch_size, num_collisions = cone_radius.shape[:2]
numerator = torch.norm(points_rel - dot_prod.unsqueeze(dim=-1) *
cone_axis.unsqueeze(dim=-2),
p=2, dim=-1)
denominator = -cone_radius / sigma * dot_prod + cone_radius
else:
batch_size, num_collisions = cone_radius.shape[:2]
numerator = torch.norm(points_rel - dot_prod.unsqueeze(-1) * cone_axis,
p=2, dim=-1)
denominator = -cone_radius.view(batch_size, num_collisions) / sigma * \
dot_prod + cone_radius.view(batch_size, num_collisions)
return numerator / (denominator + epsilon)
def conical_distance_field(triangle_points, cone_center, cone_radius,
cone_axis, sigma=0.5, vectorized=True,
penalize_outside=True, linear_max=1000):
''' Distance field calculation for a cone
Args:
- triangle_points (torch.tensor (BxCxNx3): Contains
the points whose distance from the cone we want to calculate.
- cone_center (torch.tensor (BxCx3)): The coordinates of the center
of the cone
- cone_radius (torch.tensor (BxC)): The radius of the base of the
cone
- cone_axis (torch.tensor(BxCx3)): The unit vector that represents
the axis of the cone
Keyword Arguments
- sigma (float = 0.5): The float value of the height of the cone
- vectorized (bool = True): Whether to use an iterative or a
vectorized version of the function
Returns:
- (torch.tensor BxCxN): The distance field values at the N points
for the cone
'''
if vectorized:
# Calculate the coordinates of the points relative to the center of
# the cone
points_rel = triangle_points - cone_center.unsqueeze(dim=-2)
# Calculate the dot product between the relative point coordinates and
# the axis (normal) of the cone. Essentially, it is the length of the
# projection of the relative vector on the axis of the cone
dot_prod = torch.sum(points_rel * cone_axis.unsqueeze(dim=-2), dim=-1)
# Calculate the distance of the projections of the points on the cone
# base plane to the center of cone, normalized by the height
axis_dist = dist_to_cone_axis(points_rel, dot_prod,
cone_axis, cone_radius,
sigma=sigma, vectorized=True)
circumcenter_dist = repulsion_intensity(
dot_prod, sigma=sigma, penalize_outside=penalize_outside,
linear_max=linear_max)
# Ignore the points with axis_dist > 1, since they are out of the cone
mask = axis_dist.lt(1).to(dtype=triangle_points.dtype)
distance_field = mask * ((1 - axis_dist) * circumcenter_dist).pow(2)
else:
batch_size, num_collisions, num_points = triangle_points.shape[:3]
distance_field = torch.zeros([batch_size, num_collisions, 3],
dtype=triangle_points.dtype,
device=triangle_points.device)
for idx in range(num_points):
# The relative coordinates of each point to the center of the cone
# BxCx3
points_rel = triangle_points[:, :, idx, :] - cone_center
# Calculate the dot product between the relative point coordinates
# and the axis (normal) of the cone. Essentially, it is the length
# of the projection of the relative vector on the axis of the cone
dot_prod = torch.sum(points_rel * cone_axis, dim=-1)
axis_dist = dist_to_cone_axis(points_rel, dot_prod,
cone_axis, cone_radius,
sigma=sigma,
vectorized=False)
circumcenter_dist = repulsion_intensity(
dot_prod, sigma=sigma, penalize_outside=penalize_outside)
mask = (axis_dist < 1).to(dtype=triangle_points.dtype)
distance_field[:, :, idx] = (1 - axis_dist) * mask * \
circumcenter_dist
return torch.pow(distance_field, 2)
class DistanceFieldPenetrationLoss(nn.Module):
def __init__(self, sigma=0.5, point2plane=False, vectorized=True,
penalize_outside=True, linear_max=1000):
super(DistanceFieldPenetrationLoss, self).__init__()
self.sigma = sigma
self.point2plane = point2plane
self.vectorized = vectorized
self.penalize_outside = penalize_outside
self.linear_max = linear_max
def forward(self, triangles, collision_idxs):
'''
Args:
- triangles: A torch tensor of size BxFx3x3 that contains the
coordinates of the triangle vertices
- collision_idxs: A torch tensor of size Bx(-1)x2 that contains the
indices of the colliding pairs
Returns:
A tensor with size B that contains the self penetration loss for
each mesh in the batch
'''
coll_idxs = collision_idxs[:, :, 0].ge(0).nonzero()
if len(coll_idxs) < 1:
return torch.zeros([triangles.shape[0]],
dtype=triangles.dtype,
device=triangles.device,
requires_grad=triangles.requires_grad)
receiver_faces = collision_idxs[coll_idxs[:, 0], coll_idxs[:, 1], 0]
intruder_faces = collision_idxs[coll_idxs[:, 0], coll_idxs[:, 1], 1]
batch_idxs = coll_idxs[:, 0]
num_collisions = coll_idxs.shape[0]
batch_size = triangles.shape[0]
if len(intruder_faces) < 1:
return torch.tensor(0.0, dtype=triangles.dtype,
device=triangles.device,
requires_grad=triangles.requires_grad)
# Calculate the edges of the triangles
# Size: BxFx3
edge0 = triangles[:, :, 1] - triangles[:, :, 0]
edge1 = triangles[:, :, 2] - triangles[:, :, 0]
# Compute the cross product of the edges to find the normal vector of
# the triangle
aCrossb = torch.cross(edge0, edge1, dim=2)
circumradius, circumcenter = calc_circumcircle(triangles, aCrossb)
# Normalize the result to get a unit vector
normals = aCrossb / torch.norm(aCrossb, 2, dim=2, keepdim=True)
recv_triangles = triangles[batch_idxs, receiver_faces]
intr_triangles = triangles[batch_idxs, intruder_faces]
recv_normals = normals[batch_idxs, receiver_faces]
recv_circumradius = circumradius[batch_idxs, receiver_faces]
recv_circumcenter = circumcenter[batch_idxs, receiver_faces]
intr_normals = normals[batch_idxs, intruder_faces]
intr_circumradius = circumradius[batch_idxs, intruder_faces]
intr_circumcenter = circumcenter[batch_idxs, intruder_faces]
# Compute the distance field for the intruding triangles
# B x NUM_COLLISIONS x 3
# For each batch element, for each collision pair, 3 distance values
# for the vertices of the intruding triangle
phi_receivers = conical_distance_field(
intr_triangles,
recv_circumcenter, recv_circumradius,
recv_normals,
sigma=self.sigma,
vectorized=self.vectorized,
penalize_outside=self.penalize_outside,
linear_max=self.linear_max)
# Compute the distance field for the intruding triangles
# B x NUM_COLLISIONS x 3
# For each batch element, for each collision pair, 3 distance values
# for the vertices of the intruding triangle
# Same as above, but now the receiver is the "intruder".
phi_intruders = conical_distance_field(
recv_triangles,
intr_circumcenter,
intr_circumradius,
intr_normals,
sigma=self.sigma,
vectorized=self.vectorized,
penalize_outside=self.penalize_outside,
linear_max=self.linear_max)
receiver_loss = torch.tensor(0, device=triangles.device,
dtype=torch.float32)
intruder_loss = torch.tensor(0, device=triangles.device,
dtype=torch.float32)
if self.point2plane:
receiver_loss = (-phi_receivers).pow(2).sum(dim=-1)
intruder_loss = (-phi_intruders).pow(2).sum(dim=-1)
else:
receiver_loss = torch.norm(-phi_receivers.unsqueeze(dim=-1) *
intr_normals.unsqueeze(dim=-2), p=2,
dim=-1).pow(2).sum(dim=-1)
intruder_loss = torch.norm(-phi_intruders.unsqueeze(dim=-1) *
recv_normals.unsqueeze(dim=-2), p=2,
dim=-1).pow(2).sum(dim=-1)
batch_ind = torch.arange(0, batch_size, dtype=batch_idxs.dtype,
device=triangles.device).unsqueeze(dim=1)
batch_mask = batch_ind.repeat([1, num_collisions]).eq(batch_idxs)\
.to(receiver_loss.dtype)
loss = torch.matmul(batch_mask, receiver_loss) + \
torch.matmul(batch_mask, intruder_loss)
return loss
data/SHAPY/mesh-mesh-intersection/mesh_mesh_intersection/mesh_mesh_intersection.py
+62
View File
@@ -0,0 +1,62 @@
# -*- 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 sys
import torch
import torch.nn as nn
import torch.autograd as autograd
# from loguru import logger
import mesh_mesh_intersection
import mesh_mesh_intersect_cuda
class MeshMeshIntersectionFunction(autograd.Function):
@staticmethod
@torch.no_grad()
def forward(ctx, query_triangles, target_triangles, print_timings=False,
max_collisions=32,
*args, **kwargs):
outputs = mesh_mesh_intersect_cuda.mesh_to_mesh_forward(
query_triangles, target_triangles, print_timings=print_timings,
max_collisions=max_collisions)
# ctx.save_for_backward(query_triangles, outputs)
collision_faces, collision_bcs = outputs
return collision_faces, collision_bcs
@staticmethod
def backward(ctx, grad_output, *args, **kwargs):
raise NotImplementedError
class MeshMeshIntersection(nn.Module):
def __init__(self, max_collisions=32):
super(MeshMeshIntersection, self).__init__()
self.max_collisions = max_collisions
# MeshMeshIntersectionFunction.max_collisions = self.max_collisions
def forward(self, query_triangles, target_triangles,
print_timings=False):
return MeshMeshIntersectionFunction.apply(
query_triangles, target_triangles, print_timings,
self.max_collisions)
data/SHAPY/mesh-mesh-intersection/optional-requirements.txt
+4
View File
@@ -0,0 +1,4 @@
pyrender>=0.1.23
shapely
trimesh>=2.37.6
smplx
data/SHAPY/mesh-mesh-intersection/requirements.txt
+2
View File
@@ -0,0 +1,2 @@
numpy>=1.16.2
torch>=1.0
data/SHAPY/mesh-mesh-intersection/setup.py
+100
View File
@@ -0,0 +1,100 @@
# -*- 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 and the Max Planck Institute for Biological
# Cybernetics. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.deimport io
import io
import os
import os.path as osp
from setuptools import find_packages, setup
import torch
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
# Package meta-data.
NAME = 'mesh_mesh_intersection'
DESCRIPTION = 'PyTorch module for Mesh-Mesh intersection detection'
URL = ''
EMAIL = 'vassilis.choutas@tuebingen.mpg.de'
AUTHOR = 'Vassilis Choutas'
REQUIRES_PYTHON = '>=3.6.0'
VERSION = '0.2.0'
here = os.path.abspath(os.path.dirname(__file__))
try:
FileNotFoundError
except NameError:
FileNotFoundError = IOError
# Import the README and use it as the long-description.
# Note: this will only work if 'README.md' is present in your MANIFEST.in file!
try:
with io.open(os.path.join(here, 'README.md'), encoding='utf-8') as f:
long_description = '\n' + f.read()
except FileNotFoundError:
long_description = DESCRIPTION
# Load the package's __version__.py module as a dictionary.
about = {}
if not VERSION:
with open(os.path.join(here, NAME, '__version__.py')) as f:
exec(f.read(), about)
else:
about['__version__'] = VERSION
mesh_mesh_intersect_src_files = [
'src/mesh_mesh_intersect.cpp', 'src/mesh_mesh_intersect_cuda_op.cu']
mesh_mesh_intersect_include_dirs = torch.utils.cpp_extension.include_paths() + [
osp.abspath('include'),
osp.abspath(osp.expandvars('$CUDA_SAMPLES_INC'))]
mesh_mesh_intersect_extra_compile_args = {
'nvcc': ['-DPRINT_TIMINGS=0', '-DDEBUG_PRINT=0',
'-DERROR_CHECKING=1',
'-DCOLLISION_ORDERING=1'],
'cxx': []}
mesh_mesh_intersect_extension = CUDAExtension(
'mesh_mesh_intersect_cuda', mesh_mesh_intersect_src_files,
include_dirs=mesh_mesh_intersect_include_dirs,
extra_compile_args=mesh_mesh_intersect_extra_compile_args)
render_reqs = ['pyrender>=0.1.23', 'trimesh>=2.37.6', 'shapely']
setup(name=NAME,
version=about['__version__'],
description=DESCRIPTION,
long_description=long_description,
long_description_content_type='text/markdown',
author=AUTHOR,
author_email=EMAIL,
python_requires=REQUIRES_PYTHON,
url=URL,
packages=find_packages(),
ext_modules=[mesh_mesh_intersect_extension],
classifiers=[
"License :: OSI Approved :: GNU General Public License v3 (GPLv3)",
"Environment :: Console",
"Programming Language :: Python",
"Programming Language :: Python :: 3.6",
"Programming Language :: Python :: 3.7"],
install_requires=[
'torch>=1.6.0',
],
extras_require={
'render': render_reqs,
'all': render_reqs
},
cmdclass={'build_ext': BuildExtension})
data/SHAPY/mesh-mesh-intersection/src/mesh_mesh_intersect.cpp
+64
View File
@@ -0,0 +1,64 @@
/*
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
*/
#include <torch/extension.h>
#include <vector>
#define CHECK_CUDA(x) \
TORCH_CHECK(x.device().type() == torch::kCUDA, #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) \
TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) \
CHECK_CUDA(x); \
CHECK_CONTIGUOUS(x)
void mesh_mesh_intersection_forward(const torch::Tensor &query_triangles,
const torch::Tensor &target_triangles,
torch::Tensor &collision_faces,
torch::Tensor &collision_bcs,
int max_collisions = 16,
bool print_timings = false);
std::vector<torch::Tensor>
mesh_to_mesh_intersection(torch::Tensor query_triangles,
torch::Tensor target_triangles,
int max_collisions = 16, bool print_timings = false) {
CHECK_INPUT(query_triangles);
CHECK_INPUT(target_triangles);
torch::Tensor collision_faces =
-1 * torch::ones({query_triangles.size(0),
query_triangles.size(1) * max_collisions},
torch::device(query_triangles.device())
.dtype(torch::ScalarType::Long));
torch::Tensor collision_bcs = torch::zeros(
{query_triangles.size(0), query_triangles.size(1) * max_collisions, 2, 3},
torch::device(query_triangles.device()).dtype(query_triangles.dtype()));
mesh_mesh_intersection_forward(query_triangles, target_triangles,
collision_faces, collision_bcs,
max_collisions);
return {torch::autograd::make_variable(collision_faces, false),
torch::autograd::make_variable(collision_bcs, false)};
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("mesh_to_mesh_forward", &mesh_to_mesh_intersection,
"BVH mesh-to-mesh intersection forward (CUDA)",
py::arg("query_triangles"), py::arg("target_triangles"),
py::arg("max_collisions") = 16, py::arg("print_timings") = false);
}
data/SHAPY/mesh-mesh-intersection/src/mesh_mesh_intersect_cuda_op.cu
+1079
View File
File diff suppressed because it is too large Load Diff
data/SHAPY/test_submission_format.py
+62
View File
@@ -0,0 +1,62 @@
""" Ref: https://github.com/muelea/shapy/blob/master/regressor/hbw_evaluation/test_submission_format.py """
import argparse
import numpy as np
def test_submission_file_format(
npz_file: str,
model_type: str = 'smplx'
):
submission = np.load(npz_file)
# check if keys are named correctly
keys = [x for x in submission.keys()]
assert 'image_name' in keys and 'v_shaped' in keys, \
f"Keys are not correct. Got {keys}, but expected ['image_name', 'v_shaped']"
image_names = submission['image_name']
v_shapeds = submission['v_shaped']
# check if shape and type are correct
assert type(image_names) == np.ndarray, \
f"Type of key image_name is not correct. {type(image_names)} given, but np.ndarray expected."
assert image_names.shape == (1631,), \
f"Shape of key image_name is not correct. {image_names.shape} given, but (1631,) expected."
assert type(v_shapeds) == np.ndarray, \
f"Type of key v_shaped is not correct. {type(image_names)} given, but np.ndarray expected."
if model_type == 'smplx':
assert v_shapeds.shape == (1631, 10475, 3), \
f"Shape of key v_shaped is not correct. {v_shapeds.shape} given, but (1631, 10475, 3) expected."
else:
assert v_shapeds.shape == (1631, 6890, 3), \
f"Shape of key v_shaped is not correct. {v_shapeds.shape} given, but (1631, 6890, 3) expected."
# check if each image has a prediction
hbw_images_gt = np.load('../data/SHAPY/hbw_testset_image_names.npy')
check_prediction_available = np.isin(hbw_images_gt, image_names)
assert np.all(check_prediction_available), \
f"Images without predition exist! Missing predictions: \
\n {hbw_images_gt[~check_prediction_available]}"
print(f'Your submission file passed the test.')
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--input-npz-file',
dest='input_npz_file', type=str, required=True,
help='npz containing labels and body shape parameters.')
parser.add_argument('--model-type', choices=['smpl', 'smplx'], type=str,
default='smplx',
help='The model type used for body shape prediction. ')
args = parser.parse_args()
test_submission_file_format(
npz_file=args.input_npz_file,
model_type=args.model_type
)
data/SPEC/SPEC.py
+51
View File
@@ -0,0 +1,51 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class SPEC(HumanDataset):
def __init__(self, transform, data_split):
super(SPEC, self).__init__(transform, data_split)
pre_prc_file_train = 'spec_train_smpl.npz'
pre_prc_file_test = 'spec_test_smpl.npz'
if self.data_split == 'train':
filename = getattr(cfg, 'filename', pre_prc_file_train)
else:
filename = getattr(cfg, 'filename', pre_prc_file_test)
self.img_dir = osp.join(cfg.data_dir, 'SPEC')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = (1080, 1920) # (h, w)
self.cam_param = {}
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/SSP3D/SSP3D.py
+52
View File
@@ -0,0 +1,52 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class SSP3D(HumanDataset):
def __init__(self, transform, data_split):
super(SSP3D, self).__init__(transform, data_split)
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'ssp3d_230525_311.npz')
self.img_shape = (512, 512) # (h, w)
self.cam_param = {}
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
if self.data_split == 'train':
filename = getattr(cfg, 'filename', 'ssp3d_230525_311.npz')
else:
raise ValueError('SSP3D test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'SSP3D')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data or cache
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/SynBody/SynBody.py
+47
View File
@@ -0,0 +1,47 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class SynBody(HumanDataset):
def __init__(self, transform, data_split):
super(SynBody, self).__init__(transform, data_split)
filename = 'synbody_train_230521_04000_fix_betas.npz'
self.img_dir = osp.join(cfg.data_dir, 'SynBody')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = (720, 1280) # (h, w)
self.cam_param = {
'focal': (540, 540), # (fx, fy)
'princpt': (640, 360) # (cx, cy)
}
# check image shape
img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/Talkshow/Talkshow.py
+59
View File
@@ -0,0 +1,59 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
# 'talkshow_smplx_chemistry_path.npz' zipfile.BadZipFile: File is not a zip file
# ['talkshow_smplx_conan.npz',
# 'talkshow_smplx_oliver_path.npz', 'talkshow_smplx_seth.npz']:
class Talkshow(HumanDataset):
def __init__(self, transform, data_split):
super(Talkshow, self).__init__(transform, data_split)
sample_rate = getattr(cfg, 'Talkshow_train_sample_interval', 1)
self.use_cache = getattr(cfg, 'use_cache', False)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'talkshow_smplx.npz')
self.img_shape = None # (h, w)
self.cam_param = {}
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = []
for pre_prc_file in ['talkshow_smplx_chemistry.npz', 'talkshow_smplx_conan.npz',
'talkshow_smplx_oliver.npz', 'talkshow_smplx_seth.npz']:
if self.data_split == 'train':
filename = getattr(cfg, 'filename', pre_prc_file)
else:
raise ValueError('Talkshow test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'Talkshow')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
# check image shape
# img_path = osp.join(self.img_dir, np.load(self.annot_path)['image_path'][0])
# img_shape = cv2.imread(img_path).shape[:2]
# assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(self.img_shape, img_shape)
# load data
datalist_slice = self.load_data(sample_rate)
self.datalist.extend(datalist_slice)
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)
data/UBody/UBody.py
+1147
View File
File diff suppressed because it is too large Load Diff
data/UP3D/UP3D.py
+44
View File
@@ -0,0 +1,44 @@
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config import cfg
from utils.human_models import smpl_x
from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \
get_fitting_error_3D
from utils.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class UP3D(HumanDataset):
def __init__(self, transform, data_split):
super(UP3D, self).__init__(transform, data_split)
if self.data_split == 'train':
filename = getattr(cfg, 'filename', 'up3d_trainval.npz')
else:
raise ValueError('UP3D test set is not support')
self.img_dir = osp.join(cfg.data_dir, 'UP3D')
self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename)
self.annot_path_cache = osp.join(cfg.data_dir, 'cache', filename)
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = None # (h, w)
self.cam_param = {}
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}')
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(f'[{self.__class__.__name__}] Cache not found, generating cache...')
self.datalist = self.load_data(
train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)

Some files were not shown because too many files have changed in this diff Show More