body coloring trials

data/pose_renderer.py

@@ -30,7 +30,9 @@ from pytorch3d.structures import Meshes
 class PoseRenderer(Module):
     def __init__(self, device='cuda') -> None:
         super().__init__()
-        colors = np.load('custom_data/pose_colors.npy')
+        # colors = np.load('custom_data/pose_colors.npy')
+        colors = np.load('custom_data/body_colors_interpolated_unweighted.npy')
+        
         self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
         self.cam = torch.tensor([3.0, 0.0, 0.0], device=device).cpu().numpy()
 
@@ -52,4 +54,30 @@ class PoseRenderer(Module):
 
         output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures))
 
-        return output
+        return output
+    
+if __name__ == '__main__':
+    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+    os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+    from superSMPL.smpl_torch import SMPL
+    import plotly.express as ex
+    renderer = PoseRenderer()
+    smpl = SMPL(path='superSMPL', device='cuda')
+    pose = torch.rand(72, dtype=torch.float32, device='cuda')
+    pose[:3] = 0
+    shape = torch.rand(11, dtype=torch.float32, device='cuda')
+    # pose_out[1] = 0.5 # np.random.uniform(-np.pi, np.pi)
+    # generate random number between -np.pi, np.pi
+    
+    # start = time()
+
+    body = smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False)
+    body_to_render = body[0,:,[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=body_to_render.device), body_to_render.T).T
+
+    out = renderer(body_to_render, torch.from_numpy(smpl.f_all).to(body_to_render.get_device()))
+
+    ex.imshow(out[0, ..., :3].cpu().numpy()).show()
+
+    print(out)