create pose coloring and rendering

data/pose_renderer.py

@@ -0,0 +1,55 @@
+import os
+import sys
+import numpy as np
+import torch
+from torch.nn import Module
+
+from utils.data_utils import intrinsic, extrinsic, get_camera
+from pytorch3d.renderer import (
+    RasterizationSettings, 
+    MeshRenderer, 
+    MeshRasterizer,  
+    SoftPhongShader,
+    TexturesVertex,
+    HardPhongShader,
+)
+from pytorch3d.structures import Meshes
+
+from pytorch3d.renderer import (
+    RasterizationSettings, 
+    MeshRenderer, 
+    MeshRasterizer,  
+    SoftPhongShader,
+    TexturesVertex,
+    HardPhongShader,
+    AmbientLights
+)
+from pytorch3d.structures import Meshes
+
+
+class PoseRenderer(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([5.0, 0.0, 0.0], device=device).cpu().numpy()
+
+
+    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.device, cameras=cameras, lights=AmbientLights(device=verts.device))
+        )
+
+        output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures))
+
+        return output