forgebench/code/eval/appeval/evaluate_appforce.py

"""Appearance+geometry evaluation harness CLI for single-image 3D generation.

Two modes (both driven by this one CLI):
  1. INPUT-VIEW fidelity : render the generated mesh from the INPUT camera(s)
     (K + c2w_cv from the exp .npz) and compare to the input photo/crop over
     the object mask  ->  LPIPS / SSIM / CLIP / PSNR.
  2. NOVEL-VIEW fidelity : render the generated mesh AND the GT-textured mesh
     from the locked 24-view rig, compare -> LPIPS / SSIM / CLIP / PSNR; PLUS
     geometry metrics from the meshes directly (CD_L1/L2, F@0.01/0.02/0.05,
     Normal Consistency, Volume IoU@128).

GT-texture auto-detect: datasets with UNTEXTURED GT (e.g. SA-3DAO: uniform
vertex colour / no texture) automatically SKIP novel-view APPEARANCE, but still
run geometry + input-view-vs-photo.

Meshes are assumed ALREADY canonical in [-0.5,0.5]^3 (LOCKED spec: NO ICP,
NO alignment). Pre-canonicalize predictions to the GT frame upstream.

Usage:
  python evaluate_appforce.py \
     --exp   /lp-dev/jonghoon/mv-mesh/exp_faithfulness/toys4k \
     --meshes /lp-dev/.../gen_2v/trellis \
     --out    /home/nvidia/jonghoon/mv-mesh/.debug/appeval/trellis_2v \
     --gt-mesh-dir /lp-dev/jonghoon/mv-mesh/exp_faithfulness/toys4k/renders \
     --views both --dataset toys4k --novel
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

import numpy as np
import trimesh
from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parent))

import render as R
from geometry import geometry_metrics
from appearance import appearance_metrics, average_views, composite_white


VIEW_TAGS = {"front": ["front"], "both": ["front", "side"],
             "quad": ["front", "side", "back", "oside"]}


# ----------------------------------------------------------------------------
# helpers
# ----------------------------------------------------------------------------
def load_mesh(path):
    return trimesh.load(str(path), force="mesh", process=False)


def resolve_gt(gt_dir: Path, obj: str):
    for cand in (gt_dir / f"{obj}_canon.glb",
                 gt_dir / obj / "mesh.glb",
                 gt_dir / f"{obj}.glb"):
        if cand.exists():
            return cand
    return None


def is_textured(mesh: trimesh.Trimesh) -> bool:
    """True if the mesh carries genuine surface texture (UV image or spatially
    varying vertex colour). Uniform vertex colour / no colour -> False."""
    vis = mesh.visual
    uv = getattr(vis, "uv", None)
    mat = getattr(vis, "material", None)
    if uv is not None and mat is not None:
        for a in ("baseColorTexture", "image"):
            if getattr(mat, a, None) is not None:
                return True
    vc = getattr(vis, "vertex_colors", None)
    if vc is not None:
        vc = np.asarray(vc)[:, :3].astype(np.float32)
        if vc.std(axis=0).mean() > 3.0:   # >~1% of 0..255 range varies
            return True
    return False


def read_input_png(path):
    """Return (H,W,4) float[0,1] RGBA. If no alpha, alpha=1 everywhere."""
    im = np.asarray(Image.open(path).convert("RGBA")).astype(np.float32) / 255.0
    return im


def save_sheet(imgs, path, labels=None):
    """Concatenate a list of (H,W,3) float[0,1] images horizontally and save."""
    h = min(i.shape[0] for i in imgs)
    w = min(i.shape[1] for i in imgs)
    row = np.concatenate([i[:h, :w] for i in imgs], axis=1)
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray((np.clip(row, 0, 1) * 255).astype(np.uint8)).save(path)


# ----------------------------------------------------------------------------
# per-object evaluation
# ----------------------------------------------------------------------------
def eval_object(exp: Path, meshes: Path, gt_dir: Path, obj: str, view_tags,
                do_novel: bool, out: Path, save_debug: bool, ctx):
    row = {"object": obj}
    pred_path = meshes / f"{obj}.glb"
    gt_path = resolve_gt(gt_dir, obj)
    if not pred_path.exists():
        return {"object": obj, "error": "missing pred glb"}
    if gt_path is None:
        return {"object": obj, "error": "missing gt mesh"}

    pred_mesh = load_mesh(pred_path)
    gt_mesh = load_mesh(gt_path)
    pred_gl = R.prepare_mesh(pred_mesh)
    gt_gl = R.prepare_mesh(gt_mesh)

    # ---------- INPUT-VIEW appearance (mode 1) ----------
    iv_views = []
    for tag in view_tags:
        npz = exp / "renders" / f"{obj}_{tag}.npz"
        png = exp / "inputs" / f"{obj}_{tag}.png"
        if not (npz.exists() and png.exists()):
            continue
        z = np.load(npz)
        K = {k: float(z[k]) for k in ("fx", "fy", "cx", "cy")}
        res = int(z["res"])
        c2w = z["c2w_cv"]
        bbox = z["bbox"].tolist()
        # canonical stored order is (y0, y1, x0, x1) — MUST match synth_render.crop_bbox
        y0, y1, x0, x1 = bbox
        pred_full = R.render_input_view(pred_gl, K, c2w, res, res, ctx=ctx).cpu().numpy()
        pred_crop = pred_full[y0:y1, x0:x1]
        ref = read_input_png(png)
        # match sizes (crop == input png size by construction)
        h = min(pred_crop.shape[0], ref.shape[0])
        w = min(pred_crop.shape[1], ref.shape[1])
        m = appearance_metrics(pred_crop[:h, :w], ref[:h, :w])
        m["view"] = tag
        iv_views.append(m)
        if save_debug:
            save_sheet([composite_white(ref[:h, :w]),
                        composite_white(pred_crop[:h, :w])],
                       out / "debug" / f"{obj}_inputview_{tag}.png")
    if iv_views:
        row["input_view"] = average_views(
            [{k: v for k, v in d.items() if k != "view"} for d in iv_views])
        row["input_view"]["per_view"] = iv_views

    # ---------- NOVEL-VIEW geometry + appearance (mode 2) ----------
    if do_novel:
        try:
            row["geometry"] = geometry_metrics(pred_mesh, gt_mesh)
        except Exception as e:
            row["geometry_error"] = f"{type(e).__name__}: {e}"

        gt_tex = is_textured(gt_mesh)
        row["gt_textured"] = bool(gt_tex)
        if gt_tex:
            cams = R.orbit_cameras()
            pred_r = R.render_orbit(pred_gl, cams, ctx=ctx).cpu().numpy()
            gt_r = R.render_orbit(gt_gl, cams, ctx=ctx).cpu().numpy()
            nv_views = []
            for i, cam in enumerate(cams):
                m = appearance_metrics(pred_r[i], gt_r[i])
                m["view"] = cam["name"]
                nv_views.append(m)
            row["novel_view"] = average_views(
                [{k: v for k, v in d.items() if k != "view"} for d in nv_views])
            if save_debug:
                # a 3-view proof sheet (first of each elevation)
                picks = [0, 8, 16]
                imgs = []
                for p in picks:
                    imgs.append(composite_white(gt_r[p]))
                    imgs.append(composite_white(pred_r[p]))
                save_sheet(imgs, out / "debug" / f"{obj}_novel.png")
        else:
            row["novel_view"] = None  # untextured GT -> skip novel appearance
    return row


# ----------------------------------------------------------------------------
# tables
# ----------------------------------------------------------------------------
def _fmt(x, w=8, p=4):
    return f"{x:{w}.{p}f}" if isinstance(x, (int, float)) else f"{str(x):>{w}}"


def write_tables(out: Path, method: str, results: list, do_novel: bool):
    (out / "results.json").write_text(json.dumps(results, indent=2))
    ok = [r for r in results if "error" not in r]
    lines = []
    lines.append(f"APPEARANCE+GEOMETRY EVAL  method={method}  "
                 f"n_objects={len(ok)}/{len(results)}")
    lines.append("Rankings: LPIPS(primary,lower=better)  SSIM(secondary,higher) "
                 "CLIP(tertiary,higher)  |  PSNR=NON-RANKING(info only)")
    lines.append("")

    # INPUT-VIEW table
    hdr = f"{'object':22s} | {'LPIPS':>8s} {'SSIM':>8s} {'CLIP':>8s} {'PSNR*':>8s}"
    lines.append("== INPUT-VIEW (pred render vs input photo, object-masked) ==")
    lines.append(hdr)
    lines.append("-" * len(hdr))
    iv_rows = [r for r in ok if r.get("input_view")]
    for r in sorted(iv_rows, key=lambda v: v["object"]):
        iv = r["input_view"]
        lines.append(f"{r['object'][:22]:22s} | {_fmt(iv['lpips'])} "
                     f"{_fmt(iv['ssim'])} {_fmt(iv['clip'])} {_fmt(iv['psnr'])}")
    if iv_rows:
        def mean(k): return float(np.mean([r["input_view"][k] for r in iv_rows]))
        lines.append("-" * len(hdr))
        lines.append(f"{'MEAN(' + str(len(iv_rows)) + ')':22s} | "
                     f"{_fmt(mean('lpips'))} {_fmt(mean('ssim'))} "
                     f"{_fmt(mean('clip'))} {_fmt(mean('psnr'))}")
    lines.append("")

    if do_novel:
        # NOVEL-VIEW appearance table
        lines.append("== NOVEL-VIEW APPEARANCE (24 views, pred vs GT-textured) ==")
        lines.append(hdr)
        lines.append("-" * len(hdr))
        nv_rows = [r for r in ok if r.get("novel_view")]
        for r in sorted(nv_rows, key=lambda v: v["object"]):
            nv = r["novel_view"]
            lines.append(f"{r['object'][:22]:22s} | {_fmt(nv['lpips'])} "
                         f"{_fmt(nv['ssim'])} {_fmt(nv['clip'])} {_fmt(nv['psnr'])}")
        if nv_rows:
            def mean(k): return float(np.mean([r["novel_view"][k] for r in nv_rows]))
            lines.append("-" * len(hdr))
            lines.append(f"{'MEAN(' + str(len(nv_rows)) + ')':22s} | "
                         f"{_fmt(mean('lpips'))} {_fmt(mean('ssim'))} "
                         f"{_fmt(mean('clip'))} {_fmt(mean('psnr'))}")
        else:
            lines.append("(none: GT untextured -> novel appearance skipped)")
        lines.append("")

        # GEOMETRY table
        ghdr = (f"{'object':22s} | {'CD_L1':>8s} {'CD_L2':>9s} {'F@.01':>7s} "
                f"{'F@.02':>7s} {'F@.05':>7s} {'NC':>7s} {'VolIoU':>7s}")
        lines.append("== GEOMETRY (vs GT mesh, no ICP) ==  headline: CD_L1, F@.02")
        lines.append(ghdr)
        lines.append("-" * len(ghdr))
        g_rows = [r for r in ok if r.get("geometry")]
        for r in sorted(g_rows, key=lambda v: v["object"]):
            g = r["geometry"]
            lines.append(
                f"{r['object'][:22]:22s} | {g['cd_l1']:8.4f} {g['cd_l2']:9.5f} "
                f"{g['f01']:7.4f} {g['f02']:7.4f} {g['f05']:7.4f} "
                f"{g['normal_consistency']:7.4f} {g['vol_iou']:7.4f}")
        if g_rows:
            def gm(k): return float(np.mean([r["geometry"][k] for r in g_rows]))
            lines.append("-" * len(ghdr))
            lines.append(
                f"{'MEAN(' + str(len(g_rows)) + ')':22s} | {gm('cd_l1'):8.4f} "
                f"{gm('cd_l2'):9.5f} {gm('f01'):7.4f} {gm('f02'):7.4f} "
                f"{gm('f05'):7.4f} {gm('normal_consistency'):7.4f} "
                f"{gm('vol_iou'):7.4f}")
        lines.append("")

    errs = [r for r in results if "error" in r]
    if errs:
        lines.append("ERRORS:")
        for r in errs:
            lines.append(f"  {r['object']}: {r['error']}")

    (out / "results.txt").write_text("\n".join(lines) + "\n")
    print("\n".join(lines))
    print(f"\nwrote {out}/results.json + results.txt")


# ----------------------------------------------------------------------------
def discover_objects(exp: Path, meshes: Path):
    objs = sorted(p.stem for p in meshes.glob("*.glb")
                  if not p.stem.endswith("_aligned"))
    sel = exp / "selection.json"
    if sel.exists():
        want = {s["object"] for s in json.loads(sel.read_text())["selections"]}
        objs = [o for o in objs if o in want]
    return objs


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--exp", type=Path, required=True)
    ap.add_argument("--meshes", type=Path, required=True,
                    help="DIR containing {object}.glb (one method)")
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--gt-mesh-dir", type=Path, required=True,
                    help="DIR with {object}_canon.glb or {object}/mesh.glb")
    ap.add_argument("--views", choices=["front", "both", "quad"], default="both")
    ap.add_argument("--dataset", choices=["toys4k", "sa3dao"], default="toys4k")
    ap.add_argument("--novel", dest="novel", action="store_true", default=True)
    ap.add_argument("--no-novel", dest="novel", action="store_false")
    ap.add_argument("--limit", type=int, default=0, help="cap #objects (debug)")
    ap.add_argument("--no-debug", dest="debug", action="store_false", default=True)
    ap.add_argument("--shard", type=int, default=0, help="this shard index [0,nshards)")
    ap.add_argument("--nshards", type=int, default=1, help="split objects round-robin across N parallel procs")
    args = ap.parse_args()

    out = args.out
    out.mkdir(parents=True, exist_ok=True)
    view_tags = VIEW_TAGS[args.views]
    method = args.meshes.name
    ctx = R.get_ctx()

    objs = discover_objects(args.exp, args.meshes)
    if args.limit:
        objs = objs[:args.limit]
    if args.nshards > 1:
        objs = objs[args.shard::args.nshards]
    print(f"[appeval] method={method} dataset={args.dataset} views={args.views} "
          f"novel={args.novel} n_objects={len(objs)}")

    results = []
    for i, obj in enumerate(objs):
        try:
            r = eval_object(args.exp, args.meshes, args.gt_mesh_dir, obj,
                            view_tags, args.novel, out,
                            save_debug=args.debug and i < 6, ctx=ctx)
        except Exception as e:
            import traceback
            traceback.print_exc()
            r = {"object": obj, "error": f"{type(e).__name__}: {e}"}
        results.append(r)
        tag = "OK" if "error" not in r else "ERR"
        print(f"[{i+1}/{len(objs)}] {tag} {obj}", flush=True)

    write_tables(out, method, results, args.novel)


if __name__ == "__main__":
    main()