"""ShadeNet-2 Gradio Space (CPU-friendly: int8 ONNX backend)."""
import os

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

from inference_utils import build_grid, pil_to_np, resize_pad

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

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


def decompose(image: Image.Image) -> Image.Image:
    img_rgb = resize_pad(image.convert("RGB"), IMAGE_SIZE)
    x = pil_to_np(img_rgb).astype(np.float32)
    out = _sess.run(None, {"input_rgb": x})[0]
    return build_grid(img_rgb, out)


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="Albedo | shading / depth | normal | recon"),
    title="ShadeNet-2: single-image inverse rendering (20M)",
    description=("Decomposes a photo into albedo, relative depth, surface "
                 "normals and shading (grayscale irradiance). Successor of "
                 "ShadeNet. Runs the int8 ONNX model on CPU."),
    examples=EXAMPLES,
    cache_examples=False,
)

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