---
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}
}
```