An extensive, highly customisable python package for image deblurring
Project description
DBlur: An Image Deblurring Toolkit
DBlur is an open-source python library for image deblurring. It is simple and highly versatile making it perfect for both experts and non-experts in the field.
For example, training, testing and deblurring with numerous SOTA models can be performed with just 2-3 lines of code with the default parameters mentioned in each paper. On the other hand, these methods can abe highly customised to assist existing researchers explore new avenues in the filed. For more details regarding its usage, refer to 'Usage' section below.
DBlur has a wide range functionalities which include:
- Numerous highly customisable SOTA model architectures for image deblurring. The list of supported models are mentioned below.
- Numerous popular and novel datasets which can be found here.
- Ready-to-use pipelines created according to the paper for training each of these models on a given dataset.
- Ready-to-use pipelines for testing each of these models.
- Multiple post-processing steps to improve deblurring performance.
- Deblurring images using pretrained models.
- Deblurring images by combining multiple pretrained models in different ways.
- Common evaluations metrics such as PSNR, SSIM, MSE and Brisque.
Call for Contributions
Dblur is a project which is in active development. Any helpful comments and improvements are highly encouraged. To do so, please open an issue in our official Github page.
In particular, due to the lack of GPUs on our side, the pretrained models for each of the model architectures are not yet available. Hence, if you have trained any of the models on a particular dataset, you are highly encouraged to share this pretrained model on our official Github page. This way, the pretrained models can be added to the library.
Installation
Install dblur with pip:
$ pip install dblur
The following requirements shall be installed alongside dblur:
- torch
- scipy
- numpy
- torchmetrics
- tensorboard
- piq
Supported Models
Datasets
A curated list of popular and custom datasets for general, face, and text deblurring can be found here. This datasets have been preprocessed/structured so that they can be easily used with our library.
Usage
Before diving into the usage of the library, it is important to look at a brief overview of the structure of the library. The two main classes associated with each deblurring model is a respective Trainer and Tester class. These classes for each model are all subclassed from a main BaseTrainer and BaseTester class.
Broadly speaking, each Trainer class has the following methods:
- train()
- validate()
- get_model()
- get_loss()
- get_optimizer()
- get_lr_scheduler()
- get_train_dataloader()
- get_val_dataloader()
Similarly, each Tester class has the following methods:
- test()
- deblur_imgs()
- deblur_single_img()
- get_test_dataloader()
All methods which are part of the public API have an elaborate docstring. Hence, usage of the library is made relatively straightforward. Following is a brief description on the usage of the library.
1. Training, testing, deblurring with default settings.
This section is mainly aimed at non-experts who would like to train, test a model or deblur images with a pretrained model with default settings provided in each paper. By default settings, we mean:
- optimizer
- learning rate scheduler
- loss function
- parameters associated with the model architecture, loss, scheduler, optimizer and others that are used for
- training/testing (e.g. batch size).
All the components metioned above are part of the setting. Below, we demonstrate how to train, test and deblur using the default settings for the Restormer model (can be easily generalised to other models).
from dblur.default.restormer import train, test, deblur_imgs, deblur_single_img
train_img_dir = "path_to_training_dataset"
val_img_dir = "path_to_validation_dataset"
test_img_dir = "path_to_test_dataset"
model_path = "path_for_model"
# Train model
train(model_path, train_img_dir)
# Test pretrained model
train(model_path, test_img_dir)
# Deblur images in a directory using pretrained model
deblur_imgs(model_path, "blurred_imgs_path", "sharp_imgs_path")
# Deblur single image using pretrained model
deblur_single_img(model_path, "blurred_img_path", "sharp_img_path")
2. Customize training pipeline for a model
The following code illustrates the entire training pipeline for the Restormer (can be generalised to any other model).
from dblur.trainers.restormer import RestormerTrainer
restormer_trainer = RestormerTrainer()
train_dataloader = restormer_trainer.get_train_dataloader(train_img_dir, batch_size=8)
val_dataloader = restormer_trainer.get_val_dataloader(val_img_dir, batch_size=8)
model = restormer_trainer.get_model(num_layers=4, num_refinement_blocks = 2)
optimizer = restormer_trainer.get_optimizer(model.parameters())
loss_func = restormer_trainer.get_loss()
lr_scheduler = restormer_trainer.get_lr_scheduler(optimizer)
restormer_trainer.train(model,
train_dataloader,
val_dataloader,
optimizer,
loss_func,
save_checkpoint_freq=10,
logs_folder='runs',
checkpoint_save_name="path_to_model",
val_freq=100,
write_logs=True,
lr_scheduler=lr_scheduler,
epochs=epochs)
Each method of the RestormerTrainer object has multiple parameters that can be set accordingly. This can be found in the docstring of each method. The example above only illustrates few of the parameters that can be specified. Moreover, the user is not limited to using the loss function or optimizer given by the methods of the RestormerTrainer object.
3. Customize testing pipeline for a model
The following code illustrates the entire testing pipeline for the Restormer (can be generalised to any other model).
from dblur.testers.restormer import RestormerTester
restormer_tester = RestormerTester()
test_dataloader = restormer_tester.get_test_dataloader(test_img_dir, batch_size=8)
model = restormer_tester.get_model(num_layers = 4, num_refinement_blocks = 2)
loss_func = restormer_tester.get_loss()
# test model on test dataset.
restormer_tester.test(model,
model_path,
test_dataloader,
loss_func,
is_checkpoint=True,
window_slicing=True,
window_size=256)
# deblur images in a directory using pretrained model
restormer_tester.deblur_imgs(model,
model_path,
blur_img_dir,
sharp_img_dir,
is_checkpoint=True,
batch_size=8,
window_slicing=False)
#deblur single image using pretrained model
restormer_tester.deblur_single_img(model,
model_path,
blur_img_path,
sharp_img_path,
is_checkpoint=True,
window_slicing=False)
Each method of the RestormerTester object has multiple parameters that can be set accordingly. This can be found in the docstring of each method. The example above only illustrates few of the parameters that can be specified.
4. Deblur images with multiple pretrained models
Dblur provides two ways in which you can combine multiple pretrained models to deblur images.
a) First method averages the outputs of the pretrained models provided. The code below shows how to perform this on a single image (can be done for multiple images in a directory as well).
from dblur.testers.mscnn import MSCNNTester
from dblur.testers.stack_dmphn import StackDMPHNTester
from dblur.multi_modal_deblur import multi_modal_deblur
mscnn_tester = MSCNNTester()
model1 = mscnn_tester.get_model()
dmphn_tester = StackDMPHNTester()
model2 = dmphn_tester.get_model(num_of_stacks=1)
multi_modal_deblur(models=[model1, model2],
model_names=["MSCNN", "StackDMPHN"],
model_paths=[model_path1, model_path2],
blur_img_path=blur_img_path,
sharp_img_path=sharp_img_path,
is_checkpoint=[True, True],
window_slicing=True,
window_size=256,
overlap_size=0)
b) The second method selects the output with the lowest(i.e. the best) Brisque index. Brisque is an index used to access the quality of an image. For more details regarding Brisque, refer to this. The code below shows how to perform this on a single image (can be done for multiple images in a directory as well).
from dblur.testers.mscnn import MSCNNTester
from dblur.testers.stack_dmphn import StackDMPHNTester
from dblur.multi_modal_deblur import multi_modal_deblur_by_brisque
mscnn_tester = MSCNNTester()
model1 = mscnn_tester.get_model()
dmphn_tester = StackDMPHNTester()
model2 = dmphn_tester.get_model(num_of_stacks=1)
multi_modal_deblur_by_brsique(models=[model1, model2],
model_names=["MSCNN", "StackDMPHN"],
model_paths=[model_path1, model_path2],
blur_img_path=blur_img_path,
sharp_img_path=sharp_img_path,
is_checkpoint=[True, True],
window_slicing=True,
window_size=256,
overlap_size=0)
Citation
If you use or extend this work, please consider citing it as below:
@software{Parayil_Shibu_DBlur_An_Image_2023,
author = {Parayil Shibu, Jishnu},
month = {3},
title = {{DBlur: An Image Deblurring Toolkit}},
url = {https://github.com/Jishnu8/DBlur-An-Image-Deblurring-Toolkit},
version = {1.0.0},
year = {2023}
}
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.