ShadeNet-3.2-5M (model): weights, ONNX, app + card

01f6296 verified about 1 month ago

8.71 kB

---
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**.

[![hero](assets/hero.png)](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

[![3185645793](assets/3185645793_49de805194_result.png)](assets/3185645793_49de805194_result.png)
[![241345721](assets/241345721_3f3724a7fc_result.png)](assets/241345721_3f3724a7fc_result.png)
[![252578659](assets/252578659_9e404b6430_result.png)](assets/252578659_9e404b6430_result.png)
[![307994435](assets/307994435_592f933a6d_result.png)](assets/307994435_592f933a6d_result.png)
[![2312984882](assets/2312984882_bec7849e09_result.png)](assets/2312984882_bec7849e09_result.png)
[![3440104178](assets/3440104178_6871a24e13_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%** |

[![curves](assets/training_curves.png)](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.

[![compare 1](assets/compare_154871781_ae77696b77.png)](assets/compare_154871781_ae77696b77.png)
[![compare 2](assets/compare_159712188_d530dd478c.png)](assets/compare_159712188_d530dd478c.png)
[![compare 3](assets/compare_160585932_fa6339f248.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.

[![atom dictionary](assets/atom_dictionary.png)](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}
}
```

	---
	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.

	[![hero](assets/hero.png)](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

	[![3185645793](assets/3185645793_49de805194_result.png)](assets/3185645793_49de805194_result.png)
	[![241345721](assets/241345721_3f3724a7fc_result.png)](assets/241345721_3f3724a7fc_result.png)
	[![252578659](assets/252578659_9e404b6430_result.png)](assets/252578659_9e404b6430_result.png)
	[![307994435](assets/307994435_592f933a6d_result.png)](assets/307994435_592f933a6d_result.png)
	[![2312984882](assets/2312984882_bec7849e09_result.png)](assets/2312984882_bec7849e09_result.png)
	[![3440104178](assets/3440104178_6871a24e13_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% \|

	[![curves](assets/training_curves.png)](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.

	[![compare 1](assets/compare_154871781_ae77696b77.png)](assets/compare_154871781_ae77696b77.png)
	[![compare 2](assets/compare_159712188_d530dd478c.png)](assets/compare_159712188_d530dd478c.png)
	[![compare 3](assets/compare_160585932_fa6339f248.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.

	[![atom dictionary](assets/atom_dictionary.png)](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}
	}
	```