general save

data/render.py

@@ -3,10 +3,10 @@ import sys
 import torch
 import numpy as np
 
-from deodr import Camera, Scene3D
-from deodr.triangulated_mesh import TriMesh
-from deodr.pytorch.triangulated_mesh_pytorch import ColoredTriMeshPytorch
-from deodr.pytorch.differentiable_renderer_pytorch import Scene3DPytorch, CameraPytorch
+# from deodr import Camera, Scene3D
+# from deodr.triangulated_mesh import TriMesh
+# from deodr.pytorch.triangulated_mesh_pytorch import ColoredTriMeshPytorch
+# from deodr.pytorch.differentiable_renderer_pytorch import Scene3DPytorch, CameraPytorch
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 from cloth3d.util import proj
@@ -106,7 +106,8 @@ from pytorch3d.renderer import (
     MeshRasterizer,  
     SoftPhongShader,
     TexturesUV,
-    TexturesVertex
+    TexturesVertex,
+    HardPhongShader,
 )
 from pytorch3d.structures import Meshes
 from pytorch3d.utils import cameras_from_opencv_projection
@@ -162,3 +163,33 @@ class DepthRenderer(nn.Module):
         # ex.imshow(depth.cpu().numpy()).show()
 
         return depth, mask, v2i, K, cameras, pix_to_face
+
+
+class PoseRenderer(nn.Module):
+    def __init__(self, device='cuda') -> None:
+        super().__init__()
+        colors = np.load('custom_data/pose_colors.npy')
+        self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
+        self.cam = torch.tensor([0.0, 0.0, 1.0], device=device)
+
+
+    def forward(self, verts: torch.Tensor, faces: torch.Tensor):
+        cameras, K = get_camera(self.cam, [480, 640], verts.get_device())
+
+        raster_settings = RasterizationSettings(
+                            image_size=(480, 640), 
+                            blur_radius=0.0, 
+                            faces_per_pixel=1, 
+                            bin_size=0
+                        )
+        
+        renderer = MeshRenderer(
+            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
+            shader=HardPhongShader(device=verts.get_device(), cameras=cameras)
+        )
+        # create a pytorch3d vertex texture from the colors
+
+        output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures))
+
+
+        return output