diff --git a/custom_data/pose_colors.npy b/custom_data/pose_colors.npy new file mode 100644 index 0000000..aebad44 Binary files /dev/null and b/custom_data/pose_colors.npy differ diff --git a/data/pose_renderer.py b/data/pose_renderer.py new file mode 100644 index 0000000..fafb926 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/utils/pose_coloring.py b/utils/pose_coloring.py new file mode 100644 index 0000000..7dae36e --- /dev/null +++ b/utils/pose_coloring.py @@ -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 \ No newline at end of file