<jupyter_start><jupyter_text>Image Superresolution (ADMM Plug-and-Play Priors w/ DnCNN)==========================================================This example demonstrates the use of the ADMM Plug and Play Priors(PPP) algorithm  for solvinga simple image superresolution problem.<jupyter_code>import jax

import scico
import scico.numpy as snp
import scico.random
from scico import denoiser, functional, linop, loss, metric, plot
from scico.data import kodim23
from scico.optimize.admm import ADMM, LinearSubproblemSolver
from scico.solver import cg
from scico.util import device_info
plot.config_notebook_plotting()<jupyter_output><empty_output><jupyter_text>Define downsampling function.<jupyter_code>def downsample_image(img, rate):
    img = snp.mean(snp.reshape(img, (-1, rate, img.shape[1], img.shape[2])), axis=1)
    img = snp.mean(snp.reshape(img, (img.shape[0], -1, rate, img.shape[2])), axis=2)
    return img<jupyter_output><empty_output><jupyter_text>Read a ground truth image.<jupyter_code>img = kodim23(asfloat=True)[160:416, 60:316]
img = jax.device_put(img)<jupyter_output><empty_output><jupyter_text>Create a test image by downsampling and adding Gaussian white noise.<jupyter_code>rate = 4  # downsampling rate
σ = 2e-2  # noise standard deviation

Afn = lambda x: downsample_image(x, rate=rate)
s = Afn(img)
input_shape = img.shape
output_shape = s.shape
noise, key = scico.random.randn(s.shape, seed=0)
sn = s + σ * noise<jupyter_output><empty_output><jupyter_text>Set up the PPP problem pseudo-functional. The DnCNN denoiser is used as a regularizer.<jupyter_code>A = linop.LinearOperator(input_shape=input_shape, output_shape=output_shape, eval_fn=Afn)
f = loss.SquaredL2Loss(y=sn, A=A)
C = linop.Identity(input_shape=input_shape)
g = functional.DnCNN("17M")<jupyter_output><empty_output><jupyter_text>Compute a baseline solution via denoising of the pseudo-inverse of theforward operator. This baseline solution is also used to initialize thePPP solver.<jupyter_code>xpinv, info = cg(A.T @ A, A.T @ sn, snp.zeros(input_shape))
dncnn = denoiser.DnCNN("17M")
xden = dncnn(xpinv)<jupyter_output><empty_output><jupyter_text>Set up an ADMM solver and solve.<jupyter_code>ρ = 3.4e-2  # ADMM penalty parameter
maxiter = 12  # number of ADMM iterations
solver = ADMM(
    f=f,
    g_list=[g],
    C_list=[C],
    rho_list=[ρ],
    x0=xden,
    maxiter=maxiter,
    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": 1e-3, "maxiter": 10}),
    itstat_options={"display": True},
)

print(f"Solving on {device_info()}\n")
xppp = solver.solve()
hist = solver.itstat_object.history(transpose=True)<jupyter_output>Solving on GPU (NVIDIA GeForce RTX 2080 Ti)<jupyter_text>Plot convergence statistics.<jupyter_code>plot.plot(
    snp.vstack((hist.Prml_Rsdl, hist.Dual_Rsdl)).T,
    ptyp="semilogy",
    title="Residuals",
    xlbl="Iteration",
    lgnd=("Primal", "Dual"),
)<jupyter_output><empty_output><jupyter_text>Show reference and test images.<jupyter_code>fig = plot.figure(figsize=(8, 6))
ax0 = plot.plt.subplot2grid((1, rate + 1), (0, 0), colspan=rate)
plot.imview(img, title="Reference", fig=fig, ax=ax0)
ax1 = plot.plt.subplot2grid((1, rate + 1), (0, rate))
plot.imview(sn, title="Downsampled", fig=fig, ax=ax1)
fig.show()<jupyter_output><empty_output><jupyter_text>Show recovered full-resolution images.<jupyter_code>fig, ax = plot.subplots(nrows=1, ncols=3, sharex=True, sharey=True, figsize=(21, 7))
plot.imview(xpinv, title="Pseudo-inverse: %.2f (dB)" % metric.psnr(img, xpinv), fig=fig, ax=ax[0])
plot.imview(
    xden, title="Denoised pseudo-inverse: %.2f (dB)" % metric.psnr(img, xden), fig=fig, ax=ax[1]
)
plot.imview(xppp, title="PPP solution: %.2f (dB)" % metric.psnr(img, xppp), fig=fig, ax=ax[2])
fig.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15152317880794702, 'fraction_numerical': 0.019072847682119205, 'mean_word_length': 6.506958250497018, 'pattern_counts': {'":': 3, '<': 40, '<?xml version=': 0, '>': 40, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '451099', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 1500, 'n_tokens_neox': 1420, 'n_words': 300, 'path': 'notebooks/superres_ppp_dncnn_admm.ipynb', 'repo_name': 'lanl/scico-data', 'size': '3775'}