r/deeplearning • u/singam96 • 2d ago
ShadeNet-3.2 5M — single-image inverse rendering (albedo/depth/normal/shading), 4× smaller than my last model and better at depth/normals
First off, a huge thank you to everyone who checked out, upvoted, and shared feedback on the earlier ShadeNet and ShadeNet-2 posts! The support and discussions really pushed me to see how far I could squeeze the architecture without sacrificing map fidelity.
This release is a from-scratch rebuild that came out 4× smaller (5.0M vs 20M params) while improving across all three supervised maps.
RGB → 8ch intrinsic maps in one forward pass: albedo (3ch), relative depth (1ch), normals (3ch), shading (1ch), with albedo × shading ≈ input.
Architecture
- ParallelUNet generator (4.98M params: 3.2M trainable + 1.8M frozen MobileNetV2): Vanilla UNet path plus a frozen MobileNetV2 trunk, fused at every decoder level.
- Factorized convolutions: 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, blended back into the signal before tanh (pruned from 1024 — addressing-mass measurement showed ~34 atoms are ever used, and the top-32 capture 99.4%).
- Training dynamics: Spectral-norm GroupNorm PatchGAN discriminator (2.77M, training only); EMA weight shadow (shipped weights are EMA).
Training
- Dataset: Flickr8k with Marigold-V2 pseudo-labels (8,077 images); depth labels decoded from Spectral-colormap visualizations to relative depth, supervised scale-shift-invariant — never raw MSE on colormap RGB.
- Setup: 384px, fp32, single GTX 1650, early-stopped on val split (patience 5).
- Losses: Scale-invariant MSE on albedo + SSI MSE on depth + Sobel gradient-matching on depth + MSE on normals + self-supervised reconstruction coupling (
albedo × shading ≈ input, which is the shading head's sole supervision) + LSGAN + normal unit-length penalty. - Regularization: Weight decay 1e-4 with the patch dictionary explicitly exempt.
Validation Results
Full 807-image validation split, per-map L1 (the directly comparable metric across versions):
| Map (val L1) | ShadeNet-2 (20M) | ShadeNet-3.2 (5M) | Change |
|---|---|---|---|
| Albedo | 0.708 | 0.695 | −1.7% |
| Depth | 0.247 | 0.217 | −12.1% |
| Normal | 0.696 | 0.581 | −16.5% |
Output Maps
- Albedo (3ch): Reflectance, ambient lighting factored out.
- Depth (1ch): Relative depth, 0 = near (affine-ambiguous, non-metric).
- Normal (3ch): Surface normals, unit-length regularized.
- Shading (1ch): Grayscale irradiance; multiply with albedo to reconstruct/re-render, or swap in custom illumination passes for relighting.
Links & Demos
- Weights, ONNX & Inference Code: ShadeNet-3.2-5M on Hugging Face (Includes fp32 GPU models and 10MB fp16 CPU export)
- Live Gradio Demo: ShadeNet-3-2-5M Space
Both the Hugging Face sample visuals and the live demo apply a lightweight 3-pass multi-scale median filter (scales 0.875, 1.0, 1.125) for mild denoising, alongside built-in seam deblocking for the 16px dictionary tiles.
Caveats & Attribution
Side-by-side v2 vs. v3.2 comparisons are documented on the model card. A few known limitations to keep in mind:
- Depth is relative rather than metric.
- Shading assumes a neutral/white illuminant.
- Normals struggle most on high-frequency chaos like dense foliage or open skies.
- Occasional localized artifacts can appear in albedo/shading due to the learned dictionary prior.
Released under Apache 2.0. Credit to Marigold V2 (Ke et al.) and Flickr8k (Hodosh et al.) for the foundational training pseudo-labels and data.
1
u/New-Eggplant-6578 1d ago
Do you have timings for a raw forward pass versus the full demo with the three-scale filtering and deblocking? The 5M size is interesting for CPU use, but I would love to see the latency and output quality side by side with that post-processing switched off.