55 lines
1.7 KiB
Python
55 lines
1.7 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')
|
|
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 |