BEDSR-Net

This repository is unofficial implementation of BEDSR-Net: A Deep Shadow Removal Network From a Single Document Image [Lin+, CVPR 2020] with PyTorch.

A refined version of IsHYuhi's implementation.

Fix several problems

nn.ConvTranspose2d compatible with higher version of Pytorch
gradcam uses too much vram, use pytorch-grad-cam instead
provide default correct training config
provide easy inference code

Dependencies

Pytorch, torchvision, matplotlib, wandb, albumentations, pytorch-grad-cam

Dataset Structure

The dataset should be formatted like below, train.csv and test.csv can be generated using

python make_dataset.py

.
├── csv/
│   └── Jung/
│       ├── train.csv
│       └── test.csv
└── dataset/
    └── Jung/
        ├── train/
        │   ├── input/
        │   │   ├── *.jpg
        │   │   └── ...
        │   └── target/
        │       ├── *.jpg
        │       └── ...
        └── test/
            ├── input/
            │   ├── *.jpg
            │   └── ...
            └── target/
                ├── *.jpg
                └── ...

Training

Training BE-Net

python3 train_benet.py ./configs/model\=benet/config.yaml

Training BEDSR-Net

python3 train_bedsrnet.py ./configs/model\=bedsrnet/config.yaml

You can use W&B by --use_wandb.

Infer

mask sure put all your model state_dict into pretrained directory

python infer.py