ECE 60141: Foundations of Computational Imaging

Lab 5: MAP Reconstruction with Plug-and-Play

Overview What the lab is about, the data, and the ground rules.

What This Lab Is About

Due: Friday, Oct. 23.

In Lab 4 you built a restoration that used a QGGMRF prior. Buried inside it, without being named, was a denoiser. In this lab you pull that denoiser out, hand it to a different algorithm, and then swap it for a denoiser that was never derived from a prior at all. The restoration keeps working.

That swap is the plug-and-play idea, and it is worth being careful about. Up to now every algorithm in this course minimized a cost function, and the answer was defined before the algorithm existed. After the swap there is no cost function. The algorithm still converges and the pictures still improve, but what it converges to is a new question. Step 8 asks you to answer it for your own code.

The background is Chapter 9.5 on variable splitting and ADMM, and Chapter 10 on plug-and-play. The notation here matches the book.

The Data

The same image and blur kernel as Lab 3, so that your numbers can be compared with the ones you already have.

wget https://cabouman.github.io/grad_labs/labs/resources/kodim23.pgm
wget https://cabouman.github.io/grad_labs/labs/resources/levin09_kernels.zip
unzip levin09_kernels.zip

Ground Rules

  • All code in PyTorch, in a single file named lab5.py, holding functions and nothing else. Run your experiments from a separate script.
  • You may import forward, adjoint, and load_pgm from your lab3.py. Everything else on this page is new.
  • Every function has the exact signature given here, and a docstring saying what it does, what its inputs are, and what it returns.
  • Everything in float64. No for loop over pixels.
  • Fix the random seed for the noise and report it.
Three sigmas, three jobs. This lab has \(\sigma_w\), the noise standard deviation in the data; \(\sigma_x\), the scale of the prior; and \(\sigma\), a new one. The new \(\sigma\) belongs to the algorithm, not to any model. Step 2 says what it does. Keep the three apart in your code and in your report.
Module 1: Splitting the Problem in Two Steps 1 to 3. Proximal maps, and the one that has a closed form.

Step 1: Two Models, One Optimizer

The MAP estimate you computed in Labs 3 and 4 minimizes a sum of two terms,

$$ \hat{x} = \arg\min_x \left\{ f(x) + h(x) \right\} \ , \qquad f(x) = \frac{1}{2\sigma_w^2} \| y - A x \|^2 \ , \tag{1} $$

where \(f\) is the forward model term and \(h\) is the prior term. They come from different places. The forward model came from physics, from what the camera did. The prior came from a belief about what images look like.

In Labs 3 and 4 you minimized the sum with one optimizer, which meant the two models were welded together. Changing the prior meant rederiving the gradient of the whole cost. Changing the blur meant the same. This lab takes them apart, so that each model is handled by its own piece of code and the two pieces only exchange images.

Step 2: The Proximal Map

The tool that separates them is the proximal map. For each of the two terms, define

$$ F(x) = \arg\min_v \left\{ f(v) + \frac{1}{2\sigma^2} \| v - x \|^2 \right\} \ , \qquad H(x) = \arg\min_v \left\{ h(v) + \frac{1}{2\sigma^2} \| v - x \|^2 \right\} \ , \tag{2} $$

where \(\sigma\) is a parameter you choose. Read \(F(x)\) as a question: starting from the image \(x\), where should I move to reduce \(f\), without going far? Small \(\sigma\) means stay close to \(x\). Large \(\sigma\) means go most of the way to the minimum of \(f\).

This is a change in kind, and it is the reason the chapter exists. A cost function is an objective: it scores images and says which are better. A proximal map is an action: give it an image and it hands you back another image. The proximal map converts one into the other.

Look at what \(H\) actually is. Write \(h(v) = -\log p(v)\) for the prior. Then (2) says that \(H(x)\) is the MAP estimate of an image whose prior is \(p\), given a measurement \(x\) corrupted by white Gaussian noise of standard deviation \(\sigma\). In other words, the proximal map of a prior is a MAP denoiser. You built one in Lab 4 and called it a restoration with \(A = I\).

Step 3: \(F\) Has a Closed Form

\(f\) is quadratic, so the minimization in (2) is a linear system. Set the gradient to zero:

$$ \left( \frac{1}{\sigma_w^2} A^t A + \frac{1}{\sigma^2} I \right) v = \frac{1}{\sigma_w^2} A^t y + \frac{1}{\sigma^2} x \ . \tag{3} $$

In Lab 3 you chose circular boundaries for \(A\), and now that choice pays. Circular convolution is diagonalized by the DFT, so (3) becomes a division, one entry at a time:

$$ \hat{V} = \frac{ \overline{\hat{a}} \, \hat{Y} / \sigma_w^2 + \hat{X} / \sigma^2 } { | \hat{a} |^2 / \sigma_w^2 + 1 / \sigma^2 } \ , \tag{4} $$

where a hat denotes the 2-D DFT, \(\hat{a}\) is the DFT of the blur kernel embedded in an array the size of the image, and the bar is complex conjugation. No iteration, no step size, and no convergence to worry about. \(F\) is exact.

Write these two functions.

def kernel_dft(a, shape):
    """Return the DFT of the blur kernel, sized to the image.

    Embed a in a zero array of the given shape, positioned so that this
    DFT represents the same circular convolution as forward() in Lab 3,
    then take the 2-D DFT.

    a        (Ka, Kb) float64 tensor, odd sized
    shape    (H, W), the image shape
    returns  (H, W) complex128 tensor
    """


def prox_data(x, y, a_hat, sigma_w, sigma):
    """Return F(x), the proximal map of the forward model term, using (4).

    x        (H, W) float64 tensor, the point the map is anchored at
    y        (H, W) float64 tensor, the measured data
    a_hat    (H, W) complex128 tensor from kernel_dft
    sigma_w  float, the noise standard deviation
    sigma    float, the proximal parameter
    returns  (H, W) float64 tensor
    """
Getting the centering right. The kernel has to sit in the array so that its center lands at index \((0,0)\), with the rest wrapping around. If you place it wrong, the blur shifts and every image after it is shifted too. Task 3a below catches that, and it catches it in one number, which is much easier than staring at pictures.

Task 3a. Check the DFT operator against Lab 3. For a random image \(u\), compare ifft2(fft2(u) * a_hat).real with forward(u, a) from Lab 3. Report the largest absolute difference. It should be at the level of round-off.

Task 3b. See what \(F\) does on its own. Form the blurred noisy data \( y = Ax + w \) with kernel 1 and \(\sigma_w = 0.02\). Compute \(F(y)\) for \(\sigma = 0.01\) and for \(\sigma = 1.0\), and display both next to \(y\).

HAND IN · D1

The three images \(y\), \(F(y)\) at \(\sigma = 0.01\), and \(F(y)\) at \(\sigma = 1.0\), with the number from Task 3a. Add two sentences explaining what \(\sigma\) did, and why the large \(\sigma\) result looks the way it does even though it fits the data better.

Module 2: Plugging In a Denoiser Steps 4 to 6. Two denoisers, and the algorithm that uses them.

Step 4: The Denoiser You Already Have

By Step 2, \(H\) is a MAP denoiser for the prior. For the QGGMRF prior of Lab 4 that means

$$ H(x) = \arg\min_v \left\{ \frac{1}{2\sigma^2} \| v - x \|^2 + \sum_{ \{s,r\} \in \mathcal{P} } b_{s,r} \, \rho( v_s - v_r ) \right\} \ , \tag{5} $$

with the QGGMRF potential from Lab 4, repeated here so you do not have to go looking:

$$ \rho( \Delta ) = \frac{ |\Delta|^p }{ p\,\sigma_x^p } \left( \frac{ \left| \frac{\Delta}{T \sigma_x} \right|^{q-p} } { 1 + \left| \frac{\Delta}{T \sigma_x} \right|^{q-p} } \right) \ . \tag{6} $$

This is the Lab 4 problem with the blur removed. Solve it the same way, with majorization on the outside and gradient descent on the inside. If your Lab 4 code is organized into functions you can call, call them. If it is not, this is a good moment to fix that.

def denoise_qggmrf(x, sigma, sigma_x, p, q, T, num_iters):
    """Return H(x), the MAP denoiser of equation (5).

    x          (H, W) float64 tensor, the noisy image
    sigma      float, the noise level the denoiser assumes
    sigma_x    float, the prior scale
    p, q, T    floats, the QGGMRF parameters of equation (6)
    num_iters  int, outer majorization iterations
    returns    (H, W) float64 tensor
    """

Use \(q = 2\), \(p = 1.2\), and \(T = 0.1\) throughout this lab, which are the baseline values from Lab 4.

Step 5: PnP-ADMM

Now the two agents are put in a loop. Applying ADMM to the split problem gives three lines:

PnP-ADMM

  Initialize u = 0 and v = y
  Until converged:
      x = F(v - u)
      v = H(x + u)
      u = u + (x - v)

Read the loop as a negotiation between two parties who never speak the same language. \(F\) knows about the camera and nothing about images. \(H\) knows about images and nothing about the camera. The vector \(u\) is the running record of how far apart they are, and it is what forces them to agree in the end.

When \(F\) and \(H\) are both proximal maps, this is exactly the ADMM algorithm of Chapter 9, and at convergence \(x = v\) is the minimizer of (1). Nothing has been given up yet.

def pnp_admm(y, a_hat, denoiser, sigma_w, sigma, num_iters):
    """Run the PnP-ADMM loop above.

    y          (H, W) float64 tensor, the measured data
    a_hat      (H, W) complex128 tensor from kernel_dft
    denoiser   a function taking one (H, W) tensor and returning one,
               which plays the role of H
    sigma_w    float, the noise standard deviation
    sigma      float, the proximal parameter
    num_iters  int
    returns    (x, history), the final image and a list of the values of
               ||x - v|| / ||x|| at each iteration
    """
Why the denoiser is passed in. pnp_admm never names a prior and never sees one. It calls whatever function you hand it. That is the entire mechanism of this lab, and the next step uses it.

Step 6: A Denoiser That Came From Nowhere

Non-local means denoises a pixel by averaging other pixels whose surrounding patches look similar, wherever in the image they happen to be. It was not derived from a prior. Nobody wrote down a \(p(x)\), took its logarithm, and minimized anything. It is an algorithm somebody invented because it works.

It is already installed, in the scikit-image package from Lab 2. Wrap it so that it has the same interface as your other denoiser.

from skimage.restoration import denoise_nl_means


def denoise_nlm(x, sigma):
    """Return a non-local means denoising of x, as an agent for pnp_admm.

    x        (H, W) float64 tensor
    sigma    float, the noise level the denoiser assumes
    returns  (H, W) float64 tensor
    """
    z = x.detach().numpy()
    out = denoise_nl_means(z, h=0.8 * sigma, sigma=sigma,
                           patch_size=5, patch_distance=6, fast_mode=True)
    return torch.tensor(out, dtype=torch.float64)

Task 6a. Form the noisy image \( z = x + w \) with \(\sigma_w = 0.05\), and denoise it twice at \(\sigma = 0.05\), once with denoise_qggmrf and once with denoise_nlm. Report the RMSE of each against the clean image.

HAND IN · D2

Three images displayed over \([0,1]\): the noisy \(z\), the QGGMRF denoising, and the non-local means denoising, with the two RMSE values. Add one sentence on how the two differ where the image has texture, such as the feathers.

HAND IN · D3

The specification you gave your AI assistant for prox_data and pnp_admm, quoted, with any corrections you had to make.

Module 3: Does It Work, and What Is It Doing? Steps 7 to 9. Deblur with both agents, then test whether the swap was legal.

Step 7: Deblurring Experiments

Form the data \( y = Ax + w \) with kernel 1 and \(\sigma_w = 0.02\), the same as Lab 3, Step 8. Run pnp_admm for 50 iterations twice, once with each denoiser, using \(\sigma_x = 0.05\) in the QGGMRF and \(\sigma = 0.02\) in both runs.

Produce the following, and no more.

  1. Images. The blurred noisy \(y\), the two PnP results, and your direct MAP result from Lab 4 on the same data, all over \([0,1]\).
  2. Crops. The same four, cropped to a 128 by 128 region containing an edge and some texture. State the crop coordinates.
  3. Error table. The RMSE of \(y\) and of each of the three restorations.
  4. Convergence. The history \( \|x - v\| / \|x\| \) versus iteration for both runs, on one plot with a log vertical axis.
HAND IN · D4

The four full images and the four crops, with the crop coordinates.

HAND IN · D5

The error table and the convergence plot.

Task 7a. Sweep \(\sigma\). Repeat the non-local means run for \( \sigma \in \{0.005,\ 0.02,\ 0.08\} \), keeping everything else fixed. Report the RMSE of each and show the three crops.

HAND IN · D6

The three crops from the \(\sigma\) sweep with their RMSE values, and two sentences on what \(\sigma\) controls here. Note that \(\sigma\) is a parameter of the algorithm and not of any model, and say what that means for how you would choose it in practice.

Step 8: Was the Swap Legal?

The algorithm converged and the pictures improved. That is not the same as knowing what it computed. If both agents are proximal maps of convex functions, the answer is the minimizer of (1). If they are not, there may be no cost function at all, and the fixed point is defined only by the algorithm.

Moreau settled which case you are in. An operator \(H\) is the proximal map of a proper closed convex function if and only if it is both

  1. nonexpansive, meaning \( \| H(x) - H(z) \| \leq \| x - z \| \) for all \(x\) and \(z\), and
  2. conservative, meaning \(H = \nabla \phi\) for some scalar function \(\phi\).

Both are testable numerically on a small image, and neither needs the Jacobian to be formed. Nonexpansiveness is a ratio you can sample. For conservativeness, use the fact that the Jacobian of a gradient is symmetric, so for any directions \(u\) and \(v\),

$$ u^t J v = v^t J u \ , \qquad J v \approx \frac{ H(x + \epsilon v) - H(x - \epsilon v) }{ 2 \epsilon } \ . \tag{7} $$

Task 8a. On a 64 by 64 crop, with 20 random pairs and \(\epsilon = 10^{-4}\), test both properties for denoise_qggmrf and for denoise_nlm. For nonexpansiveness report the largest value of \( \|H(x) - H(z)\| / \|x - z\| \). For conservativeness report the largest value of \( |u^t J v - v^t J u| \, / \, ( \|u\| \|v\| ) \).

def moreau_tests(denoiser, x, num_trials, epsilon):
    """Test the two conditions of Moreau's theorem numerically.

    denoiser    a function taking one (H, W) tensor and returning one
    x           (H, W) float64 tensor, the point to test around
    num_trials  int, number of random pairs
    epsilon     float, the finite difference step
    returns     (max_expansion, max_asymmetry), two floats
    """
Reading the result. A finite difference has error of its own, so neither number will be exactly zero. Judge them by size. If one denoiser gives an asymmetry a few orders of magnitude larger than the other, that difference is the answer, not the noise.
HAND IN · D7

The four numbers from Task 8a in a small table, and three sentences: which of the two denoisers is a proximal map, how you can tell from these numbers, and what that says about the restoration you produced with the other one.

Step 9: Questions to Answer

Answer each in a few sentences, using your own images and numbers as evidence.

  1. Compare the PnP result using your QGGMRF denoiser with your direct MAP result from Lab 4. They aim at the same estimate. Are they the same image? Say what would make them differ, and check your answer against the RMSE table.
  2. In the PnP-ADMM loop, \(F\) is exact and costs one pair of FFTs, while \(H\) is iterative and costs far more. Suppose you halved the work inside denoise_qggmrf. What would you expect to happen to the convergence plot, and why?
  3. The non-local means denoiser has no \(\sigma_x\). Where did the strength of its regularization come from in your runs, and what does that say about the role of \(\sigma\) in PnP?
  4. Suppose a colleague replaces \(H\) with a denoiser that sets every pixel to the image mean. Trace the three lines of PnP-ADMM and say what \(x\) converges to. What does this tell you about why convergence alone is not evidence that an answer is good?
  5. Chapter 10 goes on to consensus equilibrium, which handles more than two agents. Name a third agent you would add to this deblurring problem, say what it would know that \(F\) and \(H\) do not, and say what its proximal map would do to an image.
HAND IN · D8

Your answers to the five questions above.

Deliverables What to submit, and where each item comes from.

What to Hand In

Submit two files through Brightspace.

  1. A report, as a single PDF, labeled "Lab 5", containing the items below in this order.
  2. Your lab5.py, as a plain .py file. Do not paste the code into the report and do not send a zip file.

Also commit your code to your image-processing-labs repository under lab5 and push it before you submit.

  1. Your name, the link to your GitHub repository, and the random seed you used.
  2. D1 The three images showing what \(F\) does, with the Task 3a check number and two sentences.
  3. D2 The three denoising images with their RMSE values, and one sentence on texture.
  4. D3 The specification you gave your AI assistant.
  5. D4 The four deblurred images and the four crops.
  6. D5 The error table and the convergence plot.
  7. D6 The \(\sigma\) sweep crops with RMSE, and two sentences.
  8. D7 The Moreau test table and three sentences.
  9. D8 Your answers to the five questions.

Concise and clear beats long.

You have now built a reconstruction whose prior can be replaced without touching the algorithm, and you have checked what that replacement cost you. Every method in the rest of the course fits into the same frame: an agent that knows the physics, one or more agents that know what the answer should look like, and a rule that makes them agree.

Back to the laboratory index · ECE 60141 course page