Tidy up denoising demo.

1 files changed, 83 insertions, 50 deletions

diff --git a/Wrappers/Python/wip/demo_compare_RGLTK_TV_denoising.py b/Wrappers/Python/wip/demo_compare_RGLTK_TV_denoising.py
index fd992de..911cff4 100644
--- a/Wrappers/Python/wip/demo_compare_RGLTK_TV_denoising.py
+++ b/Wrappers/Python/wip/demo_compare_RGLTK_TV_denoising.py
@@ -1,4 +1,10 @@
 
+# This demo illustrates how the CCPi Regularisation Toolkit can be used 
+# as TV denoising for use with the FISTA algorithm of the modular 
+# optimisation framework and compares with the FBPD TV implementation as well
+# as CVXPY.
+
+# All own imports
 from ccpi.framework import ImageData, ImageGeometry, AcquisitionGeometry, DataContainer
 from ccpi.optimisation.algs import FISTA, FBPD, CGLS
 from ccpi.optimisation.funcs import Norm2sq, ZeroFun, Norm1, TV2D
@@ -6,12 +12,15 @@ from ccpi.optimisation.ops import LinearOperatorMatrix, Identity
 
 from ccpi.plugins.regularisers import _ROF_TV_, _FGP_TV_, _SB_TV_
 
+# All external imports
 import numpy as np
 import matplotlib.pyplot as plt
+
 #%%
 # Requires CVXPY, see http://www.cvxpy.org/
 # CVXPY can be installed in anaconda using
 # conda install -c cvxgrp cvxpy libgcc
+
 # Whether to use or omit CVXPY
 use_cvxpy = True
 if use_cvxpy:
@@ -19,8 +28,6 @@ if use_cvxpy:
 
 #%%
 
-# Now try 1-norm and TV denoising with FBPD, first 1-norm.
-
 # Set up phantom size NxN by creating ImageGeometry, initialising the 
 # ImageData object with this geometry and empty array and finally put some
 # data into its array, and display as image.
@@ -44,6 +51,7 @@ y = I.direct(Phantom)
 np.random.seed(0)
 y.array = y.array + 0.1*np.random.randn(N, N)
 
+# Display noisy image
 plt.imshow(y.array)
 plt.title('Noisy image')
 plt.show()
@@ -51,9 +59,8 @@ plt.show()
 #%% TV parameter
 lam_tv = 1.0
 
-#%% Do CVX as high quality ground truth
+#%% Do CVX as high quality ground truth for comparison.
 if use_cvxpy:
-    # Compare to CVXPY
     
     # Construct the problem.
     xtv_denoise = Variable(N,N)
@@ -63,17 +70,15 @@ if use_cvxpy:
     # The optimal objective is returned by prob.solve().
     resulttv_denoise = probtv_denoise.solve(verbose=False,solver=SCS,eps=1e-12)
     
-    # The optimal solution for x is stored in x.value and optimal objective value 
-    # is in result as well as in objective.value
-    print("CVXPY least squares plus TV solution and objective value:")
-    # print(xtv_denoise.value)
-    # print(objectivetv_denoise.value)
+    # The optimal solution for x is stored in x.value and optimal objective 
+    # value is in result as well as in objective.value
     
-plt.figure()
-plt.imshow(xtv_denoise.value)
-plt.title('CVX TV  with objective equal to {:.2f}'.format(objectivetv_denoise.value))
-plt.show()
-print(objectivetv_denoise.value)
+    # Display
+    plt.figure()
+    plt.imshow(xtv_denoise.value)
+    plt.title('CVX TV  with objective equal to {:.2f}'.format(objectivetv_denoise.value))
+    plt.show()
+    print(objectivetv_denoise.value)
 
 #%%
 # Data fidelity term
@@ -81,19 +86,22 @@ f_denoise = Norm2sq(I,y,c=0.5)
 
 #%%
 
-#%% THen FBPD
+#%% Then run FBPD algorithm for TV  denoising
+
 # Initial guess
 x_init_denoise = ImageData(np.zeros((N,N)))
 
+# Set up TV function
 gtv = TV2D(lam_tv)
-gtv(gtv.op.direct(x_init_denoise))
 
-opt_tv = {'tol': 1e-4, 'iter': 10000}
+# Evalutate TV of noisy image.
+gtv(gtv.op.direct(y))
 
+# Specify FBPD options and run FBPD.
+opt_tv = {'tol': 1e-4, 'iter': 10000}
 x_fbpdtv_denoise, itfbpdtv_denoise, timingfbpdtv_denoise, criterfbpdtv_denoise = FBPD(x_init_denoise, None, f_denoise, gtv,opt=opt_tv)
 
-
-print("CVXPY least squares plus TV solution and objective value:")
+print("FBPD least squares plus TV solution and objective value:")
 plt.figure()
 plt.imshow(x_fbpdtv_denoise.as_array())
 plt.title('FBPD TV with objective equal to {:.2f}'.format(criterfbpdtv_denoise[-1]))
@@ -101,16 +109,24 @@ plt.show()
 
 print(criterfbpdtv_denoise[-1])
 
-#%%
-plt.loglog([0,opt_tv['iter']], [objectivetv_denoise.value,objectivetv_denoise.value], label='CVX TV')
+# Also plot history of criterion vs. CVX
+if use_cvxpy:
+    plt.loglog([0,opt_tv['iter']], [objectivetv_denoise.value,objectivetv_denoise.value], label='CVX TV')
 plt.loglog(criterfbpdtv_denoise, label='FBPD TV')
+plt.legend()
 plt.show()
 
 #%% FISTA with ROF-TV regularisation
-g_rof = _ROF_TV_(lambdaReg = lam_tv,iterationsTV=2000,tolerance=0,time_marchstep=0.0009,device='cpu')
+g_rof = _ROF_TV_(lambdaReg = lam_tv,
+                 iterationsTV=2000,
+                 tolerance=0,
+                 time_marchstep=0.0009,
+                 device='cpu')
 
+# Evaluating the proximal operator corresponds to denoising.
 xtv_rof = g_rof.prox(y,1.0)
 
+# Display denoised image and final criterion value.
 print("CCPi-RGL TV ROF:")
 plt.figure()
 plt.imshow(xtv_rof.as_array())
@@ -120,10 +136,18 @@ plt.show()
 print(EnergytotalROF)
 
 #%% FISTA with FGP-TV regularisation
-g_fgp = _FGP_TV_(lambdaReg = lam_tv,iterationsTV=5000,tolerance=0,methodTV=0,nonnegativity=0,printing=0,device='cpu')
-
+g_fgp = _FGP_TV_(lambdaReg = lam_tv,
+                 iterationsTV=5000,
+                 tolerance=0,
+                 methodTV=0,
+                 nonnegativity=0,
+                 printing=0,
+                 device='cpu')
+
+# Evaluating the proximal operator corresponds to denoising.
 xtv_fgp = g_fgp.prox(y,1.0)
 
+# Display denoised image and final criterion value.
 print("CCPi-RGL TV FGP:")
 plt.figure()
 plt.imshow(xtv_fgp.as_array())
@@ -131,11 +155,19 @@ EnergytotalFGP = f_denoise(xtv_fgp) + g_fgp(xtv_fgp)
 plt.title('FGP TV prox with objective equal to {:.2f}'.format(EnergytotalFGP))
 plt.show()
 print(EnergytotalFGP)
-#%% Split-Bregman-TV regularisation
-g_sb = _SB_TV_(lambdaReg = lam_tv,iterationsTV=1000,tolerance=0,methodTV=0,printing=0,device='cpu')
 
+#%% Split-Bregman-TV regularisation
+g_sb = _SB_TV_(lambdaReg = lam_tv,
+               iterationsTV=1000,
+               tolerance=0,
+               methodTV=0,
+               printing=0,
+               device='cpu')
+
+# Evaluating the proximal operator corresponds to denoising.
 xtv_sb = g_sb.prox(y,1.0)
 
+# Display denoised image and final criterion value.
 print("CCPi-RGL TV SB:")
 plt.figure()
 plt.imshow(xtv_sb.as_array())
@@ -143,8 +175,8 @@ EnergytotalSB = f_denoise(xtv_sb) + g_fgp(xtv_sb)
 plt.title('SB TV prox with objective equal to {:.2f}'.format(EnergytotalSB))
 plt.show()
 print(EnergytotalSB)
-#%%
 
+#%%
 
 # Compare all reconstruction
 clims = (-0.2,1.2)
@@ -177,26 +209,27 @@ a.set_title('SB')
 imgplot = plt.imshow(xtv_sb.as_array(),vmin=clims[0],vmax=clims[1])
 plt.axis('off')
 
-current = current + 1
-a=fig.add_subplot(rows,cols,current)
-a.set_title('FBPD - CVX')
-imgplot = plt.imshow(x_fbpdtv_denoise.as_array()-xtv_denoise.value,vmin=dlims[0],vmax=dlims[1])
-plt.axis('off')
-
-current = current + 1
-a=fig.add_subplot(rows,cols,current)
-a.set_title('ROF - CVX')
-imgplot = plt.imshow(xtv_rof.as_array()-xtv_denoise.value,vmin=dlims[0],vmax=dlims[1])
-plt.axis('off')
-
-current = current + 1
-a=fig.add_subplot(rows,cols,current)
-a.set_title('FGP - CVX')
-imgplot = plt.imshow(xtv_fgp.as_array()-xtv_denoise.value,vmin=dlims[0],vmax=dlims[1])
-plt.axis('off')
-
-current = current + 1
-a=fig.add_subplot(rows,cols,current)
-a.set_title('SB - CVX')
-imgplot = plt.imshow(xtv_sb.as_array()-xtv_denoise.value,vmin=dlims[0],vmax=dlims[1])
-plt.axis('off')
+if use_cvxpy: 
+    current = current + 1
+    a=fig.add_subplot(rows,cols,current)
+    a.set_title('FBPD - CVX')
+    imgplot = plt.imshow(x_fbpdtv_denoise.as_array()-xtv_denoise.value,vmin=dlims[0],vmax=dlims[1])
+    plt.axis('off')
+    
+    current = current + 1
+    a=fig.add_subplot(rows,cols,current)
+    a.set_title('ROF - CVX')
+    imgplot = plt.imshow(xtv_rof.as_array()-xtv_denoise.value,vmin=dlims[0],vmax=dlims[1])
+    plt.axis('off')
+    
+    current = current + 1
+    a=fig.add_subplot(rows,cols,current)
+    a.set_title('FGP - CVX')
+    imgplot = plt.imshow(xtv_fgp.as_array()-xtv_denoise.value,vmin=dlims[0],vmax=dlims[1])
+    plt.axis('off')
+    
+    current = current + 1
+    a=fig.add_subplot(rows,cols,current)
+    a.set_title('SB - CVX')
+    imgplot = plt.imshow(xtv_sb.as_array()-xtv_denoise.value,vmin=dlims[0],vmax=dlims[1])
+    plt.axis('off')