SMPL model visualizerΒΆ
Visualizer for SMPL human body models. Requires a .npz model file.
- See here for download instructions:
https://github.com/vchoutas/smplx?tab=readme-ov-file#downloading-the-model
1from __future__ import annotations
2
3import time
4from dataclasses import dataclass
5from pathlib import Path
6
7import numpy as np
8import tyro
9
10import viser
11import viser.transforms as tf
12
13
14@dataclass(frozen=True)
15class SmplOutputs:
16 vertices: np.ndarray
17 faces: np.ndarray
18 T_world_joint: np.ndarray # (num_joints, 4, 4)
19 T_parent_joint: np.ndarray # (num_joints, 4, 4)
20
21
22class SmplHelper:
23 """Helper for models in the SMPL family, implemented in numpy."""
24
25 def __init__(self, model_path: Path) -> None:
26 assert model_path.suffix.lower() == ".npz", "Model should be an .npz file!"
27 body_dict = dict(**np.load(model_path, allow_pickle=True))
28
29 self.J_regressor = body_dict["J_regressor"]
30 self.weights = body_dict["weights"]
31 self.v_template = body_dict["v_template"]
32 self.posedirs = body_dict["posedirs"]
33 self.shapedirs = body_dict["shapedirs"]
34 self.faces = body_dict["f"]
35
36 self.num_joints: int = self.weights.shape[-1]
37 self.num_betas: int = self.shapedirs.shape[-1]
38 self.parent_idx: np.ndarray = body_dict["kintree_table"][0]
39
40 def get_outputs(self, betas: np.ndarray, joint_rotmats: np.ndarray) -> SmplOutputs:
41 # Get shaped vertices + joint positions, when all local poses are identity.
42 v_tpose = self.v_template + np.einsum("vxb,b->vx", self.shapedirs, betas)
43 j_tpose = np.einsum("jv,vx->jx", self.J_regressor, v_tpose)
44
45 # Local SE(3) transforms.
46 T_parent_joint = np.zeros((self.num_joints, 4, 4)) + np.eye(4)
47 T_parent_joint[:, :3, :3] = joint_rotmats
48 T_parent_joint[0, :3, 3] = j_tpose[0]
49 T_parent_joint[1:, :3, 3] = j_tpose[1:] - j_tpose[self.parent_idx[1:]]
50
51 # Forward kinematics.
52 T_world_joint = T_parent_joint.copy()
53 for i in range(1, self.num_joints):
54 T_world_joint[i] = T_world_joint[self.parent_idx[i]] @ T_parent_joint[i]
55
56 # Linear blend skinning.
57 pose_delta = (joint_rotmats[1:, ...] - np.eye(3)).flatten()
58 v_blend = v_tpose + np.einsum("byn,n->by", self.posedirs, pose_delta)
59 v_delta = np.ones((v_blend.shape[0], self.num_joints, 4))
60 v_delta[:, :, :3] = v_blend[:, None, :] - j_tpose[None, :, :]
61 v_posed = np.einsum(
62 "jxy,vj,vjy->vx", T_world_joint[:, :3, :], self.weights, v_delta
63 )
64 return SmplOutputs(v_posed, self.faces, T_world_joint, T_parent_joint)
65
66
67def main(model_path: Path) -> None:
68 server = viser.ViserServer()
69 server.scene.set_up_direction("+y")
70 server.gui.configure_theme(control_layout="collapsible")
71
72 server.scene.enable_default_lights(cast_shadow=True)
73 server.scene.add_grid("/grid", position=(0.0, -1.3, 0.0), plane="xz")
74
75 # Main loop. We'll read pose/shape from the GUI elements, compute the mesh,
76 # and then send the updated mesh in a loop.
77 model = SmplHelper(model_path)
78 gui_elements = make_gui_elements(
79 server,
80 num_betas=model.num_betas,
81 num_joints=model.num_joints,
82 parent_idx=model.parent_idx,
83 )
84 body_handle = server.scene.add_mesh_simple(
85 "/human",
86 model.v_template,
87 model.faces,
88 wireframe=gui_elements.gui_wireframe.value,
89 color=gui_elements.gui_rgb.value,
90 )
91 while True:
92 # Do nothing if no change.
93 time.sleep(0.02)
94 if not gui_elements.changed:
95 continue
96
97 gui_elements.changed = False
98
99 # If anything has changed, re-compute SMPL outputs.
100 smpl_outputs = model.get_outputs(
101 betas=np.array([x.value for x in gui_elements.gui_betas]),
102 joint_rotmats=tf.SO3.exp(
103 # (num_joints, 3)
104 np.array([x.value for x in gui_elements.gui_joints])
105 ).as_matrix(),
106 )
107
108 # Update the mesh properties based on the SMPL model output + GUI
109 # elements.
110 body_handle.vertices = smpl_outputs.vertices
111 body_handle.wireframe = gui_elements.gui_wireframe.value
112 body_handle.color = gui_elements.gui_rgb.value
113
114 # Match transform control gizmos to joint positions.
115 for i, control in enumerate(gui_elements.transform_controls):
116 control.position = smpl_outputs.T_parent_joint[i, :3, 3]
117
118
119@dataclass
120class GuiElements:
121 """Structure containing handles for reading from GUI elements."""
122
123 gui_rgb: viser.GuiInputHandle[tuple[int, int, int]]
124 gui_wireframe: viser.GuiInputHandle[bool]
125 gui_betas: list[viser.GuiInputHandle[float]]
126 gui_joints: list[viser.GuiInputHandle[tuple[float, float, float]]]
127 transform_controls: list[viser.TransformControlsHandle]
128
129 changed: bool
130 """This flag will be flipped to True whenever the mesh needs to be re-generated."""
131
132
133def make_gui_elements(
134 server: viser.ViserServer,
135 num_betas: int,
136 num_joints: int,
137 parent_idx: np.ndarray,
138) -> GuiElements:
139 """Make GUI elements for interacting with the model."""
140
141 tab_group = server.gui.add_tab_group()
142
143 def set_changed(_) -> None:
144 out.changed = True # out is define later!
145
146 # GUI elements: mesh settings + visibility.
147 with tab_group.add_tab("View", viser.Icon.VIEWFINDER):
148 gui_rgb = server.gui.add_rgb("Color", initial_value=(90, 200, 255))
149 gui_wireframe = server.gui.add_checkbox("Wireframe", initial_value=False)
150 gui_show_controls = server.gui.add_checkbox("Handles", initial_value=True)
151
152 gui_rgb.on_update(set_changed)
153 gui_wireframe.on_update(set_changed)
154
155 @gui_show_controls.on_update
156 def _(_):
157 for control in transform_controls:
158 control.visible = gui_show_controls.value
159
160 # GUI elements: shape parameters.
161 with tab_group.add_tab("Shape", viser.Icon.BOX):
162 gui_reset_shape = server.gui.add_button("Reset Shape")
163 gui_random_shape = server.gui.add_button("Random Shape")
164
165 @gui_reset_shape.on_click
166 def _(_):
167 for beta in gui_betas:
168 beta.value = 0.0
169
170 @gui_random_shape.on_click
171 def _(_):
172 for beta in gui_betas:
173 beta.value = np.random.normal(loc=0.0, scale=1.0)
174
175 gui_betas = []
176 for i in range(num_betas):
177 beta = server.gui.add_slider(
178 f"beta{i}", min=-5.0, max=5.0, step=0.01, initial_value=0.0
179 )
180 gui_betas.append(beta)
181 beta.on_update(set_changed)
182
183 # GUI elements: joint angles.
184 with tab_group.add_tab("Joints", viser.Icon.ANGLE):
185 gui_reset_joints = server.gui.add_button("Reset Joints")
186 gui_random_joints = server.gui.add_button("Random Joints")
187
188 @gui_reset_joints.on_click
189 def _(_):
190 for joint in gui_joints:
191 joint.value = (0.0, 0.0, 0.0)
192
193 @gui_random_joints.on_click
194 def _(_):
195 rng = np.random.default_rng()
196 for joint in gui_joints:
197 joint.value = tf.SO3.sample_uniform(rng).log()
198
199 gui_joints: list[viser.GuiInputHandle[tuple[float, float, float]]] = []
200 for i in range(num_joints):
201 gui_joint = server.gui.add_vector3(
202 label=f"Joint {i}",
203 initial_value=(0.0, 0.0, 0.0),
204 step=0.05,
205 )
206 gui_joints.append(gui_joint)
207
208 def set_callback_in_closure(i: int) -> None:
209 @gui_joint.on_update
210 def _(_):
211 transform_controls[i].wxyz = tf.SO3.exp(
212 np.array(gui_joints[i].value)
213 ).wxyz
214 out.changed = True
215
216 set_callback_in_closure(i)
217
218 # Transform control gizmos on joints.
219 transform_controls: list[viser.TransformControlsHandle] = []
220 prefixed_joint_names = [] # Joint names, but prefixed with parents.
221 for i in range(num_joints):
222 prefixed_joint_name = f"joint_{i}"
223 if i > 0:
224 prefixed_joint_name = (
225 prefixed_joint_names[parent_idx[i]] + "/" + prefixed_joint_name
226 )
227 prefixed_joint_names.append(prefixed_joint_name)
228 controls = server.scene.add_transform_controls(
229 f"/smpl/{prefixed_joint_name}",
230 depth_test=False,
231 scale=0.2 * (0.75 ** prefixed_joint_name.count("/")),
232 disable_axes=True,
233 disable_sliders=True,
234 visible=gui_show_controls.value,
235 )
236 transform_controls.append(controls)
237
238 def set_callback_in_closure(i: int) -> None:
239 @controls.on_update
240 def _(_) -> None:
241 axisangle = tf.SO3(transform_controls[i].wxyz).log()
242 gui_joints[i].value = (axisangle[0], axisangle[1], axisangle[2])
243
244 set_callback_in_closure(i)
245
246 out = GuiElements(
247 gui_rgb,
248 gui_wireframe,
249 gui_betas,
250 gui_joints,
251 transform_controls=transform_controls,
252 changed=True,
253 )
254 return out
255
256
257if __name__ == "__main__":
258 tyro.cli(main, description=__doc__)