Cleaning printed text using Denoising Autoencoder based on UNet architecture in PyTorch
The UNet architecture used here is borrowed from https://github.com/jvanvugt/pytorch-unet. The only modification made in the UNet architecture mentioned in the above link is the addition of dropout layers.
- torch >= 0.4
- torchvision >= 0.2.2
- opencv-python
- numpy >= 1.7.3
- matplotlib
- tqdm
Set the number of total synthetic images to be generated num_synthetic_imgs and set the percentage of training data train_percentage in config.py Then run
python generate_synthetic_dataset.py
It will generate the synthetic data in a directory named data (can be changed in the config.py) in the root dirctory.
Set the desired values of lr, epochs and batch_size in config.py
In config.py,
- set resume to False
python train.py
In config.py,
- set resume to True and
- set ckpt to the path of the model to be loaded, i.e. ckpt = 'model02.pth'
python train.py
The model was trained for 12 epochs for the configuration mentioned in config.py
In config.py,
- set ckpt to the path of the model to be loaded, i.e. ckpt = 'model02.pth'
- set test_dir to the path that contains the noisy images that you need to denoise ('data/val/noisy' by default)
- set test_bs to the desired batch size for the test set (1 by default)
python test.py
Once the testing is done, the results will be saved in a directory named results
