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 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")
    a = ap.parse_args()

    t = time.time()
    models = load_all()
    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"))
    info.pop("subs", None)
    info.pop("hr_slat", None)

    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()