# -*- 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 os.path as osp import cv2 import numpy as np import torch import torch.nn.functional as F import torchvision.transforms as transforms import trimesh from PIL import Image from lib.common.render import Render from lib.dataset.mesh_util import SMPLX, HoppeMesh, projection, rescale_smpl cape_gender = { "male": ['00032', '00096', '00122', '00127', '00145', '00215', '02474', '03284', '03375', '03394'], "female": ['00134', '00159', '03223', '03331', '03383'] } class EvalDataset: def __init__(self, cfg, device): self.root = cfg.root self.bsize = cfg.batch_size self.opt = cfg.dataset self.datasets = self.opt.types self.input_size = self.opt.input_size self.scales = self.opt.scales self.vol_res = cfg.vol_res # [(feat_name, channel_num),...] self.in_geo = [item[0] for item in cfg.net.in_geo] self.in_nml = [item[0] for item in cfg.net.in_nml] self.in_geo_dim = [item[1] for item in cfg.net.in_geo] self.in_nml_dim = [item[1] for item in cfg.net.in_nml] self.in_total = self.in_geo + self.in_nml self.in_total_dim = self.in_geo_dim + self.in_nml_dim self.rotations = range(0, 360, 120) self.datasets_dict = {} for dataset_id, dataset in enumerate(self.datasets): dataset_dir = osp.join(self.root, dataset) mesh_dir = osp.join(dataset_dir, "scans") smplx_dir = osp.join(dataset_dir, "smplx") smpl_dir = osp.join(dataset_dir, "smpl") self.datasets_dict[dataset] = { "smplx_dir": smplx_dir, "smpl_dir": smpl_dir, "mesh_dir": mesh_dir, "scale": self.scales[dataset_id], } self.datasets_dict[dataset].update({ "subjects": np.loadtxt(osp.join(dataset_dir, "all.txt"), dtype=str) }) self.subject_list = self.get_subject_list() self.smplx = SMPLX() # PIL to tensor self.image_to_tensor = transforms.Compose([ transforms.Resize(self.input_size), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) # PIL to tensor self.mask_to_tensor = transforms.Compose([ transforms.Resize(self.input_size), transforms.ToTensor(), transforms.Normalize((0.0, ), (1.0, )), ]) self.device = device self.render = Render(size=512, device=self.device) def render_normal(self, verts, faces): # render optimized mesh (normal, T_normal, image [-1,1]) self.render.load_meshes(verts, faces) return self.render.get_image() def get_subject_list(self): subject_list = [] for dataset in self.datasets: split_txt = "" if dataset == 'renderpeople': split_txt = osp.join(self.root, dataset, "loose.txt") elif dataset == 'cape': split_txt = osp.join(self.root, dataset, "pose.txt") if osp.exists(split_txt) and osp.getsize(split_txt) > 0: print(f"load from {split_txt}") subject_list += np.loadtxt(split_txt, dtype=str).tolist() return subject_list def __len__(self): return len(self.subject_list) * len(self.rotations) def __getitem__(self, index): rid = index % len(self.rotations) mid = index // len(self.rotations) rotation = self.rotations[rid] subject = self.subject_list[mid].split("/")[1] dataset = self.subject_list[mid].split("/")[0] render_folder = "/".join([dataset + f"_{self.opt.rotation_num}views", subject]) if not osp.exists(osp.join(self.root, render_folder)): render_folder = "/".join([dataset + "_36views", subject]) # setup paths data_dict = { "dataset": dataset, "subject": subject, "rotation": rotation, "scale": self.datasets_dict[dataset]["scale"], "calib_path": osp.join(self.root, render_folder, "calib", f"{rotation:03d}.txt"), "image_path": osp.join(self.root, render_folder, "render", f"{rotation:03d}.png"), } if dataset == "cape": data_dict.update({ "mesh_path": osp.join(self.datasets_dict[dataset]["mesh_dir"], f"{subject}.obj"), "smpl_path": osp.join(self.datasets_dict[dataset]["smpl_dir"], f"{subject}.obj"), }) else: data_dict.update({ "mesh_path": osp.join( self.datasets_dict[dataset]["mesh_dir"], f"{subject}.obj", ), "smplx_path": osp.join(self.datasets_dict[dataset]["smplx_dir"], f"{subject}.obj"), }) # load training data data_dict.update(self.load_calib(data_dict)) # image/normal/depth loader for name, channel in zip(self.in_total, self.in_total_dim): if f"{name}_path" not in data_dict.keys(): data_dict.update({ f"{name}_path": osp.join(self.root, render_folder, name, f"{rotation:03d}.png") }) # tensor update if os.path.exists(data_dict[f"{name}_path"]): data_dict.update({ name: self.imagepath2tensor(data_dict[f"{name}_path"], channel, inv=False) }) data_dict.update(self.load_mesh(data_dict)) data_dict.update(self.load_smpl(data_dict)) del data_dict["mesh"] return data_dict def imagepath2tensor(self, path, channel=3, inv=False): rgba = Image.open(path).convert("RGBA") # remove CAPE's noisy outliers using OpenCV's inpainting if "cape" in path and "T_" not in path: mask = cv2.imread(path.replace(path.split("/")[-2], "mask"), 0) > 1 img = np.asarray(rgba)[:, :, :3] fill_mask = ((mask & (img.sum(axis=2) == 0))).astype(np.uint8) image = Image.fromarray( cv2.inpaint(img * mask[..., None], fill_mask, 3, cv2.INPAINT_TELEA) ) mask = Image.fromarray(mask) else: mask = rgba.split()[-1] image = rgba.convert("RGB") image = self.image_to_tensor(image) mask = self.mask_to_tensor(mask) image = (image * mask)[:channel] return (image * (0.5 - inv) * 2.0).float() def load_calib(self, data_dict): calib_data = np.loadtxt(data_dict["calib_path"], dtype=float) extrinsic = calib_data[:4, :4] intrinsic = calib_data[4:8, :4] calib_mat = np.matmul(intrinsic, extrinsic) calib_mat = torch.from_numpy(calib_mat).float() return {"calib": calib_mat} def load_mesh(self, data_dict): mesh_path = data_dict["mesh_path"] scale = data_dict["scale"] scan_mesh = trimesh.load(mesh_path) verts = scan_mesh.vertices faces = scan_mesh.faces mesh = HoppeMesh(verts * scale, faces) return { "mesh": mesh, "verts": torch.as_tensor(verts * scale).float(), "faces": torch.as_tensor(faces).long(), } def load_smpl(self, data_dict): smpl_type = ("smplx" if ("smplx_path" in data_dict.keys()) else "smpl") smplx_verts = rescale_smpl(data_dict[f"{smpl_type}_path"], scale=100.0) smplx_faces = torch.as_tensor(getattr(self.smplx, f"{smpl_type}_faces")).long() smplx_verts = projection(smplx_verts, data_dict["calib"]).float() return_dict = { "smpl_verts": smplx_verts, "smpl_faces": smplx_faces, } return return_dict 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 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")