101 lines
3.7 KiB
Python
101 lines
3.7 KiB
Python
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([3.0, 0.0, 0.0], device=device).cpu().numpy()
|
|
|
|
|
|
def forward(self, verts: torch.Tensor, faces: torch.Tensor):
|
|
cameras, K = get_camera(self.cam, [512, 512], verts.get_device(), res_x=512, res_y=512)
|
|
|
|
raster_settings = RasterizationSettings(
|
|
image_size=(512, 512),
|
|
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 |