pixal3d_mrp_mlx/scripts/image_to_mesh.py

"""Image -> GLB through the full Pixal3D cascade, with a silhouette check.

Usage:  python scripts/image_to_mesh.py IMAGE [-o OUT.glb] [--fov RAD] [--seed N]

The silhouette IoU is the acceptance test: re-project the mesh through the same camera
and compare against the input matte. Pixal3D's entire claim is pixel alignment, so a
run that completes with a poor IoU has failed even though nothing raised.
"""

import argparse
import json
import sys
import time
from pathlib import Path

import mlx.core as mx
import numpy as np
import trimesh

REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO))

from pixal3d_mlx.mesh import PBR_ATTR_LAYOUT, to_camera_frame     # noqa: E402
from pixal3d_mlx.models import load_all, normalization             # noqa: E402
from pixal3d_mlx.pipeline import DEFAULT_FOV, image_to_mesh        # noqa: E402

DEFAULT_IMAGE = REPO / "upstream" / "Pixal3D" / "assets" / "images" / "0_img.png"


def silhouette_iou(vertices, image_path, fov, res=512):
    """Re-project the mesh through the generating camera; IoU against the input matte.

    Vertices MUST be rotated into the camera frame first — o_voxel returns them in the
    voxel-grid frame while ProjGrid rotates its lattice before projecting.
    """
    from PIL import Image
    from scipy.ndimage import binary_dilation

    from pixal3d_mlx.cond import preprocess_image
    from pixal3d_mlx.proj import _FRONT_VIEW, distance_from_fov, project_points

    img = Image.open(image_path)
    if img.mode != "RGBA":
        return None
    matte = np.asarray(preprocess_image(img).convert("L").resize((res, res))) > 8

    tm = _FRONT_VIEW.copy()
    tm[1, 3] = -distance_from_fov(fov, 1.0, res)
    pts = to_camera_frame(vertices).astype(np.float32)[None]
    px, _, _ = project_points(mx.array(pts), mx.array(tm[None]), fov, res)
    px = np.asarray(px)[0].astype(int)

    keep = (px[:, 0] >= 0) & (px[:, 0] < res) & (px[:, 1] >= 0) & (px[:, 1] < res)
    proj = np.zeros((res, res), bool)
    proj[px[keep, 1], px[keep, 0]] = True
    proj = binary_dilation(proj, np.ones((3, 3), bool))
    return (proj & matte).sum() / (proj | matte).sum(), proj, matte


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("image", nargs="?", default=str(DEFAULT_IMAGE))
    ap.add_argument("-o", "--output", default=str(REPO / "output.glb"))
    ap.add_argument("--fov", type=float, default=DEFAULT_FOV)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--target-faces", type=int, default=100_000,
                    help="face budget after cleanup; 0 disables decimation")
    ap.add_argument("--raw", action="store_true",
                    help="skip cleanup entirely and export the decoder output as-is")
    ap.add_argument("--min-iou", type=float, default=0.85,
                    help="fail the run below this silhouette IoU; 0 disables the gate")
    ap.add_argument("--json", help="write run metadata here")
    ap.add_argument("--texture", action="store_true",
                    help="run the texture stage and bake PBR maps through o_voxel")
    ap.add_argument("--texture-size", type=int, default=2048)
    ap.add_argument("--baker", choices=("vertex", "uv"), default="vertex",
                    help="vertex = seconds, base colour only (default); "
                         "uv = o_voxel unwrap + full PBR maps, but >20min CPU here")
    a = ap.parse_args()

    t = time.time()
    models = load_all(with_texture=a.texture)
    print(f"loaded models ({time.time() - t:.1f}s, lazy — weights fault in on first use)")

    v, f, info = image_to_mesh(a.image, models, camera_angle_x=a.fov, seed=a.seed,
                               normalization=normalization("shape"),
                               tex_normalization=normalization("tex") if a.texture else None,
                               texture=a.texture)
    subs = info.pop("subs", None)
    info.pop("hr_slat", None)
    tex_voxels = info.pop("tex_voxels", None)

    if a.texture and tex_voxels is not None:
        # CLEAN BEFORE BAKING. o_voxel's remesh+unwrap on the raw ~8M-face mesh hangs
        # (killed at 20min); the trellis2 lane hit the same wall and its operator note
        # says the uncapped bake peaks at 75GB. Welding and stripping floaters first
        # makes it tractable, and the baker samples the attribute VOLUME at mesh
        # positions, so a decimated mesh still gets correct colours.
        from pixal3d_mlx.cleanup import clean
        from pixal3d_mlx.mesh import bake_vertex_colors, to_glb

        t = time.time()
        pre, _ = clean(v.cpu().numpy(), f.cpu().numpy(),
                       target_faces=a.target_faces or 500_000)
        print(f"  pre-bake   {len(pre.faces):,} faces   {time.time() - t:6.1f}s")

        t = time.time()
        if a.baker == "vertex":
            mesh, bstat = bake_vertex_colors(pre, tex_voxels, info["output_resolution"])
            info["bake"] = bstat
            print(f"  bake       vertex colours, {bstat['hit_rate']:.1%} of vertices hit"
                  f"   {time.time() - t:6.1f}s")
        else:
            import torch
            scene = to_glb(torch.from_numpy(np.asarray(pre.vertices, np.float32)),
                           torch.from_numpy(np.asarray(pre.faces, np.int32)),
                           tex_voxels, PBR_ATTR_LAYOUT, info["output_resolution"],
                           texture_size=a.texture_size,
                           decimation_target=a.target_faces or 500_000)
            mesh = scene if isinstance(scene, trimesh.Trimesh) else scene.dump(concatenate=True)
            print(f"  bake       UV {len(mesh.faces):,} faces, {a.texture_size}px  "
                  f"{time.time() - t:6.1f}s")
        info["baker"] = a.baker
        mesh.export(a.output)
        info["faces"], info["vertices"] = len(mesh.faces), len(mesh.vertices)
        print(f"\nTOTAL {info['seconds']}s   peak {info['peak_gb']} GB   -> {a.output}")
        if a.json:
            Path(a.json).write_text(json.dumps(info, indent=2))
        return

    if a.raw:
        mesh = trimesh.Trimesh(v.cpu().numpy(), f.cpu().numpy(), process=False)
    else:
        from pixal3d_mlx.cleanup import clean
        t = time.time()
        mesh, stats = clean(v.cpu().numpy(), f.cpu().numpy(),
                            target_faces=a.target_faces or None)
        print(f"  cleanup    {time.time() - t:6.1f}s")
        info["cleanup"] = stats

    mesh.export(a.output)
    info["faces"] = len(mesh.faces)
    info["vertices"] = len(mesh.vertices)
    print(f"\nTOTAL {info['seconds']}s   peak {info['peak_gb']} GB   "
          f"{len(mesh.faces):,} faces   -> {a.output}")

    if a.min_iou > 0:
        got = silhouette_iou(mesh.vertices, a.image, a.fov)
        if got is None:
            print("(input has no alpha matte — skipping silhouette check)")
        else:
            iou = float(got[0])
            info["silhouette_iou"] = round(iou, 4)
            print(f"silhouette IoU {iou:.3f}")
            if iou < a.min_iou:
                # completing is not succeeding — a run can finish cleanly and still
                # have produced a blob that does not match the input at all
                print(f"FAIL: IoU {iou:.3f} < {a.min_iou} — not tracking the input")
                if a.json:
                    Path(a.json).write_text(json.dumps(info, indent=2))
                sys.exit(1)

    if a.json:
        Path(a.json).write_text(json.dumps(info, indent=2))


if __name__ == "__main__":
    main()