# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ MHR shape to SMPL shape converter. This class handles the conversion of MHR shape parameters to SMPL shape parameters. """ import argparse from pathlib import Path import imageio import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import trimesh from tqdm import tqdm from soma.geometry.barycentric_interp import BarycentricInterpolator from soma.geometry.batched_skinning import BatchedSkinning from soma.soma import SOMALayer from tools.vis_pyrender import MeshRenderer, look_at, set_pyopengl_platform def get_smooth_noise(T, dim, device, num_keyframes=None, mode="normal"): if num_keyframes is None: num_keyframes = max(3, T // 30) if mode == "normal": keyframes = torch.randn(1, dim, num_keyframes, device=device) elif mode == "uniform": keyframes = torch.rand(1, dim, num_keyframes, device=device) res = F.interpolate(keyframes, size=T, mode="linear", align_corners=True)[0].T return res class ShapeTransfer(nn.Module): """ MHR shape to SMPL shape converter. """ def __init__(self, data_root, device="cuda"): """ Initialize the shape transfer. """ super().__init__() self.data_root = Path(data_root) self.device = device self.mhr_soma = SOMALayer( data_root=data_root, device=device, identity_model_type="mhr", mode="warp" ) self.smpl_soma = SOMALayer( data_root=data_root, device=device, identity_model_type="smpl", mode="warp" ) self.soma_to_smpl = self.get_soma_to_smpl_interpolator() smpl_rest_mesh = trimesh.load( self.data_root / "SMPL" / "smpl_base_body.obj", maintain_order=True, process=False ) smpl_rest_shape = torch.from_numpy(smpl_rest_mesh.vertices).float().to(device).unsqueeze(0) self.smpl_rest_shape_soma = self.smpl_soma.identity_model.identity_model_to_soma( smpl_rest_shape ) self.posed_world_smpl_tpose = self.mhr_soma.skeleton_transfer.fit(self.smpl_rest_shape_soma) def get_soma_to_smpl_interpolator(self): mesh_smpl = trimesh.load( self.data_root / "SMPL" / "smpl_base_body.obj", maintain_order=True, process=False ) V_smpl = torch.from_numpy(mesh_smpl.vertices).float().to(self.device) mesh_nv = trimesh.load( self.data_root / "SMPL" / "Nova_wrap.obj", maintain_order=True, process=False ) V_nv = torch.from_numpy(mesh_nv.vertices).float().to(self.device) F_nv = torch.from_numpy(mesh_nv.faces).int().to(self.device) return BarycentricInterpolator(V_nv, F_nv, V_smpl) def forward(self, identity_coeffs, scale_params): """ Forward pass. """ batch_size = identity_coeffs.shape[0] device = identity_coeffs.device # 1. Get MHR rest shape mhr_rest_shape_soma = self.mhr_soma.identity_model(identity_coeffs, scale_params) # 2. Get MHR posed world transforms posed_world_mhr = self.mhr_soma.skeleton_transfer.fit(mhr_rest_shape_soma) # 3. Skin the MHR rest shape to get it under the SMPL rest pose batched_skinning = BatchedSkinning( self.mhr_soma.joint_parent_ids, self.mhr_soma.skinning_weights, posed_world_mhr, mhr_rest_shape_soma, joint_orient=self.posed_world_smpl_tpose[0], mode=self.mhr_soma.mode, ) pose_rotations = torch.eye(3, device=device).unsqueeze(0).expand(batch_size, 78, 3, 3) pose_translations = self.posed_world_smpl_tpose[:, 1, :3, 3].expand(batch_size, 3) vertices, T_world = batched_skinning.pose( local_rotations=pose_rotations, hips_translations=pose_translations, return_transforms=True, ) # 4. Get the SMPL topology vertices mhr_vertices_smpl = self.soma_to_smpl(vertices) # 5. Solve the betas B = ( mhr_vertices_smpl - self.smpl_soma.identity_model.identity_model.v_template[None] ).reshape(batch_size, -1) A = self.smpl_soma.identity_model.identity_model.shape_dirs.reshape(-1, 10) betas = torch.linalg.lstsq(A, B.T).solution.T return betas # Example usage if __name__ == "__main__": parser = argparse.ArgumentParser(description="Shape transfer.") parser.add_argument("--data_root", type=str, default="./assets", help="Path to the data root.") parser.add_argument("--device", default="cuda:0") parser.add_argument("--output-dir", default="out/") parser.add_argument("--image-size", type=int, default=1920) parser.add_argument("--sequence-length", type=int, default=300) parser.add_argument("--pyopengl-platform", default="osmesa") args = parser.parse_args() set_pyopengl_platform(args.pyopengl_platform) device = "cuda" shape_transfer = ShapeTransfer(args.data_root, device) T = args.sequence_length mhr_im = shape_transfer.mhr_soma.identity_model identity_coeffs = get_smooth_noise(T, mhr_im.num_identity_coeffs, device) scale_params = get_smooth_noise(T, mhr_im.num_scale_params, device, mode="normal") * 0.2 zero_pose = torch.zeros(1, 77, 3, device=device) zero_transl = torch.zeros(1, 3, device=device) betas = shape_transfer(identity_coeffs, scale_params) smpl_vertices = shape_transfer.smpl_soma(zero_pose, betas, None, zero_transl)["vertices"] mhr_vertices = shape_transfer.mhr_soma(zero_pose, identity_coeffs, scale_params, zero_transl)[ "vertices" ] smpl_vertices = smpl_vertices.detach().cpu().numpy() mhr_vertices = mhr_vertices.detach().cpu().numpy() faces = shape_transfer.mhr_soma.faces.cpu().numpy() print("Rendering videos...") colors = { "mhr": (0.98, 0.65, 0.15, 1.0), "anny": (0.25, 0.75, 1.0, 1.0), "smpl": (0.55, 0.15, 0.85, 1.0), } def save_video(frames, path, fps=30): imageio.mimsave(path, frames, fps=fps) print(f"Saved {path}") renderer = MeshRenderer(image_size=args.image_size, light_intensity=5) cam_pose = look_at( eye=np.array([0.0, 0.0, 6.0]), target=np.array([0.0, 0.0, 0.0]), up=np.array([0.0, 1.0, 0.0]), ) light_dir = np.array([0.0, -0.5, -1.0]) faces = shape_transfer.mhr_soma.faces.cpu().numpy() frames = [] for t in tqdm(range(T)): mhr_img = renderer.render( mhr_vertices[t], faces, mesh_color=colors["mhr"], cam_pose=cam_pose, light_dir=light_dir, metallic=0.0, roughness=0.5, base_color_factor=[0.9, 0.9, 0.9, 1.0], ) smpl_img = renderer.render( smpl_vertices[t], faces, mesh_color=colors["smpl"], cam_pose=cam_pose, light_dir=light_dir, metallic=0.0, roughness=0.5, base_color_factor=[0.9, 0.9, 0.9, 1.0], ) merged_img = (0.5 * mhr_img + 0.5 * smpl_img).astype(np.uint8) img = np.concatenate([mhr_img, merged_img, smpl_img], axis=1) frames.append(img[..., ::-1]) renderer.delete() save_video(frames, Path(args.output_dir) / "shape_transfer.mp4")