From 2f8bdc267177f75c663dfc6823b69d585507ecec Mon Sep 17 00:00:00 2001 From: brentyi Date: Wed, 25 Sep 2024 16:44:41 -0700 Subject: [PATCH] Example cleanup --- .../examples/07_record3d_visualizer.rst | 3 +- docs/source/examples/08_smpl_visualizer.rst | 8 ++---- .../examples/25_smpl_visualizer_skinned.rst | 28 ++++++------------- examples/08_smpl_visualizer.py | 8 ++---- examples/25_smpl_visualizer_skinned.py | 28 ++++++------------- 5 files changed, 22 insertions(+), 53 deletions(-) diff --git a/docs/source/examples/07_record3d_visualizer.rst b/docs/source/examples/07_record3d_visualizer.rst index 89b0df03..ea813c04 100644 --- a/docs/source/examples/07_record3d_visualizer.rst +++ b/docs/source/examples/07_record3d_visualizer.rst @@ -18,10 +18,11 @@ Parse and stream record3d captures. To get the demo data, see ``./assets/downloa import numpy as np import tyro + from tqdm.auto import tqdm + import viser import viser.extras import viser.transforms as tf - from tqdm.auto import tqdm def main( diff --git a/docs/source/examples/08_smpl_visualizer.rst b/docs/source/examples/08_smpl_visualizer.rst index 6d808008..dcd1ff67 100644 --- a/docs/source/examples/08_smpl_visualizer.rst +++ b/docs/source/examples/08_smpl_visualizer.rst @@ -24,7 +24,6 @@ See here for download instructions: import numpy as np import tyro - import viser import viser.transforms as tf @@ -207,12 +206,9 @@ See here for download instructions: @gui_random_joints.on_click def _(_): + rng = np.random.default_rng() for joint in gui_joints: - # It's hard to uniformly sample orientations directly in so(3), so we - # first sample on S^3 and then convert. - quat = np.random.normal(loc=0.0, scale=1.0, size=(4,)) - quat /= np.linalg.norm(quat) - joint.value = tf.SO3(wxyz=quat).log() + joint.value = tf.SO3.sample_uniform(rng).log() gui_joints: list[viser.GuiInputHandle[tuple[float, float, float]]] = [] for i in range(num_joints): diff --git a/docs/source/examples/25_smpl_visualizer_skinned.rst b/docs/source/examples/25_smpl_visualizer_skinned.rst index 1917d2ed..7d62f5c6 100644 --- a/docs/source/examples/25_smpl_visualizer_skinned.rst +++ b/docs/source/examples/25_smpl_visualizer_skinned.rst @@ -25,21 +25,18 @@ See here for download instructions: import numpy as np import tyro - import viser import viser.transforms as tf @dataclass(frozen=True) - class SmplOutputs: - vertices: np.ndarray - faces: np.ndarray + class SmplFkOutputs: T_world_joint: np.ndarray # (num_joints, 4, 4) T_parent_joint: np.ndarray # (num_joints, 4, 4) class SmplHelper: - """Helper for models in the SMPL family, implemented in numpy.""" + """Helper for models in the SMPL family, implemented in numpy. Does not include blend skinning.""" def __init__(self, model_path: Path) -> None: assert model_path.suffix.lower() == ".npz", "Model should be an .npz file!" @@ -62,7 +59,9 @@ See here for download instructions: j_tpose = np.einsum("jv,vx->jx", self.J_regressor, v_tpose) return v_tpose, j_tpose - def get_outputs(self, betas: np.ndarray, joint_rotmats: np.ndarray) -> SmplOutputs: + def get_outputs( + self, betas: np.ndarray, joint_rotmats: np.ndarray + ) -> SmplFkOutputs: # Get shaped vertices + joint positions, when all local poses are identity. v_tpose = self.v_template + np.einsum("vxb,b->vx", self.shapedirs, betas) j_tpose = np.einsum("jv,vx->jx", self.J_regressor, v_tpose) @@ -78,15 +77,7 @@ See here for download instructions: for i in range(1, self.num_joints): T_world_joint[i] = T_world_joint[self.parent_idx[i]] @ T_parent_joint[i] - # Linear blend skinning. - pose_delta = (joint_rotmats[1:, ...] - np.eye(3)).flatten() - v_blend = v_tpose + np.einsum("byn,n->by", self.posedirs, pose_delta) - v_delta = np.ones((v_blend.shape[0], self.num_joints, 4)) - v_delta[:, :, :3] = v_blend[:, None, :] - j_tpose[None, :, :] - v_posed = np.einsum( - "jxy,vj,vjy->vx", T_world_joint[:, :3, :], self.weights, v_delta - ) - return SmplOutputs(v_posed, self.faces, T_world_joint, T_parent_joint) + return SmplFkOutputs(T_world_joint, T_parent_joint) def main(model_path: Path) -> None: @@ -252,12 +243,9 @@ See here for download instructions: @gui_random_joints.on_click def _(_): + rng = np.random.default_rng() for joint in gui_joints: - # It's hard to uniformly sample orientations directly in so(3), so we - # first sample on S^3 and then convert. - quat = np.random.normal(loc=0.0, scale=1.0, size=(4,)) - quat /= np.linalg.norm(quat) - joint.value = tf.SO3(wxyz=quat).log() + joint.value = tf.SO3.sample_uniform(rng).log() gui_joints: List[viser.GuiInputHandle[Tuple[float, float, float]]] = [] for i in range(num_joints): diff --git a/examples/08_smpl_visualizer.py b/examples/08_smpl_visualizer.py index 66584605..aac2cc78 100644 --- a/examples/08_smpl_visualizer.py +++ b/examples/08_smpl_visualizer.py @@ -14,7 +14,6 @@ import numpy as np import tyro - import viser import viser.transforms as tf @@ -197,12 +196,9 @@ def _(_): @gui_random_joints.on_click def _(_): + rng = np.random.default_rng() for joint in gui_joints: - # It's hard to uniformly sample orientations directly in so(3), so we - # first sample on S^3 and then convert. - quat = np.random.normal(loc=0.0, scale=1.0, size=(4,)) - quat /= np.linalg.norm(quat) - joint.value = tf.SO3(wxyz=quat).log() + joint.value = tf.SO3.sample_uniform(rng).log() gui_joints: list[viser.GuiInputHandle[tuple[float, float, float]]] = [] for i in range(num_joints): diff --git a/examples/25_smpl_visualizer_skinned.py b/examples/25_smpl_visualizer_skinned.py index 91faa511..2b37fd71 100644 --- a/examples/25_smpl_visualizer_skinned.py +++ b/examples/25_smpl_visualizer_skinned.py @@ -20,21 +20,18 @@ import numpy as np import tyro - import viser import viser.transforms as tf @dataclass(frozen=True) -class SmplOutputs: - vertices: np.ndarray - faces: np.ndarray +class SmplFkOutputs: T_world_joint: np.ndarray # (num_joints, 4, 4) T_parent_joint: np.ndarray # (num_joints, 4, 4) class SmplHelper: - """Helper for models in the SMPL family, implemented in numpy.""" + """Helper for models in the SMPL family, implemented in numpy. Does not include blend skinning.""" def __init__(self, model_path: Path) -> None: assert model_path.suffix.lower() == ".npz", "Model should be an .npz file!" @@ -57,7 +54,9 @@ def get_tpose(self, betas: np.ndarray) -> tuple[np.ndarray, np.ndarray]: j_tpose = np.einsum("jv,vx->jx", self.J_regressor, v_tpose) return v_tpose, j_tpose - def get_outputs(self, betas: np.ndarray, joint_rotmats: np.ndarray) -> SmplOutputs: + def get_outputs( + self, betas: np.ndarray, joint_rotmats: np.ndarray + ) -> SmplFkOutputs: # Get shaped vertices + joint positions, when all local poses are identity. v_tpose = self.v_template + np.einsum("vxb,b->vx", self.shapedirs, betas) j_tpose = np.einsum("jv,vx->jx", self.J_regressor, v_tpose) @@ -73,15 +72,7 @@ def get_outputs(self, betas: np.ndarray, joint_rotmats: np.ndarray) -> SmplOutpu for i in range(1, self.num_joints): T_world_joint[i] = T_world_joint[self.parent_idx[i]] @ T_parent_joint[i] - # Linear blend skinning. - pose_delta = (joint_rotmats[1:, ...] - np.eye(3)).flatten() - v_blend = v_tpose + np.einsum("byn,n->by", self.posedirs, pose_delta) - v_delta = np.ones((v_blend.shape[0], self.num_joints, 4)) - v_delta[:, :, :3] = v_blend[:, None, :] - j_tpose[None, :, :] - v_posed = np.einsum( - "jxy,vj,vjy->vx", T_world_joint[:, :3, :], self.weights, v_delta - ) - return SmplOutputs(v_posed, self.faces, T_world_joint, T_parent_joint) + return SmplFkOutputs(T_world_joint, T_parent_joint) def main(model_path: Path) -> None: @@ -247,12 +238,9 @@ def _(_): @gui_random_joints.on_click def _(_): + rng = np.random.default_rng() for joint in gui_joints: - # It's hard to uniformly sample orientations directly in so(3), so we - # first sample on S^3 and then convert. - quat = np.random.normal(loc=0.0, scale=1.0, size=(4,)) - quat /= np.linalg.norm(quat) - joint.value = tf.SO3(wxyz=quat).log() + joint.value = tf.SO3.sample_uniform(rng).log() gui_joints: List[viser.GuiInputHandle[Tuple[float, float, float]]] = [] for i in range(num_joints):