"""ShadeNet-3.2 Gradio Space (CPU-friendly: fp16 ONNX backend)."""
import os

import gradio as gr
import numpy as np
import onnxruntime as ort
import spaces
from PIL import Image

from inference_utils import (build_grid, extract_maps, pil_to_np, resize_pad,
                             run_ensemble_onnx)

MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                           "onnx", "model_fp16.onnx")
IMAGE_SIZE = 512

_sess = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])


@spaces.GPU(duration=60)
def decompose(image: Image.Image):
    """Returns (grid, input, albedo, shading, depth, normal, recon)."""
    img_rgb = resize_pad(image.convert("RGB"), IMAGE_SIZE)
    out = run_ensemble_onnx(_sess, img_rgb, IMAGE_SIZE)
    maps = extract_maps(out, img_rgb)
    return (build_grid(img_rgb, out), maps["input"], maps["albedo"],
            maps["shading"], maps["depth"], maps["normal"], maps["recon"])


EXAMPLES = []
_exdir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                      "assets", "examples")
if os.path.isdir(_exdir):
    EXAMPLES = sorted(os.path.join(_exdir, f) for f in os.listdir(_exdir)
                      if f.lower().endswith((".png", ".jpg", ".jpeg")))

demo = gr.Interface(
    fn=decompose,
    inputs=gr.Image(type="pil", label="Input photo"),
    outputs=[
        gr.Image(type="pil", label="Overview grid"),
        gr.Image(type="pil", label="Input (resized)"),
        gr.Image(type="pil", label="Albedo"),
        gr.Image(type="pil", label="Shading"),
        gr.Image(type="pil", label="Depth (relative: dark=near)"),
        gr.Image(type="pil", label="Normal"),
        gr.Image(type="pil", label="Recon (albedo x shading)"),
    ],
    title="ShadeNet-3.2: single-image inverse rendering (7M)",
    description=("Decomposes a photo into albedo, relative depth, surface "
                 "normals and shading (grayscale irradiance). Successor of "
                 "ShadeNet-2 (20M), 3x smaller with better depth/normal "
                 "accuracy. Runs the fp16 ONNX model."),
    examples=EXAMPLES,
    cache_examples=False,
)

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