Upload smpl.py

src/spin/smpl.py

@@ -0,0 +1,436 @@
+import numpy as np
+import torch
+from smplx import SMPL as _SMPL
+from smplx import SMPLX as _SMPLX
+from smplx.body_models import SMPLOutput, SMPLXOutput
+from smplx.lbs import vertices2joints
+
+from .constants import JOINT_MAP, JOINT_NAMES
+
+# Hand joints
+SMPLX_HAND_TO_PANOPTIC = [
+    0,
+    13,
+    14,
+    15,
+    16,
+    1,
+    2,
+    3,
+    17,
+    4,
+    5,
+    6,
+    18,
+    10,
+    11,
+    12,
+    19,
+    7,
+    8,
+    9,
+    20,
+]  # Wrist Thumb to Pinky
+
+
+
+class SMPL(_SMPL):
+    """Extension of the official SMPL implementation to support more joints"""
+
+    JOINTS = (
+        "Hips",
+        "Left Upper Leg",
+        "Right Upper Leg",
+        "Spine",
+        "Left Leg",
+        "Right Leg",
+        "Spine1",
+        "Left Foot",
+        "Right Foot",
+        "Thorax",
+        "Left Toe",
+        "Right Toe",
+        "Neck",
+        "Left Shoulder",
+        "Right Shoulder",
+        "Head",
+        "Left ForeArm",
+        "Right ForeArm",
+        "Left Arm",
+        "Right Arm",
+        "Left Hand",
+        "Right Hand",
+        "Left Finger",
+        "Right Finger",
+    )
+
+    SKELETON = (
+        (0, 1),
+        (0, 2),
+        (0, 3),
+        (1, 4),
+        (2, 5),
+        (3, 6),
+        (4, 7),
+        (5, 8),
+        (6, 9),
+        (7, 10),
+        (8, 11),
+        (9, 12),
+        (12, 13),
+        (12, 14),
+        (12, 15),
+        (13, 16),
+        (14, 17),
+        (16, 18),
+        (17, 19),
+        (18, 20),
+        (19, 21),
+        (20, 22),
+        (21, 23),
+    )
+
+    def __init__(self, *args, **kwargs):
+        super(SMPL, self).__init__(*args, **kwargs)
+        joints = [JOINT_MAP[i] for i in JOINT_NAMES]
+        joint_regressor_extra = kwargs["joint_regressor_extra_path"]
+        J_regressor_extra = np.load(joint_regressor_extra)
+        self.register_buffer(
+            "J_regressor_extra", torch.tensor(J_regressor_extra, dtype=torch.float32)
+        )
+        self.joint_map = torch.tensor(joints, dtype=torch.long)
+
+    def forward(self, *args, **kwargs):
+        kwargs["get_skin"] = True
+        smpl_output = super(SMPL, self).forward(*args, **kwargs)
+        extra_joints = vertices2joints(
+            self.J_regressor_extra, smpl_output.vertices
+        )  # Additional 9 joints #Check doc/J_regressor_extra.png
+        joints = torch.cat(
+            [smpl_output.joints, extra_joints], dim=1
+        )  # [N, 24 + 21, 3]  + [N, 9, 3]
+        joints = joints[:, self.joint_map, :]
+        output = SMPLOutput(
+            vertices=smpl_output.vertices,
+            global_orient=smpl_output.global_orient,
+            body_pose=smpl_output.body_pose,
+            joints=joints,
+            betas=smpl_output.betas,
+            full_pose=smpl_output.full_pose,
+        )
+        return output
+
+
+class SMPLX(_SMPLX):
+    """Extension of the official SMPL implementation to support more joints"""
+
+    JOINTS = (
+        "Hips",
+        "Left Upper Leg",
+        "Right Upper Leg",
+        "Spine",
+        "Left Leg",
+        "Right Leg",
+        "Spine1",
+        "Left Foot",
+        "Right Foot",
+        "Thorax",
+        "Left Toe",
+        "Right Toe",
+        "Neck",
+        "Left Shoulder",
+        "Right Shoulder",
+        "Head",
+        "Left ForeArm",
+        "Right ForeArm",
+        "Left Arm",
+        "Right Arm",
+        "Left Hand",
+        "Right Hand",
+    )
+
+    SKELETON = (
+        (0, 1),
+        (0, 2),
+        (0, 3),
+        (1, 4),
+        (2, 5),
+        (3, 6),
+        (4, 7),
+        (5, 8),
+        (6, 9),
+        (7, 10),
+        (8, 11),
+        (9, 12),
+        (12, 13),
+        (12, 14),
+        (12, 15),
+        (13, 16),
+        (14, 17),
+        (16, 18),
+        (17, 19),
+        (18, 20),
+        (19, 21),
+    )
+
+    def __init__(self, *args, **kwargs):
+        kwargs["ext"] = "pkl"  # We have pkl file
+        super(SMPLX, self).__init__(*args, **kwargs)
+        joints = [JOINT_MAP[i] for i in JOINT_NAMES]
+        self.joint_map = torch.tensor(joints, dtype=torch.long)
+
+    def forward(self, *args, **kwargs):
+        kwargs["get_skin"] = True
+
+        # if pose parameter is for SMPL with 21 joints (ignoring root)
+        try:
+            if kwargs["body_pose"].shape[1] == 69:
+                kwargs["body_pose"] = kwargs["body_pose"][
+                    :, : -2 * 3
+                ]  # Ignore the last two joints (which are on the palm. Not used)
+
+            if kwargs["body_pose"].shape[1] == 23:
+                kwargs["body_pose"] = kwargs["body_pose"][
+                    :, :-2
+                ]  # Ignore the last two joints (which are on the palm. Not used)
+        except:
+            pass
+
+        smpl_output = super(SMPLX, self).forward(*args, **kwargs)
+
+        # SMPL-X Joint order: https://docs.google.com/spreadsheets/d/1_1dLdaX-sbMkCKr_JzJW_RZCpwBwd7rcKkWT_VgAQ_0/edit#gid=0
+        smplx_to_smpl = (
+            list(range(0, 22)) + [28, 43] + list(range(55, 76))
+        )  # 28 left middle finger , 43: right middle finger 1
+        smpl_joints = smpl_output.joints[
+            :, smplx_to_smpl, :
+        ]  # Convert SMPL-X to SMPL     127 ->45
+        joints = smpl_joints
+        joints = joints[:, self.joint_map, :]
+
+        smplx_lhand = (
+            [20] + list(range(25, 40)) + list(range(66, 71))
+        )  # 20 for left wrist. 20 finger joints
+        lhand_joints = smpl_output.joints[:, smplx_lhand, :]  # (N,21,3)
+        lhand_joints = lhand_joints[
+            :, SMPLX_HAND_TO_PANOPTIC, :
+        ]  # Convert SMPL-X hand order to paonptic hand order
+
+        smplx_rhand = (
+            [21] + list(range(40, 55)) + list(range(71, 76))
+        )  # 21 for right wrist. 20 finger joints
+        rhand_joints = smpl_output.joints[:, smplx_rhand, :]  # (N,21,3)
+        rhand_joints = rhand_joints[
+            :, SMPLX_HAND_TO_PANOPTIC, :
+        ]  # Convert SMPL-X hand order to paonptic hand order
+
+        output = SMPLXOutput(
+            vertices=smpl_output.vertices,
+            global_orient=smpl_output.global_orient,
+            body_pose=smpl_output.body_pose,
+            joints=joints,
+            right_hand_pose=rhand_joints,  # N,21,3
+            left_hand_pose=lhand_joints,  # N,21,3
+            betas=smpl_output.betas,
+            full_pose=smpl_output.full_pose,
+            A=smpl_output.A,
+        )
+        return output
+
+
+"""
+0	pelvis',
+1	left_hip',
+2	right_hip',
+3	spine1',
+4	left_knee',
+5	right_knee',
+6	spine2',
+7	left_ankle',
+8	right_ankle',
+9	spine3',
+10	left_foot',
+11	right_foot',
+12	neck',
+13	left_collar',
+14	right_collar',
+15	head',
+16	left_shoulder',
+17	right_shoulder',
+18	left_elbow',
+19	right_elbow',
+20	left_wrist',
+21	right_wrist',
+22	jaw',
+23	left_eye_smplhf',
+24	right_eye_smplhf',
+25	left_index1',
+26	left_index2',
+27	left_index3',
+28	left_middle1',
+29	left_middle2',
+30	left_middle3',
+31	left_pinky1',
+32	left_pinky2',
+33	left_pinky3',
+34	left_ring1',
+35	left_ring2',
+36	left_ring3',
+37	left_thumb1',
+38	left_thumb2',
+39	left_thumb3',
+40	right_index1',
+41	right_index2',
+42	right_index3',
+43	right_middle1',
+44	right_middle2',
+45	right_middle3',
+46	right_pinky1',
+47	right_pinky2',
+48	right_pinky3',
+49	right_ring1',
+50	right_ring2',
+51	right_ring3',
+52	right_thumb1',
+53	right_thumb2',
+54	right_thumb3',
+55	nose',
+56	right_eye',
+57	left_eye',
+58	right_ear',
+59	left_ear',
+60	left_big_toe',
+61	left_small_toe',
+62	left_heel',
+63	right_big_toe',
+64	right_small_toe',
+65	right_heel',
+66	left_thumb',
+67	left_index',
+68	left_middle',
+69	left_ring',
+70	left_pinky',
+71	right_thumb',
+72	right_index',
+73	right_middle',
+74	right_ring',
+75	right_pinky',
+76	right_eye_brow1',
+77	right_eye_brow2',
+78	right_eye_brow3',
+79	right_eye_brow4',
+80	right_eye_brow5',
+81	left_eye_brow5',
+82	left_eye_brow4',
+83	left_eye_brow3',
+84	left_eye_brow2',
+85	left_eye_brow1',
+86	nose1',
+87	nose2',
+88	nose3',
+89	nose4',
+90	right_nose_2',
+91	right_nose_1',
+92	nose_middle',
+93	left_nose_1',
+94	left_nose_2',
+95	right_eye1',
+96	right_eye2',
+97	right_eye3',
+98	right_eye4',
+99	right_eye5',
+100	right_eye6',
+101	left_eye4',
+102	left_eye3',
+103	left_eye2',
+104	left_eye1',
+105	left_eye6',
+106	left_eye5',
+107	right_mouth_1',
+108	right_mouth_2',
+109	right_mouth_3',
+110	mouth_top',
+111	left_mouth_3',
+112	left_mouth_2',
+113	left_mouth_1',
+114	left_mouth_5', # 59 in OpenPose output
+115	left_mouth_4', # 58 in OpenPose output
+116	mouth_bottom',
+117	right_mouth_4',
+118	right_mouth_5',
+119	right_lip_1',
+120	right_lip_2',
+121	lip_top',
+122	left_lip_2',
+123	left_lip_1',
+124	left_lip_3',
+125	lip_bottom',
+126	right_lip_3',
+127	right_contour_1',
+128	right_contour_2',
+129	right_contour_3',
+130	right_contour_4',
+131	right_contour_5',
+132	right_contour_6',
+133	right_contour_7',
+134	right_contour_8',
+135	contour_middle',
+136	left_contour_8',
+137	left_contour_7',
+138	left_contour_6',
+139	left_contour_5',
+140	left_contour_4',
+141	left_contour_3',
+142	left_contour_2',
+143	left_contour_1'
+"""
+
+
+# SMPL Joints:
+"""
+0	pelvis',
+1	left_hip',
+2	right_hip',
+3	spine1',
+4	left_knee',
+5	right_knee',
+6	spine2',
+7	left_ankle',
+8	right_ankle',
+9	spine3',
+10	left_foot',
+11	right_foot',
+12	neck',
+13	left_collar',
+14	right_collar',
+15	head',
+16	left_shoulder',
+17	right_shoulder',
+18	left_elbow',
+19	right_elbow',
+20	left_wrist',
+21	right_wrist',
+22	
+23	
+24	nose',
+25	right_eye',
+26	left_eye',
+27	right_ear',
+28	left_ear',
+29	left_big_toe',
+30	left_small_toe',
+31	left_heel',
+32	right_big_toe',
+33	right_small_toe',
+34	right_heel',
+35	left_thumb',
+36	left_index',
+37	left_middle',
+38	left_ring',
+39	left_pinky',
+40	right_thumb',
+41	right_index',
+42	right_middle',
+43	right_ring',
+44	right_pinky',
+"""