File size: 2,192 Bytes
2d1eb0d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | """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()
|