use TestData

1 files changed, 37 insertions, 18 deletions

diff --git a/Wrappers/Python/demos/PDHG_examples/PDHG_TV_Denoising_Gaussian.py b/Wrappers/Python/demos/PDHG_examples/PDHG_TV_Denoising_Gaussian.py
index afdb6a2..7966ded 100644
--- a/Wrappers/Python/demos/PDHG_examples/PDHG_TV_Denoising_Gaussian.py
+++ b/Wrappers/Python/demos/PDHG_examples/PDHG_TV_Denoising_Gaussian.py
@@ -50,42 +50,57 @@ from ccpi.optimisation.algorithms import PDHG
 from ccpi.optimisation.operators import BlockOperator, Identity, Gradient
 from ccpi.optimisation.functions import ZeroFunction, L2NormSquared, \
                       MixedL21Norm, BlockFunction
-      
+
+from ccpi.framework import TestData
+import os, sys
+loader = TestData(data_dir=os.path.join(sys.prefix, 'share','ccpi'))
+
 # Load Data                      
-N = 100
+N = 256
+M = 300
 
 data = np.zeros((N,N))
 data[round(N/4):round(3*N/4),round(N/4):round(3*N/4)] = 0.5
 data[round(N/8):round(7*N/8),round(3*N/8):round(5*N/8)] = 1
 data = ImageData(data)
-ig = ImageGeometry(voxel_num_x = N, voxel_num_y = N)
+data = loader.load(TestData.PEPPERS, size=(N,M), scale=(0,1))
+#ig = ImageGeometry(voxel_num_x = N, voxel_num_y = N)
+print (data)
+ig = data.geometry
 ag = ig
                       
 # Create Noisy data. Add Gaussian noise
 np.random.seed(10)
-noisy_data = ImageData( data.as_array() + np.random.normal(0, 0.1, size=ig.shape) )
+noisy_data = ImageData( data.as_array() + np.random.normal(0, 0.1, size=data.shape) )
+
+print ("min {} max {}".format(data.as_array().min(), data.as_array().max()))
 
 # Show Ground Truth and Noisy Data
-plt.figure(figsize=(15,15))
-plt.subplot(2,1,1)
+plt.figure()
+plt.subplot(1,3,1)
 plt.imshow(data.as_array())
 plt.title('Ground Truth')
 plt.colorbar()
-plt.subplot(2,1,2)
+plt.subplot(1,3,2)
 plt.imshow(noisy_data.as_array())
 plt.title('Noisy Data')
 plt.colorbar()
+plt.subplot(1,3,3)
+plt.imshow((data - noisy_data).as_array())
+plt.title('diff')
+plt.colorbar()
+
 plt.show()
 
 # Regularisation Parameter
-alpha = 0.2
+alpha = 2.
 
 method = '0'
 
 if method == '0':
 
     # Create operators
-    op1 = Gradient(ig)
+    op1 = Gradient(ig, correlation='SpaceChannels')
     op2 = Identity(ig, ag)
 
     # Create BlockOperator
@@ -114,28 +129,32 @@ tau = 1/(sigma*normK**2)
 
 # Setup and Run the PDHG algorithm
 pdhg = PDHG(f=f,g=g,operator=operator, tau=tau, sigma=sigma)
-pdhg.max_iteration = 3000
-pdhg.update_objective_interval = 200
-pdhg.run(3000, verbose=False)
+pdhg.max_iteration = 10000
+pdhg.update_objective_interval = 100
+pdhg.run(1000, verbose=True)
 
 # Show Results
-plt.figure(figsize=(15,15))
-plt.subplot(3,1,1)
+plt.figure()
+plt.subplot(1,3,1)
 plt.imshow(data.as_array())
 plt.title('Ground Truth')
 plt.colorbar()
-plt.subplot(3,1,2)
+plt.clim(0,1)
+plt.subplot(1,3,2)
 plt.imshow(noisy_data.as_array())
 plt.title('Noisy Data')
 plt.colorbar()
-plt.subplot(3,1,3)
+plt.clim(0,1)
+plt.subplot(1,3,3)
 plt.imshow(pdhg.get_output().as_array())
 plt.title('TV Reconstruction')
+plt.clim(0,1)
 plt.colorbar()
 plt.show()
 
-plt.plot(np.linspace(0,N,N), data.as_array()[int(N/2),:], label = 'GTruth')
-plt.plot(np.linspace(0,N,N), pdhg.get_output().as_array()[int(N/2),:], label = 'TV reconstruction')
+plt.plot(np.linspace(0,N,M), data.as_array()[int(N/2),:], label = 'GTruth')
+plt.plot(np.linspace(0,N,M), pdhg.get_output().as_array()[int(N/2),:], label = 'TV reconstruction')
+plt.plot(np.linspace(0,N,M), noisy_data.as_array()[int(N/2),:], label = 'Noisy data')
 plt.legend()
 plt.title('Middle Line Profiles')
 plt.show()