# -*- coding: utf-8 -*-
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
import os
import lib.smplx as smplx
from lib.pymaf.utils.geometry import rotation_matrix_to_angle_axis, batch_rodrigues
from lib.pymaf.utils.imutils import process_image
from lib.pymaf.core import path_config
from lib.pymaf.models import pymaf_net
from lib.common.config import cfg
from lib.common.render import Render
from lib.dataset.body_model import TetraSMPLModel
from lib.dataset.mesh_util import get_visibility, SMPLX
import os.path as osp
import torch
import numpy as np
import random
from termcolor import colored
from PIL import ImageFile
from huggingface_hub import cached_download
ImageFile.LOAD_TRUNCATED_IMAGES = True
class TestDataset():
def __init__(self, cfg, device):
random.seed(1993)
self.image_path = cfg['image_path']
self.seg_dir = cfg['seg_dir']
self.has_det = cfg['has_det']
self.hps_type = cfg['hps_type']
self.smpl_type = 'smpl' if cfg['hps_type'] != 'pixie' else 'smplx'
self.smpl_gender = 'neutral'
self.device = device
self.subject_list = [self.image_path]
# smpl related
self.smpl_data = SMPLX()
self.get_smpl_model = lambda smpl_type, smpl_gender: smplx.create(
model_path=self.smpl_data.model_dir,
gender=smpl_gender,
model_type=smpl_type,
ext='npz')
# Load SMPL model
self.smpl_model = self.get_smpl_model(
self.smpl_type, self.smpl_gender).to(self.device)
self.faces = self.smpl_model.faces
self.hps = pymaf_net(path_config.SMPL_MEAN_PARAMS,
pretrained=True).to(self.device)
self.hps.load_state_dict(torch.load(
path_config.CHECKPOINT_FILE)['model'],
strict=True)
self.hps.eval()
print(colored(f"Using {self.hps_type} as HPS Estimator\n", "green"))
self.render = Render(size=512, device=device)
def __len__(self):
return len(self.subject_list)
def compute_vis_cmap(self, smpl_verts, smpl_faces):
(xy, z) = torch.as_tensor(smpl_verts).split([2, 1], dim=1)
smpl_vis = get_visibility(xy, -z, torch.as_tensor(smpl_faces).long())
if self.smpl_type == 'smpl':
smplx_ind = self.smpl_data.smpl2smplx(np.arange(smpl_vis.shape[0]))
else:
smplx_ind = np.arange(smpl_vis.shape[0])
smpl_cmap = self.smpl_data.get_smpl_mat(smplx_ind)
return {
'smpl_vis': smpl_vis.unsqueeze(0).to(self.device),
'smpl_cmap': smpl_cmap.unsqueeze(0).to(self.device),
'smpl_verts': smpl_verts.unsqueeze(0)
}
def compute_voxel_verts(self, body_pose, global_orient, betas, trans,
scale):
smpl_path = cached_download(osp.join(self.smpl_data.model_dir, "smpl/SMPL_NEUTRAL.pkl"), use_auth_token=os.environ['ICON'])
tetra_path = cached_download(osp.join(self.smpl_data.tedra_dir,
'tetra_neutral_adult_smpl.npz'), use_auth_token=os.environ['ICON'])
smpl_model = TetraSMPLModel(smpl_path, tetra_path, 'adult')
pose = torch.cat([global_orient[0], body_pose[0]], dim=0)
smpl_model.set_params(rotation_matrix_to_angle_axis(pose),
beta=betas[0])
verts = np.concatenate(
[smpl_model.verts, smpl_model.verts_added],
axis=0) * scale.item() + trans.detach().cpu().numpy()
faces = np.loadtxt(cached_download(osp.join(self.smpl_data.tedra_dir,
'tetrahedrons_neutral_adult.txt'), use_auth_token=os.environ['ICON']),
dtype=np.int32) - 1
pad_v_num = int(8000 - verts.shape[0])
pad_f_num = int(25100 - faces.shape[0])
verts = np.pad(verts, ((0, pad_v_num), (0, 0)),
mode='constant',
constant_values=0.0).astype(np.float32) * 0.5
faces = np.pad(faces, ((0, pad_f_num), (0, 0)),
mode='constant',
constant_values=0.0).astype(np.int32)
verts[:, 2] *= -1.0
voxel_dict = {
'voxel_verts':
torch.from_numpy(verts).to(self.device).unsqueeze(0).float(),
'voxel_faces':
torch.from_numpy(faces).to(self.device).unsqueeze(0).long(),
'pad_v_num':
torch.tensor(pad_v_num).to(self.device).unsqueeze(0).long(),
'pad_f_num':
torch.tensor(pad_f_num).to(self.device).unsqueeze(0).long()
}
return voxel_dict
def __getitem__(self, index):
img_path = self.subject_list[index]
img_name = img_path.split("/")[-1].rsplit(".", 1)[0]
if self.seg_dir is None:
img_icon, img_hps, img_ori, img_mask, uncrop_param = process_image(
img_path, self.hps_type, 512, self.device)
data_dict = {
'name': img_name,
'image': img_icon.to(self.device).unsqueeze(0),
'ori_image': img_ori,
'mask': img_mask,
'uncrop_param': uncrop_param
}
else:
img_icon, img_hps, img_ori, img_mask, uncrop_param, segmentations = process_image(
img_path, self.hps_type, 512, self.device,
seg_path=os.path.join(self.seg_dir, f'{img_name}.json'))
data_dict = {
'name': img_name,
'image': img_icon.to(self.device).unsqueeze(0),
'ori_image': img_ori,
'mask': img_mask,
'uncrop_param': uncrop_param,
'segmentations': segmentations
}
with torch.no_grad():
# import ipdb; ipdb.set_trace()
preds_dict = self.hps.forward(img_hps)
data_dict['smpl_faces'] = torch.Tensor(
self.faces.astype(np.int16)).long().unsqueeze(0).to(
self.device)
if self.hps_type == 'pymaf':
output = preds_dict['smpl_out'][-1]
scale, tranX, tranY = output['theta'][0, :3]
data_dict['betas'] = output['pred_shape']
data_dict['body_pose'] = output['rotmat'][:, 1:]
data_dict['global_orient'] = output['rotmat'][:, 0:1]
data_dict['smpl_verts'] = output['verts']
elif self.hps_type == 'pare':
data_dict['body_pose'] = preds_dict['pred_pose'][:, 1:]
data_dict['global_orient'] = preds_dict['pred_pose'][:, 0:1]
data_dict['betas'] = preds_dict['pred_shape']
data_dict['smpl_verts'] = preds_dict['smpl_vertices']
scale, tranX, tranY = preds_dict['pred_cam'][0, :3]
elif self.hps_type == 'pixie':
data_dict.update(preds_dict)
data_dict['body_pose'] = preds_dict['body_pose']
data_dict['global_orient'] = preds_dict['global_pose']
data_dict['betas'] = preds_dict['shape']
data_dict['smpl_verts'] = preds_dict['vertices']
scale, tranX, tranY = preds_dict['cam'][0, :3]
elif self.hps_type == 'hybrik':
data_dict['body_pose'] = preds_dict['pred_theta_mats'][:, 1:]
data_dict['global_orient'] = preds_dict['pred_theta_mats'][:, [0]]
data_dict['betas'] = preds_dict['pred_shape']
data_dict['smpl_verts'] = preds_dict['pred_vertices']
scale, tranX, tranY = preds_dict['pred_camera'][0, :3]
scale = scale * 2
elif self.hps_type == 'bev':
data_dict['betas'] = torch.from_numpy(preds_dict['smpl_betas'])[
[0], :10].to(self.device).float()
pred_thetas = batch_rodrigues(torch.from_numpy(
preds_dict['smpl_thetas'][0]).reshape(-1, 3)).float()
data_dict['body_pose'] = pred_thetas[1:][None].to(self.device)
data_dict['global_orient'] = pred_thetas[[0]][None].to(self.device)
data_dict['smpl_verts'] = torch.from_numpy(
preds_dict['verts'][[0]]).to(self.device).float()
tranX = preds_dict['cam_trans'][0, 0]
tranY = preds_dict['cam'][0, 1] + 0.28
scale = preds_dict['cam'][0, 0] * 1.1
data_dict['scale'] = scale
data_dict['trans'] = torch.tensor(
[tranX, tranY, 0.0]).to(self.device).float()
# data_dict info (key-shape):
# scale, tranX, tranY - tensor.float
# betas - [1,10] / [1, 200]
# body_pose - [1, 23, 3, 3] / [1, 21, 3, 3]
# global_orient - [1, 1, 3, 3]
# smpl_verts - [1, 6890, 3] / [1, 10475, 3]
# from rot_mat to rot_6d for better optimization
N_body = data_dict["body_pose"].shape[1]
data_dict["body_pose"] = data_dict["body_pose"][:, :, :, :2].reshape(1, N_body,-1)
data_dict["global_orient"] = data_dict["global_orient"][:, :, :, :2].reshape(1, 1,-1)
return data_dict
def render_normal(self, verts, faces):
# render optimized mesh (normal, T_normal, image [-1,1])
self.render.load_meshes(verts, faces)
return self.render.get_rgb_image()
def render_depth(self, verts, faces):
# render optimized mesh (normal, T_normal, image [-1,1])
self.render.load_meshes(verts, faces)
return self.render.get_depth_map(cam_ids=[0, 2])