Add bound constraints in FISTA

demo_tv_bounds_reconstruction.py

@@ -0,0 +1,75 @@
+"""
+Total-variation penalization and bound constraints for tomography reconstruction
+================================================================================
+
+In this example, we reconstruct an image from its tomography projections
+with an uncomplete set of projections (l/8 angles, where l is the linear
+size of the image. For a correct reconstruction without a-priori information,
+one would usually require l or more angles). In addition, noise is added to
+the projections.
+
+In order to reconstruct the original image, we minimize a function that
+is the sum of (i) a L2 data fit term, and (ii) the total variation of the
+image, and bound constraints on the pixel values. Proximal iterations
+using the FISTA scheme are used.
+
+We compare with and without the bounds
+
+This example should take around 1mn to run and plot the results.
+"""
+
+print __doc__
+
+import numpy as np
+from reconstruction.forward_backward_tv import fista_tv
+from reconstruction.projections import build_projection_operator
+from reconstruction.util import generate_synthetic_data
+from time import time
+import matplotlib.pyplot as plt
+
+# Synthetic data
+l = 512
+np.random.seed(0)
+x = generate_synthetic_data(l)
+
+
+# Projection operator and projections data, with noise
+H = build_projection_operator(l, l / 32)
+y = H * x.ravel()[:, np.newaxis]
+y += 5 * np.random.randn(*y.shape)
+
+# Display original data
+plt.figure(figsize=(12, 5))
+plt.subplot(2, 3, 1)
+plt.imshow(x, cmap=plt.cm.gnuplot2, interpolation='nearest', vmin=-.1, vmax=1.2)
+plt.title('original data (256x256)')
+plt.axis('off')
+
+for idx, (val_min, val_max, name) in enumerate([
+                                        (None, None, 'TV'),
+                                        (0, 1, 'TV + interval'),
+                                    ]):
+    # Reconstruction
+    t1 = time()
+    res, energies = fista_tv(y, 50, 100, H, val_min=val_min,
+                             val_max=val_max)
+    t2 = time()
+
+    # Fraction of errors of segmented image wrt ground truth
+    err = np.abs(x - (res[-1] > 0.5)).mean()
+    print "%s: reconstruction done in %f s, %.3f%% segmentation error" % (
+                name, t2 - t1, 100 * err)
+
+    plt.subplot(2, 3, 2 + idx)
+    plt.imshow(res[-1], cmap=plt.cm.gnuplot2, interpolation='nearest', vmin=-.1,
+                vmax=1.2)
+    plt.title('reconstruction with %s' % name)
+    plt.axis('off')
+    ax = plt.subplot(2, 3, 5 + idx)
+    ax.yaxis.set_scale('log')
+    plt.hist(res[-1].ravel(), bins=20, normed=True)
+    plt.yticks(())
+    plt.title('Histogram of pixel intensity')
+    plt.axis('tight')
+
+plt.show()

reconstruction/forward_backward_tv.py

@@ -11,6 +11,7 @@ def tv_norm(im):
     grad_x2 = np.diff(im, axis=1)
     return np.sqrt(grad_x1[:, :-1]**2 + grad_x2[:-1, :]**2).sum()
 
+
 def tv_norm_anisotropic(im):
     """Compute the anisotropic TV norm of an image"""
     grad_x1 = np.diff(im, axis=0)
@@ -19,7 +20,8 @@ def tv_norm_anisotropic(im):
 
 # ------------------ Proximal iterators ----------------------------
 
-def fista_tv(y, beta, niter, H, verbose=0, mask=None):
+def fista_tv(y, beta, niter, H, verbose=0, mask=None,
+             val_min=None, val_max=None):
     """
     TV regression using FISTA algorithm
     (Fast Iterative Shrinkage/Thresholding Algorithm)
@@ -42,6 +44,12 @@ def fista_tv(y, beta, niter, H, verbose=0, mask=None):
 
     mask : array of bools
 
+    val_min: None or float, optional
+        an optional lower bound constraint on the reconstructed image
+
+    val_max: None or float, optional
+        an optional upper bound constraint on the reconstructed image
+
     Returns
     -------
 
@@ -102,7 +110,8 @@ def fista_tv(y, beta, niter, H, verbose=0, mask=None):
         else:
             tmp2d = tmp.reshape((l, l))
         u_n = tv_denoise_fista(tmp2d,
-                weight=beta*gamma, eps=eps)
+                weight=beta*gamma, eps=eps, val_min=val_min,
+                val_max=val_max)
         t_new = (1 + np.sqrt(1 + 4 * t_old**2))/2.
         t_old = t_new
         x = u_n + (t_old - 1)/t_new * (u_n - u_old)
@@ -218,6 +227,7 @@ def ista_tv(y, beta, niter, H=None):
         energies.append(energy)
     return res, energies
 
+
 def gfb_tv(y, beta, niter, H=None, val_min=0, val_max=1, x0=None,
            stop_tol=1.e-4):
     """
@@ -332,6 +342,7 @@ def gfb_tv(y, beta, niter, H=None, val_min=0, val_max=1, x0=None,
             break
     return res, energies
 
+
 def gfb_tv_local(y, beta, niter, mask_pix, mask_reg, H=None,
                                 val_min=0, val_max=1, x0=None):
     """

reconstruction/tv_denoising.py

@@ -11,8 +11,9 @@ def div(grad):
         this_res[-1] -= this_grad[-2]
     return res
 
+
 def gradient(img):
-    """ 
+    """
     Compute gradient of an image
 
     Parameters
@@ -26,7 +27,7 @@ def gradient(img):
         Gradient of the image: the i-th component along the first
         axis is the gradient along the i-th axis of the original
         array img
-"""
+    """
     shape = [img.ndim, ] + list(img.shape)
     gradient = np.zeros(shape, dtype=img.dtype)
     # 'Clever' code to have a view of the gradient with dimension i stop
@@ -69,9 +70,9 @@ def dual_gap(im, new, gap, weight):
     dual_gap = (gap**2).sum() + tv_new - im_norm + (new**2).sum()
     return 0.5 / im_norm * dual_gap
 
-def tv_denoise_fista(im, weight=50, eps=5.e-5, n_iter_max=200,
-                                check_gap_frequency=3):
 
+def tv_denoise_fista(im, weight=50, eps=5.e-5, n_iter_max=200,
+                     check_gap_frequency=3, val_min=None, val_max=None):
     """
     Perform total-variation denoising on 2-d and 3-d images
 
@@ -99,6 +100,12 @@ def tv_denoise_fista(im, weight=50, eps=5.e-5, n_iter_max=200,
     n_iter_max: int, optional
         maximal number of iterations used for the optimization.
 
+    val_min: None or float, optional
+        an optional lower bound constraint on the reconstructed image
+
+    val_max: None or float, optional
+        an optional upper bound constraint on the reconstructed image
+
     Returns
     -------
     out: ndarray
@@ -129,6 +136,7 @@ def tv_denoise_fista(im, weight=50, eps=5.e-5, n_iter_max=200,
     t = 1.
     i = 0
     while i < n_iter_max:
+        # error is the dual variable
         error = weight * div(grad_aux) - im
         grad_tmp = gradient(error)
         grad_tmp *= 1./ (8 * weight)
@@ -139,13 +147,26 @@ def tv_denoise_fista(im, weight=50, eps=5.e-5, n_iter_max=200,
         grad_aux = (1 + t_factor) * grad_tmp - t_factor * grad_im
         grad_im = grad_tmp
         t = t_new
-        if (i % check_gap_frequency) == 0:
+        if (val_min is not None or val_max is not None or
+                                (i % check_gap_frequency) == 0):
             gap = weight * div(grad_im)
+            # Compute the primal variable
             new = im - gap
-            dgap = dual_gap(im, new, gap, weight)
-            if dgap < eps:
-                break
+            if (val_min is not None or val_max is not None):
+                new = new.clip(val_min, val_max, out=new)
+                # Now we need to recompute the dual variable
+                grad_im = gradient(new)
+            if (i % check_gap_frequency) == 0:
+                # In the case of bound constraints, the dual gap as we
+                # computed it may not go to zero.
+                dgap = dual_gap(im, new, gap, weight)
+                if dgap < eps:
+                    break
         i += 1
+    # Compute the primal variable
+    new = im - gap
+    if (val_min is not None or val_max is not None):
+        new = new.clip(val_min, val_max, out=new)
     return new