| --- |
| license: apache-2.0 |
| pipeline_tag: image-to-image |
| library_name: pytorch |
| tags: |
| - onnx |
| - inverse-rendering |
| - image-decomposition |
| - albedo |
| - normal-map |
| - depth-estimation |
| - pbr |
| - material-estimation |
| - flickr8k |
| - shadenet |
| datasets: |
| - singam96/flickr8k_marigold_v2 |
| metrics: |
| - mse |
| model-index: |
| - name: ShadeNet-2-20M |
| results: |
| - task: |
| type: image-to-image |
| dataset: |
| name: flickr8k_marigold_v2 (val split, 807 images) |
| type: singam96/flickr8k_marigold_v2 |
| metrics: |
| - name: val/loss (weighted MSE + recon) |
| type: mse |
| value: 0.2411 |
| --- |
| |
| # ShadeNet-2 20M |
|
|
| A lightweight inverse-rendering model: one photo in, **albedo + relative depth + surface normals + shading** out (8 channels, 20M params). Successor of [ShadeNet](https://huggingface.co/singam96/ShadeNet) — smaller, sharper, and with a new self-supervised shading head. |
|
|
| [](assets/hero.png) |
|
|
| ## Examples |
|
|
| [](assets/154871781_ae77696b77_result.png) |
| [](assets/157139628_5dc483e2e4_result.png) |
| [](assets/159712188_d530dd478c_result.png) |
| [](assets/160541986_d5be2ab4c1_result.png) |
| [](assets/160566014_59528ff897_result.png) |
| [](assets/160585932_fa6339f248_result.png) |
| [](assets/160792599_6a7ec52516_result.png) |
| [](assets/161669933_3e7d8c7e2c_result.png) |
|
|
| *Each grid: input | albedo | shading / depth | normal | recon (albedo×shading). Click any image for full size.* |
|
|
| ## ShadeNet v1 vs ShadeNet-2 |
|
|
| Same inputs, both models — albedo/basecolor, normals and depth side by side (v1 top rows, v2 bottom rows). Each side uses its best ensemble: v1 runs 5 noisy passes + median (its published protocol); v2 runs a 3-pass multi-scale median tuned on the val split (noise was tested and hurts: val/loss 0.2395 vs 0.2421 single-pass). |
|
|
| [](assets/compare_154871781_ae77696b77.png) |
| [](assets/compare_159712188_d530dd478c.png) |
| [](assets/compare_160585932_fa6339f248.png) |
|
|
| ## Architecture |
|
|
| **ParallelUNet generator (20.05M params, 18.24M trainable)** + spectral-norm GroupNorm PatchGAN discriminator (2.77M, training only): |
|
|
| - Dual parallel encoders — vanilla UNet path plus a **frozen MobileNetV2** feature trunk, fused at every decoder level |
| - Single-pass RGB → 8ch output (no dual-mode round trip) |
| - Adversarial training paced 1 G-step : 1 D-step every 2nd window, EMA weight shadow (shipped weights are EMA) |
|
|
| ## Output maps |
|
|
| | Map | Channels | Range | Description | |
| |---|---|---|---| |
| | Albedo | 3 `[0:3]` | [−1, 1] | Reflectance / diffuse color, lighting factored out | |
| | Depth | 1 `[3:4]` | [−1, 1] | **Relative** depth (0=near), affine-ambiguous | |
| | Normal | 3 `[4:7]` | [−1, 1] | Surface normals, unit-length regularised | |
| | Shading | 1 `[7:8]` | [−1, 1] | Grayscale irradiance; `input ≈ albedo × shading` | |
| | Recon | — | — | `albedo × shading` re-rendering (diagnostic, not a head) | |
|
|
| ## Files |
|
|
| ``` |
| ├── app.py # Gradio Space app (int8 ONNX, CPU-only) |
| ├── inference.py # Standalone torch CLI |
| ├── inference_utils.py # Grid visualisation (numpy/PIL) |
| ├── model.py # Standalone generator architecture |
| ├── requirements.txt |
| ├── checkpoints/shadenet2.ckpt # Torch weights, EMA (333MB) |
| └── onnx/ |
| ├── model.onnx # fp32, EMA (80MB) — GPU via CUDA EP |
| └── model_quantized.onnx # int8 dynamic, EMA (28MB) — CPU |
| ``` |
|
|
| ## Usage |
|
|
| ### Gradio Space |
|
|
| Try it in your browser — no installation: **[singam96/ShadeNet-2-20M Space](https://huggingface.co/spaces/singam96/ShadeNet-2-20M)** (runs the int8 model on CPU). |
|
|
| This repo ships `app.py`, the Space entrypoint. To recreate it: New Space → Gradio SDK → point at this repo. |
|
|
| ### Torch CLI |
|
|
| ```bash |
| pip install torch torchvision pillow numpy |
| python inference.py photo.jpg --output_dir ./output |
| # --checkpoint ./checkpoints/shadenet2.ckpt --image-size 512 --no-ema to disable EMA |
| ``` |
|
|
| ### ONNX (CPU) |
|
|
| ```bash |
| pip install onnxruntime pillow numpy |
| python - <<'EOF' |
| import onnxruntime as ort, numpy as np |
| from PIL import Image |
| from inference_utils import build_grid, pil_to_np, resize_pad |
| sess = ort.InferenceSession("onnx/model_quantized.onnx", providers=["CPUExecutionProvider"]) |
| img = resize_pad(Image.open("photo.jpg"), 512) |
| out = sess.run(None, {"input_rgb": pil_to_np(img).astype(np.float32)})[0] |
| build_grid(img, out).save("result.png") |
| EOF |
| ``` |
|
|
| Input: `[1, 3, H, W]` in `[-1, 1]` (any H, W; multiples of 32 recommended). Output: `[1, 8, H, W]` in `[-1, 1]`. |
|
|
| ## Training |
|
|
| Trained from scratch on [`singam96/flickr8k_marigold_v2`](https://huggingface.co/datasets/singam96/flickr8k_marigold_v2) (8077 Flickr8k photos with Marigold-V2 pseudo-labels: albedo/depth/normal), 384px, fp32, ~17h on a single GTX 1650, early-stopped on `val/loss` (patience 5). |
|
|
| Losses: per-map weighted MSE (depth term is **scale-shift-invariant** after least-squares alignment — decoded Marigold depth is relative) + self-supervised **reconstruction coupling** (`albedo×shading ≈ input`, the shading head's only supervision) + LSGAN + normal unit-length penalty. |
|
|
| [](assets/training_curves.png) |
|
|
| | Weights | val/loss | albedo MSE | depth SSI | normal MSE | recon L1 | |
| |---|---|---|---|---|---| |
| | last.ckpt | 0.2467 | 0.2912 | 0.0954 | 0.2541 | 0.0690 | |
| | best.ckpt (raw) | 0.2464 | 0.2879 | 0.0957 | 0.2561 | 0.0688 | |
| | best.ckpt (EMA) | 0.2411 | 0.2758 | 0.0925 | 0.2542 | 0.0730 | |
| | **best-calibrated (EMA, shipped)** | **0.2384** | **0.2726** | **0.0920** | **0.2456** | 0.0787 | |
| | calibrated + 3-pass ensemble | 0.2373 | 0.2729 | 0.0917 | 0.2420 | 0.0791 | |
|
|
| Full 807-image val split. Depth is the best-learned map; normals the hardest. |
|
|
| Shipped weights carry a post-hoc Optuna calibration (inference-only, no |
| retraining): per-layer weight affines on the three parameter groups that a |
| sensitivity audit showed could matter (`head_out`/`head`/`up4` — everything |
| else is GroupNorm-invariant), tuned on val-loss. All page outputs use the |
| calibrated weights with a 3-pass multi-scale median ensemble. |
|
|
| ## Limitations |
|
|
| - Depth is **relative**, not metric — don't read meters off it. |
| - Shading assumes **white light**; strongly colored illumination (sunsets, neon) leaks into albedo. |
| - Normals are noisy in foliage/sky — those pseudo-labels were noisy too. |
| - Uncertainty is not provided: confidently-wrong pseudo-labels are fitted confidently. |
|
|
| ## Attribution |
|
|
| Supervision labels come from **Marigold V2** (Ke et al.) applied to **Flickr8k** (Hodosh et al.): |
|
|
| - Marigold: *Repurposing Diffusion-Based Image Generators for Monocular Depth Estimation* — Ke, Obukhov, Metzger, Daudt, Schindler, Schindler (CVPR 2024) |
| - Flickr8k: *Framing Image Description as a Ranking Task* — Hodosh, Young, Hockenmaier (2013) |
|
|
| This model (weights + code) is Apache-2.0; upstream dataset/model terms still apply to their artifacts. |
|
|
| ## Citation |
|
|
| ```bibtex |
| @software{shadenet2, |
| author = {Sachin}, |
| title = {ShadeNet-2: single-image inverse rendering (20M)}, |
| year = {2026}, |
| url = {https://huggingface.co/singam96/ShadeNet-2-20M} |
| } |
| ``` |
|
|
