SfPUEL: Shape from Polarization under Unknown Environment Light, NeurIPS 2024

Youwei Lyu¹, Heng Guo¹, Kailong Zhang¹, Si Li¹, and Boxin Shi^2
1Beijing University of Posts and Telecommunications, ²Peking University

This paper proposes a learning-based method to address the problem of shape from polarization (SfP) under unknown environment light.

SfPUEL is implemented with PyTorch and tested on Ubuntu 22.04.

• Python 3.10
• cv2, tqdm, pandas, einops
• PyTorch 2.0.1+cu118

conda create -n sfpuel python=3.10 -y
conda activate sfpuel

pip install -r requirements.txt
pip install --no-cache-dir torch==2.0.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
git clone https://github.com/YouweiLyu/SfPUEL.git

The real-world test dataset and synthetic test samples can be downloaded from GoogleDrive.

The test data is structured as follows. You may arrange your dataset in the same format for test.

test_data_dir/
├── pol000/       # polarization images at 0 degree
├── pol045/       # polarization images at 45 degree
├── pol090/       # polarization images at 90 degree
├── pol135/       # polarization images at 135 degree
├── normal/       # ground-truth normal maps
├── mask/         # mask

First download the trained model from GoogleDrive and unzip it in the root directory of this repository. The checkpoint is unzipped to data/checkpoints/ckpt.pth.

• Test SfPUEL without evaluation by running:

  test_data_dir="SfPUEL_test_data/real1"
  save_folder="sfpuel_real1"
  python tools/test.py --data_dir_tst $test_data_dir --suffix $save_folder

  The estimated normal maps are stored in results/$save_folder.

• Test SfPUEL and evaluate the normal predictions by running:

  test_data_dir="SfPUEL_test_data/real1"
  save_folder="sfpuel_real1"
  python tools/test.py --data_dir_tst $test_data_dir --suffix $save_folder &&\
      python tools/normal_criterion.py -d results/$save_folder --gt_data_dir $test_data_dir 
  

  The quantitative results can be found in results/${save_folder}_normal_eval.xlsx.

  Note: Due to random sampling operations in the implementation, quantitative results might slightly vary with each run.

The training dataset is provided on HuggingFace, which consists of 20K pairs of polarization images and corresponding normal maps. The training data is structured as follows.

training_data_dir/
├── pol000/       # polarization images at 0 degree
├── pol045/       # polarization images at 45 degree
├── pol090/       # polarization images at 90 degree
├── pol135/       # polarization images at 135 degree
├── normal/       # ground-truth normal maps
├── mask/         # masks
├── material_tag/         # material annotations

If you find this code useful in your research, please consider citing:

@inproceedings{lyu2024sfpuel,
    title={{SfPUEL}: Shape from Polarization under Unknown Environment Light},
    author={Youwei Lyu and Heng Guo and Kailong Zhang and Si Li and Boxin Shi},
    booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems (NeurIPS)},
    year={2024},
}