FlashNormal: Detailed Surface Normal Estimation from Flash and No-Flash Images

Ruiyang Chen^1,+, Feiran Li^2,+, Heng Guo^1,*, Zhanyu Ma¹

¹Beijing University of Posts and Telecommunications   ²Independent Researcher

⁺Equal contribution;  ^*Corresponding author

Official implementation of FlashNormal, a diffusion-based method for estimating detailed surface normals from a registered flash/no-flash image pair.

• Installation
• Download Models and Data
• Data Format
• Inference
• Training
• Citation
• Acknowledgements
• License

The code has been tested with Python 3.8, PyTorch 2.0.1, CUDA, and a Linux environment. An NVIDIA GPU is required by the current implementation.

git clone https://github.com/richardchen225/FlashNormal.git
cd FlashNormal

conda create -n flashnormal python=3.8 -y
conda activate flashnormal

pip install -r requirements.txt

Download the pretrained models, training data, evaluation benchmark, and Metric3D checkpoint:

mkdir -p checkpoints data

# Diffusion backbone used to initialize FlashNormal.
huggingface-cli download prs-eth/marigold-normals-v0-1 \
  --local-dir checkpoints/marigold-normals

# Trained FlashNormal EMA weights (model repository).
huggingface-cli download richardchencccc/FlashNormal \
  --local-dir checkpoints/flashnormal

# Flash100K training data.
huggingface-cli download richardchencccc/Flash100K \
  --repo-type dataset \
  --local-dir data/Flash100K

# Real-world evaluation benchmark.
huggingface-cli download richardchencccc/EvalFlash \
  --repo-type dataset \
  --local-dir data/EvalFlash

# Synthetic evaluation benchmark.
huggingface-cli download richardchencccc/EvalFlash-synth \
  --repo-type dataset \
  --local-dir data/EvalFlash-synth

# Metric3D ViT-Large checkpoint used to create missing coarse normal priors.
# The Metric3D source tree is included at ./Metric3D in this repository.
mkdir -p Metric3D/weights
huggingface-cli download JUGGHM/Metric3D metric_depth_vit_large_800k.pth \
  --local-dir Metric3D/weights

If a repository is gated, authenticate first with huggingface-cli login. After downloading, extract any archives while preserving the directory structures described below.

The expected model layout is:

checkpoints/
|-- marigold-normals/
|   |-- scheduler/
|   |-- text_encoder/
|   |-- tokenizer/
|   |-- unet/
|   `-- vae/
`-- flashnormal/
    `-- unet_ema/

The released real and synthetic benchmarks are hosted separately:

• EvalFlash stores real-world objects under real_benchmark/.
• EvalFlash-synth stores synthetic objects under synth_bench/.

infer/infer.py processes every object directory directly under --test_dir. The directory name and the first three file prefixes must match after conversion to lowercase. The real benchmark downloaded above has the following layout:

data/EvalFlash/
`-- real_benchmark/
    |-- Anderson/
    |   |-- anderson.png          # no-flash RGB image
    |   |-- anderson_flash.png    # registered flash RGB image
    |   |-- anderson_mask.png     # foreground mask
    |   |-- normals_m3d.png       # coarse normal prior
    |   |-- anderson_gt.png       # optional, used only for evaluation
    |   `-- ...
    `-- Egyptcat/
        |-- egyptcat.png
        |-- egyptcat_flash.png
        |-- egyptcat_mask.png
        `-- normals_m3d.png

The synthetic benchmark follows the same object-per-directory convention under data/EvalFlash-synth/synth_bench/.

Images within an object directory must be spatially registered. The current loader uses the foreground mask to crop the object, resizes the crop to 512 x 512, and uses normals_m3d.png as the coarse normal input.

normals_m3d.png may be omitted when GENERATE_COARSE_NORMALS=1; it will be generated from the no-flash image before inference.

The Flash100K training repository stores ten shards named flash_no_flash_0_png through flash_no_flash_9_png under train/:

data/Flash100K/
`-- train/
    |-- flash_no_flash_0_png/
    |   `-- object_000001/
    |       |-- 001.png           # no-flash RGB image
    |       |-- 002.png           # registered flash RGB image
    |       |-- normals.png       # 16-bit ground-truth normal map
    |       |-- normals_m3d.png   # coarse normal prior (optional with generation enabled)
    |       `-- mask.png          # foreground mask
    |-- flash_no_flash_1_png/
    |   `-- ...
    `-- flash_no_flash_9_png/
        `-- ...

Incomplete object directories are skipped. They are deleted only when training is launched with REMOVE_INCOMPLETE_SAMPLES=1; leave this disabled unless the source dataset is recoverable.

If the training or evaluation data does not include normals_m3d.png, FlashNormal can generate the missing coarse priors with the included Metric3D checkout. Download the ViT-Large checkpoint above first, then:

GENERATE_COARSE_NORMALS=1 \
METRIC3D_MODEL_PATH=/path/to/Metric3D \
METRIC3D_WEIGHTS_PATH=/path/to/Metric3D/weights/metric_depth_vit_large_800k.pth \
bash training/train.sh

Metric3D is loaded lazily and reused for all missing samples. Existing normals_m3d.png files are never overwritten. In distributed training, the main process generates missing files first and the remaining processes read the completed data.

The generator follows Metric3D's hub example without flipping normal axes or converting its output channels before cv2.imwrite. The FlashNormal data loader then applies its existing Y/Z sign conversion exactly once, for both existing and newly generated normals_m3d.png files. Do not manually flip or reorder these files.

The same behavior is available during evaluation:

GENERATE_COARSE_NORMALS=1 \
METRIC3D_MODEL_PATH=./Metric3D \
METRIC3D_WEIGHTS_PATH=./Metric3D/weights/metric_depth_vit_large_800k.pth \
bash infer/infer.sh

The inference shell script contains all paths and inference settings in one place. Its defaults match the directory layout above:

bash infer/infer.sh

By default, this evaluates data/EvalFlash/real_benchmark. To run the synthetic benchmark instead:

TEST_DIR=./data/EvalFlash-synth/synth_bench bash infer/infer.sh

Every value can be overridden without editing the script. For example:

TEST_DIR=/path/to/custom-data \
CHECKPOINT_PATH=/path/to/checkpoint-1000 \
PROCESSING_RES=768 \
ENSEMBLE_SIZE=5 \
ENABLE_XFORMERS=0 \
bash infer/infer.sh

One predicted normal map is written for each input object as:

outputs/evalflash/<object-name>_val.png

When <object-name>_gt.png and <object-name>_mask.png are available, inference also computes angular MAE, normal-vector MSE, and RMSE. Per-object values and macro averages are written to outputs/evalflash/metrics.json. Objects without ground truth are still inferred and saved, but are marked as skipped in the metrics report.

To run on custom data, arrange it in the same object-per-directory format and set TEST_DIR. Inference always loads the released EMA weights from unet_ema/.

The main inference controls exposed by infer/infer.sh are PROCESSING_RES, CROP_MARGIN, TIMESTEP, PROMPT, DOMAIN, ENSEMBLE_SIZE, DENOISING_STEPS, OUTPUT_NAME_SUFFIX, METRICS_FILENAME, and ENABLE_XFORMERS.

Training initializes the VAE, tokenizer, text encoder, scheduler, and UNet from Marigold, expands the UNet input to 12 channels, and fine-tunes it on the flash/no-flash training set. Validation uses one object directory, not the benchmark root. Batch size and spatial resolution are derived dynamically by the loss implementation.

ACCELERATE_CONFIG=./node_config/1gpu.yaml bash training/train.sh

training/train.sh uses node_config/4gpu.yaml by default. Select another supplied configuration as follows:

# Four GPUs (default)
bash training/train.sh

# Eight GPUs
ACCELERATE_CONFIG=./node_config/8gpu.yaml bash training/train.sh

Paths and hyperparameters can also be overridden through environment variables:

TRAINDATA_DIR=/path/to/Flash100K/train \
TEST_DIR=/path/to/EvalFlash/real_benchmark/Anderson \
TRAIN_BATCH_SIZE=1 \
LEARNING_RATE=1e-5 \
PROCESSING_RES=512 \
CURVATURE_LOSS_WEIGHT=10 \
DETAIL_THRESHOLD=0.002 \
bash training/train.sh

The script exposes optimization, data loading, preprocessing, loss, validation, EMA, xFormers, mixed precision, and checkpoint settings at the top of the file. Boolean settings use 0 or 1.

The effective batch size is:

train_batch_size x number_of_GPUs x gradient_accumulation_steps

Checkpoints are saved to outputs/train/checkpoint-<step>/. Each checkpoint contains Accelerate state together with unet/ and, when --use_ema is enabled, unet_ema/. These checkpoint directories can be passed directly to inference through --checkpoint_path.

To resume training from the newest checkpoint:

RESUME_FROM_CHECKPOINT=latest bash training/train.sh

To resume from a particular checkpoint, use its directory name, for example RESUME_FROM_CHECKPOINT=checkpoint-1000 bash training/train.sh.

If you find this work useful, please cite:

@article{chen2026flashnormal,
  title={{FlashNormal}: Detailed Surface Normal Estimation from Flash and No-Flash Images},
  author={Chen, Ruiyang and Li, Feiran and Guo, Heng and Ma, Zhanyu},
  journal={IEEE Transactions on Circuits and Systems for Video Technology},
  year={2026}
}

This project builds on Marigold, Diffusers, and Accelerate. We thank the authors and maintainers of these projects for releasing their work.

The original FlashNormal code is released under the Apache License 2.0. Third-party components remain subject to their respective licenses and copyright notices. In particular, the vendored Metric3D/ source is distributed under its included BSD 2-Clause License, while the Hugging Face-derived model components retain their Apache 2.0 notices.