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(mesh_or_vertices, image_path, fov, res=512, samples=3_000_000):
    """Re-project the mesh through the generating camera; IoU against the input matte.

    Points are SAMPLED UNIFORMLY OVER THE SURFACE, not taken from the vertex list.
    Projecting vertices makes the score depend on tessellation: the same shape scored
    0.965 at 227k vertices and 0.790 at 134k, purely because a sparser point cloud
    leaves holes inside its own silhouette. That would fail good assets for the crime
    of being decimated. Fixed-count surface sampling makes density a constant of the
    metric instead of a property of the mesh.

    Points MUST be rotated into the camera frame — o_voxel returns geometry 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

    if isinstance(mesh_or_vertices, trimesh.Trimesh):
        pts3, _ = trimesh.sample.sample_surface(mesh_or_vertices, samples)
    else:
        pts3 = np.asarray(mesh_or_vertices)   # bare vertices: caller accepts the bias

    tm = _FRONT_VIEW.copy()
    tm[1, 3] = -distance_from_fov(fov, 1.0, res)
    pts = to_camera_frame(pts3).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 gate(mesh, a, info):
    """Silhouette check + non-zero exit. Shared by the textured and geometry paths."""
    if a.min_iou <= 0:
        return
    got = silhouette_iou(mesh, a.image, a.fov)
    if got is None:
        print("(input has no alpha matte — skipping silhouette check)")
        return
    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)


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}")

        # The gate applies to TEXTURED runs too. Skipping it here would mean the
        # textured path — the one most likely to be used for real assets — is the only
        # one that can ship a blob silently.
        gate(mesh, a, info)
        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}")

    gate(mesh, a, info)

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



if __name__ == "__main__":
    main()