| """Standalone inference utilities (numpy/PIL only — no torch needed).""" |
| import numpy as np |
| from PIL import Image, ImageDraw |
|
|
| ALBEDO = slice(0, 3) |
| DEPTH = slice(3, 4) |
| NORMAL = slice(4, 7) |
| SHADING = slice(7, 8) |
|
|
|
|
| def resize_pad(img_rgb: Image.Image, image_size: int) -> Image.Image: |
| """Scale + center-crop to a square (RGB in, RGB out).""" |
| img_rgb = img_rgb.convert("RGB") |
| w, h = img_rgb.size |
| size = int(image_size) |
| scale = max(size / float(w), size / float(h)) |
| img_r = img_rgb.resize((max(1, int(round(w * scale))), |
| max(1, int(round(h * scale)))), Image.BICUBIC) |
| left = (img_r.size[0] - size) // 2 |
| top = (img_r.size[1] - size) // 2 |
| return img_r.crop((left, top, left + size, top + size)) |
|
|
|
|
| def pil_to_np(img_rgb: Image.Image) -> np.ndarray: |
| """PIL RGB -> [1, 3, H, W] float32 in [-1, 1].""" |
| arr = np.array(img_rgb.convert("RGB"), dtype=np.float32).transpose(2, 0, 1) |
| return arr[np.newaxis] / 255.0 * 2.0 - 1.0 |
|
|
|
|
| def _to_01(output: np.ndarray) -> np.ndarray: |
| return np.clip((output + 1.0) / 2.0, 0, 1) |
|
|
|
|
| def _tile(arr01: np.ndarray, ts: int) -> Image.Image: |
| return Image.fromarray( |
| (arr01 * 255.0 + 0.5).astype("uint8")).resize((ts, ts), Image.BICUBIC) |
|
|
|
|
| def _resize_map(out: np.ndarray, size: int) -> np.ndarray: |
| """Bilinear-resize a [1, C, H, W] float map to [1, C, size, size].""" |
| c = out.shape[1] |
| res = np.empty((1, c, size, size), dtype=np.float32) |
| for i in range(c): |
| im = Image.fromarray(out[0, i].astype(np.float32), mode="F") |
| res[0, i] = np.asarray(im.resize((size, size), Image.BILINEAR)) |
| return res |
|
|
|
|
| ENS_SCALES = (0.875, 1.0, 1.125) |
|
|
|
|
| def run_ensemble_onnx(sess, img_rgb: Image.Image, image_size: int = 512, |
| scales=ENS_SCALES) -> np.ndarray: |
| """3-pass multi-scale median (torch-free). Returns [1, 8, S, S].""" |
| outs = [] |
| for s in scales: |
| side = max(16, int(round(image_size * s)) // 16 * 16) |
| scaled = resize_pad(img_rgb, side) |
| x = pil_to_np(scaled).astype(np.float32) |
| o = sess.run(None, {"input_rgb": x})[0] |
| outs.append(o if side == image_size else _resize_map(o, image_size)) |
| return np.median(np.stack(outs, axis=0), axis=0).astype(np.float32) |
|
|
|
|
| def extract_maps(output: np.ndarray, img_rgb: Image.Image, |
| ts: int = 256) -> dict: |
| """Split a [1, 8, H, W] output ([-1, 1]) into labelled PIL tiles.""" |
| assert output.shape[1] == 8, f"want 8 channels, got {output.shape[1]}" |
| o = _to_01(output) |
| alb = o[:, ALBEDO][0].transpose(1, 2, 0) |
| sh = o[:, SHADING][0, 0] |
| d = o[:, DEPTH][0, 0] |
| span = d.max() - d.min() |
| dg = (d - d.min()) / (span if span > 1e-6 else 1.0) |
| return { |
| "input": img_rgb.resize((ts, ts), Image.BICUBIC), |
| "albedo": _tile(alb, ts), |
| "shading": _tile(np.repeat(sh[..., None], 3, axis=2), ts), |
| "depth": _tile(dg, ts).convert("RGB"), |
| "normal": _tile(o[:, NORMAL][0].transpose(1, 2, 0), ts), |
| "recon": _tile(np.clip(alb * sh[..., None], 0, 1), ts), |
| } |
|
|
|
|
| def build_grid(img_rgb: Image.Image, output: np.ndarray, |
| ts: int = 256) -> Image.Image: |
| """3x2 labelled grid: input|albedo|shading / depth|normal|recon.""" |
| maps = extract_maps(output, img_rgb, ts) |
| grid = Image.new("RGB", (ts * 3, ts * 2)) |
| grid.paste(maps["input"], (0, 0)) |
| grid.paste(maps["albedo"], (ts, 0)) |
| grid.paste(maps["shading"], (ts * 2, 0)) |
| grid.paste(maps["depth"], (0, ts)) |
| grid.paste(maps["normal"], (ts, ts)) |
| grid.paste(maps["recon"], (ts * 2, ts)) |
| draw = ImageDraw.Draw(grid) |
| for (x, y, name) in [(4, 4, "input"), (ts + 4, 4, "albedo"), |
| (ts * 2 + 4, 4, "shading"), (4, ts + 4, "depth"), |
| (ts + 4, ts + 4, "normal"), |
| (ts * 2 + 4, ts + 4, "recon")]: |
| draw.text((x, y), name, fill=(255, 255, 0)) |
| return grid |
|
|
