# -*- 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 logging import warnings warnings.filterwarnings("ignore") logging.getLogger("lightning").setLevel(logging.ERROR) logging.getLogger("trimesh").setLevel(logging.ERROR) import glob import os.path as osp import numpy as np import torch import torch.nn.functional as F from PIL import ImageFile from termcolor import colored from torchvision import transforms from torchvision.models import detection from lib.common.config import cfg from lib.common.imutils import process_image from lib.common.render import Render from lib.common.train_util import Format from lib.dataset.mesh_util import SMPLX, get_visibility from lib.pixielib.models.SMPLX import SMPLX as PIXIE_SMPLX from lib.pixielib.pixie import PIXIE from lib.pixielib.utils.config import cfg as pixie_cfg from lib.pymafx.core import path_config from lib.pymafx.models import pymaf_net ImageFile.LOAD_TRUNCATED_IMAGES = True class TestDataset: def __init__(self, cfg, device): self.image_path = cfg["image_path"] self.use_seg = cfg["use_seg"] self.hps_type = cfg["hps_type"] self.smpl_type = "smplx" self.smpl_gender = "neutral" self.vol_res = cfg["vol_res"] self.single = cfg["single"] self.device = device self.subject_list = [self.image_path] # smpl related self.smpl_data = SMPLX() if self.hps_type == "pymafx": 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() pixie_cfg.merge_from_list(["model.n_shape", 10, "model.n_exp", 10]) elif self.hps_type == "pixie": self.hps = PIXIE(config=pixie_cfg, device=self.device) self.smpl_model = PIXIE_SMPLX(pixie_cfg.model).to(self.device) self.detector = detection.maskrcnn_resnet50_fpn( weights=detection.MaskRCNN_ResNet50_FPN_V2_Weights ) self.detector.eval() print( colored( f"SMPL-X estimate with {self.hps_type.upper()}", "green" ) ) self.render = Render(size=512, device=self.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()[:, :, [0, 2, 1]]).unsqueeze(-1) smpl_cmap = self.smpl_data.cmap_smpl_vids(self.smpl_type).unsqueeze(0) return { "smpl_vis": smpl_vis.to(self.device), "smpl_cmap": smpl_cmap.to(self.device), "smpl_verts": smpl_verts, } def depth_to_voxel(self, data_dict): data_dict["depth_F"] = transforms.Resize(self.vol_res)(data_dict["depth_F"]) data_dict["depth_B"] = transforms.Resize(self.vol_res)(data_dict["depth_B"]) depth_mask = (~torch.isnan(data_dict['depth_F'])) depth_FB = torch.cat([data_dict['depth_F'], data_dict['depth_B']], dim=0) depth_FB[:, ~depth_mask[0]] = 0. # Important: index_long = depth_value - 1 index_z = (((depth_FB + 1.) * 0.5 * self.vol_res) - 1).clip(0, self.vol_res - 1).permute(1, 2, 0) index_z_ceil = torch.ceil(index_z).long() index_z_floor = torch.floor(index_z).long() index_z_frac = torch.frac(index_z) index_mask = index_z[..., 0] == torch.tensor(self.vol_res * 0.5 - 1).long() voxels = F.one_hot(index_z_ceil[..., 0], self.vol_res) * index_z_frac[..., 0] + \ F.one_hot(index_z_floor[..., 0], self.vol_res) * (1.0-index_z_frac[..., 0]) + \ F.one_hot(index_z_ceil[..., 1], self.vol_res) * index_z_frac[..., 1]+ \ F.one_hot(index_z_floor[..., 1], self.vol_res) * (1.0 - index_z_frac[..., 1]) voxels[index_mask] *= 0 voxels = torch.flip(voxels, [2]).permute(2, 0, 1).float() #[x-2, y-0, z-1] return { "depth_voxels": voxels.flip([ 0, ]).unsqueeze(0).to(self.device), } def __getitem__(self, index): img_path = self.subject_list[index] img_name = img_path.split("/")[-1].rsplit(".", 1)[0] arr_dict = process_image(img_path, self.hps_type, self.single, 512, self.detector) arr_dict.update({"name": img_name}) with torch.no_grad(): if self.hps_type == "pixie": preds_dict = self.hps.forward(arr_dict["img_hps"].to(self.device)) elif self.hps_type == 'pymafx': batch = {k: v.to(self.device) for k, v in arr_dict["img_pymafx"].items()} preds_dict, _ = self.hps.forward(batch) arr_dict["smpl_faces"] = ( torch.as_tensor(self.smpl_data.smplx_faces.astype(np.int64)).unsqueeze(0).long().to( self.device ) ) arr_dict["type"] = self.smpl_type if self.hps_type == "pymafx": output = preds_dict["mesh_out"][-1] scale, tranX, tranY = output["theta"][:, :3].split(1, dim=1) arr_dict["betas"] = output["pred_shape"] arr_dict["body_pose"] = output["rotmat"][:, 1:22] arr_dict["global_orient"] = output["rotmat"][:, 0:1] arr_dict["smpl_verts"] = output["smplx_verts"] arr_dict["left_hand_pose"] = output["pred_lhand_rotmat"] arr_dict["right_hand_pose"] = output["pred_rhand_rotmat"] arr_dict['jaw_pose'] = output['pred_face_rotmat'][:, 0:1] arr_dict["exp"] = output["pred_exp"] # 1.2009, 0.0013, 0.3954 elif self.hps_type == "pixie": arr_dict.update(preds_dict) arr_dict["global_orient"] = preds_dict["global_pose"] arr_dict["betas"] = preds_dict["shape"] #200 arr_dict["smpl_verts"] = preds_dict["vertices"] scale, tranX, tranY = preds_dict["cam"].split(1, dim=1) # 1.1435, 0.0128, 0.3520 arr_dict["scale"] = scale.unsqueeze(1) arr_dict["trans"] = ( torch.cat([tranX, tranY, torch.zeros_like(tranX)], dim=1).unsqueeze(1).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, N_pose = arr_dict["body_pose"].shape[:2] arr_dict["body_pose"] = arr_dict["body_pose"][:, :, :, :2].reshape(N_body, N_pose, -1) arr_dict["global_orient"] = arr_dict["global_orient"][:, :, :, :2].reshape(N_body, 1, -1) return arr_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_image(type="rgb") 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_image(type="depth")