singam96/ShadeNet-3.2-5M

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 — 4× smaller, better depth and normals.

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

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%

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.

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.

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.

This repo ships app.py, the Space entrypoint. To recreate it: New Space → Gradio SDK → point at this repo.

Torch CLI

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)

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

@software{shadenet32,
  author = {Sachin},
  title = {ShadeNet-3.2: single-image inverse rendering (5M)},
  year = {2026},
  url = {https://huggingface.co/singam96/ShadeNet-3.2-5M}
}