README.md

metadata

license: apache-2.0
base_model: facebook/dinov2-small
pipeline_tag: image-to-image
library_name: pytorch
tags:
  - inverse-rendering
  - intrinsic-decomposition
  - albedo
  - normal-map
  - depth-estimation
  - dinov2

ShadeNet-5

Single-image inverse rendering — one RGB photo in, albedo + relative depth + surface normal out.

Try it in your browser: ShadeNet-5 Space (free CPU, no install; single pass, single pass + refiner, or six passes).

A feed-forward network on a frozen DINOv2 ViT-S/14 backbone. In 1 forward pass it predicts a residual correction to a fixed, photo-derived prior, giving the three maps. No iteration, no time embedding, no diffusion schedule.

The images and the headline numbers on this card use a small refiner (0.91M parameters) applied to the single pass: one forward pass of the network, then a tiny UNet that cleans up the maps, then unit-length normals. It was trained to reproduce the six-pass mirror and rotation average from one pass, and recovers about two thirds of that average's gain for about a fifth of its cost (see The refiner).

Network (traced)

The diagram shows the trainable UNet, its patch dictionaries and the output heads. The frozen DINOv2 branch that feeds the decoder skips is not drawn.

parameters	34,652,431 (12,595,855 trainable)
patch dictionaries	8,582,400 — bottleneck x4, d0 x4, d1 x4, d2 x4
folded fixed patterns	in, u1, u2, u3, u4, u5, u6, d3, d4, out
model	ParallelUNet, one forward pass, residual on a fixed photo prior (mix)
input	3ch photo + 7ch state (448x448 shown)

TL;DR

Task	Single-image inverse rendering (intrinsic decomposition)
Output	7ch — albedo (3) + relative depth (1) + normal (3)
Params	34.7M total (12.6M trainable; 22.1M is the frozen DINOv2 trunk)
Backbone	DINOv2 ViT-S/14, frozen (Meta AI)
Inference	one forward pass plus a small refiner network (0.91M parameters, about 12 ms) and unit-length normals; the plain single pass and the six-pass average are also available
Input size	448x448 (multiples of 448 only: 448 or 896)
Framework	PyTorch; 4/8-bit ONNX exports of the pruned model are described under Quantised ONNX
License	Apache-2.0

What it does

Given a single RGB image it predicts three intrinsic maps: albedo (base colour), relative depth (grayscale, per-image), and a surface normal map. Output channels are [0:3]=albedo, [3:4]=depth, [4:7]=normal, all in [-1, 1].

Breaking change vs ShadeNet-3.2: this model has no shading channel (7ch, not 8ch).

Results

Per-map L1 (raw output vs label; albedo and normal masked where the label is invalid) on the 23-image held-out validation split at 448px, for the best checkpoint (epoch 19, pruned):

map	single pass (unit normals)	single pass + refiner	six passes (mirror + rotations ±4°)
albedo	0.1804	0.1762	0.1734
depth	0.1754	0.1746	0.1653
normal	0.1713	0.1509	0.1497
mean	0.1757	0.1673	0.1628
time per image (A4000, bf16)	57 ms	68 ms	353 ms

The refiner recovers 65% of the gain of the six-pass average (0.0084 of 0.0129 mean L1) for about 12 ms on top of the single pass, against 296 ms for the five extra passes. The gain is almost all in the normals and some in the albedo; the depth barely moves. Single-pass normals here are rescaled to unit length (the raw single-pass normal L1 is 0.1730).

The validation split is small (23 images) and every label is a Marigold-V2 pseudo-label (depth is decoded from Marigold's colour visualisation), so treat these as indicative, not benchmark-grade. The depth loss is scale-and-shift invariant, so raw depth L1 is not what it optimises.

Against the previous run

Same 23 validation images, same labels, same L1 metric, both as a plain single pass. The previous run used the MobileNetV2 trunk, a masked-L1 loss, 384px inputs and trained 6 epochs (best epoch shown):

map	previous run	this model
albedo	0.2474	0.1804
depth	0.3628	0.1754
normal	0.2744	0.1730
mean	0.3141	0.1763

Backbone, loss, resolution and training length all changed together, so the gain cannot be attributed to any single change.

The refiner

The six-pass average is better than a single pass but costs six forward passes. The refiner is a small UNet (0.91M parameters, 4 levels, GroupNorm) that takes the photo and the single-pass maps and adds a learned residual; its last layer starts at zero, so an untrained refiner changes nothing. It is trained by distillation: the target is what the six-pass average produces for the same photo, so it needs no labels and can use any photos.

• Data: 8,024 photos from the Open Images V7 validation split that are listed as CC BY 2.0, none of them from the training set of the main model; 7,724 for training and 300 held out. The author and source page of every photo are in refiner/attribution.csv.
• Training: 6,000 steps, batch 16, random 320x320 crops and mirrors, L1 loss, AdamW 3e-4 with a cosine decay, about 25 minutes on an RTX A4000. The teacher outputs were generated with this card's checkpoint.
• Held-out photos: against the six-pass target the mean L1 fell from 0.0567 (single pass) to 0.0398 (refined).
• Against the labels: see the Results table: normals recover about 94% of the six-pass gain, albedo about 60%, depth about 8%. The refiner does not reproduce the six-pass average's changes to the overall depth levels, so depth is essentially the single pass's.
• Cost: about 12 ms per image on top of the single pass (PyTorch, A4000) and about 0.26 s on a CPU for the 448x448 ONNX version.

Files: refiner/refiner.pt (PyTorch weights), onnx/shadenet5_refiner.onnx (the refiner alone, fp32, fixed 448x448) and onnx/shadenet5_mix_refined.onnx (35 MB: the MIX model, the refiner, the clamp and the unit-length normals in one file, photo in and final maps out; its mean L1 is identical to running the two files one after the other). Run them with python inference/run_onnx.py photos/ --refine (two files) or --merged (the single file).

On the shipped 4/8-bit MIX ONNX model (CPU, same 23 images, unit-length normals): single pass 0.1755, single pass + refiner 0.1676, six passes 0.1645. The refiner was trained on full-precision outputs, so this is the number for the file you would download.

Test-time ensembling

Which extra predictions of the same photo are worth averaging, and whether to take their mean or median? Each row combines several predictions, scored on the same 23 images, and is compared with the plain single pass image by image (the last column is the mean improvement in the average of the three L1s, plus or minus its standard error; positive is better). The normals are raw (not rescaled) in this first table.

ensemble	forwards	albedo	depth	normal	mean	gain
single pass	1	0.1804	0.1754	0.1730	0.1763	—
noise draws x16, mean or median	16	0.1804	0.1754	0.1729	0.1762	+0.0000 ± 0.0000
photo jitter x3, mean	3	0.1791	0.1753	0.1726	0.1756	+0.0006 ± 0.0002
mirror, mean	2	0.1794	0.1704	0.1549	0.1682	+0.0080 ± 0.0018
mirror + photo jitter x3, mean	6	0.1781	0.1703	0.1548	0.1677	+0.0085 ± 0.0018
shifts x5, mean	5	0.1765	0.1967	0.1748	0.1826	−0.0064 ± 0.0023
shifts x5, median	5	0.1768	0.1772	0.1738	0.1759	+0.0003 ± 0.0021
mirror + shifts x5, mean	10	0.1760	0.1941	0.1592	0.1764	−0.0001 ± 0.0031
mirror + shifts x5, median	10	0.1761	0.1755	0.1577	0.1698	+0.0065 ± 0.0025

• Mirror averaging is the clearest win in this table: two forward passes for a 4.6% lower mean L1, mostly in the normals.
• Noise draws change nothing: the noise on the depth input moves the output by about 0.0001, so averaging over them is wasted work.
• Photo-jitter ensembles add little: 0.0005 on top of mirroring, for three times the cost.
• Mean versus median: on clean predictions (mirror, photo jitter) they are within 0.0004 of each other and the mean is slightly better. With shifted copies, which likely carry wrap-around artefacts at the borders, the mean is much worse on depth (0.1967 against 0.1772 for the median at five passes). The median needs at least five passes: at three it was worse than the mean.
• Shifted ensembles improve albedo but hurt depth and normals, so they are not used.

Rotations

A second study rotated the photo by a small angle (the training augmentation covers ±12 degrees), predicted, rotated the maps back with the normals reoriented, and averaged per pixel over the passes that cover it (a rotated canvas loses the corners of the frame). Unit-length normals are applied to every row here, and each row is compared both with the plain pass and with the mirror average.

ensemble	forwards	albedo	depth	normal	mean	gain vs plain	gain vs mirror
single pass	1	0.1804	0.1754	0.1713	0.1757	—	—
mirror, mean	2	0.1794	0.1704	0.1537	0.1679	+0.0078 ± 0.0019	—
rotations ±4°, mean	3	0.1740	0.1688	0.1653	0.1693	+0.0064 ± 0.0011	−0.0015 ± 0.0018
rotations ±4°, median	3	0.1743	0.1703	0.1660	0.1702	+0.0055 ± 0.0010	−0.0024 ± 0.0018
rotations ±12°, mean	3	0.1738	0.1695	0.1650	0.1694	+0.0063 ± 0.0015	−0.0016 ± 0.0018
rotations ±4°, ±8°, ±12°, mean	7	0.1722	0.1674	0.1634	0.1677	+0.0080 ± 0.0016	+0.0002 ± 0.0019
mirror + rotations ±4°, mean	6	0.1734	0.1653	0.1497	0.1628	+0.0129 ± 0.0023	+0.0051 ± 0.0009
mirror + rotations ±4°, median	6	0.1736	0.1657	0.1500	0.1631	+0.0126 ± 0.0024	+0.0047 ± 0.0009
mirror + rotations ±4°, ±8°, mean	10	0.1721	0.1651	0.1491	0.1621	+0.0136 ± 0.0025	+0.0058 ± 0.0012
mirror + rotations ±4°, ±8°, ±12°, mean	14	0.1713	0.1649	0.1492	0.1618	+0.0139 ± 0.0027	+0.0061 ± 0.0014

• Rotation alone works: three passes (the photo and ±4°) improve all three maps and beat the plain pass on 22 of 23 images, about as well as mirroring (the difference between them is within noise).
• Mirror and rotation stack: mirror + ±4° (six passes) is 0.0051 ± 0.0009 better than the mirror average alone, on 21 of 23 images. More angles help a little (ten passes 0.0058, fourteen 0.0061), with diminishing returns, so six passes is the recipe the refiner is trained to reproduce.
• The angle hardly matters: ±4°, ±8° and ±12° give the same result with three passes (0.1693 to 0.1694), so a small tilt gets nearly all of the gain.
• Mean beats median here (0.1628 against 0.1631 with six passes, 0.1693 against 0.1702 with three): the passes are clean, so discarding values only loses information.
• The rotation sign was checked on real labels: rotating the label normals by R(+θ) keeps their correlation with the depth gradient, the opposite sign lowers it.

Gallery

Held-out images (never seen during training — all of them are in this run's validation split), rendered with a single forward pass followed by the small refiner and unit-length normals. Each panel: input | albedo / depth | normal.

snowy mountain + people

person against a wall

motocross on dirt

paraglider at sunset

metro station interior

red building street

dog with a ball

skier on a cliff

Architecture

A ParallelUNet with two parallel encoders: a depthwise-separable conv UNet (full H/32 depth) and a frozen DINOv2 ViT-S/14 (22.1M parameters, facebook/dinov2-small) that sees only the photo. At 448x448 the ViT produces a 32x32 token grid; four of its blocks (2, 5, 8, 11) are tapped, concatenated, 1x1-projected per level and resized to H/2 ... H/32, where they join the UNet's skip connections and bottleneck. The ViT never resizes its input: height and width must be multiples of 14 (the ViT) and 32 (the UNet pyramid), i.e. multiples of 224; the dictionary patch sizes narrow this to multiples of 448 (448 and 896 run, 224 and 672 do not).

Soft patch dictionaries (128 atoms each, with an addressing MLP, applied 4 times with shared weights) remain at 4 stages: bottleneck (patch 7), d0 (patch 4), d1 (patch 8), d2 (patch 8). They hold 8.6M of the 12.6M trainable parameters. The other 10 stages (in, u1, u2, u3, u4, u5, u6, d3, d4, out) were trained with a dictionary too, but each ended up picking the same atom for every tile, so after training they were folded into one fixed pattern per stage (see Pruning). The output is an unbounded residual added to the starting state, through per-map heads (albedo 3ch / depth 1ch / normal 3ch).

Inference is one forward pass: x_0 = P(photo), x_1 = clamp(x_0 + f(photo, x_0 + beta*z)). The starting prior P keeps the photo as albedo and blends 50/50 a simple photo cue (luminance for depth, the photo for normals) with a fixed prior (a vertical ramp for depth, a mean normal). beta*z perturbs only the network's input view of the depth channel (fixed noise field z, amplitude 0.3) and never enters the output. There is no time embedding and no diffusion schedule. Optionally the same pass is repeated on the mirrored photo and on photos rotated by ±4 degrees (reflected border); the maps are rotated and flipped back (normals reoriented, normal x negated for the mirror) and averaged per pixel over the passes that cover it, and the normals are rescaled to unit length. The gallery shows the refiner's output instead (see The refiner).

Pruning

The model was trained with a patch dictionary at all 14 stages. Measured on 150 training images, at 10 of them the addressing picked the same single atom for at least 95% of tiles in every loop pass, i.e. the dictionary was acting as a constant bias. After training, each of those was replaced by that atom, folded with its learned scale and loop count into one [C, p, p] pattern tiled over the feature map. No retraining was done.

	unpruned (epoch 19)	pruned (this card)
mean L1	0.1763	0.1763
model loss	0.2476	0.2476
parameters	50.34M	34.65M
time per image	70.7 ms	44.8 ms
dictionaries	14	4 + 10 fixed patterns

Same 23 validation images and noise seeds for both; time is per image on an RTX A4000, batch 1, bf16, first image excluded. The largest difference at any pixel between the two models' outputs is 0.0195 on the [-1, 1] scale. d0, d1 and bottleneck still choose between many atoms, and folding all 14 stages raised depth L1 from 0.1754 to 0.2258, so they stay dictionaries. d2 splits its tiles about half and half between two atoms; it was also left as a dictionary because folding it adds a 0.0008 depth L1 cost and a 0.0898 maximum pixel difference.

Patch dictionaries

A patch dictionary splits a feature map into non-overlapping tiles, picks a soft mixture of learned atoms for each tile from a small selector network, and adds the mixture back onto the tile. The atoms live in feature space (one [C, p, p] block per atom, C is the stage's channel count), so they are not pictures of anything: each cell below is the channel mean of one atom, blue negative and red positive. Usage is the mean address weight over every tile of every pass (looped sites counted once per loop) on 24 dataset photos, so it includes training images.

• d0: 128 atoms, the busiest takes 34.0% of the address weight, the five busiest 64%, and the average tile puts 56% on its favourite atom
• d1: 128 atoms, the busiest takes 38.7% of the address weight, the five busiest 80%, and the average tile puts 89% on its favourite atom
• d2: 128 atoms, the busiest takes 53.5% of the address weight, the five busiest 100%, and the average tile puts 100% on its favourite atom
• bottleneck: 128 atoms, the busiest takes 19.9% of the address weight, the five busiest 75%, and the average tile puts 84% on its favourite atom

Folded sites

The sites the pruning step replaced. Each is a single fixed pattern (the chosen atom times its learned scale times its loop count), tiled over the feature map: the left cell is the channel mean, then the first eight channels.

Quantised ONNX (4 and 8 bit)

Weight-only block quantisation of the exported ONNX model: each big weight is stored as int4 or int8 integers with one fp16 scale per block (32 to 128 values) in front of a DequantizeLinear, and everything else is fp16. The weights are converted back to floats when the model runs, so this shrinks the file and does not speed anything up (on CPU it is slightly slower). A per-layer sweep found that the dictionary atoms, selectors and convolutions lose almost nothing at 4 bits, and the two MLP layers of every DINOv2 block are the sensitive part, so MIX keeps those at 8 bits.

model	file size	single pass mean L1	enhanced albedo	enhanced depth	enhanced normal	enhanced mean L1
fp32 ONNX (reference)	139.4 MB	0.1764	0.1735	0.1653	0.1497	0.1629
all int4	26.7 MB	0.1900	0.1776	0.1977	0.1546	0.1766
MIX2: ViT fc2 int8, rest int4	27.9 MB	0.1809	0.1756	0.1774	0.1520	0.1683
MIX: ViT fc1/fc2 int8, rest int4	31.4 MB	0.1761	0.1729	0.1701	0.1506	0.1645

L1 on the same 23 held-out images as above; 'enhanced' is the card's recipe (mirror and rotations of 4 degrees, mean, unit normals; 6 forward passes). With so few images, differences below about 0.002 are noise. MIX is +0.0017 from the fp32 model on the mean (depth 0.1701 against 0.1653, the map the quantisation affects most); MIX2 is smaller but loses more depth accuracy, and putting every weight at 4 bits loses more again. The ONNX model takes a 448x448 photo in [-1, 1] (input input_rgb) and returns the 7 maps (maps) for a single pass; the enhanced recipe is the repository's tta_average wrapped around that call (6 passes).

Files in onnx/: shadenet5_mix_int4_int8.onnx = MIX: ViT fc1/fc2 int8, rest int4 (31.4 MB); shadenet5_mix2_int4_int8.onnx = MIX2: ViT fc2 int8, rest int4 (27.9 MB). The fp32 reference and the all-int4 model are in the table for comparison and are not uploaded.

snowy mountain + people: fp32 vs MIX vs MIX2

person against a wall: fp32 vs MIX vs MIX2

motocross on dirt: fp32 vs MIX vs MIX2

paraglider at sunset: fp32 vs MIX vs MIX2

Rows: fp32, MIX, MIX2, all with the enhanced recipe. The left tile of the lower rows is the mean absolute difference to the fp32 maps, amplified 6x.

Training

• The small refiner is trained separately, after this model, on its outputs (see The refiner); the main model's training is unchanged.
• Data: Flickr8k RGB with Marigold-V2 pseudo-labels for albedo, normal and depth. 8,054 training scenes, 23 held out; no label-quality filtering in this run.
• Warm start: weights-only transfer from the previous run's best checkpoint (itself warm-started from ShadeNet-3.2). The UNet and dictionary weights transfer; the DINOv2 projector and the two H/32 atom banks (patch size 7) start fresh.
• Per-map losses on the predicted maps: albedo = masked L1 + multiscale gradient matching; depth = scale-and-shift-invariant L1 + gradient matching + a small raw-L1 anchor (the output is clamped to [-1, 1] and SSI alone cannot tell a map from its inverse); normal = mean angle between predicted and target vectors + a small unit-length penalty. Equal weights across the three maps.
• A small PatchGAN (weight 0.02, albedo and normal maps only) is added; no perceptual loss, no EMA.
• Optimiser: schedule-free AdamW, lr 1e-4, bf16-mixed, batch 4 x gradient accumulation 2. One RTX A4000, about 28 minutes per epoch.
• Augmentation: zoom-in crop, smooth warp, in-plane roll (+-12 degrees, normals reoriented — an earlier version of the pipeline rotated them the wrong way and was fixed before this run), horizontal mirror, photometric jitter on the input photo only, and 35% of training photos replaced by a relit render (albedo x shading from the normals). No vertical mirror.
• Length: 19 epochs at 448px, stopped by hand once validation loss was improving by about 0.4% per epoch; the reported checkpoint is the one with the best mean per-map L1, pruned after training (see Pruning).

Usage

The supported way to run the model is the ONNX export (onnx/shadenet5_mix_int4_int8.onnx, see Quantised ONNX) with the scripts in the inference/ folder of this repository. They need no training code:

pip install onnxruntime numpy pillow torch huggingface_hub
python inference/run_onnx.py photos/ --refine --out out/

Each photo becomes a panel image: input, albedo, depth, normal. --refine runs the single pass and then the small refiner (about 12 ms more; the recommended mode, see The refiner). --enhanced is the six-pass post-processing described above in one switch (mirror average, rotations of ±4 degrees, unit-length normals; six forward passes); leave both flags out for the plain single pass. The whole photo is kept: its long side is scaled to 448 and the rest of the 448x448 input is filled by mirroring the photo's own edge, then the padding is cut off the maps again, so the outputs have the photo's own shape. A photo that is not square therefore uses fewer than 448 pixels on its short side. (The PyTorch model runs at 448 and 896 square; the ONNX files are fixed at 448x448.) torch is only used for the rotation resampling of --enhanced. Set --threads to the number of CPU cores you have: more threads than cores is slower, not faster. A hosted demo runs the same code: ShadeNet-5 Space.

Or call the single merged file directly:

import numpy as np, onnxruntime as ort
sess = ort.InferenceSession("shadenet5_mix_refined.onnx")
# photo: float32 [1, 3, 448, 448] in [-1, 1]
maps = sess.run(None, {"input_rgb": photo})[0]   # [1, 7, 448, 448]
albedo, depth, normal = maps[:, 0:3], maps[:, 3:4], maps[:, 4:7]

The PyTorch checkpoint (checkpoints/best-l1-pruned10-ep19.ckpt) is the full-precision reference. Loading it needs the training code, which is not published yet (the --enhanced switch of its infer.py is the same recipe), so use the ONNX file unless you have that code.

config.json in this repository lists the architecture, the input and output convention (7 maps in [-1, 1]: albedo, depth, normal), the supported sizes and the recommended inference settings.

Limitations

• Input sizes: multiples of 448 only (448 is what it was trained at). Other sizes must be resized first.

License & attribution

Code and weights: Apache-2.0.

• Backbone: DINOv2 ViT-S/14 (facebook/dinov2-small, Meta AI, Apache-2.0), kept frozen.
• Labels: the albedo, normal and depth training labels are pseudo-labels generated with Marigold V2 (huawei-bayerlab/marigold-v2-0, Apache-2.0; its base model Qwen-Image-Edit-2509 is Apache-2.0 too). They are model outputs, not ground truth; please cite the Marigold V2 paper if you build on them.
• Refiner training photos: Open Images V7 (validation split), images listed as CC BY 2.0, annotations CC BY 4.0 (Google LLC); Google makes no guarantee about each image's licence. The photos are not redistributed; the author, source page and licence of every photo used are in refiner/attribution.csv. Credit to those photographers.
• Images: the photographs come from the Flickr8k corpus (Hodosh, Young and Hockenmaier, 2013; Rashtchian et al., 2010). The corpus authors release the captions under CC BY-SA, but the photographs are individual Flickr images owned by their photographers and have no blanket licence. The example and gallery pictures on this card are held-out Flickr8k photos shown for illustration only; check the original Flickr page of an image before reusing it.