add pytorch3d rendering

data/dataset.py

@@ -13,7 +13,7 @@ import plotly.express as ex
 
 from torch.utils.data import Dataset
 
-from render import Render
+from render import Render, DepthRenderer
 
 
 class Dataset(Dataset):
@@ -22,7 +22,7 @@ class Dataset(Dataset):
         self.transforms = tv.transforms.ToTensor()
         self.frames = os.listdir(os.path.join(folder, 'mesh_posed'))
 
-        self.render = Render(max_depth=10)
+        self.render = DepthRenderer()
         self.body_faces = np.load('custom_data/smpl_faces.npy')
         
 
@@ -41,12 +41,15 @@ class Dataset(Dataset):
 
         info_folder = '/home/workspace/data/dataset/' + ('train' if self.folder.split('/')[-1] == 'train' else 'validation') + '/' + id + '/info.mat'
         info = sio.loadmat(info_folder)
-        combined = np.concatenate((body, mesh_posed[mask, :]))
-        combined_faces = np.concatenate((self.body_faces, faces + 14475))
+        combined = torch.Tensor(np.concatenate((body, mesh_posed[mask, :]))).to('cuda')
+        combined_faces = torch.Tensor(np.concatenate((self.body_faces, faces + 14475))).to('cuda')
 
         start = time.perf_counter()
-        depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0]
+        # 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]
+        depth_image = self.render(combined, combined_faces, info['camLoc'][0])
         t = time.perf_counter() - start
+        print(t)
 
         # image = self.transforms(image)
         return { 'mesh_posed': mesh_posed, 'depth_image': depth_image, 'mask': mask }