Spaces:
Sleeping
Sleeping
File size: 5,089 Bytes
15d6c34 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
import pyrender
from tqdm import tqdm
import trimesh
import numpy as np
import os
import imageio
def save_gif(gif_path, gif_frames, duration=0.01):
if gif_frames:
print(f"Saving GIF with {len(gif_frames)} frames to {gif_path}")
imageio.mimsave(uri=gif_path, ims=gif_frames, duration=duration)
else:
print("No frames to save.")
def render_meshes(model, output, should_save_gif=False, gif_path=None):
should_display = not should_save_gif
vertices_list = output.vertices.detach().cpu().numpy().squeeze()
joints_list = output.joints.detach().cpu().numpy().squeeze()
if len(vertices_list.shape) == 2:
vertices_list = [vertices_list]
joints_list = [joints_list]
scene = pyrender.Scene()
if should_display:
viewer = pyrender.Viewer(scene, run_in_thread=True)
mesh_node = None
joints_node = None
# Rotation matrix (90 degrees around the X-axis)
rot = trimesh.transformations.rotation_matrix(np.radians(90), [1, 0, 0])
gif_frames = []
if should_save_gif:
os.makedirs(os.path.dirname(gif_path), exist_ok=True)
try:
for i in tqdm(range(len(vertices_list))):
vertices = vertices_list[i]
joints = joints_list[i]
# print("Vertices shape =", vertices.shape)
# print("Joints shape =", joints.shape)
# from their demo script
plotting_module = "pyrender"
plot_joints = False
if plotting_module == "pyrender":
vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8]
tri_mesh = trimesh.Trimesh(vertices, model.faces, vertex_colors=vertex_colors)
# Apply rotation
tri_mesh.apply_transform(rot)
##### RENDER LOCK #####
if should_display:
viewer.render_lock.acquire()
if mesh_node:
scene.remove_node(mesh_node)
mesh = pyrender.Mesh.from_trimesh(tri_mesh)
mesh_node = scene.add(mesh)
camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=1.0)
min_bound, max_bound = mesh.bounds
# Calculate the center of the bounding box
center = (min_bound + max_bound) / 2
# Calculate the extents (the dimensions of the bounding box)
extents = max_bound - min_bound
# Estimate a suitable distance
distance = max(extents) * 2 # Adjust the multiplier as needed
# Create a camera pose matrix
cam_pose = np.array(
[
[1.0, 0, 0, center[0]],
[0, 1.0, 0, center[1]-1.0],
[0, 0, 1.0, center[2] + distance + 0.5],
[0, 0, 0, 1],
]
)
# Rotate around X-axis
angle = np.radians(80)
cos_angle = np.cos(angle)
sin_angle = np.sin(angle)
rot_x_10_deg = np.array([
[1, 0, 0, 0],
[0, cos_angle, -sin_angle, 0],
[0, sin_angle, cos_angle, 0],
[0, 0, 0, 1]
])
# rotate cam_pose with rot_x_10_deg
cam_pose = np.matmul(cam_pose, rot_x_10_deg)
cam_pose[:3, 3] += np.array([0, -2.2, -3.0])
scene.add(camera, pose=cam_pose)
# Add light for better visualization
light = pyrender.DirectionalLight(color=np.ones(3), intensity=2.0)
scene.add(light, pose=cam_pose)
# TODO: rotation doesn't work here, so appears sideways
if plot_joints:
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
for i, joint in enumerate(joints):
tfs[i, :3, :3] = rot[:3, :3]
tfs[i, :3, 3] = joint
joints_pcl = pyrender.Mesh.from_trimesh(sm, poses=tfs)
if joints_node:
scene.remove_node(joints_node)
joints_node = scene.add(joints_pcl)
if should_save_gif:
r = pyrender.OffscreenRenderer(viewport_width=640, viewport_height=480)
color, _ = r.render(scene)
gif_frames.append(color)
r.delete() # Free up the resources
###### RENDER LOCK RELEASE #####
if should_display:
viewer.render_lock.release()
except KeyboardInterrupt:
if should_display:
viewer.close_external()
save_gif(gif_path, gif_frames)
finally:
save_gif(gif_path, gif_frames)
|