Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,4 @@
|
||||
from . import viewer
|
||||
|
||||
# from . import render
|
||||
from . import tools
|
||||
+98
@@ -0,0 +1,98 @@
|
||||
241 91 181
|
||||
0 187 249
|
||||
254 228 64
|
||||
0 245 212
|
||||
222 158 54
|
||||
255 133 141
|
||||
93 169 233
|
||||
155 93 229
|
||||
24 242 178
|
||||
252 116 186
|
||||
138 225 252
|
||||
178 171 242
|
||||
24 231 98
|
||||
232 93 117
|
||||
207 250 231
|
||||
253 208 242
|
||||
109 157 197
|
||||
166 189 219
|
||||
253 146 207
|
||||
167 201 87
|
||||
117 68 177
|
||||
255 229 50
|
||||
251 202 239
|
||||
58 134 255
|
||||
255 0 110
|
||||
251 86 7
|
||||
188 51 209
|
||||
122 229 130
|
||||
0 48 73
|
||||
214 40 40
|
||||
229 179 179
|
||||
0 187 249
|
||||
255 190 11
|
||||
204 213 174
|
||||
0 245 212
|
||||
255 153 200
|
||||
144 251 146
|
||||
189 211 147
|
||||
230 0 86
|
||||
0 95 57
|
||||
0 174 126
|
||||
255 116 163
|
||||
189 198 255
|
||||
90 219 255
|
||||
158 0 142
|
||||
255 147 126
|
||||
164 36 0
|
||||
0 21 68
|
||||
145 208 203
|
||||
95 173 78
|
||||
107 104 130
|
||||
0 125 181
|
||||
106 130 108
|
||||
252 246 189
|
||||
208 244 222
|
||||
169 222 249
|
||||
228 193 249
|
||||
122 204 174
|
||||
194 140 159
|
||||
0 143 156
|
||||
235 0 0
|
||||
255 2 157
|
||||
104 61 59
|
||||
150 138 232
|
||||
152 255 82
|
||||
167 87 64
|
||||
1 255 254
|
||||
255 238 232
|
||||
254 137 0
|
||||
1 208 255
|
||||
187 136 0
|
||||
165 255 210
|
||||
255 166 254
|
||||
119 77 0
|
||||
122 71 130
|
||||
38 52 0
|
||||
0 71 84
|
||||
67 0 44
|
||||
181 0 255
|
||||
255 177 103
|
||||
255 219 102
|
||||
126 45 210
|
||||
229 111 254
|
||||
222 255 116
|
||||
0 255 120
|
||||
0 155 255
|
||||
0 100 1
|
||||
0 118 255
|
||||
133 169 0
|
||||
0 185 23
|
||||
120 130 49
|
||||
0 255 198
|
||||
255 110 65
|
||||
232 94 190
|
||||
1 0 103
|
||||
149 0 58
|
||||
98 14 0
|
||||
0 0 0
|
||||
+245
@@ -0,0 +1,245 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def get_seq_figure_skip(seq_name=None):
|
||||
if seq_name == "017437_mpii_test":
|
||||
return 15
|
||||
if seq_name == "012968_mpii_test":
|
||||
return 15
|
||||
if seq_name == "recording_20211002_S03_S18_04-all-100-200":
|
||||
return 11
|
||||
return 10
|
||||
|
||||
|
||||
def get_seq_static_lookat_points(seq_name=None, bounds=None):
|
||||
if seq_name == "002276_mpii_test":
|
||||
top_source = torch.tensor([-0.5, -3.0, -2.0])
|
||||
top_target = torch.tensor([0.0, 0.0, 4.0])
|
||||
|
||||
side_source = torch.tensor([6.0, -2.0, 2.0])
|
||||
side_target = torch.tensor([2.0, 0.0, 6.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "002374_mpii_test":
|
||||
top_source = torch.tensor([-4.0, 0.0, 0.0])
|
||||
top_target = torch.tensor([-2.0, 2.0, 3.0])
|
||||
|
||||
side_source = torch.tensor([-7.0, 1.0, 7.0])
|
||||
side_target = torch.tensor([0.0, 1.0, 8.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "003742_mpii_test":
|
||||
top_source = torch.tensor([3.0, -2.0, -6])
|
||||
top_target = torch.tensor([3.0, 1.0, 2.0])
|
||||
|
||||
side_source = torch.tensor([-5.0, -1.0, -2])
|
||||
side_target = torch.tensor([2.0, 0.0, 5])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "003747_mpii_test":
|
||||
top_source = torch.tensor([-5.0, -2.0, -5])
|
||||
top_target = torch.tensor([-3.0, 2.0, 1.0])
|
||||
|
||||
side_source = torch.tensor([-9.0, -1.0, -4])
|
||||
side_target = torch.tensor([-3.0, 2.0, 3])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "003943_mpii_test":
|
||||
top_source = torch.tensor([0, -1, -2])
|
||||
top_target = torch.tensor([0, 1, 3])
|
||||
|
||||
side_source = torch.tensor([6, 0, 1])
|
||||
side_target = torch.tensor([0, 0, 5])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "006537_mpii_test":
|
||||
top_source = torch.tensor([0, -4, -1])
|
||||
top_target = torch.tensor([1.0, 0.0, 5.0])
|
||||
|
||||
side_source = torch.tensor([5, -2, 0])
|
||||
side_target = torch.tensor([1.0, 1.0, 7.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "007684_mpii_test":
|
||||
top_source = torch.tensor([-2, -4, -6])
|
||||
top_target = torch.tensor([0.5, 0.0, 4.0])
|
||||
|
||||
side_source = torch.tensor([-8, -3, -2])
|
||||
side_target = torch.tensor([1.0, 0.0, 5.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "009039_mpii_test":
|
||||
top_source = torch.tensor([-2.0, -3.0, -1.5])
|
||||
top_target = torch.tensor([1.0, 0.0, 3.0])
|
||||
|
||||
side_source = torch.tensor([6, -2, -1.5])
|
||||
side_target = torch.tensor([2.0, 0.0, 6.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "009607_mpii_test":
|
||||
top_source = torch.tensor([0, -3, -1])
|
||||
top_target = torch.tensor([0, 0, 5])
|
||||
|
||||
side_source = torch.tensor([-6, -0.5, 2])
|
||||
side_target = torch.tensor([-1.0, 0.5, 7])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "012968_mpii_test":
|
||||
top_source = torch.tensor([-0.5, -3, -4])
|
||||
top_target = torch.tensor([0.0, 0.5, 2.0])
|
||||
|
||||
side_source = torch.tensor([7, -4, 0])
|
||||
side_target = torch.tensor([0, 0.5, 2.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "014531_mpii_test":
|
||||
top_source = torch.tensor([0, -2, -2])
|
||||
top_target = torch.tensor([0, 0, 2])
|
||||
|
||||
side_source = torch.tensor([-2, 0, -2])
|
||||
side_target = torch.tensor([0, 0, 2])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "015933_mpii_test":
|
||||
top_source = torch.tensor([0.0, 0.5, -5.0])
|
||||
top_target = torch.tensor([0.0, 2.0, -2.0])
|
||||
|
||||
side_source = torch.tensor([5.0, 1.5, -3.0])
|
||||
side_target = torch.tensor([0.0, 2.5, 1.5])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "016195_mpii_test":
|
||||
top_source = torch.tensor([0, -1, -2])
|
||||
top_target = torch.tensor([0, 1, 1])
|
||||
|
||||
side_source = torch.tensor([-3, 0, 1])
|
||||
side_target = torch.tensor([0, 1, 2])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "018061_mpii_test":
|
||||
top_source = torch.tensor([4, -1, 2])
|
||||
top_target = torch.tensor([1, 1, 3])
|
||||
|
||||
side_source = torch.tensor([-3, 0, 4])
|
||||
side_target = torch.tensor([1, 1, 3])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "018713_mpii_test":
|
||||
top_source = torch.tensor([0, -1, -1])
|
||||
top_target = torch.tensor([0, 0.5, 2])
|
||||
|
||||
side_source = torch.tensor([3, 0, 1])
|
||||
side_target = torch.tensor([0, 0, 3])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "022691_mpii_test":
|
||||
top_source = torch.tensor([0, -4, -4])
|
||||
top_target = torch.tensor([-2.0, 0.0, 1.0])
|
||||
|
||||
side_source = torch.tensor([-5, -2, -5])
|
||||
side_target = torch.tensor([-3.5, 0.0, 1.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "023390_mpii_test":
|
||||
top_source = torch.tensor([0.0, -1.5, -1.0])
|
||||
top_target = torch.tensor([0.0, 0.0, 2.0])
|
||||
|
||||
side_source = torch.tensor([4.0, -0.5, 0])
|
||||
side_target = torch.tensor([1.0, 0.5, 4])
|
||||
# top_source = torch.tensor([0.0, -3.0, -2.0])
|
||||
# top_target = torch.tensor([0.0, 0.5, 2.0])
|
||||
|
||||
# side_source = torch.tensor([5.0, -0.5, -3])
|
||||
# side_target = torch.tensor([1.0, 0.5, 4])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name=="024165_mpii_test":
|
||||
top_source = torch.tensor([0, -3, -3])
|
||||
top_target = torch.tensor([0, 0, 3])
|
||||
|
||||
side_source = torch.tensor([-4, -1, -2])
|
||||
side_target = torch.tensor([0, 0, 3])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "024154_mpii_test":
|
||||
top_source = torch.tensor([-1, -0.5, 1])
|
||||
top_target = torch.tensor([-1, 1, 8])
|
||||
|
||||
side_source = torch.tensor([-4, 1, 3])
|
||||
side_target = torch.tensor([-1, 1, 8])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
# if seq_name == "024159_mpii_test":
|
||||
# top_source = torch.tensor([-4.0, -2.0, -6.0])
|
||||
# top_target = torch.tensor([-4.0, 0.0, 0.0])
|
||||
|
||||
# side_source = torch.tensor([-12, -1, -2])
|
||||
# side_target = torch.tensor([-10.0, 0.0, 0.0])
|
||||
# return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "023962_mpii_test":
|
||||
top_source = torch.tensor([0, -0.5, -2])
|
||||
top_target = torch.tensor([0, 0.5, 2])
|
||||
|
||||
side_source = torch.tensor([3, 0.5, 0])
|
||||
side_target = torch.tensor([0, 0, 2])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "017437_mpii_test":
|
||||
top_source = torch.tensor([0, -8, -6])
|
||||
top_target = torch.tensor([-2.0, -1.0, 2.0])
|
||||
|
||||
side_source = torch.tensor([-5, -1, -7])
|
||||
side_target = torch.tensor([-3.0, -1.0, 1.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "recording_20211002_S03_S18_04-all-100-200":
|
||||
top_source = torch.tensor([1.0, -1, -4.5])
|
||||
top_target = torch.tensor([1.0, 0.0, 1])
|
||||
|
||||
side_source = torch.tensor([1.0, 0.5, -4.5])
|
||||
side_target = torch.tensor([1.0, 0.5, 1])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "recording_20211002_S03_S18_04-all-100-500":
|
||||
top_source = torch.tensor([0.0, -0.5, -3])
|
||||
top_target = torch.tensor([0.0, 1.0, 1])
|
||||
|
||||
side_source = torch.tensor([-2.0, 0.5, -3])
|
||||
side_target = torch.tensor([0.0, 1.0, 2])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "recording_20211002_S03_S18_01-all-700-862":
|
||||
top_source = torch.tensor([0.5, -2.0, -2])
|
||||
top_target = torch.tensor([0.5, -0.5, 2])
|
||||
|
||||
side_source = torch.tensor([2.0, -0.5, -2])
|
||||
side_target = torch.tensor([0.5, -0.5, 1])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "recording_20220315_S21_S30_02-all-0-100":
|
||||
top_source = torch.tensor([0.0, -1, -2])
|
||||
top_target = torch.tensor([0.0, 0.5, 1])
|
||||
|
||||
side_source = torch.tensor([-2.0, 0.0, -1])
|
||||
side_target = torch.tensor([0.0, 0.5, 1])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if bounds is not None:
|
||||
# (3), (3), (3)
|
||||
bb_min, bb_max, center = bounds
|
||||
print("SCENE BOUNDS", bb_min, bb_max, center)
|
||||
length = torch.abs(bb_max - bb_min).max()
|
||||
print(length)
|
||||
top_source = center + torch.tensor([0.0, -2.0, -0.9 * length])
|
||||
|
||||
side_source = center + torch.tensor([0.5 * length, -0.5, -0.7 * length])
|
||||
return (top_source, center), (side_source, center)
|
||||
|
||||
top_source = torch.tensor([0.0, -2.0, -3.0])
|
||||
top_target = torch.tensor([0.0, 0.0, 1.0])
|
||||
|
||||
side_source = torch.tensor([3.0, -1.0, -1.0])
|
||||
side_target = torch.tensor([0.0, 0.0, 1.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
+379
@@ -0,0 +1,379 @@
|
||||
import os
|
||||
import cv2
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
import subprocess
|
||||
|
||||
from body_model import run_mano
|
||||
from geometry import camera as cam_util
|
||||
from geometry.mesh import make_batch_mesh, save_mesh_scenes
|
||||
from geometry.plane import parse_floor_plane, get_plane_transform
|
||||
|
||||
from util.tensor import detach_all, to_torch, move_to
|
||||
|
||||
from .fig_specs import get_seq_figure_skip, get_seq_static_lookat_points
|
||||
from .tools import mesh_to_geometry, smooth_results
|
||||
|
||||
|
||||
def prep_result_vis(res, vis_mask, track_ids, body_model, temporal_smooth, smooth_trans=True):
|
||||
"""
|
||||
:param res (dict) with (B, T, *) tensor elements, B tracks and T frames
|
||||
:param vis_mask (B, T) with visibility of each track in each frame
|
||||
:param track_ids (B,) index of each track
|
||||
:param temporal_smooth (bool) whether to apply temporal smoothing
|
||||
:param smooth_trans (bool) whether to smooth translation (only used if temporal_smooth=True)
|
||||
"""
|
||||
print("RESULT FIELDS", res.keys())
|
||||
res = detach_all(res)
|
||||
# print(res.keys(), res["is_right"])
|
||||
with torch.no_grad():
|
||||
if temporal_smooth:
|
||||
if smooth_trans:
|
||||
print('running temporal smooth (including trans)')
|
||||
# Smooth everything including trans
|
||||
res["root_orient"], res["pose_body"], res['betas'], res["trans"] = smooth_results(res["root_orient"], res["pose_body"], res['betas'], res["is_right"], res["trans"])
|
||||
else:
|
||||
print('running temporal smooth (excluding trans)')
|
||||
# Smooth only pose and shape, not trans
|
||||
res["root_orient"], res["pose_body"], res['betas'], _ = smooth_results(res["root_orient"], res["pose_body"], res['betas'], res["is_right"])
|
||||
|
||||
world_smpl = run_mano(
|
||||
body_model,
|
||||
res["trans"],
|
||||
res["root_orient"],
|
||||
res["pose_body"],
|
||||
res["is_right"],
|
||||
res.get("betas", None),
|
||||
)
|
||||
|
||||
T_w2c = None
|
||||
floor_plane = None
|
||||
if "cam_R" in res and "cam_t" in res:
|
||||
T_w2c = cam_util.make_4x4_pose(res["cam_R"][0], res["cam_t"][0])
|
||||
if "floor_plane" in res:
|
||||
floor_plane = res["floor_plane"][0]
|
||||
return build_scene_dict(
|
||||
world_smpl,
|
||||
vis_mask,
|
||||
track_ids,
|
||||
T_w2c=T_w2c,
|
||||
floor_plane=floor_plane,
|
||||
)
|
||||
|
||||
|
||||
def build_scene_dict(
|
||||
world_smpl, vis_mask, track_ids, T_w2c=None, floor_plane=None, **kwargs
|
||||
):
|
||||
scene_dict = {}
|
||||
|
||||
# first get the geometry of the people
|
||||
# lists of length T with (B, V, 3), (F, 3), (B, 3)
|
||||
scene_dict["geometry"] = mesh_to_geometry(
|
||||
world_smpl["vertices"], world_smpl["joints"], world_smpl["l_faces"], world_smpl["r_faces"], world_smpl["is_right"], vis_mask, track_ids
|
||||
)
|
||||
|
||||
if T_w2c is None:
|
||||
T_w2c = torch.eye(4)[None]
|
||||
|
||||
T_c2w = torch.linalg.inv(T_w2c)
|
||||
# rotate the camera slightly down and translate back and up
|
||||
T = cam_util.make_4x4_pose(
|
||||
cam_util.rotx(-np.pi / 10), torch.tensor([0, -1, -2])
|
||||
).to(T_c2w.device)
|
||||
|
||||
scene_dict["cameras"] = {
|
||||
"src_cam": T_c2w,
|
||||
# "front": torch.einsum("ij,...jk->...ik", T, T_c2w),
|
||||
}
|
||||
|
||||
# Create ground plane based on hand mesh height
|
||||
# DISABLED: Ground plane visualization disabled to better see coordinate system
|
||||
# verts = scene_dict["geometry"][0] # Get vertices from geometry
|
||||
# if len(verts) > 0:
|
||||
# # Find minimum height across all frames and vertices
|
||||
# min_height = float('inf')
|
||||
# for frame_verts in verts:
|
||||
# if len(frame_verts) > 0:
|
||||
# frame_min = frame_verts[..., 1].min().item() # y-coordinate is height
|
||||
# min_height = min(min_height, frame_min)
|
||||
#
|
||||
# # Set ground plane further below minimum height
|
||||
# ground_offset = -0.5 # 20cm below minimum height
|
||||
# ground_height = min_height - ground_offset
|
||||
#
|
||||
# # Create ground plane transform
|
||||
# R = torch.eye(3) # Identity rotation (flat ground)
|
||||
# t = torch.tensor([0.0, ground_height, 0.0]) # Translate to ground height
|
||||
# scene_dict["ground"] = cam_util.make_4x4_pose(R, t)
|
||||
#
|
||||
# # Save ground mesh for Blender debugging
|
||||
# import trimesh
|
||||
# from vis.viewer import make_checkerboard
|
||||
# ground_mesh = make_checkerboard(color0=[0.9, 0.95, 1.0], color1=[0.7, 0.8, 0.85], up="y", alpha=1.0)
|
||||
# ground_mesh.apply_translation([0.0, ground_height, 0.0])
|
||||
# ground_mesh.export("ground_debug.obj")
|
||||
|
||||
# if floor_plane is not None:
|
||||
# # compute the ground transform
|
||||
# # use the first appearance of a track as the reference point
|
||||
# tid, sid = torch.where(vis_mask > 0)
|
||||
# idx = tid[torch.argmin(sid)]
|
||||
# root = world_smpl["joints"][idx, 0, 0].detach().cpu()
|
||||
# floor = parse_floor_plane(floor_plane.detach().cpu())
|
||||
# R, t = get_plane_transform(torch.tensor([0.0, 1.0, 0.0]), floor, root)
|
||||
# scene_dict["ground"] = cam_util.make_4x4_pose(R, t)
|
||||
|
||||
return scene_dict
|
||||
|
||||
|
||||
# def render_scene_dict(renderer, scene_dict, out_name, fps=30, **kwargs):
|
||||
# # lists of T (B, V, 3), (B, 3), (F, 3)
|
||||
# verts, colors, faces, bounds = scene_dict["geometry"]
|
||||
# print("NUM VERTS", len(verts))
|
||||
|
||||
# # add a top view
|
||||
# scene_dict["cameras"]["above"] = cam_util.make_4x4_pose(
|
||||
# torch.eye(3), torch.tensor([0, 0, -10])
|
||||
# )[None]
|
||||
|
||||
# for cam_name, cam_poses in scene_dict["cameras"].items():
|
||||
# print("rendering scene for", cam_name)
|
||||
# # cam_poses are (T, 4, 4)
|
||||
# render_bg = cam_name == "src_cam"
|
||||
# ground_pose = scene_dict.get("ground", None)
|
||||
# print(cam_name, ground_pose)
|
||||
# frames = renderer.render_video(
|
||||
# cam_poses[None], verts, faces, colors, render_bg, ground_pose=ground_pose
|
||||
# )
|
||||
# os.makedirs(f"{out_name}_{cam_name}/", exist_ok=True)
|
||||
# for idx, i in enumerate(frames):
|
||||
# imageio.imwrite(f"{out_name}_{cam_name}/" + f'{str(idx).zfill(6)}.jpg', i)
|
||||
# imageio.mimwrite(f"{out_name}_{cam_name}.mp4", frames, fps=fps)
|
||||
|
||||
|
||||
def animate_scene(
|
||||
vis,
|
||||
scene,
|
||||
out_name,
|
||||
seq_name=None,
|
||||
accumulate=False,
|
||||
render_views=["src_cam", "front", "above", "side"],
|
||||
render_bg=True,
|
||||
render_cam=True,
|
||||
render_ground=True,
|
||||
debug=False,
|
||||
render_keypoints=False,
|
||||
**kwargs,
|
||||
):
|
||||
if len(render_views) < 1:
|
||||
return
|
||||
|
||||
scene = build_pyrender_scene(
|
||||
vis,
|
||||
scene,
|
||||
seq_name,
|
||||
render_views=render_views,
|
||||
render_cam=render_cam,
|
||||
accumulate=accumulate,
|
||||
debug=debug,
|
||||
)
|
||||
|
||||
print("RENDERING VIEWS", scene["cameras"].keys())
|
||||
render_ground = render_ground and "ground" in scene
|
||||
save_paths = []
|
||||
for cam_name, cam_poses in scene["cameras"].items():
|
||||
is_src = cam_name == "src_cam"
|
||||
show_bg = is_src and render_bg
|
||||
show_ground = render_ground and not is_src
|
||||
show_cam = render_cam and not is_src
|
||||
show_keypoints = is_src and render_keypoints # Only show keypoints on source view
|
||||
vis_name = f"{out_name}_{cam_name}"
|
||||
print(f"{cam_name} has {len(cam_poses)} poses")
|
||||
skip = 10 if debug else 1
|
||||
vis.set_camera_seq(cam_poses[::skip])
|
||||
save_path = vis.animate(
|
||||
vis_name,
|
||||
render_bg=show_bg,
|
||||
render_ground=show_ground,
|
||||
render_cam=show_cam,
|
||||
render_keypoints=show_keypoints,
|
||||
**kwargs,
|
||||
)
|
||||
save_paths.append(save_path)
|
||||
|
||||
return save_paths
|
||||
|
||||
|
||||
def build_pyrender_scene(
|
||||
vis,
|
||||
scene,
|
||||
seq_name,
|
||||
render_views=["src_cam", "front", "above", "side"],
|
||||
render_cam=True,
|
||||
accumulate=False,
|
||||
debug=False,
|
||||
):
|
||||
"""
|
||||
:param vis (viewer object)
|
||||
:param scene (scene dict with geometry, cameras, etc)
|
||||
:param accumulate (optional bool, default False) whether to render entire trajectory together
|
||||
:param render_views (list str) camera views to render
|
||||
"""
|
||||
if len(render_views) < 1:
|
||||
return
|
||||
|
||||
assert all(view in ["src_cam", "front", "above", "side"] for view in render_views)
|
||||
|
||||
scene = move_to(detach_all(scene), "cpu")
|
||||
src_cams = scene["cameras"]["src_cam"]
|
||||
verts, _, colors, l_faces, r_faces, is_right, bounds = scene["geometry"]
|
||||
T = len(verts)
|
||||
print(f"{T} mesh frames for {seq_name}, {len(verts)}")
|
||||
|
||||
# set camera views
|
||||
if not "cameras" in scene:
|
||||
scene["cameras"] = {}
|
||||
|
||||
# remove default views from source camera perspective if desired
|
||||
if "src_cam" not in render_views:
|
||||
scene["cameras"].pop("src_cam", None)
|
||||
if "front" not in render_views:
|
||||
scene["cameras"].pop("front", None)
|
||||
|
||||
# add static viewpoints if desired
|
||||
top_pose, side_pose, _skip = get_static_views(seq_name, bounds)
|
||||
if "above" in render_views:
|
||||
scene["cameras"]["above"] = top_pose[None]
|
||||
if "side" in render_views:
|
||||
scene["cameras"]["side"] = side_pose[None]
|
||||
if "front" in render_views:
|
||||
# Use a static front camera: place it in front of the hand mesh, looking at the center
|
||||
# Place the camera at a fixed distance along the -z axis from the center
|
||||
if bounds is not None:
|
||||
bb_min, bb_max, center = bounds
|
||||
length = torch.abs(bb_max - bb_min).max()
|
||||
front_offset = 1.0 # 1 closer, 2 farther
|
||||
front_source = center + torch.tensor([0.0, -0.5, -front_offset])
|
||||
front_target = center
|
||||
up = torch.tensor([0.0, 1, 0.0])
|
||||
front_pose = cam_util.lookat_matrix(front_source, front_target, up)
|
||||
scene["cameras"]["front"] = front_pose[None]
|
||||
else:
|
||||
# fallback to previous behavior if bounds not available
|
||||
pass
|
||||
|
||||
# accumulate meshes if possible (can only accumulate for static camera)
|
||||
moving_cam = "src_cam" in render_views or "front" in render_views
|
||||
accumulate = accumulate and not moving_cam
|
||||
skip = _skip if accumulate else 1
|
||||
|
||||
vis.clear_meshes()
|
||||
|
||||
if "ground" in scene:
|
||||
vis.set_ground(scene["ground"])
|
||||
|
||||
if debug:
|
||||
skip = 10
|
||||
times = list(range(0, T, skip))
|
||||
# print(times, verts[0].shape)
|
||||
# raise ValueError
|
||||
for t in times:
|
||||
if len(is_right[t]) > 1:
|
||||
assert (is_right[t].cpu().numpy().tolist() == [0,1])
|
||||
l_meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
r_meshes = make_batch_mesh(verts[t][1][None], r_faces[t], colors[t][1][None])
|
||||
assert len(l_meshes) == 1
|
||||
assert len(r_meshes) == 1
|
||||
meshes = [l_meshes[0], r_meshes[0]]
|
||||
else:
|
||||
assert len(is_right[t]) == 1
|
||||
if is_right[t] == 0:
|
||||
meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
elif is_right[t] == 1:
|
||||
meshes = make_batch_mesh(verts[t][0][None], r_faces[t], colors[t][0][None])
|
||||
|
||||
if accumulate:
|
||||
vis.add_static_meshes(meshes)
|
||||
else:
|
||||
vis.add_mesh_frame(meshes, debug=debug)
|
||||
|
||||
# add camera markers
|
||||
if render_cam:
|
||||
if accumulate:
|
||||
vis.add_camera_markers_static(src_cams[::skip])
|
||||
else:
|
||||
vis.add_camera_markers(src_cams[::skip])
|
||||
|
||||
return scene
|
||||
|
||||
|
||||
def get_static_views(seq_name=None, bounds=None):
|
||||
print("STATIC VIEWS FOR SEQ NAME", seq_name)
|
||||
up = torch.tensor([0.0, 1.0, 0.0])
|
||||
|
||||
skip = get_seq_figure_skip(seq_name)
|
||||
top_vp, side_vp = get_seq_static_lookat_points(seq_name, bounds)
|
||||
top_source, top_target = top_vp
|
||||
side_source, side_target = side_vp
|
||||
top_pose = cam_util.lookat_matrix(top_source, top_target, up)
|
||||
side_pose = cam_util.lookat_matrix(side_source, side_target, up)
|
||||
return top_pose, side_pose, skip
|
||||
|
||||
|
||||
def make_video_grid_2x2(out_path, vid_paths, fps, overwrite=False):
|
||||
if os.path.isfile(out_path) and not overwrite:
|
||||
print(f"{out_path} already exists, skipping.")
|
||||
return
|
||||
|
||||
if any(not os.path.isfile(v) for v in vid_paths):
|
||||
print("not all inputs exist!", vid_paths)
|
||||
return
|
||||
|
||||
# Open all videos for reading
|
||||
readers = [imageio.get_reader(v, 'ffmpeg') for v in vid_paths]
|
||||
try:
|
||||
# Check min frame count to avoid index errors; also get shape/fps
|
||||
lengths = [reader.count_frames() for reader in readers]
|
||||
min_len = min(lengths)
|
||||
if min_len == 0:
|
||||
print("Some videos have 0 frames!")
|
||||
for reader in readers:
|
||||
reader.close()
|
||||
return
|
||||
|
||||
# We assume the first video dictates the output fps and (half)size
|
||||
meta = readers[0].get_meta_data()
|
||||
# fps = meta.get('fps', 30)
|
||||
orig_size = meta['size'] # (width, height)
|
||||
width = orig_size[0] // 2
|
||||
height = orig_size[1] // 2
|
||||
out_size = (width * 2, height * 2)
|
||||
|
||||
writer = imageio.get_writer(out_path, fps=fps, codec='libx264', ffmpeg_log_level="quiet", macro_block_size=None)
|
||||
|
||||
for t in range(min_len):
|
||||
frames = []
|
||||
for r in readers:
|
||||
try:
|
||||
frame = r.get_data(t)
|
||||
except Exception:
|
||||
frame = None
|
||||
frames.append(frame)
|
||||
if any(f is None for f in frames):
|
||||
break
|
||||
# Resize all to (width, height)
|
||||
frames = [cv2.resize(f, (width, height)) for f in frames]
|
||||
row1 = np.hstack(frames[:2])
|
||||
row2 = np.hstack(frames[2:])
|
||||
grid = np.vstack([row1, row2])
|
||||
# Ensure uint8
|
||||
if grid.dtype != np.uint8:
|
||||
grid = np.clip(grid, 0, 255).astype(np.uint8)
|
||||
writer.append_data(grid)
|
||||
|
||||
writer.close()
|
||||
print(f"Video grid created: {out_path}")
|
||||
finally:
|
||||
for reader in readers:
|
||||
reader.close()
|
||||
+254
@@ -0,0 +1,254 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
|
||||
from pytorch3d.renderer import (
|
||||
RasterizationSettings,
|
||||
MeshRenderer,
|
||||
MeshRasterizer,
|
||||
HardPhongShader,
|
||||
PointLights,
|
||||
PerspectiveCameras,
|
||||
)
|
||||
from pytorch3d.structures import Meshes
|
||||
from pytorch3d.structures.meshes import join_meshes_as_scene
|
||||
from pytorch3d.renderer import TexturesAtlas, TexturesVertex
|
||||
from pytorch3d.transforms import euler_angles_to_matrix
|
||||
|
||||
from .tools import get_colors, read_image, transform_torch3d, checkerboard_geometry
|
||||
|
||||
|
||||
def init_renderer(img_size, intrins, device, vis_scale=1.0, bg_paths=None):
|
||||
img_size = int(vis_scale * img_size[0]), int(vis_scale * img_size[1])
|
||||
renderer = RenderBase(device, img_size, intrins=vis_scale * intrins)
|
||||
if bg_paths is not None:
|
||||
bg_imgs = [
|
||||
torch.from_numpy(read_image(p, scale=vis_scale)).float() / 255
|
||||
for p in bg_paths
|
||||
]
|
||||
print(bg_imgs[0].shape)
|
||||
renderer.set_bg_seq(bg_imgs)
|
||||
print("RENDERER", renderer)
|
||||
return renderer
|
||||
|
||||
|
||||
class RenderBase(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
device,
|
||||
img_size,
|
||||
intrins=None,
|
||||
blur_radius: float = 0.0, # rasterizer settings
|
||||
faces_per_pixel: int = 3,
|
||||
):
|
||||
"""
|
||||
Uses Pytorch3D to render results.
|
||||
We use the PyTorch3D coordinate system, which assumes:
|
||||
+X:left, +Y: up and +Z: from us to scene (right-handed)
|
||||
https://pytorch3d.org/docs/cameras
|
||||
"""
|
||||
super().__init__()
|
||||
self.device = device
|
||||
W, H = img_size
|
||||
self.img_size = (H, W)
|
||||
if intrins is not None:
|
||||
intrins = intrins.to(device)
|
||||
self.cam_f, self.cam_center = intrins[:2], intrins[2:]
|
||||
cx, cy = self.cam_center.detach().cpu()
|
||||
else:
|
||||
self.cam_f = torch.tensor([0.5 * (H + W), 0.5 * (H + W)], device=device)
|
||||
self.cam_center = torch.tensor([0.5 * W, 0.5 * H], device=device)
|
||||
print("cam_f", self.cam_f, "cam_center", self.cam_center)
|
||||
print("img_size", self.img_size)
|
||||
self.blur_radius = blur_radius
|
||||
self.faces_per_pixel = faces_per_pixel
|
||||
|
||||
self.light_loc = [[0, 0, -0.5]]
|
||||
self.cameras = self.make_cameras()
|
||||
self.bg_imgs = None
|
||||
self.set_ground()
|
||||
self.build_renderer()
|
||||
|
||||
def build_renderer(self):
|
||||
self.lights = PointLights(device=self.device, location=self.light_loc)
|
||||
|
||||
self.rasterizer = RasterizationSettings(
|
||||
image_size=self.img_size,
|
||||
blur_radius=self.blur_radius,
|
||||
faces_per_pixel=self.faces_per_pixel,
|
||||
max_faces_per_bin=100000,
|
||||
)
|
||||
self.renderer = MeshRenderer(
|
||||
rasterizer=MeshRasterizer(
|
||||
cameras=self.cameras, raster_settings=self.rasterizer
|
||||
),
|
||||
shader=HardPhongShader(
|
||||
cameras=self.cameras, lights=self.lights, device=self.device
|
||||
),
|
||||
)
|
||||
|
||||
def make_cameras(self, cam_poses=None):
|
||||
device = self.device
|
||||
if cam_poses is None:
|
||||
cam_poses = torch.eye(4, device=device)[None]
|
||||
cam_poses = cam_poses.to(device)
|
||||
cam_R, cam_t = transform_torch3d(cam_poses)
|
||||
Nc = len(cam_poses)
|
||||
focal = self.cam_f[None].repeat(Nc, 1)
|
||||
ppt = self.cam_center[None].repeat(Nc, 1)
|
||||
img_size = [self.img_size for _ in range(Nc)]
|
||||
return PerspectiveCameras(
|
||||
device=self.device,
|
||||
in_ndc=False,
|
||||
focal_length=focal,
|
||||
principal_point=ppt,
|
||||
image_size=img_size,
|
||||
R=cam_R,
|
||||
T=cam_t,
|
||||
)
|
||||
|
||||
def render_frame(self, cam_poses, verts, faces, colors):
|
||||
"""
|
||||
:param cam_poses (C, 4, 4)
|
||||
:param verts (B, V, 3)
|
||||
:param faces (F, 3)
|
||||
:param colors (B, 4)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
cameras = self.make_cameras(cam_poses)
|
||||
mesh = build_meshes(verts, faces, colors)
|
||||
image = self.renderer(mesh, cameras=cameras)
|
||||
return image
|
||||
|
||||
def render_with_ground(self, cam_poses, verts, faces, colors, ground_pose):
|
||||
"""
|
||||
:param cam_poses (C, 4, 4)
|
||||
:param verts (B, V, 3)
|
||||
:param faces (F, 3)
|
||||
:param colors (B, 3)
|
||||
:param ground_pose (4, 4)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
self.set_ground_pose(ground_pose)
|
||||
cameras = self.make_cameras(cam_poses)
|
||||
# (B, V, 3), (B, F, 3), (B, V, 3)
|
||||
verts, faces, colors = prep_shared_geometry(verts, faces, colors)
|
||||
# (V, 3), (F, 3), (V, 3)
|
||||
gv, gf, gc = self.ground_geometry
|
||||
verts = list(torch.unbind(verts, dim=0)) + [gv]
|
||||
faces = list(torch.unbind(faces, dim=0)) + [gf]
|
||||
colors = list(torch.unbind(colors, dim=0)) + [gc[..., :3]]
|
||||
mesh = create_meshes(verts, faces, colors)
|
||||
|
||||
image = self.renderer(mesh, cameras=cameras)
|
||||
return image
|
||||
|
||||
def set_bg_seq(self, bg_imgs):
|
||||
self.bg_imgs = bg_imgs
|
||||
|
||||
def set_ground(self, pose=None):
|
||||
device = self.device
|
||||
v, f, vc, fc = map(torch.from_numpy, checkerboard_geometry(up="y"))
|
||||
v, f, vc = v.to(device), f.to(device), vc.to(device)
|
||||
self.ground_geometry = [v, f, vc]
|
||||
if pose is not None:
|
||||
self.set_ground_pose(pose)
|
||||
|
||||
def set_ground_pose(self, pose):
|
||||
"""
|
||||
:param pose (4, 4)
|
||||
"""
|
||||
if self.ground_geometry is None:
|
||||
return
|
||||
v = self.ground_geometry[0]
|
||||
pose = pose.to(self.device)
|
||||
# R, t = pose[:3, :3], pose[:3, 3]
|
||||
R, t = transform_torch3d(pose)
|
||||
self.ground_geometry[0] = torch.einsum("ij,vj->vi", R, v) + t[None]
|
||||
|
||||
def composite_bg(self, frames, fac=1.0):
|
||||
T = len(frames)
|
||||
if self.bg_imgs is None:
|
||||
return frames
|
||||
|
||||
out_frames = []
|
||||
for bg, frame in zip(self.bg_imgs, frames):
|
||||
alpha = fac * frame[..., 3:]
|
||||
out = alpha * frame[..., :3] + (1 - alpha) * bg
|
||||
out_frames.append(out)
|
||||
return out_frames
|
||||
|
||||
def render_video(
|
||||
self,
|
||||
cam_poses,
|
||||
verts,
|
||||
faces,
|
||||
colors,
|
||||
render_bg=False,
|
||||
fac=1.0,
|
||||
ground_pose=None,
|
||||
):
|
||||
"""
|
||||
:param cam_poses (C, T, 4, 4)
|
||||
:param verts list of T (B, V, 3)
|
||||
:param faces list of T (F, 3)
|
||||
:param colors list of T (B, 4)
|
||||
:param render_bg (optional default False)
|
||||
:param fac (optional float, default 1)
|
||||
:param ground_pose (optional, default None) (4, 4)
|
||||
returns list of T (H, W, 4)
|
||||
"""
|
||||
T = len(verts)
|
||||
print(f"rendering video with {T} frames")
|
||||
if cam_poses.shape[1] == 1:
|
||||
cam_poses = cam_poses.expand(-1, T, -1, -1)
|
||||
|
||||
if render_bg: # don't render ground on source frame
|
||||
ground_pose = None
|
||||
|
||||
frames = []
|
||||
for t in range(T):
|
||||
if ground_pose is None:
|
||||
# (H, W, 4)
|
||||
frame = self.render_frame(
|
||||
cam_poses[:, t], verts[t], faces[t], colors[t]
|
||||
)
|
||||
else:
|
||||
# (H, W, 4)
|
||||
frame = self.render_with_ground(
|
||||
cam_poses[:, t], verts[t], faces[t], colors[t], ground_pose
|
||||
)
|
||||
frames.append(frame.cpu()[0])
|
||||
|
||||
if render_bg and self.bg_imgs is not None:
|
||||
frames = self.composite_bg(frames, fac=fac)
|
||||
|
||||
return [(255 * f).byte() for f in frames]
|
||||
|
||||
|
||||
def build_meshes(verts, faces, colors):
|
||||
return create_meshes(*prep_shared_geometry(verts, faces, colors))
|
||||
|
||||
|
||||
def prep_shared_geometry(verts, faces, colors):
|
||||
"""
|
||||
:param verts (B, V, 3)
|
||||
:param faces (F, 3)
|
||||
:param colors (B, 4)
|
||||
"""
|
||||
B, V, _ = verts.shape
|
||||
F, _ = faces.shape
|
||||
colors = colors.unsqueeze(1).expand(B, V, -1)[..., :3]
|
||||
faces = faces.unsqueeze(0).expand(B, F, -1)
|
||||
return verts, faces, colors
|
||||
|
||||
|
||||
def create_meshes(verts, faces, colors):
|
||||
"""
|
||||
:param verts (B, V, 3)
|
||||
:param faces (B, F, 3)
|
||||
:param colors (B, V, 3)
|
||||
"""
|
||||
textures = TexturesVertex(verts_features=colors)
|
||||
meshes = Meshes(verts=verts, faces=faces, textures=textures)
|
||||
return join_meshes_as_scene(meshes)
|
||||
+1109
File diff suppressed because it is too large
Load Diff
+672
@@ -0,0 +1,672 @@
|
||||
import os
|
||||
import imageio
|
||||
import numpy as np
|
||||
|
||||
import time
|
||||
import torch
|
||||
import trimesh
|
||||
|
||||
os.environ["PYOPENGL_PLATFORM"] = "egl"
|
||||
|
||||
import pyrender
|
||||
from pyrender.constants import RenderFlags
|
||||
from pyrender.light import DirectionalLight
|
||||
from pyrender.node import Node
|
||||
|
||||
from geometry.camera import make_rotation, make_translation
|
||||
from .tools import (
|
||||
read_image,
|
||||
checkerboard_geometry,
|
||||
camera_marker_geometry,
|
||||
transform_pyrender,
|
||||
)
|
||||
|
||||
|
||||
def init_viewer(
|
||||
img_size, intrins, vis_scale=1.0, bg_paths=None, fps=30,
|
||||
):
|
||||
img_size = int(vis_scale * img_size[0]), int(vis_scale * img_size[1])
|
||||
intrins = vis_scale * intrins
|
||||
|
||||
platform = os.environ.get("PYOPENGL_PLATFORM", "pyglet")
|
||||
if platform == "pyglet":
|
||||
vis = AnimationViewer(img_size, intrins=intrins, fps=fps)
|
||||
print("VIS", vis)
|
||||
return vis
|
||||
|
||||
if platform != "egl":
|
||||
raise NotImplementedError
|
||||
|
||||
vis = OffscreenAnimation(img_size, intrins=intrins, fps=fps)
|
||||
if bg_paths is not None:
|
||||
bg_imgs = [
|
||||
read_image(p, scale=vis_scale).astype(np.float32) / 255 for p in bg_paths
|
||||
]
|
||||
vis.set_bg_seq(bg_imgs)
|
||||
print("VIS", vis)
|
||||
return vis
|
||||
|
||||
|
||||
class AnimationBase(object):
|
||||
"""
|
||||
Render an animation
|
||||
IMPORTANT: pyrender uses coordinates with Y UP (i.e. xyz = right up back),
|
||||
we must keep this in mind because the human prior frame has Z UP (i.e.
|
||||
xyz = right forward up)
|
||||
"""
|
||||
|
||||
def __init__(self, img_size, intrins=None):
|
||||
self.scene = pyrender.Scene(
|
||||
ambient_light=[0.3, 0.3, 0.3], bg_color=[1.0, 1.0, 1.0, 0.0]
|
||||
)
|
||||
self.scene.bg_color = np.array([1.0, 1.0, 1.0, 0.0])
|
||||
|
||||
self.viewport_size = img_size
|
||||
self.camera = make_pyrender_camera(img_size, intrins)
|
||||
|
||||
# set the ground with Z up
|
||||
self.ground_pose = np.eye(4)
|
||||
ground = pyrender.Mesh.from_trimesh(
|
||||
make_checkerboard(color0=[0.9, 0.95, 1.0], color1=[0.7, 0.8, 0.85], up="y", alpha=1.0), smooth=False
|
||||
)
|
||||
self.ground_node = self.scene.add(ground, name="ground", pose=self.ground_pose)
|
||||
self.cam_node = self.scene.add(self.camera, name="camera")
|
||||
|
||||
self.cam_marker_node = None
|
||||
self.cam_marker_poses = []
|
||||
|
||||
self.light_nodes = []
|
||||
self.light_poses = []
|
||||
|
||||
self.bg_seq = []
|
||||
|
||||
# pyrender nodes, reused between frames
|
||||
self.anim_nodes = []
|
||||
# list of meshes per frame of animation
|
||||
self.anim_meshes = []
|
||||
# list of camera poses per frame of animation
|
||||
self.anim_cameras = None
|
||||
# current timestep of animation
|
||||
self.anim_idx = 0
|
||||
|
||||
def acquire_lock(self):
|
||||
pass
|
||||
|
||||
def release_lock(self):
|
||||
pass
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
def delete(self):
|
||||
pass
|
||||
|
||||
def clear_meshes(self):
|
||||
self.anim_meshes = []
|
||||
self.anim_idx = 0
|
||||
|
||||
@property
|
||||
def anim_len(self):
|
||||
return len(self.anim_meshes)
|
||||
|
||||
def add_lighting(self, color=np.ones(3), intensity=1.0):
|
||||
self.light_poses = get_light_poses()
|
||||
self.light_poses.append(np.eye(4))
|
||||
cam_pose = self.scene.get_pose(self.cam_node)
|
||||
for i, pose in enumerate(self.light_poses):
|
||||
matrix = cam_pose @ pose
|
||||
node = Node(
|
||||
name=f"light-{i:02d}",
|
||||
light=DirectionalLight(color=color, intensity=intensity),
|
||||
matrix=matrix,
|
||||
)
|
||||
if self.scene.has_node(node):
|
||||
continue
|
||||
self.scene.add_node(node)
|
||||
self.light_nodes.append(node)
|
||||
|
||||
def add_mesh_frame(self, meshes, debug=False):
|
||||
"""
|
||||
:param meshes (list) trimesh meshes for this frame timestep
|
||||
"""
|
||||
meshes = [pyrender.Mesh.from_trimesh(m) for m in meshes]
|
||||
num_nodes = len(meshes)
|
||||
self.anim_meshes.append(meshes)
|
||||
|
||||
if num_nodes > len(self.anim_nodes):
|
||||
# need to add a node in the pyrender scene
|
||||
self.acquire_lock()
|
||||
for i in range(len(self.anim_nodes), num_nodes):
|
||||
self.anim_nodes.append(self.scene.add(meshes[i], name=f"mesh_{i:03d}"))
|
||||
self.release_lock()
|
||||
|
||||
def add_static_meshes(self, meshes, smooth=True):
|
||||
"""
|
||||
add all meshes to a single frame (timestep)
|
||||
"""
|
||||
meshes = [pyrender.Mesh.from_trimesh(m, smooth=smooth) for m in meshes]
|
||||
if len(self.anim_meshes) < 1:
|
||||
self.anim_meshes.append([])
|
||||
self.anim_meshes[0].extend(meshes)
|
||||
num_meshes = len(self.anim_meshes[0])
|
||||
num_nodes = len(self.anim_nodes)
|
||||
if num_meshes > num_nodes:
|
||||
# need to add a node in the pyrender scene
|
||||
self.acquire_lock()
|
||||
self.anim_nodes.extend(
|
||||
[
|
||||
self.scene.add(self.anim_meshes[0][i])
|
||||
for i in range(num_nodes, num_meshes)
|
||||
]
|
||||
)
|
||||
self.release_lock()
|
||||
|
||||
def set_camera_seq(self, poses):
|
||||
"""
|
||||
:param poses (*, 4, 4)
|
||||
"""
|
||||
poses = transform_pyrender(poses)
|
||||
poses = np.asarray(poses)
|
||||
assert poses.ndim >= 2 and poses.shape[-2:] == (4, 4)
|
||||
if poses.ndim < 3:
|
||||
poses = poses[None]
|
||||
print(f"Adding camera sequence length {len(poses)}")
|
||||
self.anim_cameras = poses
|
||||
|
||||
def add_camera_markers(self, poses):
|
||||
"""
|
||||
Add a single camera marker node that we move around when we animate
|
||||
"""
|
||||
print("ADDING CAMERA MARKERS")
|
||||
if self.cam_marker_node is None:
|
||||
cam_marker = make_camera_marker(up="y")
|
||||
self.cam_marker_node = self.scene.add(
|
||||
pyrender.Mesh.from_trimesh(cam_marker, smooth=False)
|
||||
)
|
||||
poses = transform_pyrender(poses)
|
||||
poses = np.asarray(poses)
|
||||
self.cam_marker_poses = poses
|
||||
|
||||
def add_camera_markers_static(self, poses):
|
||||
"""
|
||||
Add all camera markers as separate mesh nodes to render in one pass
|
||||
"""
|
||||
print("ADDING STATIC CAMERA MARKERS")
|
||||
poses = transform_pyrender(poses)
|
||||
poses = np.asarray(poses)
|
||||
cam_meshes = [make_camera_marker(up="y", transform=pose) for pose in poses]
|
||||
self.add_static_meshes(cam_meshes, smooth=False)
|
||||
|
||||
def update_camera(self, cam_pose):
|
||||
self.acquire_lock()
|
||||
self.scene.set_pose(self.cam_node, pose=cam_pose)
|
||||
self.release_lock()
|
||||
|
||||
def set_ground(self, pose):
|
||||
pose = transform_pyrender(pose)
|
||||
print("Setting ground pose to be", pose)
|
||||
self.ground_pose = np.asarray(pose)
|
||||
self.scene.set_pose(self.ground_node, pose=self.ground_pose)
|
||||
|
||||
def set_bg_seq(self, bg_imgs):
|
||||
pass
|
||||
|
||||
def set_mesh_visibility(self, vis):
|
||||
self.acquire_lock()
|
||||
for node in self.anim_nodes:
|
||||
node.mesh.is_visible = vis
|
||||
self.release_lock()
|
||||
|
||||
def check_mesh_visibility(self):
|
||||
for node in self.anim_nodes:
|
||||
if node.mesh.is_visible:
|
||||
return True
|
||||
|
||||
if self.ground_node.mesh.is_visible:
|
||||
return True
|
||||
|
||||
if self.cam_marker_node is not None and self.cam_marker_node.mesh.is_visible:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def update_frame(self):
|
||||
t = self.anim_idx % self.anim_len
|
||||
|
||||
# update camera
|
||||
if self.anim_cameras is not None:
|
||||
cam_t = min(t, len(self.anim_cameras) - 1)
|
||||
cam_pose = self.anim_cameras[cam_t]
|
||||
self.update_camera(cam_pose)
|
||||
if self.cam_marker_node is not None and t < len(self.cam_marker_poses):
|
||||
self.acquire_lock()
|
||||
self.scene.set_pose(self.cam_marker_node, self.cam_marker_poses[t])
|
||||
self.release_lock()
|
||||
|
||||
# update meshes
|
||||
meshes = self.anim_meshes[t]
|
||||
self.acquire_lock()
|
||||
for i, node in enumerate(self.anim_nodes):
|
||||
if i < len(meshes):
|
||||
node.mesh = meshes[i]
|
||||
node.mesh.is_visible = True
|
||||
else:
|
||||
node.mesh.is_visible = False
|
||||
self.release_lock()
|
||||
|
||||
|
||||
class OffscreenAnimation(AnimationBase):
|
||||
def __init__(self, img_size, intrins=None, fps=30, ext="mp4"):
|
||||
super().__init__(img_size, intrins=intrins)
|
||||
self.add_lighting(0.9)
|
||||
|
||||
self.fps = fps
|
||||
self.ext = ext
|
||||
self.viewer = pyrender.OffscreenRenderer(*self.viewport_size)
|
||||
|
||||
self.bg_seq = []
|
||||
self.keypoints_seq = None # Store GT 2D keypoints for overlay
|
||||
self.pred_keypoints_seq = None # Store predicted 2D keypoints
|
||||
|
||||
def set_bg_seq(self, bg_imgs):
|
||||
assert isinstance(bg_imgs, list)
|
||||
print(f"setting length {len(bg_imgs)} bg sequence")
|
||||
self.bg_seq = bg_imgs
|
||||
|
||||
def set_keypoints_seq(self, keypoints_seq):
|
||||
"""
|
||||
Set GT 2D keypoints sequence for overlay visualization
|
||||
Args:
|
||||
keypoints_seq: List of (J, 3) arrays with [x, y, confidence] for each frame
|
||||
"""
|
||||
self.keypoints_seq = keypoints_seq
|
||||
print(f"setting length {len(keypoints_seq)} GT keypoints sequence")
|
||||
|
||||
def set_pred_keypoints_seq(self, pred_keypoints_seq):
|
||||
"""
|
||||
Set predicted 2D keypoints sequence
|
||||
Args:
|
||||
pred_keypoints_seq: List of (J, 2) arrays with [x, y] for each frame
|
||||
"""
|
||||
self.pred_keypoints_seq = pred_keypoints_seq
|
||||
print(f"setting length {len(pred_keypoints_seq)} predicted keypoints sequence")
|
||||
|
||||
def delete(self):
|
||||
self.viewer.delete()
|
||||
|
||||
def close(self):
|
||||
self.delete()
|
||||
|
||||
def draw_keypoints_on_image(self, img, keypoints, radius=4, color=(0, 255, 0), thickness=-1, conf_threshold=0.3):
|
||||
"""
|
||||
Draw 2D keypoints on an image
|
||||
Args:
|
||||
img: numpy array (H, W, 3)
|
||||
keypoints: (J, 3) array with [x, y, confidence]
|
||||
radius: circle radius for keypoints
|
||||
color: BGR color tuple
|
||||
thickness: -1 for filled circles
|
||||
conf_threshold: minimum confidence to draw keypoint
|
||||
Returns:
|
||||
img with keypoints drawn
|
||||
"""
|
||||
import cv2
|
||||
img = img.copy()
|
||||
for j in range(keypoints.shape[0]):
|
||||
x, y, conf = keypoints[j]
|
||||
if conf > conf_threshold:
|
||||
# Draw circle at keypoint location
|
||||
cv2.circle(img, (int(x), int(y)), radius, color, thickness)
|
||||
return img
|
||||
|
||||
def render(
|
||||
self, render_bg=True, render_ground=True, render_cam=True, fac=1.0, render_keypoints=False, **kwargs
|
||||
):
|
||||
flags = RenderFlags.RGBA | RenderFlags.SHADOWS_DIRECTIONAL
|
||||
self.ground_node.mesh.is_visible = render_ground and (not render_bg)
|
||||
if self.cam_marker_node is not None:
|
||||
self.cam_marker_node.mesh.is_visible = render_cam
|
||||
|
||||
rgba = np.zeros((*self.viewport_size[::-1], 4))
|
||||
if self.check_mesh_visibility(): # if any meshes are visible
|
||||
rgba, _ = self.viewer.render(self.scene, flags=flags)
|
||||
|
||||
img = rgba
|
||||
if render_bg and len(self.bg_seq) > 0:
|
||||
t = min(self.anim_idx, len(self.bg_seq) - 1)
|
||||
bg = self.bg_seq[t]
|
||||
rgba = rgba.astype(np.float32) / 255
|
||||
alpha = fac * rgba[..., 3:]
|
||||
img = alpha * rgba[..., :3] + (1 - alpha) * bg
|
||||
img = (255 * img).astype(np.uint8)
|
||||
|
||||
# Overlay 2D keypoints if available and requested
|
||||
if render_keypoints and render_bg:
|
||||
# Draw GT keypoints in GREEN (larger)
|
||||
if self.keypoints_seq is not None:
|
||||
t = min(self.anim_idx, len(self.keypoints_seq) - 1)
|
||||
keypoints = self.keypoints_seq[t]
|
||||
if keypoints is not None and len(keypoints) > 0:
|
||||
img = self.draw_keypoints_on_image(img, keypoints, color=(0, 255, 0), radius=5)
|
||||
|
||||
# Draw predicted keypoints in RED (smaller)
|
||||
if self.pred_keypoints_seq is not None:
|
||||
t = min(self.anim_idx, len(self.pred_keypoints_seq) - 1)
|
||||
pred_kp = self.pred_keypoints_seq[t]
|
||||
if pred_kp is not None and len(pred_kp) > 0:
|
||||
# Add confidence=1.0 if only (J, 2)
|
||||
if pred_kp.ndim == 2 and pred_kp.shape[1] == 2:
|
||||
pred_kp = np.concatenate([pred_kp, np.ones((len(pred_kp), 1))], axis=1)
|
||||
img = self.draw_keypoints_on_image(img, pred_kp, color=(0, 0, 255), radius=3, conf_threshold=0.0)
|
||||
|
||||
return img
|
||||
|
||||
def render_mesh_layers(self, render_ground=True, render_cam=True, **kwargs):
|
||||
# render each mesh in its own layer
|
||||
flags = RenderFlags.RGBA | RenderFlags.SHADOWS_DIRECTIONAL
|
||||
render_cam = render_cam and self.cam_marker_node is not None
|
||||
|
||||
# current timestep
|
||||
t = self.anim_idx % self.anim_len
|
||||
meshes = self.anim_meshes[t]
|
||||
layers = []
|
||||
|
||||
# toggle ground and camera
|
||||
self.ground_node.mesh.is_visible = False
|
||||
if self.cam_marker_node is not None:
|
||||
self.cam_marker_node.mesh.is_visible = False
|
||||
|
||||
# render animation meshes for current timestep
|
||||
for i in range(len(meshes)):
|
||||
for j, node in enumerate(self.anim_nodes):
|
||||
if i == j:
|
||||
node.mesh.is_visible = True
|
||||
else:
|
||||
node.mesh.is_visible = False
|
||||
rgba, _ = self.viewer.render(self.scene, flags=flags)
|
||||
layers.append(rgba)
|
||||
|
||||
if not render_ground and not render_cam:
|
||||
return layers
|
||||
|
||||
# render the ground and camera
|
||||
for node in self.anim_nodes:
|
||||
node.mesh.is_visible = False
|
||||
|
||||
if render_cam:
|
||||
self.cam_marker_node.mesh.is_visible = render_cam
|
||||
rgba, _ = self.viewer.render(self.scene, flags=flags)
|
||||
layers.append(rgba)
|
||||
self.cam_marker_node.mesh.is_visible = False
|
||||
|
||||
if render_ground:
|
||||
self.ground_node.mesh.is_visible = render_ground
|
||||
rgba, _ = self.viewer.render(self.scene, flags=flags)
|
||||
layers.append(rgba)
|
||||
|
||||
return layers
|
||||
|
||||
def render_frames_layers(self):
|
||||
frames_layers = []
|
||||
for t in range(self.anim_len):
|
||||
self.anim_idx = t
|
||||
self.update_frame()
|
||||
layers = self.render_mesh_layers()
|
||||
frames_layers.append(layers)
|
||||
return frames_layers
|
||||
|
||||
def animate(self, save_name, save_frames=False, render_layers=False, **kwargs):
|
||||
if render_layers:
|
||||
frames_layers = self.render_frames_layers()
|
||||
os.makedirs(save_name, exist_ok=True)
|
||||
for t, layers in enumerate(frames_layers):
|
||||
for l, layer in enumerate(layers):
|
||||
path = f"{save_name}/{t:06d}_{l:03d}.png"
|
||||
imageio.imwrite(path, layer)
|
||||
print(f"saved frame layers to {save_name}")
|
||||
return save_name
|
||||
|
||||
frames = self.render_frames(**kwargs)
|
||||
if len(frames) > 1:
|
||||
if save_frames:
|
||||
os.makedirs(save_name, exist_ok=True)
|
||||
for i, frame in enumerate(frames):
|
||||
path = f"{save_name}/{i:06d}.png"
|
||||
imageio.imwrite(path, frame)
|
||||
return save_name
|
||||
save_path = f"{save_name}.{self.ext}"
|
||||
os.makedirs(save_name, exist_ok=True)
|
||||
for idx, i in enumerate(frames):
|
||||
imageio.imwrite(f"{save_name}/" + f'{str(idx).zfill(6)}.jpg', i[:, :, :3])
|
||||
imageio.mimwrite(save_path, frames, fps=self.fps)
|
||||
print("wrote video to", save_path)
|
||||
return save_path
|
||||
|
||||
save_path = f"{save_name}.png"
|
||||
imageio.imwrite(save_path, frames[0])
|
||||
print("wrote image to", save_path)
|
||||
return save_path
|
||||
|
||||
def render_frames(self, **kwargs):
|
||||
print("ANIMATION LENGTH", self.anim_len)
|
||||
frames = []
|
||||
for t in range(self.anim_len):
|
||||
self.anim_idx = t
|
||||
self.update_frame()
|
||||
img = self.render(**kwargs)
|
||||
frames.append(img)
|
||||
return frames
|
||||
|
||||
def render_layers(self, save_dir, composite=False, fac=0.4):
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
# render all frames
|
||||
frames = self.render_frames(render_bg=False, render_ground=False)
|
||||
|
||||
# render background
|
||||
self.set_mesh_visibility(False)
|
||||
bg_img = self.render(render_bg=False, render_ground=True) / 255
|
||||
|
||||
# save frames
|
||||
for t, frame in enumerate(frames):
|
||||
imageio.imwrite(f"{save_dir}/{t:06d}.png", frame)
|
||||
imageio.imwrite(f"{save_dir}/background.png", bg_img)
|
||||
|
||||
if composite:
|
||||
rgba = composite_layers(frames, bg_img, fac=fac)
|
||||
imageio.imwrite(f"{save_dir}/composite.png", (255 * rgba).astype(np.uint8))
|
||||
|
||||
|
||||
def composite_layers(frames, bg_img, fac=0.4):
|
||||
# composite front to back
|
||||
H, W = frames[0].shape[:2]
|
||||
comp = np.zeros((H, W, 3))
|
||||
vis = np.ones((H, W, 1))
|
||||
for frame in frames:
|
||||
frame = frame / 255
|
||||
alpha = fac * frame[..., 3:]
|
||||
comp += vis * alpha * frame[..., :3]
|
||||
vis *= 1 - alpha
|
||||
|
||||
comp += vis * bg_img[..., 3:] * bg_img[..., :3]
|
||||
vis *= 1 - bg_img[..., 3:]
|
||||
return comp + vis
|
||||
rgba = np.concatenate([comp, 1 - vis], axis=-1)
|
||||
return rgba
|
||||
|
||||
|
||||
class AnimationViewer(AnimationBase):
|
||||
def __init__(self, img_size, intrins=None, fps=30, num_repeats=1):
|
||||
super().__init__(img_size, intrins=intrins)
|
||||
|
||||
self.fps = fps
|
||||
self.ext = "gif"
|
||||
|
||||
def pause_play_callback(_, viewer):
|
||||
viewer.is_paused = not viewer.is_paused
|
||||
|
||||
def step_callback(_, viewer, step):
|
||||
viewer.anim_idx = (viewer.anim_idx + step) % viewer.anim_len
|
||||
|
||||
def stop_callback(_, viewer):
|
||||
viewer.do_animate = False
|
||||
|
||||
def loop_callback(_, viewer):
|
||||
viewer.num_repeats *= -1
|
||||
|
||||
registered_keys = {
|
||||
"p": (pause_play_callback, [self]),
|
||||
".": (step_callback, [self, 1]),
|
||||
",": (step_callback, [self, -1]),
|
||||
"l": (loop_callback, [self]),
|
||||
"s": (stop_callback, [self]),
|
||||
}
|
||||
|
||||
self.viewer = pyrender.Viewer(
|
||||
self.scene,
|
||||
run_in_thread=True,
|
||||
viewport_size=self.viewport_size,
|
||||
use_raymond_lighting=True,
|
||||
registered_keys=registered_keys,
|
||||
)
|
||||
|
||||
self.num_repeats = num_repeats
|
||||
self.do_animate = False
|
||||
self.is_paused = False
|
||||
|
||||
def acquire_lock(self):
|
||||
self.viewer.render_lock.acquire()
|
||||
|
||||
def release_lock(self):
|
||||
self.viewer.render_lock.release()
|
||||
|
||||
def close(self):
|
||||
self.viewer.close_external()
|
||||
|
||||
def animate(self, save_name=None, **kwargs):
|
||||
print("===================")
|
||||
print(f"PLAYING ANIMATION {self.num_repeats} TIMES")
|
||||
print("VIEWER CONTROLS:")
|
||||
print("p\tpause/play")
|
||||
print("l\tloop/unloop")
|
||||
print(".\tstep forward")
|
||||
print(",\tstep backward")
|
||||
print("s\tstop")
|
||||
print("===================")
|
||||
|
||||
self.anim_idx = 0
|
||||
self.do_animate = True
|
||||
frame_dur = 1 / self.fps
|
||||
if save_name is not None:
|
||||
self.viewer.viewer_flags["record"] = True
|
||||
|
||||
# set up initial frame
|
||||
self.update_frame()
|
||||
|
||||
while self.do_animate:
|
||||
if self.is_paused:
|
||||
self.update_frame()
|
||||
continue
|
||||
|
||||
if (
|
||||
self.num_repeats > 0
|
||||
and self.anim_idx >= self.num_repeats * self.anim_len
|
||||
):
|
||||
break
|
||||
|
||||
if save_name is not None and self.anim_idx >= self.anim_len:
|
||||
self.viewer.viewer_flags["record"] = False
|
||||
|
||||
start_time = time.time()
|
||||
self.update_frame()
|
||||
render_time = time.time() - start_time
|
||||
|
||||
self.anim_idx += 1
|
||||
sleep_time = frame_dur - render_time
|
||||
if sleep_time > 0:
|
||||
time.sleep(sleep_time)
|
||||
|
||||
if save_name is not None:
|
||||
print(f"SAVING RECORDING TO {save_name}")
|
||||
self.viewer.save_gif(f"{save_name}.{self.ext}")
|
||||
|
||||
|
||||
def make_checkerboard(
|
||||
length=25.0,
|
||||
color0=[0.8, 0.9, 0.9],
|
||||
color1=[0.6, 0.7, 0.7],
|
||||
tile_width=0.5,
|
||||
alpha=1.0,
|
||||
up="y",
|
||||
):
|
||||
v, f, _, fc = checkerboard_geometry(length, color0, color1, tile_width, alpha, up)
|
||||
return trimesh.Trimesh(v, f, face_colors=fc, process=False)
|
||||
|
||||
|
||||
def make_pyrender_camera(img_size, intrins=None):
|
||||
if intrins is not None:
|
||||
print("USING INTRINSICS CAMERA", intrins)
|
||||
fx, fy, cx, cy = intrins
|
||||
return pyrender.IntrinsicsCamera(fx, fy, cx, cy)
|
||||
|
||||
W, H = img_size
|
||||
focal = 0.5 * (H + W)
|
||||
yfov = 2 * np.arctan(0.5 * H / focal)
|
||||
print("USING PERSPECTIVE CAMERA", H, W, focal)
|
||||
return pyrender.PerspectiveCamera(yfov=yfov, aspectRatio=W / H)
|
||||
|
||||
|
||||
def make_camera_marker(radius=0.1, height=0.2, up="y", transform=None):
|
||||
"""
|
||||
:param radius (default 0.1) radius of pyramid base, diagonal of image plane
|
||||
:param height (default 0.2) height of pyramid, focal length
|
||||
:param up (default y) camera up vector
|
||||
:param transform (default None) (4, 4) cam to world transform
|
||||
"""
|
||||
verts, faces, face_colors = camera_marker_geometry(radius, height, up)
|
||||
if transform is not None:
|
||||
assert transform.shape == (4, 4)
|
||||
verts = (
|
||||
np.einsum("ij,nj->ni", transform[:3, :3], verts) + transform[None, :3, 3]
|
||||
)
|
||||
return trimesh.Trimesh(verts, faces, face_colors=face_colors, process=False)
|
||||
|
||||
|
||||
def get_light_poses(n_lights=5, elevation=np.pi / 3, dist=12):
|
||||
# get lights in a circle around origin at elevation
|
||||
thetas = elevation * np.ones(n_lights)
|
||||
phis = 2 * np.pi * np.arange(n_lights) / n_lights
|
||||
poses = []
|
||||
trans = make_translation(torch.tensor([0, 0, dist]))
|
||||
for phi, theta in zip(phis, thetas):
|
||||
rot = make_rotation(rx=-theta, ry=phi, order="xyz")
|
||||
poses.append((rot @ trans).numpy())
|
||||
return poses
|
||||
|
||||
|
||||
def get_raymond_light_poses(up="z"):
|
||||
thetas = np.pi * np.ones(3) / 6
|
||||
phis = 2 * np.pi * np.arange(3) / 3
|
||||
|
||||
poses = []
|
||||
|
||||
for phi, theta in zip(phis, thetas):
|
||||
xp = np.sin(theta) * np.cos(phi)
|
||||
yp = np.sin(theta) * np.sin(phi)
|
||||
zp = np.cos(theta)
|
||||
|
||||
if up == "y":
|
||||
z = np.array([xp, zp, -yp])
|
||||
else:
|
||||
z = np.array([xp, yp, zp])
|
||||
z = z / np.linalg.norm(z)
|
||||
x = np.array([-z[1], z[0], 0.0])
|
||||
if np.linalg.norm(x) == 0:
|
||||
x = np.array([1.0, 0.0, 0.0])
|
||||
x = x / np.linalg.norm(x)
|
||||
y = np.cross(z, x)
|
||||
|
||||
matrix = np.eye(4)
|
||||
matrix[:3, :3] = np.c_[x, y, z]
|
||||
poses.append(matrix)
|
||||
return poses
|
||||
Reference in New Issue
Block a user