Upload pipeline/tpose_smpl.py with huggingface_hub

pipeline/tpose_smpl.py

@@ -0,0 +1,413 @@
+"""
+tpose_smpl.py -- T-pose a humanoid GLB via inverse Linear Blend Skinning.
+
+Pipeline:
+  1. Render front view and run HMR2 -> SMPL body_pose + betas
+  2. Read rigged.glb: mesh verts (rig world space), skinning weights, T-pose joints
+  3. Compute FK transforms in rig world space using HMR2 body_pose
+  4. Apply inverse LBS: v_tpose = (Sum_j W_j * A_j)^-1 * v_posed
+  5. Map T-posed verts back to original mesh coordinate space, preserve UV/texture
+  6. Optionally export SKEL bone mesh in T-pose
+
+Usage:
+    python tpose_smpl.py --body /tmp/triposg_textured.glb \
+                         --rig  /tmp/rig_out/rigged.glb \
+                         --out  /tmp/tposed_surface.glb \
+                         [--skel_out /tmp/tposed_bones.glb] \
+                         [--debug_dir /tmp/tpose_debug]
+"""
+
+import os, sys, argparse, struct, json, warnings
+warnings.filterwarnings('ignore')
+
+import numpy as np
+import cv2
+import torch
+import trimesh
+from trimesh.visual.texture import TextureVisuals
+from trimesh.visual.material import PBRMaterial
+from scipy.spatial.transform import Rotation as R
+
+sys.path.insert(0, '/root/MV-Adapter')
+SMPL_NEUTRAL = '/root/body_models/smpl/SMPL_NEUTRAL.pkl'
+SKEL_DIR     = '/root/body_models/skel'
+
+SMPL_PARENTS = [-1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9,
+                12, 13, 14, 16, 17, 18, 19, 20, 21]
+
+
+# ---- Step 1: Render front view -----------------------------------------------
+
+def render_front(body_glb, H=1024, W=768, device='cuda'):
+    from mvadapter.utils.mesh_utils import (
+        NVDiffRastContextWrapper, load_mesh, get_orthogonal_camera, render,
+    )
+    ctx     = NVDiffRastContextWrapper(device=device, context_type='cuda')
+    mesh_mv = load_mesh(body_glb, rescale=True, device=device)
+    camera  = get_orthogonal_camera(
+        elevation_deg=[0], distance=[1.8],
+        left=-0.55, right=0.55, bottom=-0.55, top=0.55,
+        azimuth_deg=[-90], device=device,
+    )
+    out = render(ctx, mesh_mv, camera, height=H, width=W,
+                 render_attr=True, render_depth=False, render_normal=False,
+                 attr_background=0.5)
+    img_np = (out.attr[0].cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
+    return cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+
+
+# ---- Step 2: HMR2 pose estimation --------------------------------------------
+
+def run_hmr2(img_bgr, device='cuda'):
+    from pathlib import Path
+    from hmr2.configs import CACHE_DIR_4DHUMANS
+    from hmr2.models import load_hmr2, DEFAULT_CHECKPOINT, download_models
+    from hmr2.utils import recursive_to
+    from hmr2.datasets.vitdet_dataset import ViTDetDataset
+    from hmr2.utils.utils_detectron2 import DefaultPredictor_Lazy
+    from detectron2.config import LazyConfig
+    import hmr2 as hmr2_pkg
+
+    download_models(CACHE_DIR_4DHUMANS)
+    model, model_cfg = load_hmr2(DEFAULT_CHECKPOINT)
+    model = model.to(device).eval()
+
+    cfg_path = Path(hmr2_pkg.__file__).parent / 'configs' / 'cascade_mask_rcnn_vitdet_h_75ep.py'
+    det_cfg  = LazyConfig.load(str(cfg_path))
+    det_cfg.train.init_checkpoint = (
+        'https://dl.fbaipublicfiles.com/detectron2/ViTDet/COCO/cascade_mask_rcnn_vitdet_h'
+        '/f328730692/model_final_f05665.pkl'
+    )
+    for i in range(3):
+        det_cfg.model.roi_heads.box_predictors[i].test_score_thresh = 0.25
+    detector = DefaultPredictor_Lazy(det_cfg)
+
+    det_out   = detector(img_bgr)
+    instances = det_out['instances']
+    valid     = (instances.pred_classes == 0) & (instances.scores > 0.5)
+    boxes     = instances.pred_boxes.tensor[valid].cpu().numpy()
+    if len(boxes) == 0:
+        raise RuntimeError('HMR2: no person detected in render')
+
+    areas = (boxes[:,2]-boxes[:,0]) * (boxes[:,3]-boxes[:,1])
+    boxes = boxes[areas.argmax():areas.argmax()+1]
+
+    dataset    = ViTDetDataset(model_cfg, img_bgr, boxes)
+    dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)
+    for batch in dataloader:
+        batch = recursive_to(batch, device)
+        with torch.no_grad():
+            out = model(batch)
+        sp = out['pred_smpl_params']
+        return {
+            'body_pose': sp['body_pose'][0].cpu(),    # (23, 3, 3)
+            'betas':     sp['betas'][0].cpu(),         # (10,)
+        }
+
+
+# ---- Step 3: Read all data from rigged.glb -----------------------------------
+
+def read_rigged_glb(rig_glb):
+    """
+    Returns dict with:
+      verts        : (N, 3) mesh vertices in rig world space
+      j_idx        : (N, 4) joint indices
+      w_arr        : (N, 4) skinning weights
+      J_bind       : (24, 3) T-pose joint world positions
+    """
+    with open(rig_glb, 'rb') as fh:
+        raw = fh.read()
+    ch_len, _ = struct.unpack_from('<II', raw, 12)
+    gltf = json.loads(raw[20:20+ch_len])
+    bin_data = raw[20+ch_len+8:]
+
+    def _read(acc_i):
+        acc = gltf['accessors'][acc_i]
+        bv  = gltf['bufferViews'][acc['bufferView']]
+        off = bv.get('byteOffset', 0) + acc.get('byteOffset', 0)
+        cnt = acc['count']
+        n   = {'SCALAR':1,'VEC2':2,'VEC3':3,'VEC4':4,'MAT4':16}[acc['type']]
+        fmt = {5121:'B',5123:'H',5125:'I',5126:'f'}[acc['componentType']]
+        nb  = {'B':1,'H':2,'I':4,'f':4}[fmt]
+        return np.frombuffer(bin_data[off:off+cnt*n*nb],
+                             dtype=np.dtype(fmt)).reshape(cnt, n)
+
+    prim  = gltf['meshes'][0]['primitives'][0]['attributes']
+    verts = _read(prim['POSITION']).astype(np.float64)   # (N, 3)
+    j_idx = _read(prim['JOINTS_0']).astype(int)          # (N, 4)
+    w_arr = _read(prim['WEIGHTS_0']).astype(np.float64)  # (N, 4)
+    row_sum = w_arr.sum(axis=1, keepdims=True)
+    w_arr /= np.where(row_sum > 0, row_sum, 1.0)
+
+    # Read T-pose joint world positions by accumulating node translations
+    nodes   = gltf['nodes']
+    skin    = gltf['skins'][0]
+    j_nodes = skin['joints']                             # [0, 1, ..., 23]
+    J_bind  = np.zeros((24, 3), dtype=np.float64)
+    for ji, ni in enumerate(j_nodes):
+        t_local = np.array(nodes[ni].get('translation', [0, 0, 0]))
+        p = SMPL_PARENTS[ji]
+        J_bind[ji] = (J_bind[p] if p >= 0 else np.zeros(3)) + t_local
+
+    print('  Rig verts: %d  Y: [%.3f, %.3f]  X: [%.3f, %.3f]' % (
+        len(verts),
+        verts[:,1].min(), verts[:,1].max(),
+        verts[:,0].min(), verts[:,0].max()))
+    print('  J_bind pelvis: (%.3f, %.3f, %.3f)  L_shoulder: (%.3f, %.3f, %.3f)' % (
+        *J_bind[0], *J_bind[16]))
+    return {'verts': verts, 'j_idx': j_idx, 'w_arr': w_arr, 'J_bind': J_bind}
+
+
+# ---- Step 4: FK in rig world space -> A matrices -----------------------------
+
+_FLIP_X = np.diag([-1.0, 1.0, 1.0])   # X-axis mirror matrix
+
+
+def _adapt_rotmat_to_flipped_x(R_smpl):
+    """
+    Convert an SO(3) rotation matrix from SMPL convention (left=+X)
+    to rig convention (left=-X).  F @ R @ F  where F = diag(-1,1,1).
+    """
+    return _FLIP_X @ R_smpl @ _FLIP_X
+
+
+def compute_rig_fk_transforms(J_bind, body_pose_rotmats):
+    """
+    Compute A_j = G_j_posed * IBM_j in rig world space.
+    A_j maps T-pose -> posed, so A_j^{-1} maps posed -> T-pose.
+
+    HMR2 returns rotations in SMPL convention (left shoulder at +X).
+    The rig uses the opposite convention (left shoulder at -X).
+    We convert by conjugating with the X-flip matrix before building FK.
+
+    J_bind          : (24, 3) T-pose joint world positions from rig
+    body_pose_rotmats: (23, 3, 3) HMR2 body pose rotation matrices (joints 1-23)
+    Returns A: (24, 4, 4)
+    """
+    G = [None] * 24
+    for j in range(24):
+        p = SMPL_PARENTS[j]
+        # Convert rotation from SMPL (+X=left) to rig (-X=left) convention
+        R_smpl = body_pose_rotmats[j-1].numpy() if j >= 1 else np.eye(3)
+        R_j    = _adapt_rotmat_to_flipped_x(R_smpl)
+
+        if p < 0:
+            t_j = J_bind[j]           # root: absolute world position
+        else:
+            t_j = J_bind[j] - J_bind[p]
+
+        L = np.eye(4, dtype=np.float64)
+        L[:3, :3] = R_j
+        L[:3, 3]  = t_j
+
+        G[j] = L if p < 0 else G[p] @ L
+
+    G = np.stack(G)
+
+    A = np.zeros((24, 4, 4), dtype=np.float64)
+    for j in range(24):
+        IBM = np.eye(4, dtype=np.float64)
+        IBM[:3, 3] = -J_bind[j]
+        A[j] = G[j] @ IBM
+
+    return A
+
+
+# ---- Step 5: Inverse LBS -----------------------------------------------------
+
+def inverse_lbs(verts, j_idx, w_arr, A):
+    """
+    v_tpose = (Sum_j W_j * A_j)^{-1} * v_posed
+    All inputs in rig world space.
+    Returns (N, 3) T-posed vertices.
+    """
+    N = len(verts)
+    # Blend forward transforms
+    T_fwd = np.zeros((N, 4, 4), dtype=np.float64)
+    for k in range(4):
+        ji   = j_idx[:, k]
+        w    = w_arr[:, k]
+        mask = w > 1e-6
+        if mask.any():
+            T_fwd[mask] += w[mask, None, None] * A[ji[mask]]
+
+    T_inv = np.linalg.inv(T_fwd)
+    v_h   = np.concatenate([verts, np.ones((N, 1))], axis=1)
+    v_tp  = np.einsum('nij,nj->ni', T_inv, v_h)[:, :3]
+
+    disp  = np.linalg.norm(v_tp - verts, axis=1)
+    print('  inverse LBS: mean_disp=%.4f  max_disp=%.4f' % (disp.mean(), disp.max()))
+    return v_tp
+
+
+# ---- Step 6: Map T-posed rig verts back to original mesh space ---------------
+
+def rig_to_original_space(rig_verts_tposed, rig_verts_original, orig_mesh_verts):
+    """
+    Rig verts are a scaled + translated version of the original mesh verts.
+    Recover the (scale, offset) from the mapping:
+      rig_vert = orig_vert * scale + offset
+
+    Estimates scale from height ratio, offset from floor alignment.
+    Returns T-posed vertices in original mesh coordinate space.
+    """
+    rig_h  = rig_verts_original[:, 1].max() - rig_verts_original[:, 1].min()
+    orig_h = orig_mesh_verts[:, 1].max()    - orig_mesh_verts[:, 1].min()
+    scale  = rig_h / max(orig_h, 1e-6)
+
+    # The rig aligns: orig * scale, then v[:,1] -= v[:,1].min() (floor at 0)
+    # and v[:,0] += smpl_joints[0,0] - cx; v[:,2] += smpl_joints[0,2] - cz
+    # We can recover offset from comparing means/floors
+    # offset = rig_floor_Y - (orig_floor_Y * scale)
+    rig_floor  = rig_verts_original[:, 1].min()
+    orig_floor = orig_mesh_verts[:, 1].min()
+    y_offset   = rig_floor - orig_floor * scale
+
+    # X, Z: center offset
+    rig_cx  = (rig_verts_original[:, 0].max() + rig_verts_original[:, 0].min()) * 0.5
+    orig_cx = (orig_mesh_verts[:, 0].max()    + orig_mesh_verts[:, 0].min())    * 0.5
+    x_offset = rig_cx - orig_cx * scale
+
+    rig_cz  = (rig_verts_original[:, 2].max() + rig_verts_original[:, 2].min()) * 0.5
+    orig_cz = (orig_mesh_verts[:, 2].max()    + orig_mesh_verts[:, 2].min())    * 0.5
+    z_offset = rig_cz - orig_cz * scale
+
+    print('  rig->orig: scale=%.4f  offset=[%.3f, %.3f, %.3f]' % (scale, x_offset, y_offset, z_offset))
+
+    # Invert: orig_vert = (rig_vert - offset) / scale
+    # For T-posed verts: they're in rig space but T-posed, so same inversion
+    tposed_orig = np.zeros_like(rig_verts_tposed)
+    tposed_orig[:, 0] = (rig_verts_tposed[:, 0] - x_offset) / scale
+    tposed_orig[:, 1] = (rig_verts_tposed[:, 1] - y_offset) / scale
+    tposed_orig[:, 2] = (rig_verts_tposed[:, 2] - z_offset) / scale
+    return tposed_orig
+
+
+# ---- SKEL bone geometry ------------------------------------------------------
+
+def export_skel_bones(betas, out_path, gender='male'):
+    try:
+        from skel.skel_model import SKEL
+    except ImportError:
+        print('  [skel] Not installed')
+        return None
+    skel_file = os.path.join(SKEL_DIR, 'skel_%s.pkl' % gender)
+    if not os.path.exists(skel_file):
+        print('  [skel] Weights not found: %s' % skel_file)
+        return None
+    try:
+        skel_model = SKEL(gender=gender, model_path=SKEL_DIR)
+        betas_t    = betas.unsqueeze(0)[:, :10]
+        poses_zero = torch.zeros(1, 46)
+        trans_zero = torch.zeros(1, 3)
+        with torch.no_grad():
+            out = skel_model(poses=poses_zero, betas=betas_t, trans=trans_zero, skelmesh=True)
+        bone_verts = out.skel_verts[0].numpy()
+        bone_faces = skel_model.skel_f.numpy()
+        mesh = trimesh.Trimesh(vertices=bone_verts, faces=bone_faces, process=False)
+        mesh.export(out_path)
+        print('  [skel] Bone mesh -> %s  (%d verts)' % (out_path, len(bone_verts)))
+        return out_path
+    except Exception as e:
+        print('  [skel] Export failed: %s' % e)
+        return None
+
+
+# ---- Main --------------------------------------------------------------------
+
+def tpose_smpl(body_glb, out_glb, rig_glb=None, debug_dir=None, skel_out=None):
+    device = 'cuda'
+
+    if not rig_glb or not os.path.exists(rig_glb):
+        raise RuntimeError('--rig is required: provide the rigged.glb from the Rig step.')
+
+    print('[tpose_smpl] Rendering front view ...')
+    img_bgr = render_front(body_glb, device=device)
+    if debug_dir:
+        cv2.imwrite(os.path.join(debug_dir, 'tpose_render.png'), img_bgr)
+
+    print('[tpose_smpl] Running HMR2 pose estimation ...')
+    hmr2_out = run_hmr2(img_bgr, device=device)
+    print('  betas: %s' % hmr2_out['betas'].numpy().round(3))
+
+    print('[tpose_smpl] Reading rigged GLB (rig world space) ...')
+    rig_data = read_rigged_glb(rig_glb)
+
+    print('[tpose_smpl] Loading original mesh for UV/texture ...')
+    scene = trimesh.load(body_glb)
+    if isinstance(scene, trimesh.Scene):
+        geom_name = list(scene.geometry.keys())[0]
+        orig_mesh  = scene.geometry[geom_name]
+    else:
+        orig_mesh = scene; geom_name = None
+
+    orig_verts = np.array(orig_mesh.vertices, dtype=np.float64)
+    uvs        = np.array(orig_mesh.visual.uv, dtype=np.float64)
+    orig_tex   = orig_mesh.visual.material.baseColorTexture
+    print('  Orig mesh: %d verts  Y: [%.3f, %.3f]  X: [%.3f, %.3f]' % (
+        len(orig_verts),
+        orig_verts[:,1].min(), orig_verts[:,1].max(),
+        orig_verts[:,0].min(), orig_verts[:,0].max()))
+
+    print('[tpose_smpl] Computing FK transforms in rig world space ...')
+    body_pose_rotmats = hmr2_out['body_pose']   # (23, 3, 3)
+    A = compute_rig_fk_transforms(rig_data['J_bind'], body_pose_rotmats)
+
+    # Verify zero-pose gives identity (sanity check)
+    A_zero = compute_rig_fk_transforms(rig_data['J_bind'],
+                                        torch.zeros(23, 3, 3) + torch.eye(3))
+    v_test = rig_data['verts'][:3]
+    v_h = np.concatenate([v_test, np.ones((3,1))], axis=1)
+    T_fwd_test = np.zeros((3, 4, 4))
+    for k in range(4):
+        ji = rig_data['j_idx'][:3, k]; w = rig_data['w_arr'][:3, k]
+        T_fwd_test += w[:, None, None] * A_zero[ji]
+    identity_err = np.abs(T_fwd_test - np.eye(4)).max()
+    print('  zero-pose identity check: max_err=%.6f (expect ~0)' % identity_err)
+
+    print('[tpose_smpl] Applying inverse LBS ...')
+    rig_verts_tposed = inverse_lbs(
+        rig_data['verts'], rig_data['j_idx'], rig_data['w_arr'], A)
+
+    print('[tpose_smpl] T-posed rig verts: Y: [%.3f, %.3f]  X: [%.3f, %.3f]' % (
+        rig_verts_tposed[:,1].min(), rig_verts_tposed[:,1].max(),
+        rig_verts_tposed[:,0].min(), rig_verts_tposed[:,0].max()))
+
+    print('[tpose_smpl] Mapping back to original mesh coordinate space ...')
+    tposed_orig = rig_to_original_space(
+        rig_verts_tposed, rig_data['verts'], orig_verts)
+
+    print('[tpose_smpl] T-posed orig: Y: [%.3f, %.3f]  X: [%.3f, %.3f]' % (
+        tposed_orig[:,1].min(), tposed_orig[:,1].max(),
+        tposed_orig[:,0].min(), tposed_orig[:,0].max()))
+
+    orig_mesh.vertices = tposed_orig
+    orig_mesh.visual = TextureVisuals(uv=uvs,
+                                      material=PBRMaterial(baseColorTexture=orig_tex))
+
+    if geom_name and isinstance(scene, trimesh.Scene):
+        scene.geometry[geom_name] = orig_mesh
+        scene.export(out_glb)
+    else:
+        orig_mesh.export(out_glb)
+
+    print('[tpose_smpl] Saved: %s  (%d KB)' % (out_glb, os.path.getsize(out_glb)//1024))
+
+    if skel_out:
+        print('[tpose_smpl] Exporting SKEL bone geometry ...')
+        export_skel_bones(hmr2_out['betas'], skel_out)
+
+    return out_glb
+
+
+if __name__ == '__main__':
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--body',      required=True)
+    ap.add_argument('--out',       required=True)
+    ap.add_argument('--rig',       required=True, help='Rigged GLB from rig step')
+    ap.add_argument('--skel_out',  default=None,  help='SKEL BSM bone mesh output')
+    ap.add_argument('--debug_dir', default=None)
+    args = ap.parse_args()
+    os.makedirs(args.debug_dir, exist_ok=True) if args.debug_dir else None
+    tpose_smpl(args.body, args.out, rig_glb=args.rig,
+               debug_dir=args.debug_dir, skel_out=args.skel_out)