add new view definition to dataset

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,