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, 

test.py

@@ -43,13 +43,15 @@ def chamfer_loss(depth_image, vertices, depth_mask, mesh_mask):
     K = torch.tensor(intrinsic()[None, :, :], device='cuda')
     pc = kornia.geometry.depth.depth_to_3d(depth_image, K, normalize_points=True)
     pc = [torch.permute(pc[i,:,depth_mask[i]], [1,0]) for i in range(depth_image.shape[0])]
-    pc_pc = Pointclouds(pc)
-    mesh_gt = [vertices[i, mesh_mask[i], :] for i in range(vertices.shape[0])]
-    mesh_gt_pc = Pointclouds(mesh_gt)
     pc_norm = get_normals_from_depth(depth_image, depth_mask, K)
-    mesh_norm = estimate_pointcloud_normals(mesh_gt_pc, neighborhood_size=50)
+    pc_norm = [pc_norm[i,depth_mask[i],:] for i in range(pc_norm.shape[0])]
+    mesh_gt = [vertices[i, mesh_mask[i], :] for i in range(vertices.shape[0])]
+    pc_pc = Pointclouds(pc, normals=pc_norm)
+    mesh_gt_pc = Pointclouds(mesh_gt)
+    mesh_norm = list(torch.unbind(estimate_pointcloud_normals(mesh_gt_pc, neighborhood_size=50)))
+    mesh_gt_pc = Pointclouds(mesh_gt, normals=mesh_norm)
 
-    chamfer = chamfer_distance(pc_pc, mesh_gt_pc, x_normals=pc_norm, y_normals=mesh_norm)
+    chamfer = chamfer_distance(pc_pc, mesh_gt_pc)
 
     return chamfer
 
@@ -73,8 +75,8 @@ def train_step(batch):
 
     # 2. Calculate loss
     mse = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
-    chamfer, _ = chamfer_loss(y_pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
+    chamfer, normal = chamfer_loss(y_pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
-    loss = mse + 1e-1 * chamfer
+    loss = mse + 1e-1 * (chamfer + normal)
 
     # 3. Zero grad optimizer
     optimizer.zero_grad()
@@ -90,15 +92,15 @@ def train_step(batch):
 
 def valid_step(batch):
     ### Testing
-    model.eval() # put the model in evaluation mode for testing (inference)
+    model.eval() # put the model in evaluation mode for testing (inference) 
     # 1. Forward pass
     with torch.inference_mode():
         pred = model([batch['rgb_image'], batch['view']])
 
         # 2. Calculate the loss
         mse = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
-        chamfer, _ = chamfer_loss(pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
+        chamfer, normal = chamfer_loss(pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
-        loss = mse + 1e-1 * chamfer
+        loss = mse + 1e-1 * (chamfer + normal)
 
     return { 'valid/loss': loss.detach().cpu().numpy() }