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