add pytorch3d rendering

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