Merge pull request #11 from vais-ral/RGLTK_TV_denoising_demo

Rgltk tv denoising demo

2 files changed, 117 insertions, 53 deletions

diff --git a/Wrappers/Python/ccpi/plugins/regularisers.py b/Wrappers/Python/ccpi/plugins/regularisers.py
index e9c88a4..46464a9 100644
--- a/Wrappers/Python/ccpi/plugins/regularisers.py
+++ b/Wrappers/Python/ccpi/plugins/regularisers.py
@@ -25,7 +25,6 @@ from ccpi.optimisation.ops import Operator
 import numpy as np
 
 
-
 class _ROF_TV_(Operator):
     def __init__(self,lambdaReg,iterationsTV,tolerance,time_marchstep,device):
         # set parameters
@@ -36,7 +35,7 @@ class _ROF_TV_(Operator):
     def __call__(self,x):
         # evaluate objective function of TV gradient
         EnergyValTV = TV_ENERGY(np.asarray(x.as_array(), dtype=np.float32), np.asarray(x.as_array(), dtype=np.float32), self.lambdaReg, 2)
-        return EnergyValTV
+        return 0.5*EnergyValTV[0]
     def prox(self,x,Lipshitz):
         pars = {'algorithm' : ROF_TV, \
                'input' : np.asarray(x.as_array(), dtype=np.float32),\
@@ -63,7 +62,7 @@ class _FGP_TV_(Operator):
     def __call__(self,x):
         # evaluate objective function of TV gradient
         EnergyValTV = TV_ENERGY(np.asarray(x.as_array(), dtype=np.float32), np.asarray(x.as_array(), dtype=np.float32), self.lambdaReg, 2)
-        return EnergyValTV
+        return 0.5*EnergyValTV[0]
     def prox(self,x,Lipshitz):
         pars = {'algorithm' : FGP_TV, \
                'input' : np.asarray(x.as_array(), dtype=np.float32),\
@@ -96,7 +95,7 @@ class _SB_TV_(Operator):
     def __call__(self,x):
         # evaluate objective function of TV gradient
         EnergyValTV = TV_ENERGY(np.asarray(x.as_array(), dtype=np.float32), np.asarray(x.as_array(), dtype=np.float32), self.lambdaReg, 2)
-        return EnergyValTV
+        return 0.5*EnergyValTV[0]
     def prox(self,x,Lipshitz):
         pars = {'algorithm' : SB_TV, \
                'input' : np.asarray(x.as_array(), dtype=np.float32),\
diff --git a/Wrappers/Python/wip/demo_compare_RGLTK_TV_denoising.py b/Wrappers/Python/wip/demo_compare_RGLTK_TV_denoising.py
index 559679e..911cff4 100644
--- a/Wrappers/Python/wip/demo_compare_RGLTK_TV_denoising.py
+++ b/Wrappers/Python/wip/demo_compare_RGLTK_TV_denoising.py
@@ -1,12 +1,22 @@
 
+# 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
-
 from ccpi.optimisation.ops import LinearOperatorMatrix, Identity
 
-from ccpi.plugins.regularisers import _ROF_TV_, _FGP_TV_
+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
@@ -16,11 +26,7 @@ use_cvxpy = True
 if use_cvxpy:
     from cvxpy import *
 
-import numpy as np
-import matplotlib.pyplot as plt
-
-
-# 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
@@ -45,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()
@@ -52,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)
@@ -64,72 +70,118 @@ 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.imshow(xtv_denoise.value)
-plt.title('CVX TV')
-plt.show()
-print(objectivetv_denoise.value)
-
-
-#%% THen FBPD
+    # 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
 f_denoise = Norm2sq(I,y,c=0.5)
 
+#%%
+
+#%% 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')
+plt.title('FBPD TV with objective equal to {:.2f}'.format(criterfbpdtv_denoise[-1]))
 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=5000,tolerance=1e-5,time_marchstep=0.01,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())
-plt.title('ROF TV prox')
+EnergytotalROF = f_denoise(xtv_rof) + g_rof(xtv_rof)
+plt.title('ROF TV prox with objective equal to {:.2f}'.format(EnergytotalROF))
 plt.show()
-print(g_rof(xtv_rof))
+print(EnergytotalROF)
 
 #%% FISTA with FGP-TV regularisation
-g_fgp = _FGP_TV_(lambdaReg = lam_tv,iterationsTV=5000,tolerance=1e-5,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())
-plt.title('FGP TV prox')
+EnergytotalFGP = f_denoise(xtv_fgp) + g_fgp(xtv_fgp)
+plt.title('FGP TV prox with objective equal to {:.2f}'.format(EnergytotalFGP))
 plt.show()
-print(g_fgp(xtv_fgp))
+print(EnergytotalFGP)
+
+#%% 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())
+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)
 dlims = (-0.2,0.2)
-cols = 3
+cols = 4
 rows = 2
 current = 1
 
@@ -153,18 +205,31 @@ 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 - TV')
-imgplot = plt.imshow(xtv_rof.as_array()-xtv_denoise.value,vmin=dlims[0],vmax=dlims[1])
+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('FGP - TV')
-imgplot = plt.imshow(xtv_fgp.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')