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