Merge pull request #137 from vais-ral/pdtv

Adds primal-dual TV version for CPU/GPU

4 files changed, 183 insertions, 19 deletions

diff --git a/src/Python/ccpi/filters/regularisers.py b/src/Python/ccpi/filters/regularisers.py
index 0b5b2ee..5f4001a 100644
--- a/src/Python/ccpi/filters/regularisers.py
+++ b/src/Python/ccpi/filters/regularisers.py
@@ -2,9 +2,9 @@
 script which assigns a proper device core function based on a flag ('cpu' or 'gpu')
 """
 
-from ccpi.filters.cpu_regularisers import TV_ROF_CPU, TV_FGP_CPU, TV_SB_CPU, dTV_FGP_CPU, TNV_CPU, NDF_CPU, Diff4th_CPU, TGV_CPU, LLT_ROF_CPU, PATCHSEL_CPU, NLTV_CPU
+from ccpi.filters.cpu_regularisers import TV_ROF_CPU, TV_FGP_CPU, TV_PD_CPU, TV_SB_CPU, dTV_FGP_CPU, TNV_CPU, NDF_CPU, Diff4th_CPU, TGV_CPU, LLT_ROF_CPU, PATCHSEL_CPU, NLTV_CPU
 try:
-    from ccpi.filters.gpu_regularisers import TV_ROF_GPU, TV_FGP_GPU, TV_SB_GPU, dTV_FGP_GPU, NDF_GPU, Diff4th_GPU, TGV_GPU, LLT_ROF_GPU, PATCHSEL_GPU
+    from ccpi.filters.gpu_regularisers import TV_ROF_GPU, TV_FGP_GPU, TV_PD_GPU, TV_SB_GPU, dTV_FGP_GPU, NDF_GPU, Diff4th_GPU, TGV_GPU, LLT_ROF_GPU, PATCHSEL_GPU
     gpu_enabled = True
 except ImportError:
     gpu_enabled = False
@@ -51,6 +51,33 @@ def FGP_TV(inputData, regularisation_parameter,iterations,
             raise ValueError ('GPU is not available')
         raise ValueError('Unknown device {0}. Expecting gpu or cpu'\
                          .format(device))
+
+def PD_TV(inputData, regularisation_parameter, iterations,
+                     tolerance_param, methodTV, nonneg, lipschitz_const, tau, device='cpu'):
+    if device == 'cpu':
+        return TV_PD_CPU(inputData,
+                     regularisation_parameter,
+                     iterations,
+                     tolerance_param,
+                     methodTV,
+                     nonneg,
+                     lipschitz_const,
+                     tau)
+    elif device == 'gpu' and gpu_enabled:
+        return TV_PD_GPU(inputData,
+                     regularisation_parameter,
+                     iterations,
+                     tolerance_param,
+                     methodTV,
+                     nonneg,
+                     lipschitz_const,
+                     tau)
+    else:
+        if not gpu_enabled and device == 'gpu':
+            raise ValueError ('GPU is not available')
+        raise ValueError('Unknown device {0}. Expecting gpu or cpu'\
+                         .format(device))
+
 def SB_TV(inputData, regularisation_parameter, iterations,
                      tolerance_param, methodTV, device='cpu'):
     if device == 'cpu':
@@ -212,4 +239,3 @@ def NDF_INP(inputData, maskData, regularisation_parameter, edge_parameter, itera
 
 def NVM_INP(inputData, maskData, SW_increment, iterations):
         return NVM_INPAINT_CPU(inputData, maskData, SW_increment, iterations)
-
diff --git a/src/Python/setup-regularisers.py.in b/src/Python/setup-regularisers.py.in
index 4c578e3..c4ad143 100644
--- a/src/Python/setup-regularisers.py.in
+++ b/src/Python/setup-regularisers.py.in
@@ -8,13 +8,13 @@ from Cython.Distutils import build_ext
 import os
 import sys
 import numpy
-import platform	
+import platform
 
 cil_version=os.environ['CIL_VERSION']
 if  cil_version == '':
     print("Please set the environmental variable CIL_VERSION")
     sys.exit(1)
-	
+
 library_include_path = ""
 library_lib_path = ""
 try:
@@ -23,7 +23,7 @@ try:
 except:
     library_include_path = os.environ['PREFIX']+'/include'
     pass
-    
+
 extra_include_dirs = [numpy.get_include(), library_include_path]
 #extra_library_dirs = [os.path.join(library_include_path, "..", "lib")]
 extra_compile_args = []
@@ -39,6 +39,7 @@ extra_include_dirs += [os.path.join(".." , "Core"),
                        os.path.join(".." , "Core",  "inpainters_CPU"),
                        os.path.join(".." , "Core",  "regularisers_GPU" , "TV_FGP" ) ,
                        os.path.join(".." , "Core",  "regularisers_GPU" , "TV_ROF" ) ,
+                       os.path.join(".." , "Core",  "regularisers_GPU" , "TV_PD" ) ,
                        os.path.join(".." , "Core",  "regularisers_GPU" , "TV_SB" ) ,
                        os.path.join(".." , "Core",  "regularisers_GPU" , "TGV" ) ,
                        os.path.join(".." , "Core",  "regularisers_GPU" , "LLTROF" ) ,
@@ -48,12 +49,12 @@ extra_include_dirs += [os.path.join(".." , "Core"),
                        os.path.join(".." , "Core",  "regularisers_GPU" , "PatchSelect" ) ,
 						   "."]
 
-if platform.system() == 'Windows':				   
-    extra_compile_args[0:] = ['/DWIN32','/EHsc','/DBOOST_ALL_NO_LIB' , '/openmp' ]   
+if platform.system() == 'Windows':
+    extra_compile_args[0:] = ['/DWIN32','/EHsc','/DBOOST_ALL_NO_LIB' , '/openmp' ]
 else:
     extra_compile_args = ['-fopenmp','-O2', '-funsigned-char', '-Wall', '-std=c++0x']
     extra_libraries += [@EXTRA_OMP_LIB@]
-    
+
 setup(
     name='ccpi',
 	description='CCPi Core Imaging Library - Image regularisers',
@@ -61,13 +62,13 @@ setup(
     cmdclass = {'build_ext': build_ext},
     ext_modules = [Extension("ccpi.filters.cpu_regularisers",
                              sources=[os.path.join("." , "src", "cpu_regularisers.pyx" ) ],
-                             include_dirs=extra_include_dirs, 
-							 library_dirs=extra_library_dirs, 
-							 extra_compile_args=extra_compile_args, 
-							 libraries=extra_libraries ), 
-    
+                             include_dirs=extra_include_dirs,
+							 library_dirs=extra_library_dirs,
+							 extra_compile_args=extra_compile_args,
+							 libraries=extra_libraries ),
+
     ],
-	zip_safe = False,	
+	zip_safe = False,
 	packages = {'ccpi', 'ccpi.filters', 'ccpi.supp'},
 )
 
diff --git a/src/Python/src/cpu_regularisers.pyx b/src/Python/src/cpu_regularisers.pyx
index 4917d06..8de6aea 100644
--- a/src/Python/src/cpu_regularisers.pyx
+++ b/src/Python/src/cpu_regularisers.pyx
@@ -20,6 +20,7 @@ cimport numpy as np
 
 cdef extern float TV_ROF_CPU_main(float *Input, float *Output, float *infovector, float *lambdaPar, int lambda_is_arr, int iterationsNumb, float tau, float epsil, int dimX, int dimY, int dimZ);
 cdef extern float TV_FGP_CPU_main(float *Input, float *Output, float *infovector, float lambdaPar, int iterationsNumb, float epsil, int methodTV, int nonneg, int dimX, int dimY, int dimZ);
+cdef extern float PDTV_CPU_main(float *Input, float *U, float *infovector, float lambdaPar, int iterationsNumb, float epsil, float lipschitz_const, int methodTV, int nonneg, float tau, int dimX, int dimY, int dimZ);
 cdef extern float SB_TV_CPU_main(float *Input, float *Output, float *infovector, float mu, int iter, float epsil, int methodTV, int dimX, int dimY, int dimZ);
 cdef extern float LLT_ROF_CPU_main(float *Input, float *Output, float *infovector, float lambdaROF, float lambdaLLT, int iterationsNumb, float tau, float epsil, int dimX, int dimY, int dimZ);
 cdef extern float TGV_main(float *Input, float *Output, float *infovector, float lambdaPar, float alpha1, float alpha0, int iterationsNumb, float L2, float epsil, int dimX, int dimY, int dimZ);
@@ -58,9 +59,6 @@ def TV_ROF_2D(np.ndarray[np.float32_t, ndim=2, mode="c"] inputData,
     cdef np.ndarray[np.float32_t, ndim=1, mode="c"] infovec = \
             np.ones([2], dtype='float32')
 
-    # Run ROF iterations for 2D data
-    # TV_ROF_CPU_main(&inputData[0,0], &outputData[0,0], &infovec[0], regularisation_parameter, iterationsNumb, marching_step_parameter, tolerance_param, dims[1], dims[0], 1)
-     # Run ROF iterations for 2D data
     if isinstance (regularisation_parameter, np.ndarray):
         reg = regularisation_parameter.copy()
         TV_ROF_CPU_main(&inputData[0,0], &outputData[0,0], &infovec[0], &reg[0,0],  1, iterationsNumb, marching_step_parameter, tolerance_param, dims[1], dims[0], 1)
@@ -158,6 +156,75 @@ def TV_FGP_3D(np.ndarray[np.float32_t, ndim=3, mode="c"] inputData,
                        dims[2], dims[1], dims[0])
     return (outputData,infovec)
 
+#****************************************************************#
+#****************** Total-variation Primal-dual *****************#
+#****************************************************************#
+def TV_PD_CPU(inputData, regularisation_parameter, iterationsNumb, tolerance_param, methodTV, nonneg, lipschitz_const, tau):
+    if inputData.ndim == 2:
+        return TV_PD_2D(inputData, regularisation_parameter, iterationsNumb, tolerance_param, methodTV, nonneg, lipschitz_const, tau)
+    elif inputData.ndim == 3:
+        return TV_PD_3D(inputData, regularisation_parameter, iterationsNumb, tolerance_param, methodTV, nonneg, lipschitz_const, tau)
+
+def TV_PD_2D(np.ndarray[np.float32_t, ndim=2, mode="c"] inputData,
+                     float regularisation_parameter,
+                     int iterationsNumb,
+                     float tolerance_param,
+                     int methodTV,
+                     int nonneg,
+                     float lipschitz_const,
+                     float tau):
+
+    cdef long dims[2]
+    dims[0] = inputData.shape[0]
+    dims[1] = inputData.shape[1]
+
+    cdef np.ndarray[np.float32_t, ndim=2, mode="c"] outputData = \
+            np.zeros([dims[0],dims[1]], dtype='float32')
+
+    cdef np.ndarray[np.float32_t, ndim=1, mode="c"] infovec = \
+            np.ones([2], dtype='float32')
+
+    #/* Run FGP-TV iterations for 2D data */
+    PDTV_CPU_main(&inputData[0,0], &outputData[0,0], &infovec[0], regularisation_parameter,
+                       iterationsNumb,
+                       tolerance_param,
+                       lipschitz_const,
+                       methodTV,
+                       nonneg,
+                       tau,
+                       dims[1],dims[0], 1)
+    return (outputData,infovec)
+
+def TV_PD_3D(np.ndarray[np.float32_t, ndim=3, mode="c"] inputData,
+                     float regularisation_parameter,
+                     int iterationsNumb,
+                     float tolerance_param,
+                     int methodTV,
+                     int nonneg,                     
+                     float lipschitz_const,
+                     float tau):
+
+    cdef long dims[3]
+    dims[0] = inputData.shape[0]
+    dims[1] = inputData.shape[1]
+    dims[2] = inputData.shape[2]
+
+    cdef np.ndarray[np.float32_t, ndim=3, mode="c"] outputData = \
+            np.zeros([dims[0], dims[1], dims[2]], dtype='float32')
+    cdef np.ndarray[np.float32_t, ndim=1, mode="c"] infovec = \
+            np.zeros([2], dtype='float32')
+
+    #/* Run FGP-TV iterations for 3D data */
+    PDTV_CPU_main(&inputData[0,0,0], &outputData[0,0,0], &infovec[0], regularisation_parameter,
+                       iterationsNumb,
+                       tolerance_param,
+                       lipschitz_const,
+                       methodTV,
+                       nonneg,
+                       tau,
+                       dims[2], dims[1], dims[0])
+    return (outputData,infovec)
+
 #***************************************************************#
 #********************** Total-variation SB *********************#
 #***************************************************************#
diff --git a/src/Python/src/gpu_regularisers.pyx b/src/Python/src/gpu_regularisers.pyx
index 8cd8c93..b22d15e 100644
--- a/src/Python/src/gpu_regularisers.pyx
+++ b/src/Python/src/gpu_regularisers.pyx
@@ -22,6 +22,7 @@ CUDAErrorMessage = 'CUDA error'
 
 cdef extern int TV_ROF_GPU_main(float* Input, float* Output, float *infovector, float *lambdaPar, int lambda_is_arr, int iter, float tau, float epsil, int N, int M, int Z);
 cdef extern int TV_FGP_GPU_main(float *Input, float *Output, float *infovector, float lambdaPar, int iter, float epsil, int methodTV, int nonneg, int N, int M, int Z);
+cdef extern int TV_PD_GPU_main(float *Input, float *Output, float *infovector, float lambdaPar, int iter, float epsil, float lipschitz_const, int methodTV, int nonneg, float tau, int dimX, int dimY, int dimZ);
 cdef extern int TV_SB_GPU_main(float *Input, float *Output, float *infovector, float lambdaPar, int iter, float epsil, int methodTV, int N, int M, int Z);
 cdef extern int LLT_ROF_GPU_main(float *Input, float *Output, float *infovector, float lambdaROF, float lambdaLLT, int iterationsNumb, float tau,  float epsil, int N, int M, int Z);
 cdef extern int TGV_GPU_main(float *Input, float *Output, float *infovector, float lambdaPar, float alpha1, float alpha0, int iterationsNumb, float L2, float epsil, int dimX, int dimY, int dimZ);
@@ -70,6 +71,75 @@ def TV_FGP_GPU(inputData,
                      tolerance_param,
                      methodTV,
                      nonneg)
+# Total-variation Primal-Dual (PD)
+def TV_PD_GPU(inputData, regularisation_parameter, iterationsNumb, tolerance_param, methodTV, nonneg, lipschitz_const, tau):
+    if inputData.ndim == 2:
+        return TVPD2D(inputData, regularisation_parameter, iterationsNumb, tolerance_param, methodTV, nonneg, lipschitz_const, tau)
+    elif inputData.ndim == 3:
+        return TVPD3D(inputData, regularisation_parameter, iterationsNumb, tolerance_param, methodTV, nonneg, lipschitz_const, tau)
+
+def TVPD2D(np.ndarray[np.float32_t, ndim=2, mode="c"] inputData,
+                     float regularisation_parameter,
+                     int iterationsNumb,
+                     float tolerance_param,
+                     int methodTV,
+                     int nonneg,
+                     float lipschitz_const,
+                     float tau):
+
+    cdef long dims[2]
+    dims[0] = inputData.shape[0]
+    dims[1] = inputData.shape[1]
+
+    cdef np.ndarray[np.float32_t, ndim=2, mode="c"] outputData = \
+            np.zeros([dims[0],dims[1]], dtype='float32')
+
+    cdef np.ndarray[np.float32_t, ndim=1, mode="c"] infovec = \
+            np.ones([2], dtype='float32')
+
+    if (TV_PD_GPU_main(&inputData[0,0], &outputData[0,0], &infovec[0], regularisation_parameter,
+                       iterationsNumb,
+                       tolerance_param,
+                       lipschitz_const,
+                       methodTV,
+                       nonneg,
+                       tau,
+                       dims[1],dims[0], 1) ==0):
+        return (outputData,infovec)
+    else:
+        raise ValueError(CUDAErrorMessage);
+
+def TVPD3D(np.ndarray[np.float32_t, ndim=3, mode="c"] inputData,
+                     float regularisation_parameter,
+                     int iterationsNumb,
+                     float tolerance_param,
+                     int methodTV,
+                     int nonneg,
+                     float lipschitz_const,
+                     float tau):
+
+    cdef long dims[3]
+    dims[0] = inputData.shape[0]
+    dims[1] = inputData.shape[1]
+    dims[2] = inputData.shape[2]
+
+    cdef np.ndarray[np.float32_t, ndim=3, mode="c"] outputData = \
+            np.zeros([dims[0], dims[1], dims[2]], dtype='float32')
+    cdef np.ndarray[np.float32_t, ndim=1, mode="c"] infovec = \
+            np.zeros([2], dtype='float32')
+
+    if (TV_PD_GPU_main(&inputData[0,0,0], &outputData[0,0,0], &infovec[0], regularisation_parameter,
+                       iterationsNumb,
+                       tolerance_param,
+                       lipschitz_const,
+                       methodTV,
+                       nonneg,
+                       tau,
+                       dims[2], dims[1], dims[0]) ==0):
+        return (outputData,infovec)
+    else:
+        raise ValueError(CUDAErrorMessage);
+
 # Total-variation Split Bregman (SB)
 def TV_SB_GPU(inputData,
                      regularisation_parameter,
@@ -195,7 +265,7 @@ def ROFTV2D(np.ndarray[np.float32_t, ndim=2, mode="c"] inputData,
     if isinstance (regularisation_parameter, np.ndarray):
         reg = regularisation_parameter.copy()
         # Running CUDA code here
-        if (TV_ROF_GPU_main(&inputData[0,0], &outputData[0,0], &infovec[0], 
+        if (TV_ROF_GPU_main(&inputData[0,0], &outputData[0,0], &infovec[0],
                        &reg[0,0], 1,
                        iterations,
                        time_marching_parameter,