Skip to content

leodr/hdri-estimation

Repository files navigation

Face2HDRI: HDRI Environment Map Estimation from a Single Portrait

Leonhard Driesch (M.Sc. Informatics, TUM)  ·  Henri Kiesel (M.Sc. Robotics, Cognition, Intelligence, TUM)

Advanced Deep Learning for Computer Vision (ADL4CV) — Prof. Matthias Nießner, TU Munich

📄 Read the paper (paper/main.pdf)

Qualitative results: ground-truth vs. predicted environment maps, probe spheres, and relit portraits

Abstract

Estimating high-dynamic-range illumination from a single portrait is ambiguous because a mostly diffuse face reveals only a coarse view of its surroundings. We present Face2HDRI, a two-stage approach that separates observable lighting from uncertain scene detail. A spherical-query transformer first predicts a coarse environment map from synthetic portraits, obtained by relighting one-light-at-a-time face captures with a bank of HDR environment maps. A conditional flow-matching prior then adds plausible detail and can be trained on the environment maps alone. On held-out test maps, our spherical-query model reduces the mirror-ball scale-invariant RMSE of the closest prior work, SPLiT, from 0.243 to 0.164 indoor and from 0.236 to 0.115 outdoor. The detail prior improves FID from 120 to 50, while scale-invariant RMSE increases from 0.164 to 0.186.

Method

A human face is a natural light probe: its shading, highlights, and shadows encode the surrounding illumination. We exploit this with a two-stage pipeline that factors p(HDRI | portrait) = p(coarse | portrait) · p(detail | coarse).

Two-stage architecture: spherical-query transformer + flow-matching detail prior

  • Stage 1 — deterministic lighting estimator. A ViT-Ti (trained from scratch) and a frozen DINOv3-S tokenize the portrait; an 8×16 grid of spherical-query tokens attends over them through a sphere-query decoder and is seam-wrapped into a coarse log-YCoCg environment map.
  • Stage 2 — flow-matching detail prior. A wrap-convolution U-Net, trained on the HDRI bank alone and conditioned on the blurred coarse map plus a peak map, hallucinates plausible high-frequency detail.

Stage 1 coarse regression vs. Stage 2 with the flow-matching prior vs. ground truth

Training data is fully synthetic: one-light-at-a-time (OLAT) face captures are relit under a bank of HDR environment maps, giving exact portrait–illumination pairs.

Repository layout

paper/              # LaTeX source and compiled PDF of the report (CVPR template)
project_proposal/   # LaTeX project proposal
src/                # Python package
  data/             #   dataset loading, preprocessing, relighting, pre-rendering
  models/           #   ConvEncDec, ClsBottleneck, SphereQueryViT, flow prior
  training/         #   shared trainer and scheduler utilities
  scripts/          #   preprocess / train / sampling / evaluation entry points
data_studio/        # Gradio app: browse the HDRI bank + relighting studio + inference
figures/            # Figure generation scripts and selected outputs
notebooks/          # Exploration notebooks
related_works/      # Literature PDFs and posters
datasets/           # Symlinks to raw and processed local/server data (not committed)
third_party/
  FaceOLAT/         #   vendored capture/relighting code (pinned snapshot)
  SPLiT/            #   closest related work, git submodule (eval protocol + models)

Generated outputs are ignored by Git (runs/, outputs/, logs, Python caches, and .venv/). Dataset paths under datasets/ are symlinks into local/server storage and are not part of the repository.

Pretrained weights

The trained checkpoints are not in Git (they are large binaries); they are attached to the weights-v1 release. Two files are all that is needed for end-to-end inference — the frozen DINOv3-S backbone is fetched automatically by timm on first load:

File Stage Size
stage1_sphere_query.pt Stage 1 — spherical-query transformer 134 MB
stage2_flow_prior.pt Stage 2 — flow-matching detail prior 213 MB
gh release download weights-v1 --repo leodr/hdri-estimation \
  -p stage1_sphere_query.pt -p stage2_flow_prior.pt

Running

The project uses uv. Run everything as a module from the repo root.

# Preprocess: portrait masks, OLAT stacks, HDRI packing, pre-rendered relit portraits
uv run python -m src.scripts.preprocess_dataset

# Train stage 1 (spherical-query transformer is the main model)
uv run python -m src.scripts.train --model sphere-query

# Train the stage-2 flow-matching detail prior on the HDRI bank
uv run python -m src.scripts.train_hdri_prior

# End-to-end sampling: portrait -> coarse map -> detailed HDRI
# (using the pretrained weights from the weights-v1 release, see above)
uv run python -m src.scripts.sample_hdri_prior \
    --stage1 stage1_sphere_query.pt --stage1-kwargs '{"use_dino": true}' \
    --prior stage2_flow_prior.pt --guidance-scale 2.0

# Interactive data + relighting + inference studio (Gradio)
uv run python -m data_studio

Training progress, checkpoints, and TensorBoard logs are written under runs/. The submodule holding the SPLiT baseline is fetched with git clone --recursive (or git submodule update --init); its pretrained model_files/ are downloaded separately (see third_party/README.md).

About

Estimating a High Dynamic Range (HDR) environment map from a single portrait image of a person.

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages