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

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):

Checkpoint variants (full val, pruned top-32 dictionary):

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

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