ICON / lib /dataset /PIFuDataset.py
Yuliang's picture
done
2d5f249
raw
history blame
22.5 kB
from lib.renderer.mesh import load_fit_body
from lib.dataset.hoppeMesh import HoppeMesh
from lib.dataset.body_model import TetraSMPLModel
from lib.common.render import Render
from lib.dataset.mesh_util import SMPLX, projection, cal_sdf_batch, get_visibility
from lib.pare.pare.utils.geometry import rotation_matrix_to_angle_axis
from termcolor import colored
import os.path as osp
import numpy as np
from PIL import Image
import random
import os
import trimesh
import torch
from kaolin.ops.mesh import check_sign
import torchvision.transforms as transforms
from huggingface_hub import hf_hub_download, cached_download
class PIFuDataset():
def __init__(self, cfg, split='train', vis=False):
self.split = split
self.root = cfg.root
self.bsize = cfg.batch_size
self.overfit = cfg.overfit
# for debug, only used in visualize_sampling3D
self.vis = vis
self.opt = cfg.dataset
self.datasets = self.opt.types
self.input_size = self.opt.input_size
self.scales = self.opt.scales
self.workers = cfg.num_threads
self.prior_type = cfg.net.prior_type
self.noise_type = self.opt.noise_type
self.noise_scale = self.opt.noise_scale
noise_joints = [4, 5, 7, 8, 13, 14, 16, 17, 18, 19, 20, 21]
self.noise_smpl_idx = []
self.noise_smplx_idx = []
for idx in noise_joints:
self.noise_smpl_idx.append(idx * 3)
self.noise_smpl_idx.append(idx * 3 + 1)
self.noise_smpl_idx.append(idx * 3 + 2)
self.noise_smplx_idx.append((idx-1) * 3)
self.noise_smplx_idx.append((idx-1) * 3 + 1)
self.noise_smplx_idx.append((idx-1) * 3 + 2)
self.use_sdf = cfg.sdf
self.sdf_clip = cfg.sdf_clip
# [(feat_name, channel_num),...]
self.in_geo = [item[0] for item in cfg.net.in_geo]
self.in_nml = [item[0] for item in cfg.net.in_nml]
self.in_geo_dim = [item[1] for item in cfg.net.in_geo]
self.in_nml_dim = [item[1] for item in cfg.net.in_nml]
self.in_total = self.in_geo + self.in_nml
self.in_total_dim = self.in_geo_dim + self.in_nml_dim
if self.split == 'train':
self.rotations = np.arange(
0, 360, 360 / self.opt.rotation_num).astype(np.int32)
else:
self.rotations = range(0, 360, 120)
self.datasets_dict = {}
for dataset_id, dataset in enumerate(self.datasets):
mesh_dir = None
smplx_dir = None
dataset_dir = osp.join(self.root, dataset)
if dataset in ['thuman2']:
mesh_dir = osp.join(dataset_dir, "scans")
smplx_dir = osp.join(dataset_dir, "fits")
smpl_dir = osp.join(dataset_dir, "smpl")
self.datasets_dict[dataset] = {
"subjects": np.loadtxt(osp.join(dataset_dir, "all.txt"), dtype=str),
"smplx_dir": smplx_dir,
"smpl_dir": smpl_dir,
"mesh_dir": mesh_dir,
"scale": self.scales[dataset_id]
}
self.subject_list = self.get_subject_list(split)
self.smplx = SMPLX()
# PIL to tensor
self.image_to_tensor = transforms.Compose([
transforms.Resize(self.input_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# PIL to tensor
self.mask_to_tensor = transforms.Compose([
transforms.Resize(self.input_size),
transforms.ToTensor(),
transforms.Normalize((0.0, ), (1.0, ))
])
self.device = torch.device(f"cuda:{cfg.gpus[0]}")
self.render = Render(size=512, device=self.device)
def render_normal(self, verts, faces):
# render optimized mesh (normal, T_normal, image [-1,1])
self.render.load_meshes(verts, faces)
return self.render.get_rgb_image()
def get_subject_list(self, split):
subject_list = []
for dataset in self.datasets:
split_txt = osp.join(self.root, dataset, f'{split}.txt')
if osp.exists(split_txt):
print(f"load from {split_txt}")
subject_list += np.loadtxt(split_txt, dtype=str).tolist()
else:
full_txt = osp.join(self.root, dataset, 'all.txt')
print(f"split {full_txt} into train/val/test")
full_lst = np.loadtxt(full_txt, dtype=str)
full_lst = [dataset+"/"+item for item in full_lst]
[train_lst, test_lst, val_lst] = np.split(
full_lst, [500, 500+5, ])
np.savetxt(full_txt.replace(
"all", "train"), train_lst, fmt="%s")
np.savetxt(full_txt.replace("all", "test"), test_lst, fmt="%s")
np.savetxt(full_txt.replace("all", "val"), val_lst, fmt="%s")
print(f"load from {split_txt}")
subject_list += np.loadtxt(split_txt, dtype=str).tolist()
if self.split != 'test':
subject_list += subject_list[:self.bsize -
len(subject_list) % self.bsize]
print(colored(f"total: {len(subject_list)}", "yellow"))
random.shuffle(subject_list)
# subject_list = ["thuman2/0008"]
return subject_list
def __len__(self):
return len(self.subject_list) * len(self.rotations)
def __getitem__(self, index):
# only pick the first data if overfitting
if self.overfit:
index = 0
rid = index % len(self.rotations)
mid = index // len(self.rotations)
rotation = self.rotations[rid]
subject = self.subject_list[mid].split("/")[1]
dataset = self.subject_list[mid].split("/")[0]
render_folder = "/".join([dataset +
f"_{self.opt.rotation_num}views", subject])
# setup paths
data_dict = {
'dataset': dataset,
'subject': subject,
'rotation': rotation,
'scale': self.datasets_dict[dataset]["scale"],
'mesh_path': osp.join(self.datasets_dict[dataset]["mesh_dir"], f"{subject}/{subject}.obj"),
'smplx_path': osp.join(self.datasets_dict[dataset]["smplx_dir"], f"{subject}/smplx_param.pkl"),
'smpl_path': osp.join(self.datasets_dict[dataset]["smpl_dir"], f"{subject}.pkl"),
'calib_path': osp.join(self.root, render_folder, 'calib', f'{rotation:03d}.txt'),
'vis_path': osp.join(self.root, render_folder, 'vis', f'{rotation:03d}.pt'),
'image_path': osp.join(self.root, render_folder, 'render', f'{rotation:03d}.png')
}
# load training data
data_dict.update(self.load_calib(data_dict))
# image/normal/depth loader
for name, channel in zip(self.in_total, self.in_total_dim):
if f'{name}_path' not in data_dict.keys():
data_dict.update({
f'{name}_path': osp.join(self.root, render_folder, name, f'{rotation:03d}.png')
})
# tensor update
data_dict.update({
name: self.imagepath2tensor(
data_dict[f'{name}_path'], channel, inv=False)
})
data_dict.update(self.load_mesh(data_dict))
data_dict.update(self.get_sampling_geo(
data_dict, is_valid=self.split == "val", is_sdf=self.use_sdf))
data_dict.update(self.load_smpl(data_dict, self.vis))
if self.prior_type == 'pamir':
data_dict.update(self.load_smpl_voxel(data_dict))
if (self.split != 'test') and (not self.vis):
del data_dict['verts']
del data_dict['faces']
if not self.vis:
del data_dict['mesh']
path_keys = [
key for key in data_dict.keys() if '_path' in key or '_dir' in key
]
for key in path_keys:
del data_dict[key]
return data_dict
def imagepath2tensor(self, path, channel=3, inv=False):
rgba = Image.open(path).convert('RGBA')
mask = rgba.split()[-1]
image = rgba.convert('RGB')
image = self.image_to_tensor(image)
mask = self.mask_to_tensor(mask)
image = (image * mask)[:channel]
return (image * (0.5 - inv) * 2.0).float()
def load_calib(self, data_dict):
calib_data = np.loadtxt(data_dict['calib_path'], dtype=float)
extrinsic = calib_data[:4, :4]
intrinsic = calib_data[4:8, :4]
calib_mat = np.matmul(intrinsic, extrinsic)
calib_mat = torch.from_numpy(calib_mat).float()
return {'calib': calib_mat}
def load_mesh(self, data_dict):
mesh_path = data_dict['mesh_path']
scale = data_dict['scale']
mesh_ori = trimesh.load(mesh_path,
skip_materials=True,
process=False,
maintain_order=True)
verts = mesh_ori.vertices * scale
faces = mesh_ori.faces
vert_normals = np.array(mesh_ori.vertex_normals)
face_normals = np.array(mesh_ori.face_normals)
mesh = HoppeMesh(verts, faces, vert_normals, face_normals)
return {
'mesh': mesh,
'verts': torch.as_tensor(mesh.verts).float(),
'faces': torch.as_tensor(mesh.faces).long()
}
def add_noise(self,
beta_num,
smpl_pose,
smpl_betas,
noise_type,
noise_scale,
type,
hashcode):
np.random.seed(hashcode)
if type == 'smplx':
noise_idx = self.noise_smplx_idx
else:
noise_idx = self.noise_smpl_idx
if 'beta' in noise_type and noise_scale[noise_type.index("beta")] > 0.0:
smpl_betas += (np.random.rand(beta_num) -
0.5) * 2.0 * noise_scale[noise_type.index("beta")]
smpl_betas = smpl_betas.astype(np.float32)
if 'pose' in noise_type and noise_scale[noise_type.index("pose")] > 0.0:
smpl_pose[noise_idx] += (
np.random.rand(len(noise_idx)) -
0.5) * 2.0 * np.pi * noise_scale[noise_type.index("pose")]
smpl_pose = smpl_pose.astype(np.float32)
if type == 'smplx':
return torch.as_tensor(smpl_pose[None, ...]), torch.as_tensor(smpl_betas[None, ...])
else:
return smpl_pose, smpl_betas
def compute_smpl_verts(self, data_dict, noise_type=None, noise_scale=None):
dataset = data_dict['dataset']
smplx_dict = {}
smplx_param = np.load(data_dict['smplx_path'], allow_pickle=True)
smplx_pose = smplx_param["body_pose"] # [1,63]
smplx_betas = smplx_param["betas"] # [1,10]
smplx_pose, smplx_betas = self.add_noise(
smplx_betas.shape[1],
smplx_pose[0],
smplx_betas[0],
noise_type,
noise_scale,
type='smplx',
hashcode=(hash(f"{data_dict['subject']}_{data_dict['rotation']}")) % (10**8))
smplx_out, _ = load_fit_body(fitted_path=data_dict['smplx_path'],
scale=self.datasets_dict[dataset]['scale'],
smpl_type='smplx',
smpl_gender='male',
noise_dict=dict(betas=smplx_betas, body_pose=smplx_pose))
smplx_dict.update({"type": "smplx",
"gender": 'male',
"body_pose": torch.as_tensor(smplx_pose),
"betas": torch.as_tensor(smplx_betas)})
return smplx_out.vertices, smplx_dict
def compute_voxel_verts(self,
data_dict,
noise_type=None,
noise_scale=None):
smpl_param = np.load(data_dict['smpl_path'], allow_pickle=True)
smplx_param = np.load(data_dict['smplx_path'], allow_pickle=True)
smpl_pose = rotation_matrix_to_angle_axis(
torch.as_tensor(smpl_param['full_pose'][0])).numpy()
smpl_betas = smpl_param["betas"]
smpl_path = cached_download(osp.join(self.smplx.model_dir, "smpl/SMPL_MALE.pkl"), use_auth_token=os.environ['ICON'])
tetra_path = cached_download(osp.join(self.smplx.tedra_dir,
"tetra_male_adult_smpl.npz"), use_auth_token=os.environ['ICON'])
smpl_model = TetraSMPLModel(smpl_path, tetra_path, 'adult')
smpl_pose, smpl_betas = self.add_noise(
smpl_model.beta_shape[0],
smpl_pose.flatten(),
smpl_betas[0],
noise_type,
noise_scale,
type='smpl',
hashcode=(hash(f"{data_dict['subject']}_{data_dict['rotation']}")) % (10**8))
smpl_model.set_params(pose=smpl_pose.reshape(-1, 3),
beta=smpl_betas,
trans=smpl_param["transl"])
verts = (np.concatenate([smpl_model.verts, smpl_model.verts_added],
axis=0) * smplx_param["scale"] + smplx_param["translation"]
) * self.datasets_dict[data_dict['dataset']]['scale']
faces = np.loadtxt(cached_download(osp.join(self.smplx.tedra_dir, "tetrahedrons_male_adult.txt"), use_auth_token=os.environ['ICON']),
dtype=np.int32) - 1
pad_v_num = int(8000 - verts.shape[0])
pad_f_num = int(25100 - faces.shape[0])
verts = np.pad(verts, ((0, pad_v_num), (0, 0)),
mode='constant',
constant_values=0.0).astype(np.float32)
faces = np.pad(faces, ((0, pad_f_num), (0, 0)),
mode='constant',
constant_values=0.0).astype(np.int32)
return verts, faces, pad_v_num, pad_f_num
def load_smpl(self, data_dict, vis=False):
smplx_verts, smplx_dict = self.compute_smpl_verts(
data_dict, self.noise_type,
self.noise_scale) # compute using smpl model
smplx_verts = projection(smplx_verts, data_dict['calib']).float()
smplx_faces = torch.as_tensor(self.smplx.faces).long()
smplx_vis = torch.load(data_dict['vis_path']).float()
smplx_cmap = torch.as_tensor(
np.load(self.smplx.cmap_vert_path)).float()
# get smpl_signs
query_points = projection(data_dict['samples_geo'],
data_dict['calib']).float()
pts_signs = 2.0 * (check_sign(smplx_verts.unsqueeze(0),
smplx_faces,
query_points.unsqueeze(0)).float() - 0.5).squeeze(0)
return_dict = {
'smpl_verts': smplx_verts,
'smpl_faces': smplx_faces,
'smpl_vis': smplx_vis,
'smpl_cmap': smplx_cmap,
'pts_signs': pts_signs
}
if smplx_dict is not None:
return_dict.update(smplx_dict)
if vis:
(xy, z) = torch.as_tensor(smplx_verts).to(
self.device).split([2, 1], dim=1)
smplx_vis = get_visibility(xy, z, torch.as_tensor(
smplx_faces).to(self.device).long())
T_normal_F, T_normal_B = self.render_normal(
(smplx_verts*torch.tensor([1.0, -1.0, 1.0])).to(self.device),
smplx_faces.to(self.device))
return_dict.update({"T_normal_F": T_normal_F.squeeze(0),
"T_normal_B": T_normal_B.squeeze(0)})
query_points = projection(data_dict['samples_geo'],
data_dict['calib']).float()
smplx_sdf, smplx_norm, smplx_cmap, smplx_vis = cal_sdf_batch(
smplx_verts.unsqueeze(0).to(self.device),
smplx_faces.unsqueeze(0).to(self.device),
smplx_cmap.unsqueeze(0).to(self.device),
smplx_vis.unsqueeze(0).to(self.device),
query_points.unsqueeze(0).contiguous().to(self.device))
return_dict.update({
'smpl_feat':
torch.cat(
(smplx_sdf[0].detach().cpu(),
smplx_cmap[0].detach().cpu(),
smplx_norm[0].detach().cpu(),
smplx_vis[0].detach().cpu()),
dim=1)
})
return return_dict
def load_smpl_voxel(self, data_dict):
smpl_verts, smpl_faces, pad_v_num, pad_f_num = self.compute_voxel_verts(
data_dict, self.noise_type,
self.noise_scale) # compute using smpl model
smpl_verts = projection(smpl_verts, data_dict['calib'])
smpl_verts *= 0.5
return {
'voxel_verts': smpl_verts,
'voxel_faces': smpl_faces,
'pad_v_num': pad_v_num,
'pad_f_num': pad_f_num
}
def get_sampling_geo(self, data_dict, is_valid=False, is_sdf=False):
mesh = data_dict['mesh']
calib = data_dict['calib']
# Samples are around the true surface with an offset
n_samples_surface = 4 * self.opt.num_sample_geo
vert_ids = np.arange(mesh.verts.shape[0])
thickness_sample_ratio = np.ones_like(vert_ids).astype(np.float32)
thickness_sample_ratio /= thickness_sample_ratio.sum()
samples_surface_ids = np.random.choice(vert_ids,
n_samples_surface,
replace=True,
p=thickness_sample_ratio)
samples_normal_ids = np.random.choice(vert_ids,
self.opt.num_sample_geo // 2,
replace=False,
p=thickness_sample_ratio)
surf_samples = mesh.verts[samples_normal_ids, :]
surf_normals = mesh.vert_normals[samples_normal_ids, :]
samples_surface = mesh.verts[samples_surface_ids, :]
# Sampling offsets are random noise with constant scale (15cm - 20cm)
offset = np.random.normal(scale=self.opt.sigma_geo,
size=(n_samples_surface, 1))
samples_surface += mesh.vert_normals[samples_surface_ids, :] * offset
# Uniform samples in [-1, 1]
calib_inv = np.linalg.inv(calib)
n_samples_space = self.opt.num_sample_geo // 4
samples_space_img = 2.0 * np.random.rand(n_samples_space, 3) - 1.0
samples_space = projection(samples_space_img, calib_inv)
# z-ray direction samples
if self.opt.zray_type and not is_valid:
n_samples_rayz = self.opt.ray_sample_num
samples_surface_cube = projection(samples_surface, calib)
samples_surface_cube_repeat = np.repeat(samples_surface_cube,
n_samples_rayz,
axis=0)
thickness_repeat = np.repeat(0.5 *
np.ones_like(samples_surface_ids),
n_samples_rayz,
axis=0)
noise_repeat = np.random.normal(scale=0.40,
size=(n_samples_surface *
n_samples_rayz, ))
samples_surface_cube_repeat[:,
-1] += thickness_repeat * noise_repeat
samples_surface_rayz = projection(samples_surface_cube_repeat,
calib_inv)
samples = np.concatenate(
[samples_surface, samples_space, samples_surface_rayz], 0)
else:
samples = np.concatenate([samples_surface, samples_space], 0)
np.random.shuffle(samples)
# labels: in->1.0; out->0.0.
if is_sdf:
sdfs = mesh.get_sdf(samples)
inside_samples = samples[sdfs < 0]
outside_samples = samples[sdfs >= 0]
inside_sdfs = sdfs[sdfs < 0]
outside_sdfs = sdfs[sdfs >= 0]
else:
inside = mesh.contains(samples)
inside_samples = samples[inside >= 0.5]
outside_samples = samples[inside < 0.5]
nin = inside_samples.shape[0]
if nin > self.opt.num_sample_geo // 2:
inside_samples = inside_samples[:self.opt.num_sample_geo // 2]
outside_samples = outside_samples[:self.opt.num_sample_geo // 2]
if is_sdf:
inside_sdfs = inside_sdfs[:self.opt.num_sample_geo // 2]
outside_sdfs = outside_sdfs[:self.opt.num_sample_geo // 2]
else:
outside_samples = outside_samples[:(self.opt.num_sample_geo - nin)]
if is_sdf:
outside_sdfs = outside_sdfs[:(self.opt.num_sample_geo - nin)]
if is_sdf:
samples = np.concatenate(
[inside_samples, outside_samples, surf_samples], 0)
labels = np.concatenate([
inside_sdfs, outside_sdfs, 0.0 * np.ones(surf_samples.shape[0])
])
normals = np.zeros_like(samples)
normals[-self.opt.num_sample_geo // 2:, :] = surf_normals
# convert sdf from [-14, 130] to [0, 1]
# outside: 0, inside: 1
# Note: Marching cubes is defined on occupancy space (inside=1.0, outside=0.0)
labels = -labels.clip(min=-self.sdf_clip, max=self.sdf_clip)
labels += self.sdf_clip
labels /= (self.sdf_clip * 2)
else:
samples = np.concatenate([inside_samples, outside_samples])
labels = np.concatenate([
np.ones(inside_samples.shape[0]),
np.zeros(outside_samples.shape[0])
])
normals = np.zeros_like(samples)
samples = torch.from_numpy(samples).float()
labels = torch.from_numpy(labels).float()
normals = torch.from_numpy(normals).float()
return {'samples_geo': samples, 'labels_geo': labels}