Spaces:
Running on Zero
Running on Zero
| """Shared prediction visualization and output helpers.""" | |
| from pathlib import Path | |
| import numpy as np | |
| from PIL import Image | |
| # ColorBrewer Spectral (reversed by default for disparity-style predictions). | |
| _SPECTRAL = np.asarray( | |
| [ | |
| [158, 1, 66], [213, 62, 79], [244, 109, 67], [253, 174, 97], | |
| [254, 224, 139], [255, 255, 191], [230, 245, 152], [171, 221, 164], | |
| [102, 194, 165], [50, 136, 189], [94, 79, 162], | |
| ], | |
| dtype=np.float32, | |
| ) | |
| def depth_to_vis(depth, reverse_color=True): | |
| """Convert a 2-D depth/disparity array to an RGB visualization.""" | |
| values = np.asarray(depth, dtype=np.float32) | |
| if values.ndim != 2: | |
| raise ValueError("depth must have shape [H, W]") | |
| finite = np.isfinite(values) | |
| normalized = np.zeros_like(values) | |
| if finite.any(): | |
| low = float(values[finite].min()) | |
| high = float(values[finite].max()) | |
| if high > low: | |
| normalized[finite] = (values[finite] - low) / (high - low) | |
| if reverse_color: | |
| normalized = 1.0 - normalized | |
| position = normalized * (len(_SPECTRAL) - 1) | |
| lower = np.floor(position).astype(np.intp) | |
| upper = np.minimum(lower + 1, len(_SPECTRAL) - 1) | |
| weight = (position - lower)[..., None] | |
| rgb = _SPECTRAL[lower] * (1.0 - weight) + _SPECTRAL[upper] * weight | |
| rgb[~finite] = 0 | |
| return Image.fromarray(np.rint(rgb).astype(np.uint8), mode="RGB") | |
| def normal_to_vis(normal): | |
| """Convert normal vectors in [-1, 1] to an RGB visualization.""" | |
| values = np.asarray(normal, dtype=np.float32) | |
| if values.ndim != 3 or values.shape[-1] != 3: | |
| raise ValueError("normal must have shape [H, W, 3]") | |
| rgb = np.rint((np.clip(values, -1.0, 1.0) + 1.0) * 127.5).astype(np.uint8) | |
| return Image.fromarray(rgb, mode="RGB") | |
| def save_prediction_outputs(output_root, stem, rgb, depth, normal): | |
| """Save raw predictions and their visualizations using the release layout.""" | |
| root = Path(output_root) | |
| directories = { | |
| name: root / name | |
| for name in ("depth_raw", "depth_vis", "normal_raw", "normal_vis", "rgb_recon") | |
| } | |
| for directory in directories.values(): | |
| directory.mkdir(parents=True, exist_ok=True) | |
| # Both public backends store normalized disparity in [0, 1]. | |
| depth = np.clip(np.asarray(depth, dtype=np.float32), 0.0, 1.0) | |
| normal = np.asarray(normal, dtype=np.float32) | |
| np.save(directories["depth_raw"] / (stem + ".npy"), depth) | |
| np.save(directories["normal_raw"] / (stem + ".npy"), normal) | |
| depth_to_vis(depth, reverse_color=True).save(directories["depth_vis"] / (stem + ".png")) | |
| normal_to_vis(normal).save(directories["normal_vis"] / (stem + ".png")) | |
| rgb.convert("RGB").save(directories["rgb_recon"] / (stem + ".png")) | |