diff --git a/data/dataset.py b/data/dataset.py index fe5f83d..c8eeb1a 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -7,6 +7,7 @@ import math import time import torch import kornia +import pytorch3d import torchvision as tv import numpy as np import scipy.io as sio @@ -53,9 +54,12 @@ class Dataset(Dataset): path = os.path.join(self.folder, frame.id) info = self.reader.read_info(frame.id) - zrot = info['zrot'] if not self.random_view else (((np.random.random(1) * 2) -1) * np.pi)[0] + zrot = info['zrot'] # if not self.random_view else (((np.random.random(1) * 2) -1) * np.pi)[0] + trans = info['trans'][:,frame.frame_nr] # mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True) - mesh_posed = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, info['trans'][:,frame.frame_nr], zrot) + + mesh_posed = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, trans, zrot) + mesh_posed = torch.tensor(mesh_posed, device=self.device) # mesh_posed_no_trans = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, absolute=False) mesh_mask = np.load(os.path.join(path, frame.garment + '_mask.npy')) @@ -63,14 +67,28 @@ class Dataset(Dataset): faces = mask_faces(self.body_faces, mesh_mask) # faces = quads2tris(self.reader.read_garment_topology(frame.id, frame.garment)) # body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True) - body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, info['trans'][:,frame.frame_nr], zrot) + body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, trans, zrot) # body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, absolute=False) body_faces = quads2tris(body_faces) + if self.random_view: + x_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi) / 10 + y_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi) / 10 + z_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi) + view = pytorch3d.transforms.axis_angle_to_matrix(torch.concat((x_angle, y_angle, z_angle))) + + transform = pytorch3d.transforms.Rotate(view) + mesh_posed = transform.transform_points(mesh_posed) + body = transform.transform_points(body) + else: + view = torch.eye(3) + combined = torch.Tensor(np.concatenate((body, mesh_posed[mesh_mask, :]))).to(self.device) # combined_no_trans = torch.Tensor(np.concatenate((body_no_trans, mesh_posed_no_trans[mesh_mask, :]))).to('cuda') combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to(self.device) + + start = time.perf_counter() # depth_image = self.render.render(mesh_posed[mask, :], faces, info['camLoc'][0])[:,:,0] # depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0] @@ -194,10 +212,11 @@ class Dataset(Dataset): rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(torch.float32) - view = torch.tensor(np.concatenate((np.array([zrot], dtype=np.float32), info['camLoc'].astype(np.float32), info['trans'][:,frame.frame_nr].astype(np.float32)), axis=0), device=self.device) + # view = torch.tensor(np.concatenate((np.array([zrot], dtype=np.float32), info['camLoc'].astype(np.float32), trans.astype(np.float32)), axis=0), device=self.device) + # image = self.transforms(image) - data = { 'mesh_posed': torch.tensor(mesh_posed, device=self.device), + data = { 'mesh_posed': mesh_posed, 'depth_image': depth_image, 'mesh_mask': mesh_mask, 'depth_mask': depth_mask,