2 files changed, 57 insertions, 17 deletions
diff --git a/Wrappers/Python/src/cpu_regularisers.pyx b/Wrappers/Python/src/cpu_regularisers.pyx
index 7d57ed1..11a0617 100644
--- a/Wrappers/Python/src/cpu_regularisers.pyx
+++ b/Wrappers/Python/src/cpu_regularisers.pyx
@@ -22,7 +22,7 @@ cdef extern float TV_ROF_CPU_main(float *Input, float *Output, float lambdaPar,
 cdef extern float TV_FGP_CPU_main(float *Input, float *Output, float lambdaPar, int iterationsNumb, float epsil, int methodTV, int nonneg, int printM, int dimX, int dimY, int dimZ);
 cdef extern float SB_TV_CPU_main(float *Input, float *Output, float lambdaPar, int iterationsNumb, float epsil, int methodTV, int printM, int dimX, int dimY, int dimZ);
 cdef extern float LLT_ROF_CPU_main(float *Input, float *Output, float lambdaROF, float lambdaLLT, int iterationsNumb, float tau, int dimX, int dimY, int dimZ);
-cdef extern float TGV_main(float *Input, float *Output, float lambdaPar, float alpha1, float alpha0, int iterationsNumb, float L2, int dimX, int dimY);
+cdef extern float TGV_main(float *Input, float *Output, float lambdaPar, float alpha1, float alpha0, int iterationsNumb, float L2, int dimX, int dimY, int dimZ);
 cdef extern float Diffusion_CPU_main(float *Input, float *Output, float lambdaPar, float sigmaPar, int iterationsNumb, float tau, int penaltytype, int dimX, int dimY, int dimZ);
 cdef extern float Diffus4th_CPU_main(float *Input, float *Output, float lambdaPar, float sigmaPar, int iterationsNumb, float tau, int dimX, int dimY, int dimZ);
 cdef extern float TNV_CPU_main(float *Input, float *u, float lambdaPar, int maxIter, float tol, int dimX, int dimY, int dimZ);
@@ -202,12 +202,8 @@ def TGV_CPU(inputData, regularisation_parameter, alpha1, alpha0, iterations, Lip
         return TGV_2D(inputData, regularisation_parameter, alpha1, alpha0, 
                       iterations, LipshitzConst)
     elif inputData.ndim == 3:
-        shape = inputData.shape
-        out = inputData.copy()
-        for i in range(shape[0]):
-            out[i,:,:] = TGV_2D(inputData[i,:,:], regularisation_parameter, 
-               alpha1, alpha0, iterations, LipshitzConst)
-        return out
+        return TGV_3D(inputData, regularisation_parameter, alpha1, alpha0, 
+                      iterations, LipshitzConst)
 
 def TGV_2D(np.ndarray[np.float32_t, ndim=2, mode="c"] inputData, 
                      float regularisation_parameter,
@@ -229,7 +225,30 @@ def TGV_2D(np.ndarray[np.float32_t, ndim=2, mode="c"] inputData,
                        alpha0,
                        iterationsNumb, 
                        LipshitzConst,
-                       dims[1],dims[0])                           
+                       dims[1],dims[0],1)
+    return outputData
+def TGV_3D(np.ndarray[np.float32_t, ndim=3, mode="c"] inputData, 
+                     float regularisation_parameter,
+                     float alpha1,
+                     float alpha0,
+                     int iterationsNumb, 
+                     float LipshitzConst):
+                         
+    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')
+                   
+    #/* Run TGV iterations for 3D data */
+    TGV_main(&inputData[0,0,0], &outputData[0,0,0], regularisation_parameter, 
+                       alpha1,
+                       alpha0,
+                       iterationsNumb, 
+                       LipshitzConst,
+                       dims[2], dims[1], dims[0])
     return outputData
 
 #***************************************************************#
diff --git a/Wrappers/Python/src/gpu_regularisers.pyx b/Wrappers/Python/src/gpu_regularisers.pyx
index 47a6149..b52f669 100644
--- a/Wrappers/Python/src/gpu_regularisers.pyx
+++ b/Wrappers/Python/src/gpu_regularisers.pyx
@@ -23,7 +23,7 @@ CUDAErrorMessage = 'CUDA error'
 cdef extern int TV_ROF_GPU_main(float* Input, float* Output, float lambdaPar, int iter, float tau, int N, int M, int Z);
 cdef extern int TV_FGP_GPU_main(float *Input, float *Output, float lambdaPar, int iter, float epsil, int methodTV, int nonneg, int printM, int N, int M, int Z);
 cdef extern int TV_SB_GPU_main(float *Input, float *Output, float lambdaPar, int iter, float epsil, int methodTV, int printM, int N, int M, int Z);
-cdef extern int TGV_GPU_main(float *Input, float *Output, float lambdaPar, float alpha1, float alpha0, int iterationsNumb, float L2, int dimX, int dimY);
+cdef extern int TGV_GPU_main(float *Input, float *Output, float lambdaPar, float alpha1, float alpha0, int iterationsNumb, float L2, int dimX, int dimY, int dimZ);
 cdef extern int LLT_ROF_GPU_main(float *Input, float *Output, float lambdaROF, float lambdaLLT, int iterationsNumb, float tau, int N, int M, int Z);
 cdef extern int NonlDiff_GPU_main(float *Input, float *Output, float lambdaPar, float sigmaPar, int iterationsNumb, float tau, int penaltytype, int N, int M, int Z);
 cdef extern int dTV_FGP_GPU_main(float *Input, float *InputRef, float *Output, float lambdaPar, int iterationsNumb, float epsil, float eta, int methodTV, int nonneg, int printM, int N, int M, int Z);
@@ -102,12 +102,7 @@ def TGV_GPU(inputData, regularisation_parameter, alpha1, alpha0, iterations, Lip
     if inputData.ndim == 2:
         return TGV2D(inputData, regularisation_parameter, alpha1, alpha0, iterations, LipshitzConst)
     elif inputData.ndim == 3:
-        shape = inputData.shape
-        out = inputData.copy()
-        for i in range(shape[0]):
-            out[i,:,:] = TGV2D(inputData[i,:,:], regularisation_parameter, 
-               alpha1, alpha0, iterations, LipshitzConst)
-        return out
+        return TGV3D(inputData, regularisation_parameter, alpha1, alpha0, iterations, LipshitzConst)
 # Directional Total-variation Fast-Gradient-Projection (FGP)
 def dTV_FGP_GPU(inputData,
                      refdata,
@@ -393,7 +388,6 @@ def LLT_ROF_GPU3D(np.ndarray[np.float32_t, ndim=3, mode="c"] inputData,
         raise ValueError(CUDAErrorMessage);
 
 
-
 #***************************************************************#
 #***************** Total Generalised Variation *****************#
 #***************************************************************#
@@ -417,11 +411,38 @@ def TGV2D(np.ndarray[np.float32_t, ndim=2, mode="c"] inputData,
                        alpha0,
                        iterationsNumb, 
                        LipshitzConst,
-                       dims[1],dims[0])==0):
+                       dims[1],dims[0], 1)==0):
         return outputData
     else:
         raise ValueError(CUDAErrorMessage);
 
+def TGV3D(np.ndarray[np.float32_t, ndim=3, mode="c"] inputData, 
+                     float regularisation_parameter,
+                     float alpha1,
+                     float alpha0,
+                     int iterationsNumb, 
+                     float LipshitzConst):
+    
+    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')
+          
+    # Running CUDA code here    
+    if (TGV_GPU_main(
+            &inputData[0,0,0], &outputData[0,0,0], regularisation_parameter,
+                       alpha1,
+                       alpha0,
+                       iterationsNumb, 
+                       LipshitzConst,
+                       dims[2], dims[1], dims[0])==0):
+        return outputData;
+    else:
+        raise ValueError(CUDAErrorMessage);
+
 
 #****************************************************************#
 #**************Directional Total-variation FGP ******************#