Upload pipeline/face_inswap_bake.py with huggingface_hub

pipeline/face_inswap_bake.py

@@ -0,0 +1,302 @@
+"""
+face_inswap_bake.py — Proper face swap on rendered views, then UV-bake.
+
+Pipeline:
+  1. Render the mesh from multiple views (front + L/R 3-quarter)
+  2. Run inswapper_128 to swap reference face onto each rendered view
+  3. uv_render_attr() bakes each swapped render directly into UV texture
+     (render-space coords shared with UV lookup — no coordinate transforms)
+  4. Composite multiple views (front takes priority, sides fill gaps)
+  5. Save updated GLB
+
+Usage:
+    python face_inswap_bake.py \
+        --body   /tmp/triposg_textured.glb \
+        --face   /tmp/triposg_face_ref.png \
+        --out    /tmp/face_swapped.glb \
+        [--uv_size 4096] [--debug_dir /tmp]
+"""
+
+import os, sys, argparse, warnings
+warnings.filterwarnings('ignore')
+
+import numpy as np
+import cv2
+import torch
+import torch.nn.functional as F
+from PIL import Image
+import trimesh
+from trimesh.visual.texture import TextureVisuals
+from trimesh.visual.material import PBRMaterial
+
+sys.path.insert(0, '/root/MV-Adapter')
+from mvadapter.utils.mesh_utils import (
+    NVDiffRastContextWrapper, load_mesh, get_orthogonal_camera, render,
+)
+from mvadapter.utils.mesh_utils.uv import (
+    uv_precompute, uv_render_geometry, uv_render_attr,
+)
+from insightface.app import FaceAnalysis
+import insightface
+from gfpgan import GFPGANer
+
+
+GFPGAN_PATH = '/root/MV-Adapter/checkpoints/GFPGANv1.4.pth'
+
+
+# ── helpers ───────────────────────────────────────────────────────────────────
+
+def _build_front_face_uv_mask(mesh_t, tex_H, tex_W, neck_frac=0.76):
+    """UV-space mask covering only front-facing head triangles (no back-of-head)."""
+    verts = np.array(mesh_t.vertices, dtype=np.float64)
+    faces = np.array(mesh_t.faces,    dtype=np.int32)
+    uvs   = np.array(mesh_t.visual.uv, dtype=np.float64)
+
+    y_min, y_max = verts[:, 1].min(), verts[:, 1].max()
+    neck_y   = float(y_min + (y_max - y_min) * neck_frac)
+    head_idx = np.where(verts[:, 1] > neck_y)[0]
+    hv = verts[head_idx]
+
+    z_thresh = float(np.percentile(hv[:, 2], 40))
+    front    = hv[:, 2] >= z_thresh
+    if front.sum() < 30:
+        front = np.ones(len(hv), bool)
+
+    face_vert_idx  = head_idx[front]
+    face_vert_mask = np.zeros(len(verts), bool)
+    face_vert_mask[face_vert_idx] = True
+    face_tri_mask  = face_vert_mask[faces].all(axis=1)
+    face_tris      = faces[face_tri_mask]
+    print(f'  Geometry mask: {face_tri_mask.sum()} front-face triangles '
+          f'(neck_y={neck_y:.3f}, z_thresh={z_thresh:.3f})')
+
+    geom_mask = np.zeros((tex_H, tex_W), dtype=np.float32)
+    pts_list = []
+    for tri in face_tris:
+        uv = uvs[tri]
+        px = uv[:, 0] * tex_W
+        py = (1.0 - uv[:, 1]) * tex_H
+        pts_list.append(np.column_stack([px, py]).astype(np.int32))
+    if pts_list:
+        cv2.fillPoly(geom_mask, pts_list, 1.0)
+
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    geom_mask = cv2.dilate(geom_mask, kernel, iterations=2)
+    geom_mask = cv2.erode(geom_mask,  kernel, iterations=1)
+    geom_mask = cv2.GaussianBlur(geom_mask, (31, 31), 8)
+    return geom_mask
+
+
+def _detect_largest_face(img_bgr, app):
+    faces = app.get(img_bgr)
+    if not faces:
+        return None
+    return max(faces, key=lambda f: (f.bbox[2]-f.bbox[0])*(f.bbox[3]-f.bbox[1]))
+
+
+def _render_view(ctx, mesh_mv, uv_pre, azimuth_deg, H, W, device):
+    """Render the mesh from a given azimuth; return (camera, uv_geom)."""
+    camera = get_orthogonal_camera(
+        elevation_deg=[0], distance=[1.8],
+        left=-0.55, right=0.55, bottom=-0.55, top=0.55,
+        azimuth_deg=[azimuth_deg], device=device,
+    )
+    uv_geom = uv_render_geometry(
+        ctx, mesh_mv, camera,
+        view_height=H, view_width=W,
+        uv_precompute_output=uv_pre,
+        compute_depth_grad=False,
+    )
+    return camera, uv_geom
+
+
+def face_inswap_bake(body_glb, face_img_path, out_glb,
+                     uv_size=4096, debug_dir=None):
+
+    device = 'cuda'
+    INSWAPPER_PATH = '/root/MV-Adapter/checkpoints/inswapper_128.onnx'
+
+    # ── Load GFPGAN enhancer ──────────────────────────────────────────────────
+    print('[fib] Loading GFPGANv1.4 ...')
+    enhancer = GFPGANer(
+        model_path=GFPGAN_PATH,
+        upscale=1,
+        arch='clean',
+        channel_multiplier=2,
+        bg_upsampler=None,
+    )
+
+    # ── Load mesh ─────────────────────────────────────────────────────────────
+    print(f'[fib] Loading mesh: {body_glb}')
+    ctx = NVDiffRastContextWrapper(device=device, context_type='cuda')
+    mesh_mv = load_mesh(body_glb, rescale=True, device=device)
+
+    scene_t = trimesh.load(body_glb)
+    if isinstance(scene_t, trimesh.Scene):
+        geom_name = list(scene_t.geometry.keys())[0]
+        mesh_t = scene_t.geometry[geom_name]
+    else:
+        mesh_t = scene_t; geom_name = None
+
+    orig_tex_np = np.array(mesh_t.visual.material.baseColorTexture, dtype=np.float32) / 255.0
+    uvs = np.array(mesh_t.visual.uv, dtype=np.float64)
+    tex_H, tex_W = orig_tex_np.shape[:2]
+    print(f'  Texture: {tex_W}×{tex_H}')
+
+    # Build geometry mask (front-face head triangles only) at UV resolution
+    print('[fib] Building front-face geometry UV mask ...')
+    geom_uv_mask = _build_front_face_uv_mask(mesh_t, uv_size, uv_size)
+
+    # Render dimensions (match triposg_app.py)
+    H_r, W_r = 1024, 768
+
+    # ── Precompute UV geometry ─────────────────────────────────────────────────
+    print(f'[fib] Precomputing UV geometry ({uv_size}×{uv_size}) ...')
+    uv_pre = uv_precompute(ctx, mesh_mv, height=uv_size, width=uv_size)
+
+    # ── Load face swap model + face detector ──────────────────────────────────
+    print('[fib] Loading inswapper_128 ...')
+    swapper = insightface.model_zoo.get_model(
+        INSWAPPER_PATH, download=False,
+        providers=['CUDAExecutionProvider', 'CPUExecutionProvider'],
+    )
+
+    app = FaceAnalysis(providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+    app.prepare(ctx_id=0, det_size=(640, 640))
+
+    ref_bgr = cv2.imread(face_img_path)
+    ref_face = _detect_largest_face(ref_bgr, app)
+    if ref_face is None:
+        raise RuntimeError(f'No face detected in reference: {face_img_path}')
+    print(f'  Reference face detected: bbox={ref_face.bbox.astype(int).tolist()}')
+
+    # ── Process each view ─────────────────────────────────────────────────────
+    # Views: front (azimuth=-90), slight left (-60), slight right (-120)
+    # Azimuth convention from MV-Adapter: -90 = front-facing
+    views = [
+        ('front',       -90,   1.0),   # (name, azimuth_deg, priority_weight)
+        ('threequarter_r', -60,  0.7),
+        ('threequarter_l', -120, 0.7),
+    ]
+
+    # Accumulators for weighted UV compositing
+    uv_colour_acc = np.zeros((uv_size, uv_size, 3), dtype=np.float32)
+    uv_weight_acc = np.zeros((uv_size, uv_size),    dtype=np.float32)
+
+    for view_name, azimuth, weight in views:
+        print(f'\n[fib] View: {view_name} (azimuth={azimuth}°)')
+
+        # Create camera + UV geometry for this view
+        camera, uv_geom = _render_view(ctx, mesh_mv, uv_pre, azimuth, H_r, W_r, device)
+
+        # Render textured mesh from this view
+        render_out = render(ctx, mesh_mv, camera, height=H_r, width=W_r,
+                            render_attr=True, render_depth=False, render_normal=False,
+                            attr_background=0.0)
+        # render_out.attr: (1, H, W, 3) float in [0,1]
+        rendered_np = (render_out.attr[0].cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
+        rendered_bgr = cv2.cvtColor(rendered_np, cv2.COLOR_RGB2BGR)
+
+        if debug_dir:
+            cv2.imwrite(os.path.join(debug_dir, f'fib_render_{view_name}.png'), rendered_bgr)
+
+        # Detect face in this rendered view
+        tgt_face = _detect_largest_face(rendered_bgr, app)
+        if tgt_face is None:
+            print(f'  No face in {view_name} render — skipping')
+            continue
+        print(f'  Target face: bbox={tgt_face.bbox.astype(int).tolist()}')
+
+        # Swap face
+        swapped_bgr = swapper.get(rendered_bgr.copy(), tgt_face, ref_face, paste_back=True)
+
+        # Enhance face detail with GFPGAN
+        _, _, enhanced_bgr = enhancer.enhance(
+            swapped_bgr, has_aligned=False, only_center_face=False, paste_back=True)
+        if enhanced_bgr is not None:
+            swapped_bgr = enhanced_bgr
+            print(f'  GFPGAN enhanced')
+
+        if debug_dir:
+            cv2.imwrite(os.path.join(debug_dir, f'fib_swapped_{view_name}.png'), swapped_bgr)
+
+        swapped_rgb = cv2.cvtColor(swapped_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
+
+        # Build render-space face hull mask
+        kps = tgt_face.kps
+        hull_pts = cv2.convexHull(kps.astype(np.float32)).squeeze(1)
+        hull_cx, hull_cy = hull_pts.mean(axis=0)
+        hull_exp = (hull_pts - [hull_cx, hull_cy]) * 3.5 + [hull_cx, hull_cy]
+        face_mask = np.zeros((H_r, W_r), dtype=np.float32)
+        cv2.fillPoly(face_mask, [hull_exp.astype(np.int32)], 1.0)
+        face_mask = cv2.GaussianBlur(face_mask, (61, 61), 20)
+
+        # Bake swapped render into UV space
+        swapped_t = torch.tensor(swapped_rgb, device=device).unsqueeze(0)  # (1,H,W,3)
+        mask_t    = torch.tensor(face_mask[None], device=device)
+
+        uv_out = uv_render_attr(
+            images=swapped_t,
+            masks=mask_t,
+            uv_render_geometry_output=uv_geom,
+        )
+        uv_img  = uv_out.uv_attr_proj[0].cpu().numpy()   # (uv, uv, 3)
+        uv_mask = uv_out.uv_mask_proj[0].cpu().numpy()   # (uv, uv)
+
+        # Kill back-of-head UV islands
+        uv_mask = uv_mask * geom_uv_mask
+
+        # Weighted accumulate
+        w = uv_mask * weight
+        uv_colour_acc += uv_img * w[..., None]
+        uv_weight_acc += w
+        print(f'  Painted texels: {(uv_mask > 0.05).sum()}')
+
+    # ── Composite ──────────────────────────────────────────────────────────────
+    print('\n[fib] Compositing views ...')
+    valid = uv_weight_acc > 0.01
+    uv_final = np.where(valid[..., None],
+                        uv_colour_acc / np.maximum(uv_weight_acc[..., None], 1e-6),
+                        orig_tex_np[:uv_size, :uv_size] if uv_size <= tex_H else orig_tex_np)
+
+    # Resize to texture resolution if needed
+    if uv_size != tex_H or uv_size != tex_W:
+        uv_final_rs = cv2.resize(uv_final,         (tex_W, tex_H), interpolation=cv2.INTER_LINEAR)
+        weight_rs   = cv2.resize(uv_weight_acc,    (tex_W, tex_H), interpolation=cv2.INTER_LINEAR)
+    else:
+        uv_final_rs = uv_final
+        weight_rs   = uv_weight_acc
+
+    # Blend with original texture: use face-swap result where painted, orig elsewhere
+    alpha = np.clip(weight_rs, 0, 1)[..., None]
+    new_tex = uv_final_rs * alpha + orig_tex_np * (1.0 - alpha)
+    print(f'  Total painted texels (tex res): {(weight_rs > 0.05).sum()}')
+
+    if debug_dir:
+        Image.fromarray((uv_final_rs * 255).clip(0,255).astype(np.uint8)).save(
+            os.path.join(debug_dir, 'fib_uv_composite.png'))
+
+    # ── Save GLB ──────────────────────────────────────────────────────────────
+    new_pil = Image.fromarray((new_tex * 255).clip(0, 255).astype(np.uint8))
+    mesh_t.visual = TextureVisuals(uv=uvs, material=PBRMaterial(baseColorTexture=new_pil))
+
+    if geom_name and isinstance(scene_t, trimesh.Scene):
+        scene_t.geometry[geom_name] = mesh_t
+        scene_t.export(out_glb)
+    else:
+        mesh_t.export(out_glb)
+
+    print(f'[fib] Saved: {out_glb}  ({os.path.getsize(out_glb)//1024} KB)')
+    return out_glb
+
+
+if __name__ == '__main__':
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--body',      required=True)
+    ap.add_argument('--face',      required=True)
+    ap.add_argument('--out',       required=True)
+    ap.add_argument('--uv_size',   type=int, default=4096)
+    ap.add_argument('--debug_dir', default=None)
+    args = ap.parse_args()
+    face_inswap_bake(args.body, args.face, args.out,
+                     uv_size=args.uv_size, debug_dir=args.debug_dir)