File size: 8,706 Bytes
e8d4d21 2d1eb0d e8d4d21 2d1eb0d 516582d 01f6296 516582d 2d1eb0d 516582d 2d1eb0d | 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 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 | ---
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-3.2-5M
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.1671
---
# ShadeNet-3.2 5M
A lightweight inverse-rendering model: one photo in, **albedo + relative depth + surface normals + shading** out (8 channels, 5.0M params). Successor of [ShadeNet-2](https://huggingface.co/singam96/ShadeNet-2-20M) — **4× smaller, better depth and normals**.
[](assets/hero.png)
## TL;DR
| | |
|---|---|
| Params | **5.0M** (3.2M trainable + 1.8M frozen MobileNetV2 trunk) |
| Input | RGB `[1, 3, H, W]` in `[-1, 1]` (384px trained; any multiple of 16) |
| Output | 8ch `[1, 8, H, W]` in `[-1, 1]`: albedo, relative depth, normals, shading |
| Val L1 (807 imgs) | albedo 0.695 · depth 0.217 · normal 0.581 |
| Formats | fp32 ONNX (20MB), fp16 ONNX (10MB), torch checkpoint (44MB) |
## Examples
[](assets/3185645793_49de805194_result.png)
[](assets/241345721_3f3724a7fc_result.png)
[](assets/252578659_9e404b6430_result.png)
[](assets/307994435_592f933a6d_result.png)
[](assets/2312984882_bec7849e09_result.png)
[](assets/3440104178_6871a24e13_result.png)
*Each grid: input | albedo | shading / depth | normal | recon (albedo×shading). Click any image for full size.*
## Results
Full 807-image val split, per-map L1 (the comparable metric across versions —
the headline `val/loss` formula changed between v2 and v3):
| Map (val L1) | ShadeNet-2 (20M) | **ShadeNet-3.2 (5M)** | change |
|---|---|---|---|
| Albedo | 0.708 | **0.695** | −1.7% |
| Depth (SSI-aligned) | 0.247 | **0.217** | **−12%** |
| Normal | 0.696 | **0.581** | **−16%** |
[](assets/training_curves.png)
Checkpoint variants (full val, pruned top-32 dictionary):
| Weights | val/loss | albedo L1 | depth L1 | normal L1 |
|---|---|---|---|---|
| best.ckpt (raw) | **0.1669** | 0.7013 | 0.2241 | **0.5709** |
| best.ckpt (EMA) | 0.1671 | **0.6952** | **0.2174** | 0.5807 |
Shipped ONNX uses the **EMA** weights. All page outputs and the Space run a
**3-pass multi-scale median** (scales 0.875/1.0/1.125) — a mild denoise; the
model's raw single-pass albedo is sharper than its pseudo-labels, so this
trades a little detail for lower variance.
## ShadeNet-2 vs ShadeNet-3.2
Same inputs (ShadeNet-2 top rows, ShadeNet-3.2 bottom rows), maps only. Both
use their shipped weights; ShadeNet-3.2 runs the 3-pass multi-scale median.
[](assets/compare_154871781_ae77696b77.png)
[](assets/compare_159712188_d530dd478c.png)
[](assets/compare_160585932_fa6339f248.png)
## Architecture
**ParallelUNet generator (4.98M params)** + 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
- **Depthwise-separable factorized convs** (1×3 + 3×1) throughout; full H/32 bottleneck; reflect padding
- **Patch-dictionary output tail**: 16×16 tiles softmax-addressed over **32 learned per-channel atoms** (pruned from 1024 — the top-32 hold 99.4% of addressing mass), blended back into the signal before tanh
- Single-pass RGB → 8ch output; EMA weight shadow (shipped weights are EMA)
## Patch dictionary
The tail softmax-addresses 32 learned 16×16 atoms per tile (kept from 1024
after measuring per-atom selection: only ~34 atoms are ever used, ~31 cover
99% of the mass). Shown below per output channel — each panel is the 8×4 atom
grid, shared grayscale scale.
[](assets/atom_dictionary.png)
## 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 (fp16 ONNX)
├── inference.py # Standalone torch CLI
├── inference_utils.py # Grid visualisation (numpy/PIL)
├── model.py # Standalone generator architecture
├── requirements.txt
├── checkpoints/shadenet32.ckpt # Torch weights, EMA (44MB)
└── onnx/
├── model.onnx # fp32, EMA (20MB) — GPU via CUDA EP
└── model_fp16.onnx # fp16, EMA (10MB) — CPU
```
## Usage
### Gradio Space
Try it in your browser — no installation: **[singam96/ShadeNet-3-2-5M Space](https://huggingface.co/spaces/singam96/ShadeNet-3-2-5M)**.
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/shadenet32.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_fp16.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 16 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), 384px, fp32, single GTX 1650, early-stopped on `val/loss` (patience 5) at epoch 12. The 1024-atom dictionary was pruned to its top-32 by addressing mass (no retraining; output error vs full ~5e-4 mean) for the release.
Losses: scale-invariant MSE on albedo (per-channel std alignment) + **scale-shift-invariant** MSE on depth after least-squares alignment (decoded Marigold depth is relative) + **Sobel gradient-matching** on depth (edge crispness) + MSE on normals + self-supervised **reconstruction coupling** (`albedo×shading ≈ input`, the shading head's only supervision) + LSGAN + normal unit-length penalty. Weight decay 1e-4 (patch dictionary exempt).
## 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.
- Occasional localized artifacts in albedo/shading (learned prior pockets).
- No shadows/global illumination — relighting-style use is approximate.
- 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{shadenet32,
author = {Sachin},
title = {ShadeNet-3.2: single-image inverse rendering (5M)},
year = {2026},
url = {https://huggingface.co/singam96/ShadeNet-3.2-5M}
}
```
|