88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
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')
|
|
colors_ = np.load('custom_data/body_colors_interpolated.npy').astype(np.float32) / 255.0
|
|
color_order = np.load('custom_data/color_ordered_indices.npy')
|
|
colors = np.zeros_like(colors_)
|
|
colors[color_order] = colors_[np.arange(27578)]
|
|
|
|
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=None,
|
|
max_faces_per_bin=None,
|
|
)
|
|
|
|
renderer = MeshRenderer(
|
|
rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
|
|
shader=HardPhongShader(device=verts.device, cameras=cameras, lights=AmbientLights(device=verts.device))
|
|
# shader=SoftPhongShader(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__':
|
|
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) |