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 }

data/render.py

@@ -81,6 +81,7 @@ def intrinsic():
 					 [0,     fy_px, v],
 					 [0,     0,     1]], np.float32)
 				
+
 def extrinsic(camLoc):
 	R_w2bc = np.array([[0, 1, 0],
 					   [0, 0, 1],
@@ -121,4 +122,57 @@ class Render:
         depth[~mask] = 0.0
         mask = mask.astype(np.float32)
 
-        return depth
+        return depth
+
+
+from torch import nn
+from pytorch3d.renderer import (
+    look_at_view_transform,
+    FoVPerspectiveCameras, 
+    PerspectiveCameras,
+    PointLights, 
+    DirectionalLights, 
+    Materials, 
+    RasterizationSettings, 
+    MeshRenderer, 
+    MeshRasterizer,  
+    SoftPhongShader,
+    TexturesUV,
+    TexturesVertex
+)
+from pytorch3d.structures import Meshes
+
+class DepthRenderer(nn.Module):
+    def __init__(self) -> None:
+         super().__init__()
+
+
+
+    def forward(self, verts, faces, cam):
+        R, T = look_at_view_transform(2.7, 0, 0) 
+        cameras = FoVPerspectiveCameras(device='cuda', R=R, T=T)
+        ext = extrinsic(cam)
+        cameras = PerspectiveCameras(device='cuda', R=ext[None, :, :3], T=ext[None, :, 3], focal_length=50)
+
+        raster_settings = RasterizationSettings(
+                            image_size=512, 
+                            blur_radius=0.0, 
+                            faces_per_pixel=1, 
+                            bin_size=0
+                        )
+
+        rasterizer = MeshRasterizer(
+                    cameras=cameras, 
+                    raster_settings=raster_settings
+                )
+        depth = rasterizer(Meshes(verts[None,:,:], faces[None,:,:])).zbuf[0,:,:,0]
+
+        mask = depth != -1
+        # dmin = depth.min()
+        dmean = depth[mask].mean()
+        # dmax = depth[mask].max()
+        depth[mask] = depth[mask] - dmean
+        depth[~mask] = 0.0
+        # mask = mask.astype(np.float32)
+
+        return depth