create pose coloring and rendering

2023-05-09 11:17:33 +00:00
parent cb779f38d9
commit 4081496c24
3 changed files with 156 additions and 0 deletions
@@ -0,0 +1,101 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+import torch
+from torch.nn import Module
+import numpy as np
+from superSMPL.smpl_torch import SMPL
+import trimesh
+
+import plotly.express as ex
+import plotly.graph_objects as go
+
+
+from data_utils import intrinsic, extrinsic, get_camera
+
+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
+    
+
+if __name__ == '__main__':
+
+    smpl = SMPL(path='superSMPL', device='cuda')
+    pose = torch.zeros((1,72), device='cuda')
+    # generate random pose
+    pose[:,3:] = torch.randn((1,69), device='cuda') * 0.1
+    body, _ = smpl(pose, torch.zeros((1,11), device='cuda'), remove=False)
+    colors = np.zeros((body.shape[1], 3))
+
+
+    renderer = PoseRenderer()
+    body_to_render = body[0]
+    # swap the y,z axes of the body_to_render 
+    body_to_render = body_to_render[:,[0,2,1]]
+    # and rotate it 90 degress along the z axis
+    body_to_render = torch.matmul(torch.tensor([[0,1,0],[-1,0,0],[0,0,1]], dtype=torch.float32, device='cuda'), body_to_render.T).T
+
+    out = renderer(body_to_render, torch.from_numpy(smpl.f_all).to(body.get_device()))
+    ex.imshow(out[0,:,:,:3].cpu().numpy()).show()
+
+    body = body.cpu().numpy()[0]
+    faces = smpl.f_all
+
+    # define a gradinet color cube for the body that goes from green to red based on the coordinates of the vertices and not their indices
+    # colors = np.zeros((body.shape[0], 3))
+    # colors[:,0] = (body[:,0] - np.min(body[:,0])) / (np.max(body[:,0]) - np.min(body[:,0]))
+    # colors[:,1] = (body[:,1] - np.min(body[:,1])) / (np.max(body[:,1]) - np.min(body[:,1]))
+    # colors[:,2] = (body[:,2] - np.min(body[:,2])) / (np.max(body[:,2]) - np.min(body[:,2]))
+    # fig = go.Figure(data=[go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], vertexcolor=colors)])
+    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+    # fig.show()  
+    # np.save('custom_data/pose_colors.npy', colors.astype(np.float32))
+
+    # same as a scatter plot
+    # fig = go.Figure(data=[go.Scatter3d(x=body[:,0], y=body[:,2], z=body[:,1], mode='markers', marker=dict(size=2, color=colors, opacity=0.8))])
+    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+    # fig.show()
+
+
+    
+    pass