DuRP

Official PyTorch implementation for DuRP, a polarization-based image reconstruction framework. This repository provides dataset synthesis utilities, staged training code, inference scripts, and evaluation tools for reconstructing polarization properties.

Highlights

• End-to-end PyTorch implementation of the DuRP reconstruction pipeline.
• Three training stages:
  • PRNet: polarization reconstruction stage.
  • IRNet: image/radiance reconstruction stage.
  • DuRP: full-stage model that combines PRNet and IRNet.
• Synthetic dataset generation from clean full-polarization images and depth maps.
• Inference with FP32, FP16, and optional torch.compile acceleration.
• Evaluation scripts for S0/R, DoLP, and AoLP using PSNR and SSIM.

Repository Structure

DuRP/
|-- base/                 # Base trainer/checkpoint utilities
|-- dataset/              # Dataset loaders for train/val/test
|-- model/                # PRNet, IRNet, and full DuRP model definitions
|-- scheduler/            # Warmup and learning-rate schedulers
|-- scripts/
|   |-- make_dataset.py   # Synthetic data generation
|   |-- test_speed.py     # Inference, result saving, latency/FPS reporting
|   `-- calc_metric.py    # PSNR/SSIM evaluation
|-- train/                # Training entry point, options, losses
|-- trainer/              # Stage-specific trainer implementations
|-- utils/                # Logging, metrics, checkpoint, visualization helpers
|-- train_img_reconstruct.sh
|-- train_pol_reconstruct.sh
|-- train_full_stage.sh
`-- test.sh

Installation

The code was developed for CUDA-enabled PyTorch. A minimal environment can be created as follows:

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

# Install PyTorch following your CUDA version:
# https://pytorch.org/get-started/locally/
pip install torch torchvision

pip install numpy opencv-python tqdm timm tensorboardX albumentations scikit-image polanalyser einops

Optional but recommended:

pip install tensorboard

Data Preparation

DuRP expects generated .npy data with the following layout:

datas/SRPG/
|-- train/
|   |-- input/
|   |   |-- I_alpha/      # Four polarization-angle RGB
|   |   `-- m_I/          # Complex observed polarization 
|   `-- target/
|       |-- R/            # Ground-truth scene radiance / S0 image
|       |-- m_A/          # Airlight polarization representation
|       |-- A_inf/        # Atmospheric light
|       |-- m_D/          # Direct-transmission polarization 
|       `-- K/            # Transmission-related target
`-- test/
    `-- ...

To synthesize this structure from raw clean polarization data and depth maps, prepare:

raw_data/
|-- train/
|   |-- full_pol/*.npy
|   `-- depth/*.npy
`-- test/
    |-- full_pol/*.npy
    `-- depth/*.npy

Then run:

cd DuRP

python scripts/make_dataset.py \
  --input_dir raw_data \
  --output_dir datas/SRPG \
  --mode train \
  --enlarge_factor 3

python scripts/make_dataset.py \
  --input_dir raw_data \
  --output_dir datas/SRPG \
  --mode test \
  --enlarge_factor 1

Training

All training commands should be run from the DuRP/ directory.

Stage 1: Polarization Reconstruction

python train/train.py \
  --description polarization_reconstruct \
  --arch PRNet \
  --stage pol \
  --loss pol_loss \
  --train_dir ./datas/SRPG/train \
  --val_dir ./datas/SRPG/test \
  --batch_size 1 \
  --train_ps 512 \
  --val_ps 512 \
  --nepoch 400 \
  --lr_initial 1e-4 \
  --weight_decay 1e-4 \
  --gpu_device 0 \
  --save_dir ./experiment/

Checkpoints are saved to:

experiment/PRNet_0/checkpoints/

Stage 2: Image Reconstruction

python train/train.py \
  --description image_reconstruct \
  --arch IRNet \
  --stage img \
  --loss img_loss \
  --train_dir ./datas/SRPG/train \
  --val_dir ./datas/SRPG/test \
  --batch_size 4 \
  --train_ps 512 \
  --val_ps 512 \
  --nepoch 400 \
  --lr_initial 1e-4 \
  --weight_decay 1e-4 \
  --gpu_device 0 \
  --save_dir ./experiment/

Checkpoints are saved to:

experiment/IRNet_0/checkpoints/

Stage 3: Full DuRP Fine-Tuning

The full-stage trainer initializes DuRP from pretrained PRNet and IRNet weights:

python train/train.py \
  --description full_stage \
  --arch DuRP \
  --stage full \
  --loss full_loss \
  --train_dir ./datas/SRPG/train \
  --val_dir ./datas/SRPG/test \
  --PRNet_weights ./experiment/PRNet_0/checkpoints/model_best.pth \
  --IRNet_weights ./experiment/IRNet_0/checkpoints/model_best.pth \
  --batch_size 1 \
  --train_ps 512 \
  --val_ps 512 \
  --nepoch 400 \
  --T_max 32 \
  --lr_initial 1e-5 \
  --eta_min 1e-6 \
  --weight_decay 1e-4 \
  --gpu_device 0 \
  --save_dir ./experiment/

Inference

Run inference and save reconstructed outputs:

python scripts/test_speed.py \
  --val_dir ./datas/SRPG/test \
  --weights ./experiment/DuRP_0/checkpoints/model_latest.pth \
  --result_dir ./results/SRPG/test \
  --arch DuRP \
  --stage test \
  --gpu_device 0 \
  --mode fp16 \
  --save

Supported inference modes:

• fp32: standard full-precision inference.
• fp16: mixed-precision inference with autocast.
• fp16_compile: mixed precision with torch.compile enabled; requires PyTorch 2.0+.

Saved results follow this layout:

results/SRPG/test/
|-- S0_R/       # Recovered radiance images, PNG
|-- dolp_R/     # Recovered DoLP, NPY
`-- aolp_R/     # Recovered AoLP, NPY

Evaluation

After inference, compute PSNR and SSIM:

python scripts/calc_metric.py \
  --gt_dir ./datas/SRPG/test/GT \
  --pred_dir ./results/SRPG/test

The script reports metrics for:

• S0: reconstructed scene radiance image.
• DoLP: degree of linear polarization.
• AoLP: angle of linear polarization, using periodic PSNR.

Reproducibility

Logs and TensorBoard summaries are written under:

experiment/<ARCH>_<ENV>/logs/
experiment/<ARCH>_<ENV>/runs/

To resume from a checkpoint:

python train/train.py \
  --resume \
  --pretrain_weights ./experiment/PRNet_0/checkpoints/model_latest.pth \
  ...

Checkpoints and Datasets

Pretrained checkpoints and processed datasets will be released at:

• Checkpoints: https://drive.google.com/drive/folders/1ytiAFX0Y3nIlrZdntMB15I7BrAgGudoo?usp=sharing
• dataset: https://drive.google.com/drive/folders/1Tfy8q8aYtpMnYO2QcIF53azJ9xfpiDtm?usp=sharing

Citation

If you find this repository useful, please cite:

@inproceedings{yang2026durp,
  title={DuRP: Dual-Stage Physics-Embedded Learning for Joint Radiance and Polarization Restoration},
  author={Yang, Zhenshuo and He, Qian and Liu, Zhiyuan and Fan, Baojie and Tian, Jiandong},
  booktitle={Proceedings of the 43rd International Conference on Machine Learning},
  year={2026}
}

License

This project is released for research use. Please add the final license file before public release.

Acknowledgements

This implementation uses PyTorch and builds on common open-source utilities for image restoration, polarization analysis, and experiment logging. Please also cite the relevant datasets and dependencies when using this code.