File size: 7,706 Bytes
21e2ec3 2f93ecf 21e2ec3 b166b71 21e2ec3 feadaf9 21e2ec3 feadaf9 21e2ec3 360aac9 21e2ec3 360aac9 21e2ec3 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | ---
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}
}
```
|