src/infer.py

"""Reusable PBR inference helpers shared by the CLI and the Gradio demo."""

import json
import os

import numpy as np
import torch
from PIL import Image

from src.model import PBRUNet, category_to_index


def default_ckpt_for_run(run_dir: str) -> str:
    """Prefer the EMA checkpoint; fall back to best.pt."""
    ema = os.path.join(run_dir, "best_ema.pt")
    best = os.path.join(run_dir, "best.pt")
    if os.path.isfile(ema):
        return ema
    if os.path.isfile(best):
        return best
    raise FileNotFoundError(f"No best_ema.pt or best.pt in {run_dir}")


def _build_from_args_json(run_dir: str, device) -> PBRUNet:
    with open(os.path.join(run_dir, "args.json")) as f:
        a = json.load(f)
    kwargs = dict(
        encoder_name=a.get("encoder", "resnet34"),
        encoder_weights=None,
        use_category=a.get("use_category", False),
        normal_xy_only=a.get("normal_xy", False),
        separate_normal_decoder=a.get("separate_normal_decoder", False),
        predict_height=a.get("predict_height", False),
    )
    return PBRUNet(**kwargs).to(device)


def load_pbr_model(run_dir: str, ckpt_path, device) -> PBRUNet:
    """Build a PBRUNet from run_dir/args.json and load its checkpoint."""
    if ckpt_path is None:
        ckpt_path = default_ckpt_for_run(run_dir)
    model = _build_from_args_json(run_dir, device)
    state = torch.load(ckpt_path, weights_only=False, map_location=device)
    if isinstance(state, dict) and "model" in state:
        state = state["model"]
    model.load_state_dict(state)
    model.eval()
    return model


def predict_maps(model: PBRUNet, image, category: str, device, size: int = 256) -> dict:
    """Run inference on a PIL image, returning CPU map tensors in [0,1]."""
    img = image.convert("RGB").resize((size, size), Image.Resampling.BICUBIC)
    arr = np.asarray(img, dtype=np.float32) / 255.0
    basecolor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(device)
    cat = torch.tensor([category_to_index(category)], dtype=torch.long, device=device)
    with torch.no_grad():
        preds = model(basecolor, category=cat)
    return {
        name: preds[name][0].clamp(0, 1).cpu()
        for name in ("normal", "roughness", "metallic")
    }


def available_runs(models_root: str, allow=None) -> list[str]:
    """Sorted run names under models_root with args.json + a checkpoint."""
    out = []
    if not os.path.isdir(models_root):
        return out
    for name in os.listdir(models_root):
        run = os.path.join(models_root, name)
        if not os.path.isdir(run):
            continue
        if not os.path.isfile(os.path.join(run, "args.json")):
            continue
        has_ckpt = os.path.isfile(os.path.join(run, "best_ema.pt")) or os.path.isfile(
            os.path.join(run, "best.pt")
        )
        if not has_ckpt:
            continue
        if allow is not None and name not in allow:
            continue
        out.append(name)
    return sorted(out)