Trash or Treasure? An Interactive Dual-Stream Strategy for Single Image Reflection Separation (NeurIPS 2021)

by Qiming Hu, Xiaojie Guo.

Dependencies

• Python3
• PyTorch>=1.0
• OpenCV-Python, TensorboardX, Visdom
• NVIDIA GPU+CUDA

Network Architecture

🚀 1. Single Image Reflection Separation

Data Preparation

Training dataset

• 7,643 images from the Pascal VOC dataset, center-cropped as 224 x 224 slices to synthesize training pairs.
• 90 real-world training pairs provided by Zhang et al.

Testing dataset

• 45 real-world testing images from CEILNet dataset.
• 20 real testing pairs provided by Zhang et al.
• 454 real testing pairs from SIR^2 dataset, containing three subsets (i.e., Objects (200), Postcard (199), Wild (55)).

Usage

Training

• For stage 1: python train_sirs.py --inet ytmt_ucs --model ytmt_model_sirs --name ytmt_ucs_sirs --hyper --if_align
• For stage 2: python train_twostage_sirs.py --inet ytmt_ucs --model twostage_ytmt_model --name ytmt_uct_sirs --hyper --if_align --resume --resume_epoch xx --checkpoints_dir xxx

Testing

python test_sirs.py --inet ytmt_ucs_old --model twostage_ytmt_model --name ytmt_uct_sirs_test --hyper --if_align --resume --icnn_path ./checkpoints/ytmt_uct_sirs/ytmt_uct_sirs_68_077_00595364.pt

Note: "ytmt_ucs_old" is only for our provided checkpoint, and please change it as "ytmt_ucs" when you train our model by yourself, since it is a refactorized verison for a better view.

Trained weights

Google Drive

Visual comparison on real20 and SIR^2

Visual comparison on real45

🚀 2. Single Image Denoising

Data Preparation

Training datasets

400 images from the Berkeley segmentation dataset, following DnCNN.

Tesing datasets

BSD68 dataset and Set12.

Usage

Training

python train_denoising.py --inet ytmt_pas --name ytmt_pas_denoising --preprocess True --num_of_layers 9 --mode B --preprocess True

Testing

python test_denoising.py --inet ytmt_pas --name ytmt_pas_denoising_blindtest_25 --test_noiseL 25 --num_of_layers 9 --test_data Set68 --icnn_path ./checkpoints/ytmt_pas_denoising_49_157500.pt

Trained weights

Google Drive

Visual comparison on a sample from BSD68

🚀 3. Single Image Demoireing

Data Preparation

Training dataset

AIM 2019 Demoireing Challenge

Tesing dataset

100 moireing and clean pairs from AIM 2019 Demoireing Challenge.

Usage

Training

python train_demoire.py --inet ytmt_ucs --model ytmt_model_demoire --name ytmt_uas_demoire --hyper --if_align

Testing

python test_demoire.py --inet ytmt_ucs --model ytmt_model_demoire --name ytmt_uas_demoire_test --hyper --if_align --resume --icnn_path ./checkpoints/ytmt_ucs_demoire/ytmt_ucs_opt_086_00860000.pt

Trained weights

Google Drive

Visual comparison on the validation set of LCDMoire

🚀 4. Intrinsic Image Decomposition

Data Preparation

MIT-intrinsic dataset, pre-processed following Direct Intrinsics

Usage

Trained weights

Google Drive

Visual comparison on the validation split of MIT-Intrinsic

Training

python train_intrinsic.py --inet ytmt_ucs --model ytmt_model_intrinsic_decomp --name ytmt_ucs_intrinsic

Testing

python test_intrinsic.py --inet ytmt_ucs --model ytmt_model_intrinsic_decomp --name ytmt_ucs_intrinsic --resume --icnn_path [Path to your weight]