""" Thanks to the author of this API Tomoya Akiyama: https://research.cyberagent.ai/people/tomoya_akiyama/ """ import math import cv2 import numpy as np import torch import smplx from pytorch3d.renderer import PerspectiveCameras from torchvision.io import write_video from torchvision.transforms.functional import convert_image_dtype SMPLX_BODY_JOINT_EDGES = [ {"indices": [12, 17], "color": [255, 0, 0]}, {"indices": [12, 16], "color": [255, 85, 0]}, {"indices": [17, 19], "color": [255, 170, 0]}, {"indices": [19, 21], "color": [255, 255, 0]}, {"indices": [16, 18], "color": [170, 255, 0]}, {"indices": [18, 20], "color": [85, 255, 0]}, {"indices": [2, 12], "color": [0, 255, 0]}, {"indices": [2, 5], "color": [0, 255, 85]}, {"indices": [5, 8], "color": [0, 255, 170]}, {"indices": [1, 12], "color": [0, 255, 255]}, {"indices": [1, 4], "color": [0, 170, 255]}, {"indices": [4, 7], "color": [0, 85, 255]}, {"indices": [12, 55], "color": [0, 0, 255]}, {"indices": [55, 56], "color": [85, 0, 255]}, {"indices": [56, 58], "color": [170, 0, 255]}, {"indices": [55, 57], "color": [255, 0, 255]}, {"indices": [57, 59], "color": [255, 0, 170]}, ] SMPLX_BODY_JOINTS = [ {"index": 55, "color": [255, 0, 0]}, {"index": 12, "color": [255, 85, 0]}, {"index": 17, "color": [255, 170, 0]}, {"index": 19, "color": [255, 255, 0]}, {"index": 21, "color": [170, 255, 0]}, {"index": 16, "color": [85, 255, 0]}, {"index": 18, "color": [0, 255, 0]}, {"index": 20, "color": [0, 255, 85]}, {"index": 2, "color": [0, 255, 170]}, {"index": 5, "color": [0, 255, 255]}, {"index": 8, "color": [0, 170, 255]}, {"index": 1, "color": [0, 85, 255]}, {"index": 4, "color": [0, 0, 255]}, {"index": 7, "color": [85, 0, 255]}, {"index": 56, "color": [170, 0, 255]}, {"index": 57, "color": [255, 0, 255]}, {"index": 58, "color": [255, 0, 170]}, {"index": 59, "color": [255, 0, 85]}, ] SMPLX_HAND_JOINT_EDGES = [ {"indices": [21, 52], "color": [255, 0, 0]}, {"indices": [52, 53], "color": [255, 76, 0]}, {"indices": [53, 54], "color": [255, 153, 0]}, {"indices": [54, 71], "color": [255, 229, 0]}, {"indices": [21, 40], "color": [204, 255, 0]}, {"indices": [40, 41], "color": [128, 255, 0]}, {"indices": [41, 42], "color": [51, 255, 0]}, {"indices": [42, 72], "color": [0, 255, 26]}, {"indices": [21, 43], "color": [0, 255, 102]}, {"indices": [43, 44], "color": [0, 255, 179]}, {"indices": [44, 45], "color": [0, 255, 255]}, {"indices": [45, 73], "color": [0, 179, 255]}, {"indices": [21, 49], "color": [0, 102, 255]}, {"indices": [49, 50], "color": [0, 26, 255]}, {"indices": [50, 51], "color": [51, 0, 255]}, {"indices": [51, 74], "color": [128, 0, 255]}, {"indices": [21, 46], "color": [204, 0, 255]}, {"indices": [46, 47], "color": [255, 0, 230]}, {"indices": [47, 48], "color": [255, 0, 153]}, {"indices": [48, 75], "color": [255, 0, 77]}, {"indices": [20, 37], "color": [255, 0, 0]}, {"indices": [37, 38], "color": [255, 76, 0]}, {"indices": [38, 39], "color": [255, 153, 0]}, {"indices": [39, 66], "color": [255, 229, 0]}, {"indices": [20, 25], "color": [204, 255, 0]}, {"indices": [25, 26], "color": [128, 255, 0]}, {"indices": [26, 27], "color": [51, 255, 0]}, {"indices": [27, 67], "color": [0, 255, 26]}, {"indices": [20, 28], "color": [0, 255, 102]}, {"indices": [28, 29], "color": [0, 255, 179]}, {"indices": [29, 30], "color": [0, 255, 255]}, {"indices": [30, 68], "color": [0, 179, 255]}, {"indices": [20, 34], "color": [0, 102, 255]}, {"indices": [34, 35], "color": [0, 26, 255]}, {"indices": [35, 36], "color": [51, 0, 255]}, {"indices": [36, 69], "color": [128, 0, 255]}, {"indices": [20, 31], "color": [204, 0, 255]}, {"indices": [31, 32], "color": [255, 0, 230]}, {"indices": [32, 33], "color": [255, 0, 153]}, {"indices": [33, 70], "color": [255, 0, 77]}, ] SMPLX_HAND_JOINTS = [20, 21] + list(range(25, 55)) + list(range(66, 76)) SMPLX_FACE_LANDMARKS = list(range(76, 144)) def _draw_bodypose(canvas, joints_np): c = canvas.copy() for edge_dict in SMPLX_BODY_JOINT_EDGES: i = edge_dict["indices"] color = edge_dict["color"] xy = joints_np[i] center = np.mean(xy, axis=0).astype(int) length = np.linalg.norm(xy[0] - xy[1]) angle = math.degrees(math.atan2(xy[0, 1] - xy[1, 1], xy[0, 0] - xy[1, 0])) polygon = cv2.ellipse2Poly(center, (int(length / 2), 4), int(angle), 0, 360, 1) cv2.fillConvexPoly(c, polygon, color) c = (c * 0.6).astype(np.uint8) for j_info in SMPLX_BODY_JOINTS: center = joints_np[j_info["index"]].astype(int) cv2.circle(c, tuple(center), 4, (255, 255, 255), -1) return c def _draw_handpose(canvas, joints_np): c = canvas.copy() for edge_dict in SMPLX_HAND_JOINT_EDGES: i = edge_dict["indices"] color = edge_dict["color"] xy = joints_np[i].astype(int) if xy.min() > 0: cv2.line(c, tuple(xy[0]), tuple(xy[1]), color, 2) for j_idx in SMPLX_HAND_JOINTS: center = joints_np[j_idx].astype(int) if center.min() > 0: cv2.circle(c, tuple(center), 4, (0, 0, 255), -1) return c def _draw_facepose(canvas, joints_np): c = canvas.copy() for j_idx in SMPLX_FACE_LANDMARKS: center = joints_np[j_idx].astype(int) if center.min() > 0: cv2.circle(c, tuple(center), 3, (255, 255, 255), -1) return c def _draw_joints_2d(joints_2d, height, width, face_only): outputs = [] for j2d in joints_2d: # Convert each frame's joints to NumPy j2d_np = j2d.detach().cpu().numpy() c = np.zeros((height, width, 3), dtype=np.uint8) if face_only: c = _draw_facepose(c, j2d_np) else: c = _draw_bodypose(c, j2d_np) c = _draw_handpose(c, j2d_np) c = _draw_facepose(c, j2d_np) outputs.append(convert_image_dtype(torch.tensor(c, dtype=torch.uint8), torch.uint8)) return torch.stack(outputs).permute(0, 3, 1, 2) def _draw_joints_3d(joints_3d, height, width, face_only): outputs = [] for j3d in joints_3d: xy = j3d[:, :2].detach().cpu().numpy().copy() z = j3d[:, 2].detach().cpu().numpy().copy() z_min, z_max = z.min(), z.max() z_norm = (z - z_min) / (z_max - z_min + 1e-8) # Normalize XY to fit in the image xy[:, 0] = (xy[:, 0] - xy[:, 0].min()) / (xy[:, 0].max() - xy[:, 0].min() + 1e-8) * (width - 1) xy[:, 1] = (xy[:, 1] - xy[:, 1].min()) / (xy[:, 1].max() - xy[:, 1].min() + 1e-8) * (height - 1) c = np.zeros((height, width, 3), dtype=np.uint8) # j2d: [num_joints, 3], last dim is the normalized z j2d = np.hstack([xy, z_norm.reshape(-1, 1)]) if face_only: c = _draw_facepose(c, j2d) else: c = _draw_bodypose(c, j2d) c = _draw_handpose(c, j2d) c = _draw_facepose(c, j2d) outputs.append(convert_image_dtype(torch.tensor(c, dtype=torch.uint8), torch.uint8)) return torch.stack(outputs).permute(0, 3, 1, 2) def _load_motion_dict( motion_dict, device, remove_global=False, face_only=False ): n = motion_dict["poses"].shape[0] smplx_inputs = { "betas": torch.tensor(motion_dict["betas"]).view(1, -1), "global_orient": torch.tensor(motion_dict["poses"][:, :3]), "body_pose": torch.tensor(motion_dict["poses"][:, 3 : 22 * 3]), "left_hand_pose": torch.tensor(motion_dict["poses"][:, 25 * 3 : 40 * 3]), "right_hand_pose": torch.tensor(motion_dict["poses"][:, 40 * 3 : 55 * 3]), "transl": torch.tensor(motion_dict["trans"]), "expression": torch.tensor(motion_dict["expressions"]), "jaw_pose": torch.tensor(motion_dict["poses"][:, 22 * 3 : 23 * 3]), "leye_pose": torch.tensor(motion_dict["poses"][:, 23 * 3 : 24 * 3]), "reye_pose": torch.tensor(motion_dict["poses"][:, 24 * 3 : 25 * 3]), } # Move everything to device for k, v in smplx_inputs.items(): smplx_inputs[k] = v.to(device=device, dtype=torch.float32) # 1) If remove_global == True, keep 'transl' at the first frame's value for all frames if remove_global: first_frame_trans = smplx_inputs["transl"][0].clone() smplx_inputs["transl"][:] = first_frame_trans # 2) If face_only == True, zero out everything but the jaw pose if face_only: smplx_inputs["global_orient"][:] = 0.0 smplx_inputs["body_pose"][:] = 0.0 smplx_inputs["left_hand_pose"][:] = 0.0 smplx_inputs["right_hand_pose"][:] = 0.0 smplx_inputs["leye_pose"][:] = 0.0 smplx_inputs["reye_pose"][:] = 0.0 # The jaw_pose and expression remain as is (allowing mouth movements), # so the head is "frozen" in place except for jaw animation. return n, smplx_inputs def _get_smplx_model(smplx_folder, batch_size, device): smplx_model = smplx.create( model_path=smplx_folder, model_type="smplx", gender="NEUTRAL_2020", create_global_orient=True, create_body_pose=True, create_betas=True, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=True, create_jaw_pose=True, create_leye_pose=True, create_reye_pose=True, create_transl=True, use_face_contour=True, use_pca=False, flat_hand_mean=False, use_hands=True, use_face=True, num_betas=300, num_expression_coeffs=100, batch_size=batch_size, dtype=torch.float32, ).to(device) return smplx_model.eval() def _get_cameras( batch_size, height, width, focal_length, camera_transl, device ): r = torch.tensor( [[-1, 0, 0], [ 0, 1, 0], [ 0, 0, 1]], device=device, dtype=torch.float32 ) t = torch.tensor(camera_transl, device=device, dtype=torch.float32) cameras = PerspectiveCameras( focal_length=focal_length, principal_point=((width / 2, height / 2),), in_ndc=False, R=r.expand(batch_size, -1, -1), T=t.expand(batch_size, -1), image_size=((height, width),), device=device, ) return cameras # New fix code snippet (inside render2d or render3d): def render2d( motion_dict, resolution=(512, 512), face_only=False, remove_global=False, smplx_folder="./emage_evaltools/smplx_models/", focal_length=5000.0, camera_transl=(0.0, -0.8, 16.0), device=torch.device("cuda"), ): h, w = resolution # for face-only, override camera to zoom in if face_only: camera_transl = (0.0, -1.55, 6.0) focal_length = 10000.0 n, smplx_inputs = _load_motion_dict( motion_dict, device, remove_global=remove_global, face_only=face_only ) model = _get_smplx_model(smplx_folder, n, device) out = model(**smplx_inputs) cams = _get_cameras(n, h, w, focal_length, camera_transl, device) j2d = cams.transform_points_screen(out.joints)[:, :, :2] frames_2d = _draw_joints_2d(j2d, h, w, face_only) return frames_2d def render3d( motion_dict, resolution=(512, 512), face_only=False, remove_global=False, smplx_folder="./emage_evaltools/smplx_models/", device=torch.device("cuda"), ): h, w = resolution n, smplx_inputs = _load_motion_dict( motion_dict, device, remove_global=remove_global, face_only=face_only ) model = _get_smplx_model(smplx_folder, n, device) out = model(**smplx_inputs) frames_3d = _draw_joints_3d(out.joints, h, w, face_only) return frames_3d def example_usage(): # Suppose we have an NPZ with "poses", "trans", "betas", "expressions", etc. motion_dict = np.load("/result_motion.npz", allow_pickle=True) # 2D face (freeze body, remove global motion) v2d_face = render2d( motion_dict, resolution=(512, 512), face_only=True, remove_global=True ) write_video("/save_path_face_2d.mp4", v2d_face.permute(0, 2, 3, 1), fps=30) # 2D body (show entire body, keep global motion) v2d_body = render2d( motion_dict, resolution=(1080, 1920), face_only=False, remove_global=False ) write_video("/save_path_body_2d.mp4", v2d_body.permute(0, 2, 3, 1), fps=30) # 3D face (freeze body, remove global motion) v3d_face = render3d( motion_dict, resolution=(512, 512), face_only=True, remove_global=True ) write_video("/save_path_face_3d.mp4", v3d_face.permute(0, 2, 3, 1), fps=30) # 3D body (show entire body, keep global motion) v3d_body = render3d( motion_dict, resolution=(1080, 1920), face_only=False, remove_global=False ) write_video("/save_path_body_3d.mp4", v3d_body.permute(0, 2, 3, 1), fps=30) if __name__ == "__main__": example_usage()