singam96/ShadeNet-2-20M

ShadeNet-2 20M

A lightweight inverse-rendering model: one photo in, albedo + relative depth + surface normals + shading out (8 channels, 20M params). Successor of ShadeNet — smaller, sharper, and with a new self-supervised shading head.

Examples

Each grid: input | albedo | shading / depth | normal | recon (albedo×shading). Click any image for full size.

ShadeNet v1 vs ShadeNet-2

Same inputs, both models — albedo/basecolor, normals and depth side by side (v1 top rows, v2 bottom rows). Each side uses its best ensemble: v1 runs 5 noisy passes + median (its published protocol); v2 runs a 3-pass multi-scale median tuned on the val split (noise was tested and hurts: val/loss 0.2395 vs 0.2421 single-pass).

Architecture

ParallelUNet generator (20.05M params, 18.24M trainable) + 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
• Single-pass RGB → 8ch output (no dual-mode round trip)
• Adversarial training paced 1 G-step : 1 D-step every 2nd window, EMA weight shadow (shipped weights are EMA)

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 (int8 ONNX, CPU-only)
├── inference.py              # Standalone torch CLI
├── inference_utils.py        # Grid visualisation (numpy/PIL)
├── model.py                  # Standalone generator architecture
├── requirements.txt
├── checkpoints/shadenet2.ckpt      # Torch weights, EMA (333MB)
└── onnx/
    ├── model.onnx                  # fp32, EMA (80MB) — GPU via CUDA EP
    └── model_quantized.onnx        # int8 dynamic, EMA (28MB) — CPU

Usage

Gradio Space

Try it in your browser — no installation: singam96/ShadeNet-2-20M Space (runs the int8 model on CPU).

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/shadenet2.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_quantized.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 32 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: albedo/depth/normal), 384px, fp32, ~17h on a single GTX 1650, early-stopped on val/loss (patience 5).

Losses: per-map weighted MSE (depth term is scale-shift-invariant after least-squares alignment — decoded Marigold depth is relative) + self-supervised reconstruction coupling (albedo×shading ≈ input, the shading head's only supervision) + LSGAN + normal unit-length penalty.

Weights	val/loss	albedo MSE	depth SSI	normal MSE	recon L1
last.ckpt	0.2467	0.2912	0.0954	0.2541	0.0690
best.ckpt (raw)	0.2464	0.2879	0.0957	0.2561	0.0688
best.ckpt (EMA)	0.2411	0.2758	0.0925	0.2542	0.0730
best-calibrated (EMA, shipped)	0.2384	0.2726	0.0920	0.2456	0.0787
calibrated + 3-pass ensemble	0.2373	0.2729	0.0917	0.2420	0.0791

Full 807-image val split. Depth is the best-learned map; normals the hardest.

Shipped weights carry a post-hoc Optuna calibration (inference-only, no retraining): per-layer weight affines on the three parameter groups that a sensitivity audit showed could matter (head_out/head/up4 — everything else is GroupNorm-invariant), tuned on val-loss. All page outputs use the calibrated weights with a 3-pass multi-scale median ensemble.

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.
• 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{shadenet2,
  author = {Sachin},
  title = {ShadeNet-2: single-image inverse rendering (20M)},
  year = {2026},
  url = {https://huggingface.co/singam96/ShadeNet-2-20M}
}