app.py

# app.py
"""Gradio demo: basecolor image -> predicted PBR maps + a 3D material-ball preview.

For the bundled example textures (which ship their ground-truth maps), the demo
shows the model's prediction next to the ground-truth render on the same lit
ball — an honest side-by-side, not a polished illusion.

Runs locally (reads ./outputs) and on a Hugging Face Space (set PBR_MODELS_ROOT
to wherever the curated checkpoints live).
"""

import hashlib
import json
import os

import numpy as np
import torch
from PIL import Image

from src.infer import available_runs, load_pbr_model, predict_maps
from src.render_preview import render_sphere_preview
from src.model import CATEGORIES

MODELS_ROOT = os.environ.get("PBR_MODELS_ROOT", "outputs")
EXAMPLES_DIR = "examples"
_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Default light/detail for the material-ball preview.
_DEF_AZ, _DEF_EL, _DEF_DETAIL = -35.0, 35.0, 1.3

_DEFAULT_KEEPERS = {
    "S1_bce", "S1B_bce_gan", "S1B_bce_long", "S2_dual_w10", "S3_rw1",
    "S4_baseline", "S4_gan_light", "S4_gan_mid", "S4_gan_heavy",
}


def _load_keepers() -> set[str]:
    path = os.path.join(MODELS_ROOT, "keepers.json")
    if os.path.isfile(path):
        with open(path) as f:
            data = json.load(f)
        runs = data.get("runs", {})
        chosen = {k for k, v in runs.items() if v}
        if chosen:
            return chosen
    return set(_DEFAULT_KEEPERS)


KEEPERS = _load_keepers()
_MODEL_CACHE: dict = {}


def _get_model(run_name: str):
    if run_name not in _MODEL_CACHE:
        run_dir = os.path.join(MODELS_ROOT, run_name)
        _MODEL_CACHE[run_name] = load_pbr_model(run_dir, None, _DEVICE)
    return _MODEL_CACHE[run_name]


def _to_img(tensor) -> Image.Image:
    arr = (tensor.clamp(0, 1) * 255).byte().cpu().numpy()
    if arr.shape[0] == 1:
        return Image.fromarray(arr[0], mode="L")
    return Image.fromarray(arr.transpose(1, 2, 0), mode="RGB")


def _basecolor_tensor(pil: Image.Image) -> torch.Tensor:
    arr = np.asarray(
        pil.convert("RGB").resize((256, 256), Image.Resampling.BICUBIC),
        dtype=np.float32,
    ) / 255.0
    return torch.from_numpy(arr).permute(2, 0, 1)


def _img_key(pil: Image.Image) -> str:
    """Content hash of the 256x256 basecolor, to match example textures to GT."""
    arr = np.asarray(
        pil.convert("RGB").resize((256, 256), Image.Resampling.BICUBIC),
        dtype=np.uint8,
    )
    return hashlib.md5(arr.tobytes()).hexdigest()


def _load_map_png(path: str, channels: int) -> torch.Tensor:
    arr = np.asarray(Image.open(path), dtype=np.float32) / 255.0
    if channels == 3:
        if arr.ndim == 2:
            arr = np.repeat(arr[..., None], 3, axis=2)
        return torch.from_numpy(arr).permute(2, 0, 1)
    if arr.ndim == 3:
        arr = arr[..., 0]
    return torch.from_numpy(arr).unsqueeze(0)


def _gt_registry() -> dict:
    """Map example-basecolor hash -> ground-truth maps loaded from shipped PNGs."""
    reg = {}
    manifest = os.path.join(EXAMPLES_DIR, "manifest.json")
    if not os.path.isfile(manifest):
        return reg
    with open(manifest) as f:
        items = json.load(f)
    for it in items:
        if "gt_normal" not in it:
            continue
        try:
            key = _img_key(Image.open(os.path.join(EXAMPLES_DIR, it["file"])))
            reg[key] = {
                "normal": _load_map_png(os.path.join(EXAMPLES_DIR, it["gt_normal"]), 3),
                "roughness": _load_map_png(os.path.join(EXAMPLES_DIR, it["gt_roughness"]), 1),
                "metallic": _load_map_png(os.path.join(EXAMPLES_DIR, it["gt_metallic"]), 1),
            }
        except Exception:
            pass
    return reg


_GT = _gt_registry()


def _render_ball(maps: dict, basecolor, az, el, detail) -> Image.Image:
    sphere = render_sphere_preview(
        maps["normal"], maps["roughness"], maps["metallic"], basecolor,
        light_az_deg=az, light_el_deg=el, detail_strength=detail,
    )
    return _to_img(sphere)


def run_inference(image, run_name: str, category: str,
                  az: float = _DEF_AZ, el: float = _DEF_EL, detail: float = _DEF_DETAIL):
    """Predict maps and render the prediction ball + (if known) the GT ball.

    Returns (normal_img, roughness_img, metallic_img, pred_ball, gt_ball, state).
    gt_ball is None for arbitrary uploads (no ground truth available).
    """
    if image is None:
        raise ValueError("Please provide an input image.")
    pil = image if isinstance(image, Image.Image) else Image.fromarray(np.asarray(image))
    model = _get_model(run_name)
    maps = predict_maps(model, pil, category, _DEVICE, size=256)
    basecolor = _basecolor_tensor(pil)
    gt = _GT.get(_img_key(pil))

    state = {
        "normal": maps["normal"],
        "roughness": maps["roughness"],
        "metallic": maps["metallic"],
        "basecolor": basecolor,
        "gt": gt,
    }
    pred_ball = _render_ball(maps, basecolor, az, el, detail)
    gt_ball = _render_ball(gt, basecolor, az, el, detail) if gt else None
    return (
        _to_img(maps["normal"]),
        _to_img(maps["roughness"]),
        _to_img(maps["metallic"]),
        pred_ball,
        gt_ball,
        state,
    )


def relight(state, az: float, el: float, detail: float):
    """Re-render both balls from stored maps when a slider moves (no model run)."""
    if not state:
        return None, None
    bc = state["basecolor"]
    pred_ball = _render_ball(state, bc, az, el, detail)
    gt_ball = _render_ball(state["gt"], bc, az, el, detail) if state.get("gt") else None
    return pred_ball, gt_ball


def _run_choices() -> list[str]:
    runs = available_runs(MODELS_ROOT, allow=KEEPERS)
    return runs or sorted(KEEPERS)


def _examples(default_run: str):
    """Build [path, run, category] rows from examples/manifest.json (or bare PNGs)."""
    if not os.path.isdir(EXAMPLES_DIR):
        return []
    manifest = os.path.join(EXAMPLES_DIR, "manifest.json")
    if os.path.isfile(manifest):
        with open(manifest) as f:
            items = json.load(f)
        return [[os.path.join(EXAMPLES_DIR, it["file"]), default_run, it["category"]]
                for it in items if os.path.isfile(os.path.join(EXAMPLES_DIR, it["file"]))]
    return [[os.path.join(EXAMPLES_DIR, f), default_run, "unknown"]
            for f in sorted(os.listdir(EXAMPLES_DIR))
            if f.lower().endswith((".png", ".jpg", ".jpeg"))]


def build_ui():
    import gradio as gr

    runs = _run_choices()
    default_run = "S4_gan_light" if "S4_gan_light" in runs else (runs[0] if runs else None)

    with gr.Blocks(title="PBR Material Predictor") as demo:
        gr.Markdown(
            "# PBR Material Predictor\n"
            "Upload a basecolor texture to predict its **normal**, **roughness**, and "
            "**metallic** maps, then see the material rendered on a 3D ball you can "
            "relight. Pick which trained run to use.\n\n"
            "_The bundled examples are **held-out test materials the model never "
            "trained on**, shown with the **ground-truth** render beside the "
            "prediction. The gap — mostly missing high-frequency normal detail — is "
            "the model's main open limitation, shown honestly._"
        )
        state = gr.State(None)
        with gr.Row():
            with gr.Column(scale=1):
                inp = gr.Image(type="pil", label="Basecolor input")
                run_dd = gr.Dropdown(choices=runs, value=default_run, label="Model (run)")
                cat_dd = gr.Dropdown(choices=CATEGORIES, value="unknown", label="Category")
                btn = gr.Button("Predict", variant="primary")
                az = gr.Slider(-180, 180, value=_DEF_AZ, step=5, label="Light azimuth")
                el = gr.Slider(5, 85, value=_DEF_EL, step=5, label="Light elevation")
                detail = gr.Slider(0.3, 2.5, value=_DEF_DETAIL, step=0.1,
                                   label="Surface detail")
            with gr.Column(scale=2):
                with gr.Row():
                    out_ball = gr.Image(label="Model prediction (drag sliders to relight)")
                    out_gt = gr.Image(label="Ground truth (example textures only)")
                with gr.Row():
                    out_normal = gr.Image(label="Pred normal")
                    out_rough = gr.Image(label="Pred roughness")
                    out_metal = gr.Image(label="Pred metallic")

        rows = _examples(default_run)
        if rows:
            gr.Examples(examples=rows, inputs=[inp, run_dd, cat_dd])

        btn.click(
            run_inference,
            inputs=[inp, run_dd, cat_dd, az, el, detail],
            outputs=[out_normal, out_rough, out_metal, out_ball, out_gt, state],
        )
        for slider in (az, el, detail):
            slider.release(relight, inputs=[state, az, el, detail],
                           outputs=[out_ball, out_gt])
    return demo


if __name__ == "__main__":
    build_ui().launch()