Spaces:
Running on Zero
Running on Zero
| """ShadeNet-5 Gradio Space: photo -> albedo, relative depth, surface normals. | |
| Runs the 4/8-bit "MIX" ONNX export of the model on the free CPU tier | |
| (onnxruntime only; torch is used just for the rotation resampling of the | |
| optional enhanced mode). | |
| """ | |
| import json | |
| import os | |
| import time | |
| import gradio as gr | |
| import numpy as np | |
| import onnxruntime as ort | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from PIL import Image, ImageOps | |
| from ensemble import mean_of, median_of, tta_passes, unit_normals | |
| from framing import crop_to_box, display_size, letterbox, resize_maps | |
| try: # present on ZeroGPU Spaces only | |
| import spaces | |
| except ImportError: # local tests, plain CPU hosts | |
| spaces = None | |
| if spaces is not None: | |
| # ZeroGPU hosting (the free tier for Gradio Spaces) wants one function | |
| # decorated with @spaces.GPU. It is never called: all inference below runs | |
| # on the Space's CPU through ONNX Runtime, so no GPU quota is used. | |
| def _zerogpu_placeholder(): | |
| return None | |
| MODEL_REPO = "singam96/ShadeNet-5-35M" | |
| ONNX_FILE = "onnx/shadenet5_mix_int4_int8.onnx" | |
| REFINER_FILE = "onnx/shadenet5_refiner.onnx" | |
| ROT_DEG = 4.0 # the card's enhanced recipe: mirror + +-4 deg | |
| # config.json also tells us the input size (and is what the Hub counts as a | |
| # download of the model). | |
| with open(hf_hub_download(MODEL_REPO, "config.json")) as f: | |
| CFG = json.load(f) | |
| SIZE = int(CFG["image_size"]) | |
| def cpu_budget() -> int: | |
| """CPUs this container may use (its cgroup quota), not the host's core | |
| count: asking ONNX Runtime for more threads than the quota made it about | |
| 9x slower in a measurement (64 host cores, 12 CPUs of quota).""" | |
| try: | |
| quota, period = open("/sys/fs/cgroup/cpu.max").read().split() | |
| if quota != "max": | |
| return max(1, min(8, int(int(quota) / int(period)))) | |
| except Exception: | |
| pass | |
| try: | |
| return max(1, min(8, len(os.sched_getaffinity(0)))) | |
| except Exception: | |
| return 2 | |
| THREADS = int(os.environ.get("ORT_THREADS", cpu_budget())) | |
| torch.set_num_threads(THREADS) | |
| _so = ort.SessionOptions() | |
| _so.log_severity_level = 3 | |
| _so.intra_op_num_threads = THREADS | |
| SESS = ort.InferenceSession(hf_hub_download(MODEL_REPO, ONNX_FILE), _so, | |
| providers=["CPUExecutionProvider"]) | |
| RSESS = ort.InferenceSession(hf_hub_download(MODEL_REPO, REFINER_FILE), _so, | |
| providers=["CPUExecutionProvider"]) | |
| MODE_SINGLE = "Single pass" | |
| MODE_REFINE = "Single pass + refiner (recommended)" | |
| MODE_SIX = "Six passes (mirror + rotations, slower)" | |
| def to_tensor(img: Image.Image) -> torch.Tensor: | |
| arr = np.asarray(img, dtype=np.float32) / 255.0 | |
| return torch.from_numpy(arr.transpose(2, 0, 1)[None] * 2.0 - 1.0) | |
| def refine(c: torch.Tensor) -> torch.Tensor: | |
| out = SESS.run(None, {"input_rgb": np.ascontiguousarray(c.numpy())})[0] | |
| return torch.from_numpy(out) | |
| def run_refiner(photo: torch.Tensor, maps: torch.Tensor) -> torch.Tensor: | |
| """The small refiner (trained to turn a single pass into the six-pass result).""" | |
| out = RSESS.run(None, {"photo": np.ascontiguousarray(photo.numpy()), | |
| "maps": np.ascontiguousarray(maps.numpy())})[0] | |
| return torch.from_numpy(out) | |
| def to_img(x01: torch.Tensor) -> Image.Image: | |
| if x01.shape[1] == 1: | |
| x01 = x01.expand(-1, 3, -1, -1) | |
| arr = (x01[0].permute(1, 2, 0).clamp(0, 1).numpy() * 255).round() | |
| return Image.fromarray(arr.astype("uint8")) | |
| def predict(image, mode, median, progress=gr.Progress()): | |
| if image is None: | |
| raise gr.Error("Upload a photo first.") | |
| photo = ImageOps.exif_transpose(image).convert("RGB") | |
| # Fit the WHOLE photo into the square input (long side -> 448, the rest | |
| # mirrored from the photo's own edge) and cut the padding off the maps again. | |
| square, box = letterbox(photo, SIZE) | |
| x = to_tensor(square) | |
| ow, oh = display_size(*photo.size) | |
| enhanced = mode == MODE_SIX | |
| use_median = bool(enhanced and median) | |
| diff = None | |
| t0 = time.time() | |
| with torch.no_grad(): | |
| if enhanced: | |
| passes = tta_passes( | |
| refine, x, flip=True, rot_deg=ROT_DEG, | |
| progress=lambda i, n: progress((i, n), desc="forward pass")) | |
| mean_maps = mean_of(passes) | |
| maps = median_of(passes) if use_median else mean_maps | |
| elif mode == MODE_REFINE: | |
| single = refine(x) | |
| maps = run_refiner(x, single).float().clamp(-1.0, 1.0) | |
| else: | |
| maps = refine(x) | |
| maps = resize_maps(crop_to_box(unit_normals(maps.float()), box), ow, oh) | |
| if use_median: # where the two reductions disagree | |
| ref = resize_maps(crop_to_box(unit_normals(mean_maps.float()), box), ow, oh) | |
| d = ((maps - ref).abs().mean(dim=1, keepdim=True) * 8).clamp(0, 1) | |
| diff = to_img(d) | |
| dt = time.time() - t0 | |
| albedo = to_img((maps[:, 0:3] + 1) / 2) | |
| depth = to_img((maps[:, 3:4] + 1) / 2) | |
| normal = to_img((maps[:, 4:7] + 1) / 2) | |
| how = {MODE_SINGLE: "1 forward pass", | |
| MODE_REFINE: "1 forward pass + the small refiner", | |
| MODE_SIX: "6 forward passes (" + ("median" if use_median else "mean") + ")"}[mode] | |
| note = (f"{how} in {dt:.1f} s on CPU; the whole {photo.size[0]}x{photo.size[1]} " | |
| f"photo is kept (model input {SIZE}x{SIZE}, long side scaled to {SIZE}).") | |
| if median and not enhanced: | |
| note += " Median only applies to the six-pass mode." | |
| return photo.resize((ow, oh), Image.BICUBIC), albedo, depth, normal, diff, note | |
| EXAMPLES = [[os.path.join("examples", f), MODE_REFINE, False] | |
| for f in sorted(os.listdir("examples")) | |
| if f.lower().endswith((".jpg", ".png"))] | |
| with gr.Blocks(title="ShadeNet-5") as demo: | |
| gr.Markdown( | |
| "# ShadeNet-5: albedo, depth and normals from one photo\n" | |
| "A 35M-parameter feed-forward network on a frozen DINOv2 backbone. " | |
| "This Space runs the 4/8-bit **MIX** ONNX export on a free CPU. " | |
| "Depth is relative (per-image scale and shift); normals are " | |
| "x right, y up, z toward the viewer.\n\n" | |
| "**Single pass** is one forward pass. **Single pass + refiner** adds a " | |
| "small trained UNet (0.91M parameters) that cleans up the maps, mostly " | |
| "the normals, for a fraction of a second more: it recovers about " | |
| "three quarters of what the slower mode gains. **Six passes** averages the " | |
| "mirrored and rotated (+-4 degrees) photo and is the best but about six " | |
| "times slower; tick **Median** to take the median of the six instead of " | |
| "the mean (marginally worse in the card's measurements) and see the " | |
| "difference image. " | |
| f"[Model card and files](https://huggingface.co/{MODEL_REPO}). " | |
| "The example photos are Flickr8k photos shown for illustration only; " | |
| "see the model card for their licence status.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| inp = gr.Image(type="pil", label="Photo", sources=["upload", "clipboard"]) | |
| mode = gr.Radio([MODE_SINGLE, MODE_REFINE, MODE_SIX], value=MODE_REFINE, | |
| label="Mode") | |
| med = gr.Checkbox(False, label="Median instead of mean (six passes only)", | |
| interactive=False) | |
| btn = gr.Button("Decompose", variant="primary") | |
| msg = gr.Markdown() | |
| with gr.Column(): | |
| with gr.Row(): | |
| o_in = gr.Image(label="Input") | |
| o_alb = gr.Image(label="Albedo") | |
| with gr.Row(): | |
| o_dep = gr.Image(label="Depth (lighter = farther)") | |
| o_nor = gr.Image(label="Normal") | |
| o_diff = gr.Image(label="Median minus mean (x8, all maps; black = identical)") | |
| mode.change(lambda m: gr.update(interactive=(m == MODE_SIX), value=False), mode, med) | |
| outs = [o_in, o_alb, o_dep, o_nor, o_diff, msg] | |
| btn.click(predict, [inp, mode, med], outs) | |
| gr.Examples(EXAMPLES, [inp, mode, med], outs, fn=predict, cache_examples=False) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=8).launch() | |