#### NOTE
New to DeepInverse? Get started with the basics with the
[5 minute quickstart tutorial](https://deepinv.org/auto_examples/basics/demo_quickstart.html.md#sphx-glr-auto-examples-basics-demo-quickstart-py)..
# Image deblurring with custom deep explicit prior.
In this example, we show how to solve a deblurring inverse problem using an explicit prior.
Here we use the simple L2 prior that penalizes the squared norm of the reconstruction, with an ADMM algorithm.
```Python
import deepinv as dinv
from pathlib import Path
import torch
from torch.utils.data import DataLoader
from deepinv.optim.prior import Prior
from deepinv.optim.data_fidelity import L2
from deepinv.optim import ADMM
from deepinv.training import test
from torchvision import transforms
from deepinv.utils import load_dataset
```
## Setup paths for data loading and results.
```Python
# Setup paths for data loading, results and checkpoints.
BASE_DIR = Path(".")
DATA_DIR = BASE_DIR / "measurements"
RESULTS_DIR = BASE_DIR / "results"
DEG_DIR = BASE_DIR / "degradations"
# Set the global random seed from pytorch to ensure reproducibility of the example.
torch.manual_seed(0)
device = dinv.utils.get_device()
```
```none
Selected GPU 0 with 28319.25 MiB free memory
```
## Load base image datasets and degradation operators.
In this example, we use the CBSD68 dataset from the paper of Zhang et al. (2017) and the motion blur kernels from Levin *et al.*[1](#footcite-levin2009understanding).
```Python
# Set up the variable to fetch dataset and operators.
method = "L2_prior"
dataset_name = "set3c"
operation = "deblur"
img_size = 256
val_transform = transforms.Compose(
[transforms.CenterCrop(img_size), transforms.ToTensor()]
)
dataset = load_dataset(dataset_name, transform=val_transform)
```
```none
/local/jtachell/deepinv/deepinv/deepinv/datasets/base.py:495: DeprecationWarning: The tuple format for dataset outputs is deprecated and will be removed in a future version.It is recommended to set `use_dict_output=True` for better readability and flexibility in returned outputs.The default is currently `False` for backward compatibility, but will be switched to `True` in a future version.
super().__init__(use_dict_output=use_dict_output)
```
## Define physics operator
We use the [`deepinv.physics.BlurFFT`](https://deepinv.org/api/stubs/deepinv.physics.BlurFFT.html.md#deepinv.physics.BlurFFT) operator from the physics module to generate a dataset of blurred images.
The BlurFFT class performs the convolutions via the Fourier transform.
In this example, we choose a gaussian kernel with standard deviation 3, and we add a Gaussian noise with standard
deviation 0.03.
```Python
# Generate a Gaussian blur filter.
filter_torch = dinv.physics.functional.gaussian_blur(sigma=(3, 3))
noise_level_img = 0.03 # Gaussian Noise standard deviation for the degradation
n_channels = 3 # 3 for color images, 1 for gray-scale images
# The BlurFFT instance from physics enables to compute efficiently backward operators with Fourier transform.
p = dinv.physics.BlurFFT(
img_size=(n_channels, img_size, img_size),
filter=filter_torch,
device=device,
noise_model=dinv.physics.GaussianNoise(sigma=noise_level_img),
)
```
## Generate a dataset of blurred images
```Python
# Use parallel dataloader if using a GPU to speed up training, otherwise, as all computes are on CPU, use synchronous
# data loading.
num_workers = 4 if torch.cuda.is_available() else 0
n_images_max = 3 # Maximal number of images to restore from the input dataset
measurement_dir = DATA_DIR / dataset_name / operation
deepinv_dataset_path = dinv.datasets.generate_dataset(
train_dataset=dataset,
test_dataset=None,
physics=p,
device=device,
save_dir=measurement_dir,
train_datapoints=n_images_max,
num_workers=num_workers,
)
```
```none
/local/jtachell/deepinv/deepinv/.pixi/envs/docs/lib/python3.12/multiprocessing/popen_fork.py:66: DeprecationWarning: This process (pid=4088353) is multi-threaded, use of fork() may lead to deadlocks in the child.
self.pid = os.fork()
Dataset has been saved at measurements/set3c/deblur/dinv_dataset0.h5
```
## Set up the optimization algorithm to solve the inverse problem.
We use the [`deepinv.optim.optimizers.ADMM`](https://deepinv.org/api/stubs/deepinv.optim.ADMM.html.md#deepinv.optim.ADMM) to instantiate the ADMM optimization algorithm.
The optimization algorithm is a proximal gradient descent algorithm that solves the following optimization problem:
$$
\min_{x} \frac{1}{2} \|y - Ax\|_2^2 + \lambda \|x\|_2^2
$$
where $A$ is the forward blurring operator, $y$ is the measurement
and $\lambda$ is a regularization parameter.
```Python
# Create a custom prior which inherits from the base Prior class.
class L2Prior(Prior):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.explicit_prior = True
def fn(self, x, args, **kwargs):
return (
0.5 * torch.linalg.vector_norm(x, dim=tuple(range(1, x.dim())), ord=2) ** 2
)
# Specify the custom prior
prior = L2Prior()
# Select the data fidelity term
data_fidelity = L2()
# Logging parameters
verbose = True
# Parameters of the algorithm to solve the inverse problem
early_stop = True # Stop algorithm when convergence criteria is reached
crit_conv = "cost" # Convergence is reached when the difference of cost function between consecutive iterates is
# smaller than thres_conv
thres_conv = 1e-5
max_iter = 500 # Maximum number of iterations
# Instantiate the algorithm class to solve the IP problem.
model = ADMM(
prior=prior,
g_first=False,
data_fidelity=data_fidelity,
stepsize=1.0,
lambda_reg=0.1,
early_stop=early_stop,
max_iter=max_iter,
crit_conv=crit_conv,
thres_conv=thres_conv,
verbose=verbose,
)
```
## Evaluate the reconstruction algorithm on the problem.
We can use the [`deepinv.test()`](https://deepinv.org/api/stubs/deepinv.test.html.md#deepinv.test) function to evaluate the reconstruction algorithm on a test set.
```Python
batch_size = 1
plot_images = True # plot results
plot_convergence_metrics = True # compute performance and convergence metrics along the algorithm, curves saved in RESULTS_DIR
dataset = dinv.datasets.HDF5Dataset(
path=deepinv_dataset_path, train=True, use_dict_output=True
)
dataloader = DataLoader(
dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False
)
test(
model=model,
test_dataloader=dataloader,
physics=p,
device=device,
plot_images=plot_images,
save_folder=RESULTS_DIR / method / operation / dataset_name,
plot_convergence_metrics=plot_convergence_metrics,
verbose=verbose,
)
```
```none
/local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:593: UserWarning: Update progress bar frequency of 1 may slow down training on GPU. Consider setting freq_update_progress_bar > 1.
warnings.warn(
0%| | 0/3 [00:00
* **References:**
* **[1]** Anat Levin, Yair Weiss, Fredo Durand, and William T Freeman. Understanding and evaluating blind deconvolution algorithms. In *2009 IEEE conference on computer vision and pattern recognition*, 1964–1971. IEEE, 2009.
**Total running time of the script:** (0 minutes 1.058 seconds)
🏷 Tags: Deblurring
[`Download Jupyter notebook: demo_custom_prior.ipynb`](https://deepinv.org/auto_examples/optimization/_downloads/1f000b13523fc415ca34330757b9bc85/demo_custom_prior.ipynb)
[`Download Python source code: demo_custom_prior.py`](https://deepinv.org/auto_examples/optimization/_downloads/441b496817b69d18efaa1721eeb1d232/demo_custom_prior.py)
[`Download zipped: demo_custom_prior.zip`](https://deepinv.org/auto_examples/optimization/_downloads/2de8d5a9408304463cdb0c5c1c05bcee/demo_custom_prior.zip)
[Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)