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