SMPL human model¶
Visualize SMPL human body models with pose and shape parameter controls.
- Requires a .npz model file. See here for download instructions:
https://github.com/vchoutas/smplx?tab=readme-ov-file#downloading-the-model
Features:
SMPL model loading and mesh generation
Interactive pose parameter sliders
Real-time body shape and pose updates
3D mesh visualization with
viser.SceneApi.add_mesh_simple()
Note
This example requires external assets. To download them, run:
cd /path/to/viser/examples/assets
./download_assets.sh
Source: examples/04_demos/03_smpl_visualizer.py

Code¶
1from __future__ import annotations
2
3import time
4from dataclasses import dataclass
5from pathlib import Path
6
7import numpy as np
8import trimesh
9import tyro
10
11import viser
12import viser.transforms as tf
13
14
15@dataclass(frozen=True)
16class SmplOutputs:
17 vertices: np.ndarray
18 faces: np.ndarray
19 T_world_joint: np.ndarray # (num_joints, 4, 4)
20 T_parent_joint: np.ndarray # (num_joints, 4, 4)
21
22
23class SmplHelper:
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(
68 model_path: Path = Path(__file__).parent / "../assets/SMPLH_NEUTRAL.npz",
69) -> None:
70 server = viser.ViserServer()
71 server.scene.set_up_direction("+y")
72 server.scene.add_grid("/grid", position=(0.0, -1.3, 0.0), plane="xz")
73
74 # Main loop. We'll read pose/shape from the GUI elements, compute the mesh,
75 # and then send the updated mesh in a loop.
76 model = SmplHelper(model_path)
77 gui_elements = make_gui_elements(
78 server,
79 num_betas=model.num_betas,
80 num_joints=model.num_joints,
81 parent_idx=model.parent_idx,
82 )
83 body_handle = server.scene.add_mesh_simple(
84 "/human",
85 model.v_template,
86 model.faces,
87 wireframe=gui_elements.gui_wireframe.value,
88 color=gui_elements.gui_rgb.value,
89 )
90
91 # Add a vertex selector to the mesh. This will allow us to click on
92 # vertices to get indices.
93 red_sphere = trimesh.creation.icosphere(radius=0.001, subdivisions=1)
94 red_sphere.visual.vertex_colors = (255, 0, 0, 255) # type: ignore
95 vertex_selector = server.scene.add_batched_meshes_trimesh(
96 "/selector",
97 red_sphere,
98 batched_positions=model.v_template,
99 batched_wxyzs=((1.0, 0.0, 0.0, 0.0),) * model.v_template.shape[0],
100 )
101
102 @vertex_selector.on_click
103 def _(event: viser.SceneNodePointerEvent) -> None:
104 event.client.add_notification(
105 f"Clicked on vertex {event.instance_index}",
106 body="",
107 auto_close=3000,
108 )
109
110 while True:
111 # Do nothing if no change.
112 time.sleep(0.02)
113 if not gui_elements.changed:
114 continue
115
116 gui_elements.changed = False
117
118 # If anything has changed, re-compute SMPL outputs.
119 smpl_outputs = model.get_outputs(
120 betas=np.array([x.value for x in gui_elements.gui_betas]),
121 joint_rotmats=tf.SO3.exp(
122 # (num_joints, 3)
123 np.array([x.value for x in gui_elements.gui_joints])
124 ).as_matrix(),
125 )
126
127 # Update the mesh properties based on the SMPL model output + GUI
128 # elements.
129 body_handle.vertices = smpl_outputs.vertices
130 body_handle.wireframe = gui_elements.gui_wireframe.value
131 body_handle.color = gui_elements.gui_rgb.value
132 vertex_selector.batched_positions = smpl_outputs.vertices
133
134 # Match transform control gizmos to joint positions.
135 for i, control in enumerate(gui_elements.transform_controls):
136 control.position = smpl_outputs.T_parent_joint[i, :3, 3]
137
138
139@dataclass
140class GuiElements:
141
142 gui_rgb: viser.GuiInputHandle[tuple[int, int, int]]
143 gui_wireframe: viser.GuiInputHandle[bool]
144 gui_betas: list[viser.GuiInputHandle[float]]
145 gui_joints: list[viser.GuiInputHandle[tuple[float, float, float]]]
146 transform_controls: list[viser.TransformControlsHandle]
147
148 changed: bool
149
150
151def make_gui_elements(
152 server: viser.ViserServer,
153 num_betas: int,
154 num_joints: int,
155 parent_idx: np.ndarray,
156) -> GuiElements:
157
158 tab_group = server.gui.add_tab_group()
159
160 def set_changed(_) -> None:
161 out.changed = True # out is define later!
162
163 # GUI elements: mesh settings + visibility.
164 with tab_group.add_tab("View", viser.Icon.VIEWFINDER):
165 gui_rgb = server.gui.add_rgb("Color", initial_value=(90, 200, 255))
166 gui_wireframe = server.gui.add_checkbox("Wireframe", initial_value=False)
167 gui_show_controls = server.gui.add_checkbox("Handles", initial_value=True)
168
169 gui_rgb.on_update(set_changed)
170 gui_wireframe.on_update(set_changed)
171
172 @gui_show_controls.on_update
173 def _(_):
174 for control in transform_controls:
175 control.visible = gui_show_controls.value
176
177 # GUI elements: shape parameters.
178 with tab_group.add_tab("Shape", viser.Icon.BOX):
179 gui_reset_shape = server.gui.add_button("Reset Shape")
180 gui_random_shape = server.gui.add_button("Random Shape")
181
182 @gui_reset_shape.on_click
183 def _(_):
184 for beta in gui_betas:
185 beta.value = 0.0
186
187 @gui_random_shape.on_click
188 def _(_):
189 for beta in gui_betas:
190 beta.value = np.random.normal(loc=0.0, scale=1.0)
191
192 gui_betas = []
193 for i in range(num_betas):
194 beta = server.gui.add_slider(
195 f"beta{i}", min=-5.0, max=5.0, step=0.01, initial_value=0.0
196 )
197 gui_betas.append(beta)
198 beta.on_update(set_changed)
199
200 # GUI elements: joint angles.
201 with tab_group.add_tab("Joints", viser.Icon.ANGLE):
202 gui_reset_joints = server.gui.add_button("Reset Joints")
203 gui_random_joints = server.gui.add_button("Random Joints")
204
205 @gui_reset_joints.on_click
206 def _(_):
207 for joint in gui_joints:
208 joint.value = (0.0, 0.0, 0.0)
209
210 @gui_random_joints.on_click
211 def _(_):
212 rng = np.random.default_rng()
213 for joint in gui_joints:
214 joint.value = tf.SO3.sample_uniform(rng).log()
215
216 gui_joints: list[viser.GuiInputHandle[tuple[float, float, float]]] = []
217 for i in range(num_joints):
218 gui_joint = server.gui.add_vector3(
219 label=f"Joint {i}",
220 initial_value=(0.0, 0.0, 0.0),
221 step=0.05,
222 )
223 gui_joints.append(gui_joint)
224
225 def set_callback_in_closure(i: int) -> None:
226 @gui_joint.on_update
227 def _(_):
228 transform_controls[i].wxyz = tf.SO3.exp(
229 np.array(gui_joints[i].value)
230 ).wxyz
231 out.changed = True
232
233 set_callback_in_closure(i)
234
235 # Transform control gizmos on joints.
236 transform_controls: list[viser.TransformControlsHandle] = []
237 prefixed_joint_names = [] # Joint names, but prefixed with parents.
238 for i in range(num_joints):
239 prefixed_joint_name = f"joint_{i}"
240 if i > 0:
241 prefixed_joint_name = (
242 prefixed_joint_names[parent_idx[i]] + "/" + prefixed_joint_name
243 )
244 prefixed_joint_names.append(prefixed_joint_name)
245 controls = server.scene.add_transform_controls(
246 f"/smpl/{prefixed_joint_name}",
247 depth_test=False,
248 scale=0.2 * (0.75 ** prefixed_joint_name.count("/")),
249 disable_axes=True,
250 disable_sliders=True,
251 visible=gui_show_controls.value,
252 )
253 transform_controls.append(controls)
254
255 def set_callback_in_closure(i: int) -> None:
256 @controls.on_update
257 def _(_) -> None:
258 axisangle = tf.SO3(transform_controls[i].wxyz).log()
259 gui_joints[i].value = (axisangle[0], axisangle[1], axisangle[2])
260
261 set_callback_in_closure(i)
262
263 out = GuiElements(
264 gui_rgb,
265 gui_wireframe,
266 gui_betas,
267 gui_joints,
268 transform_controls=transform_controls,
269 changed=True,
270 )
271 return out
272
273
274if __name__ == "__main__":
275 tyro.cli(main, description=__doc__)