Upload pipeline/tpose.py with huggingface_hub

pipeline/tpose.py

@@ -0,0 +1,332 @@
+"""
+tpose.py  —  T-pose a humanoid GLB using YOLO pose estimation.
+
+Pipeline:
+  1. Render the mesh from front view (azimuth=-90)
+  2. Run YOLOv8-pose to get 17 COCO keypoints in render-space
+  3. Unproject keypoints through the orthographic camera to 3D
+  4. Build Blender armature with bones at detected 3D joint positions (current pose)
+  5. Auto-weight skin the mesh to this armature
+  6. Rotate arm/leg bones to T-pose, apply deformation, export
+
+Usage:
+    blender --background --python tpose.py -- <input.glb> <output.glb>
+"""
+
+import bpy, sys, math, mathutils, os, tempfile
+import numpy as np
+
+# ── Args ─────────────────────────────────────────────────────────────────────
+argv = sys.argv
+argv = argv[argv.index("--") + 1:] if "--" in argv else []
+if len(argv) < 2:
+    print("Usage: blender --background --python tpose.py -- input.glb output.glb")
+    sys.exit(1)
+input_glb  = argv[0]
+output_glb = argv[1]
+
+# ── Step 1: Render front view using nvdiffrast (outside Blender) ───────────────
+# We do this via a subprocess call before Blender scene setup,
+# using the triposg Python env which has MV-Adapter + nvdiffrast.
+import subprocess, json
+
+TRIPOSG_PYTHON = '/root/miniconda/envs/triposg/bin/python'
+RENDER_SCRIPT = '/tmp/_tpose_render.py'
+RENDER_OUT    = '/tmp/_tpose_front.png'
+KP_OUT        = '/tmp/_tpose_kp.json'
+
+render_code = r"""
+import sys, json
+sys.path.insert(0, '/root/MV-Adapter')
+import numpy as np, cv2, torch
+from mvadapter.utils.mesh_utils import (
+    NVDiffRastContextWrapper, load_mesh, get_orthogonal_camera, render,
+)
+
+body_glb = sys.argv[1]
+out_png   = sys.argv[2]
+
+device = 'cuda'
+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=1024, width=768,
+             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('uint8')
+cv2.imwrite(out_png, cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR))
+print(f"Rendered to {out_png}")
+"""
+
+with open(RENDER_SCRIPT, 'w') as f:
+    f.write(render_code)
+
+print("[tpose] Rendering front view ...")
+r = subprocess.run([TRIPOSG_PYTHON, RENDER_SCRIPT, input_glb, RENDER_OUT],
+                   capture_output=True, text=True)
+print(r.stdout.strip()); print(r.stderr[-500:] if r.stderr else '')
+
+# ── Step 2: YOLO pose estimation ──────────────────────────────────────────────
+YOLO_SCRIPT = '/tmp/_tpose_yolo.py'
+yolo_code = r"""
+import sys, json
+import cv2
+from ultralytics import YOLO
+import numpy as np
+
+img_path = sys.argv[1]
+kp_path  = sys.argv[2]
+
+model = YOLO('yolov8n-pose.pt')
+img   = cv2.imread(img_path)
+H, W  = img.shape[:2]
+
+results = model(img, verbose=False)
+if not results or results[0].keypoints is None:
+    print("ERROR: no person detected"); sys.exit(1)
+
+# Pick detection with highest confidence
+kps_all = results[0].keypoints.data.cpu().numpy()  # (N, 17, 3)
+confs   = kps_all[:, :, 2].mean(axis=1)
+best    = kps_all[confs.argmax()]  # (17, 3): x, y, conf
+
+# COCO 17 keypoints:
+# 0=nose 1=left_eye 2=right_eye 3=left_ear 4=right_ear
+# 5=left_shoulder 6=right_shoulder 7=left_elbow 8=right_elbow
+# 9=left_wrist 10=right_wrist 11=left_hip 12=right_hip
+# 13=left_knee 14=right_knee 15=left_ankle 16=right_ankle
+
+names = ['nose','left_eye','right_eye','left_ear','right_ear',
+         'left_shoulder','right_shoulder','left_elbow','right_elbow',
+         'left_wrist','right_wrist','left_hip','right_hip',
+         'left_knee','right_knee','left_ankle','right_ankle']
+
+kp_dict = {}
+for i, name in enumerate(names):
+    x, y, c = best[i]
+    kp_dict[name] = {'x': float(x)/W, 'y': float(y)/H, 'conf': float(c)}
+    print(f"  {name}: ({x:.1f},{y:.1f}) conf={c:.2f}")
+
+kp_dict['img_hw'] = [int(H), int(W)]
+with open(kp_path, 'w') as f:
+    json.dump(kp_dict, f)
+print(f"Keypoints saved to {kp_path}")
+"""
+
+with open(YOLO_SCRIPT, 'w') as f:
+    f.write(yolo_code)
+
+print("[tpose] Running YOLO pose estimation ...")
+r2 = subprocess.run([TRIPOSG_PYTHON, YOLO_SCRIPT, RENDER_OUT, KP_OUT],
+                    capture_output=True, text=True)
+print(r2.stdout.strip()); print(r2.stderr[-300:] if r2.stderr else '')
+
+if not os.path.exists(KP_OUT):
+    print("ERROR: YOLO failed — falling back to heuristic")
+    kp_data = None
+else:
+    with open(KP_OUT) as f:
+        kp_data = json.load(f)
+
+# ── Step 3: Unproject render-space keypoints to 3D ────────────────────────────
+# Orthographic camera: left=-0.55, right=0.55, bottom=-0.55, top=0.55
+# Render: 768×1024.  NDC x = 2*(px/W)-1, ndc y = 1-2*(py/H)
+# World X = ndc_x * 0.55, World Y (mesh up) = ndc_y * 0.55
+# We need 3D positions in the ORIGINAL mesh coordinate space.
+# After Blender GLB import, original mesh Y → Blender Z, original Z → Blender -Y
+
+def kp_to_3d(name, z_default=0.0):
+    """Convert YOLO keypoint (image fraction) → Blender 3D coords."""
+    if kp_data is None or name not in kp_data:
+        return None
+    k = kp_data[name]
+    if k['conf'] < 0.3:
+        return None
+    # Image coords (fractions) → NDC
+    ndc_x =  2 * k['x'] - 1.0      # left→right  = mesh X
+    ndc_y = -(2 * k['y'] - 1.0)    # top→bottom  = mesh Y (up)
+    # Orthographic: frustum ±0.55
+    mesh_x = ndc_x * 0.55
+    mesh_y = ndc_y * 0.55           # this is mesh-space Y (vertical)
+    # After GLB import: mesh Y → Blender Z, mesh Z → Blender -Y
+    bl_x = mesh_x
+    bl_z = mesh_y           # height
+    bl_y = z_default        # depth (not observable from front view)
+    return (bl_x, bl_y, bl_z)
+
+# Key joint positions in Blender space
+J = {}
+for name in ['nose','left_shoulder','right_shoulder','left_elbow','right_elbow',
+             'left_wrist','right_wrist','left_hip','right_hip',
+             'left_knee','right_knee','left_ankle','right_ankle']:
+    p = kp_to_3d(name)
+    if p: J[name] = p
+
+print(f"[tpose] Detected joints: {list(J.keys())}")
+
+# ── Step 4: Set up Blender scene ──────────────────────────────────────────────
+bpy.ops.wm.read_factory_settings(use_empty=True)
+bpy.ops.import_scene.gltf(filepath=input_glb)
+bpy.context.view_layer.update()
+
+mesh_obj = next((o for o in bpy.data.objects if o.type == 'MESH'), None)
+if not mesh_obj:
+    print("ERROR: no mesh"); sys.exit(1)
+
+verts_w = np.array([mesh_obj.matrix_world @ v.co for v in mesh_obj.data.vertices])
+z_min, z_max = verts_w[:,2].min(), verts_w[:,2].max()
+x_c = (verts_w[:,0].min() + verts_w[:,0].max()) / 2
+y_c = (verts_w[:,1].min() + verts_w[:,1].max()) / 2
+H_mesh = z_max - z_min
+
+def zh(frac): return z_min + frac * H_mesh
+def jv(name, fallback_frac=None, fallback_x=0.0):
+    """Get joint position from YOLO or use fallback."""
+    if name in J:
+        x, y, z = J[name]
+        return (x, y_c, z)  # use mesh y_c for depth
+    if fallback_frac is not None:
+        return (x_c + fallback_x, y_c, zh(fallback_frac))
+    return None
+
+# ── Step 5: Build armature in CURRENT pose ────────────────────────────────────
+bpy.ops.object.armature_add(location=(x_c, y_c, zh(0.5)))
+arm_obj = bpy.context.object
+arm_obj.name = 'PoseRig'
+arm = arm_obj.data
+
+bpy.ops.object.mode_set(mode='EDIT')
+eb = arm.edit_bones
+
+def V(xyz): return mathutils.Vector(xyz)
+
+def add_bone(name, head, tail, parent=None, connect=False):
+    b = eb.new(name)
+    b.head = V(head)
+    b.tail = V(tail)
+    if parent and parent in eb:
+        b.parent = eb[parent]
+        b.use_connect = connect
+    return b
+
+# Helper: midpoint
+def mid(a, b): return tuple((a[i]+b[i])/2 for i in range(3))
+def offset(p, dx=0, dy=0, dz=0): return (p[0]+dx, p[1]+dy, p[2]+dz)
+
+# ── Spine / hips ─────────────────────────────────────────────────────────────
+hip_L  = jv('left_hip',  0.48, -0.07)
+hip_R  = jv('right_hip', 0.48,  0.07)
+sh_L   = jv('left_shoulder',  0.77, -0.20)
+sh_R   = jv('right_shoulder', 0.77,  0.20)
+nose   = jv('nose', 0.92)
+
+hips_c = mid(hip_L, hip_R) if (hip_L and hip_R) else (x_c, y_c, zh(0.48))
+sh_c   = mid(sh_L, sh_R)   if (sh_L and sh_R)   else (x_c, y_c, zh(0.77))
+
+add_bone('Hips',  hips_c, offset(hips_c, dz=H_mesh*0.08))
+add_bone('Spine', hips_c, offset(hips_c, dz=(sh_c[2]-hips_c[2])*0.5), 'Hips')
+add_bone('Chest', offset(hips_c, dz=(sh_c[2]-hips_c[2])*0.5), sh_c, 'Spine', True)
+if nose:
+    neck_z = sh_c[2] + (nose[2]-sh_c[2])*0.35
+    head_z = sh_c[2] + (nose[2]-sh_c[2])*0.65
+    add_bone('Neck', (x_c, y_c, neck_z), (x_c, y_c, head_z),  'Chest')
+    add_bone('Head', (x_c, y_c, head_z), (x_c, y_c, nose[2]+H_mesh*0.05), 'Neck', True)
+else:
+    add_bone('Neck', sh_c, offset(sh_c, dz=H_mesh*0.06), 'Chest')
+    add_bone('Head', offset(sh_c, dz=H_mesh*0.06), offset(sh_c, dz=H_mesh*0.14), 'Neck', True)
+
+# ── Arms (placed at DETECTED current pose positions) ─────────────────────────
+el_L  = jv('left_elbow',  0.60, -0.30)
+el_R  = jv('right_elbow', 0.60,  0.30)
+wr_L  = jv('left_wrist',  0.45, -0.25)
+wr_R  = jv('right_wrist', 0.45,  0.25)
+
+for side, sh, el, wr in (('L', sh_L, el_L, wr_L), ('R', sh_R, el_R, wr_R)):
+    if not sh: continue
+    el_pos = el if el else offset(sh, dz=-H_mesh*0.15)
+    wr_pos = wr if wr else offset(el_pos, dz=-H_mesh*0.15)
+    hand   = offset(wr_pos, dz=(wr_pos[2]-el_pos[2])*0.4)
+    add_bone(f'UpperArm.{side}', sh,     el_pos, 'Chest')
+    add_bone(f'ForeArm.{side}',  el_pos, wr_pos, f'UpperArm.{side}', True)
+    add_bone(f'Hand.{side}',     wr_pos, hand,   f'ForeArm.{side}',  True)
+
+# ── Legs ─────────────────────────────────────────────────────────────────────
+kn_L  = jv('left_knee',   0.25, -0.07)
+kn_R  = jv('right_knee',  0.25,  0.07)
+an_L  = jv('left_ankle',  0.04, -0.06)
+an_R  = jv('right_ankle', 0.04,  0.06)
+
+for side, hp, kn, an in (('L', hip_L, kn_L, an_L), ('R', hip_R, kn_R, an_R)):
+    if not hp: continue
+    kn_pos = kn if kn else offset(hp, dz=-H_mesh*0.23)
+    an_pos = an if an else offset(kn_pos, dz=-H_mesh*0.22)
+    toe    = offset(an_pos, dy=-H_mesh*0.06, dz=-H_mesh*0.02)
+    add_bone(f'UpperLeg.{side}', hp,     kn_pos, 'Hips')
+    add_bone(f'LowerLeg.{side}', kn_pos, an_pos, f'UpperLeg.{side}', True)
+    add_bone(f'Foot.{side}',     an_pos, toe,    f'LowerLeg.{side}', True)
+
+bpy.ops.object.mode_set(mode='OBJECT')
+
+# ── Step 6: Skin mesh to armature ────────────────────────────────────────────
+bpy.context.view_layer.objects.active = arm_obj
+mesh_obj.select_set(True)
+arm_obj.select_set(True)
+bpy.ops.object.parent_set(type='ARMATURE_AUTO')
+print("[tpose] Auto-weights applied")
+
+# ── Step 7: Pose arms to T-pose ───────────────────────────────────────────────
+# Compute per-arm rotation: from (current elbow - shoulder) direction → horizontal ±X
+bpy.context.view_layer.objects.active = arm_obj
+bpy.ops.object.mode_set(mode='POSE')
+
+pb = arm_obj.pose.bones
+
+def set_tpose_arm(side, sh_pos, el_pos):
+    if not sh_pos or not el_pos:
+        return
+    if f'UpperArm.{side}' not in pb:
+        return
+    # Current upper-arm direction in armature local space
+    sx = -1 if side == 'L' else 1
+    # T-pose direction: ±X horizontal
+    tpose_dir = mathutils.Vector((sx, 0, 0))
+    # Current bone direction (head→tail) in world space
+    bone = arm_obj.data.bones[f'UpperArm.{side}']
+    cur_dir = (bone.tail_local - bone.head_local).normalized()
+    # Rotation needed in bone's local space
+    rot_quat = cur_dir.rotation_difference(tpose_dir)
+    pb[f'UpperArm.{side}'].rotation_mode = 'QUATERNION'
+    pb[f'UpperArm.{side}'].rotation_quaternion = rot_quat
+
+    # Straighten forearm along the same axis
+    if f'ForeArm.{side}' in pb:
+        pb[f'ForeArm.{side}'].rotation_mode = 'QUATERNION'
+        pb[f'ForeArm.{side}'].rotation_quaternion = mathutils.Quaternion((1,0,0,0))
+
+set_tpose_arm('L', sh_L, el_L)
+set_tpose_arm('R', sh_R, el_R)
+
+bpy.context.view_layer.update()
+bpy.ops.object.mode_set(mode='OBJECT')
+
+# ── Step 8: Apply armature modifier ──────────────────────────────────────────
+bpy.context.view_layer.objects.active = mesh_obj
+mesh_obj.select_set(True)
+for mod in mesh_obj.modifiers:
+    if mod.type == 'ARMATURE':
+        bpy.ops.object.modifier_apply(modifier=mod.name)
+        print(f"[tpose] Applied modifier: {mod.name}")
+        break
+
+bpy.data.objects.remove(arm_obj, do_unlink=True)
+
+# ── Step 9: Export ────────────────────────────────────────────────────────────
+bpy.ops.export_scene.gltf(
+    filepath=output_glb, export_format='GLB',
+    export_texcoords=True, export_normals=True,
+    export_materials='EXPORT', use_selection=False)
+print(f"[tpose] Done → {output_glb}")