mirror of
https://github.com/gosticks/body-pose-animation.git
synced 2025-10-16 11:45:42 +00:00
128 lines
3.8 KiB
Python
128 lines
3.8 KiB
Python
|
|
import smplx
|
|
import torch
|
|
import numpy as np
|
|
import trimesh
|
|
import pyrender
|
|
|
|
|
|
class SMPLyRenderer():
|
|
def __init__(
|
|
self,
|
|
):
|
|
print("TODO: do some preparation here")
|
|
|
|
# TODO: use __call__ for this
|
|
def render_model(
|
|
self,
|
|
model,
|
|
betas: torch.TensorType,
|
|
body_pose: torch.TensorType,
|
|
):
|
|
model_out = model()
|
|
# TODO: check if this also works with CUDA
|
|
vertices = model_out.vertices.detach().cpu().numpy().squeeze()
|
|
joints = model_out.joints.detach().cpu().numpy().squeeze()
|
|
|
|
# set vertex colors, maybe use this to highlight accuracies
|
|
vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8]
|
|
|
|
# triangulate vertex mesh
|
|
tri_mesh = trimesh.Trimesh(vertices, model.faces,
|
|
vertex_colors=vertex_colors)
|
|
|
|
return (tri_mesh, joints, vertices)
|
|
|
|
def set_keypoints(self, keypoints):
|
|
scene = self.scene
|
|
self.viewer.render_lock.acquire()
|
|
|
|
if self.keypoints_node is not None:
|
|
scene.remove(self.keypoints_node)
|
|
|
|
sm = trimesh.creation.uv_sphere(radius=0.01)
|
|
sm.visual.vertex_colors = [0.0, 0.0, 1.0, 1.0]
|
|
tfs = np.tile(np.eye(4), (len(keypoints), 1, 1))
|
|
tfs[:, :3, 3] = keypoints
|
|
keypoints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
|
|
self.keypoints_node = scene.add(keypoints_pcl, name="keypoints")
|
|
|
|
self.viewer.render_lock.release()
|
|
|
|
def set_model(self, model):
|
|
scene = self.scene
|
|
self.viewer.render_lock.acquire()
|
|
|
|
if self.keypoints_node is not None:
|
|
scene.remove(self.keypoints_node)
|
|
|
|
sm = trimesh.creation.uv_sphere(radius=0.01)
|
|
sm.visual.vertex_colors = [0.0, 0.0, 1.0, 1.0]
|
|
tfs = np.tile(np.eye(4), (len(keypoints), 1, 1))
|
|
tfs[:, :3, 3] = keypoints
|
|
keypoints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
|
|
self.keypoints_node = scene.add(keypoints_pcl, name="keypoints")
|
|
|
|
self.viewer.render_lock.release()
|
|
|
|
def display_mesh(
|
|
self,
|
|
tri_mesh: trimesh.Trimesh,
|
|
joints=None,
|
|
keypoints=None,
|
|
render_openpose_wireframe=True,
|
|
):
|
|
mesh = pyrender.Mesh.from_trimesh(tri_mesh)
|
|
|
|
scene = pyrender.Scene()
|
|
self.body_node = scene.add(mesh, name="body_mesh")
|
|
if joints is not None:
|
|
sm = trimesh.creation.uv_sphere(radius=0.005)
|
|
sm.visual.vertex_colors = [0.9, 0.1, 0.1, 1.0]
|
|
tfs = np.tile(np.eye(4), (len(joints), 1, 1))
|
|
tfs[:, :3, 3] = joints
|
|
joints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
|
|
self.joints_node = scene.add(joints_pcl, name="joints")
|
|
|
|
if keypoints is not None:
|
|
|
|
self.start_render(scene)
|
|
|
|
def start_render(
|
|
self,
|
|
scene,
|
|
):
|
|
self.scene = scene
|
|
self.viewer = pyrender.Viewer(scene,
|
|
use_raymond_lighting=True,
|
|
# show_world_axis=True
|
|
run_in_thread=True
|
|
)
|
|
# self.update()
|
|
while True:
|
|
pass
|
|
|
|
def update(self):
|
|
pose = self.body_node.get_pose()
|
|
print(pose)
|
|
self.viewer.render_lock.acquire()
|
|
self.scene.set_pose(self.body_node, pose)
|
|
self.viewer.render_lock.release()
|
|
|
|
def end_scene(self):
|
|
self.viewer.close_external()
|
|
while self.viewer.is_active:
|
|
pass
|
|
|
|
def display_model(
|
|
self,
|
|
model,
|
|
betas: torch.TensorType,
|
|
body_pose: torch.TensorType,
|
|
keypoints=None,
|
|
):
|
|
(tri_mesh, joints, vertices) = self.render_model(
|
|
model, betas, body_pose)
|
|
|
|
self.display_mesh(tri_mesh, joints, keypoints)
|