Snapshot Polarimetric Display Inverse Rendering

An LCD displays a linearly polarized RGB binary pattern; an RGB polarization camera with a quarter-wave plate in front of the lens records the scene in a single exposure. That one capture is decomposed into nine measurements — three spectral channels × three polarization states (unpolarized, linearly polarized, circularly polarized) — and a feed-forward transformer turns them into per-pixel normal, albedo, roughness and metallicity.

This repository holds the training and inference code. The project page lives on the project_page branch.

git clone https://github.com/MichaelCSJ/PDIR.git
cd PDIR
conda create -n pdir python=3.10 -y && conda activate pdir
pip install -r requirements.txt

The code was developed against PyTorch 2.4 / CUDA 12.4 on NVIDIA A6000 GPUs.

Pretrained weights and five sample captures live on Hugging Face, because the checkpoint is 329 MB:

python download_assets.py

That fills two gitignored directories:

checkpoints/pdir_best.ckpt
sample_data/{cat,bowl,case,foil,owl}/
    mask.png
    quad{0,1,2,3}_main_hdr_{s0,s1,s2}.npy

The sample directories hold exactly what the network reads: the four per-quadrant Stokes captures and the object mask. Then:

python inference.py --ckpt checkpoints/pdir_best.ckpt \
    --dataset real --data-root sample_data --out-dir results \
    --exposure-norm p95_0.95

python render_relight.py --maps-dir results --out-dir relight

--exposure-norm p95_0.95 matters: it is what pdir_best was trained with, and inference has to match (see Inference).

Two loaders are provided. Both hand the network raw per-quadrant Stokes components; the pattern weighting and the unpolarized/LP/CP decomposition happen inside the model, so the same checkpoint reads synthetic and real captures.

Synthetic (--dataset synth), one directory per scene:

<synth-root>/
  0000/
    obs_original.npz    diffuse_{1_0,0_0,1_1,0_1}, specular_{...}, cop_{...}   each [H, W, 3]
    gt_maps.npz         basecolor [H,W,3], roughness [H,W,1], metallic [H,W,1],
                        normal [H,W,3], mask [H,W]
  0001/
  ...

Real captures (--dataset real), one directory per scene. Any directory holding these files counts as a scene, so readable names work as well as the capture ids:

<real-root>/
  cat/
    quad{0,1,2,3}_main_hdr_s0.npy    [H, W, 3]  total intensity, >= 0
    quad{0,1,2,3}_main_hdr_s1.npy    [H, W, 3]  linear polarization, signed
    quad{0,1,2,3}_main_hdr_s2.npy    [H, W, 3]  circular polarization, signed
    mask.png                         uint8, 0/255
  bowl/
  ...

Without split files, every --val-stride-th scene goes to validation. Pass --synth-split-train / --synth-split-val (and the --real- equivalents) to control the split explicitly; each line is one scene directory name.

Supervised training on synthetic scenes:

python train.py --dataset synth \
    --synth-root /path/to/synth_scenes \
    --session pdir_synth --devices 4

Joint synthetic + real co-training, which is what the paper uses. The real branch contributes a render loss only — it needs no ground-truth maps:

python train.py --dataset joint \
    --synth-root /path/to/synth_scenes \
    --real-root  /path/to/real_captures \
    --session pdir_joint --devices 4

Training writes to outputs/<timestamp>_<session>/: checkpoints/ (top-3 by validation loss, plus last.ckpt) and logs/ (CSV scalars).

Things worth knowing before a long run:

• --image-size must be exactly twice --canonical-resolution. Defaults are 384 / 192.
• Unless you pass --no-pretrained, training initialises from the LINO-PBR weights, downloaded once into checkpoints/. Use --init-ckpt to start from your own checkpoint instead.
• --pattern-color-strength defaults to 1.0, which is what the released pdir_best was trained with. 0.80 is the value calibrated against the real captures; whichever you train with, pass the same one at inference time.
• The real loss ramps in as a staircase: zero for the first --real-warmup-stair-steps steps, then one step up per block until it reaches full weight after --real-warmup-stairs blocks. While the weight is zero the real forward pass is skipped entirely, which keeps early training stable.
• Joint mode runs under manual optimization, so --accumulate-grad-batches must stay at 1.
• Multi-GPU uses DDP with bf16 mixed precision.

python train.py --help lists everything else.

python inference.py --ckpt outputs/<run>/checkpoints/last.ckpt \
    --dataset real --data-root /path/to/real_captures \
    --out-dir results

For every scene this writes <out-dir>/<scene>/:

Use --dataset synth to run the same checkpoint over synthetic scenes. Images are processed in one tile by default; --patch-size splits them instead, and must divide --image-size.

Two flags must match the values the checkpoint was trained with, or the encoder sees a different input distribution than it learned on:

The estimated maps are directly renderable. render_relight.py orbits a point light around a scene and renders one frame per step through the same principled BRDF the training render loss uses:

python render_relight.py --maps-dir results --out-dir relight

That writes relight/<scene>/frame####.png and relight/<scene>.mp4 for every scene under --maps-dir. Useful options:

--brightness is a single global exposure, so scenes with dark albedo or high metallicity come out dimmer than others; raise it per scene if a render looks too dark. H.264 video needs imageio-ffmpeg (see requirements.txt); without it the script falls back to OpenCV's MPEG-4 encoder, which most browsers cannot play.

download_assets.py          fetches the weights and sample captures
upload_assets.py            publishes them (maintainers only)
train.py                    training entrypoint
inference.py                inference entrypoint
render_relight.py           rotating-light relighting of the estimated maps
calibration/                calibrated LCD emitter positions used by the renderer
src/io/                     synthetic and real dataset loaders
src/lightning/module.py     LightningModule: losses, joint co-training, optimizer
src/model/net_train_module.py   the network: pattern simulation, encoder, heads, losses
src/model/module/utils.py       image encoder, GLC aggregation, regressor
src/model/aggregator/           DINOv2 backbone and the alternating-attention aggregator
src/model/renderer/             principled BRDF used by the render loss

@article{choi2026snapshot,
  title     = {Snapshot Polarimetric Display Inverse Rendering},
  author    = {Choi, Seokjun and Moon, Yunseong and Kang, Kaizhang and Chung, Hoon-Gyu
               and Kim, Jin-Nyeong and Nam, Giljoo and Baek, Seung-Hwan},
  journal   = {ACM Transactions on Graphics},
  volume    = {45},
  number    = {6},
  articleno = {201},
  year      = {2026},
  doi       = {10.1145/3842531}
}

The transformer backbone builds on DINOv2 and LINO-UniPS.