further progress

2 files changed, 14 insertions, 7 deletions

diff --git a/src/Python/ccpi/filters/regularisers.py b/src/Python/ccpi/filters/regularisers.py
index 00afcc2..610907d 100644
--- a/src/Python/ccpi/filters/regularisers.py
+++ b/src/Python/ccpi/filters/regularisers.py
@@ -127,11 +127,13 @@ def NDF(inputData, regularisation_parameter, edge_parameter, iterations,
     	    raise ValueError ('GPU is not available')
         raise ValueError('Unknown device {0}. Expecting gpu or cpu'\
                          .format(device))
-def NDF_MASK(inputData, maskdata, diffuswindow, regularisation_parameter, edge_parameter, iterations,
+def NDF_MASK(inputData, maskdata, select_classes, total_classesNum, diffuswindow, regularisation_parameter, edge_parameter, iterations,
                      time_marching_parameter, penalty_type, tolerance_param, device='cpu'):
     if device == 'cpu':
         return NDF_MASK_CPU(inputData,
                      maskdata,
+                     select_classes,
+                     total_classesNum,
                      diffuswindow,
                      regularisation_parameter,
                      edge_parameter,
@@ -142,6 +144,8 @@ def NDF_MASK(inputData, maskdata, diffuswindow, regularisation_parameter, edge_p
     elif device == 'gpu' and gpu_enabled:
         return NDF_MASK_CPU(inputData,
                      maskdata,
+                     select_classes,
+                     total_classesNum,
                      diffuswindow,
                      regularisation_parameter,
                      edge_parameter,
diff --git a/src/Python/src/cpu_regularisers.pyx b/src/Python/src/cpu_regularisers.pyx
index ca402c1..78c46e2 100644
--- a/src/Python/src/cpu_regularisers.pyx
+++ b/src/Python/src/cpu_regularisers.pyx
@@ -24,7 +24,7 @@ cdef extern float SB_TV_CPU_main(float *Input, float *Output, float *infovector,
 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);
 cdef extern float Diffusion_CPU_main(float *Input, float *Output, float *infovector, float lambdaPar, float sigmaPar, int iterationsNumb, float tau, int penaltytype, float epsil, int dimX, int dimY, int dimZ);
-cdef extern float DiffusionMASK_CPU_main(float *Input, unsigned char *MASK, unsigned char *MASK_upd, float *Output, float *infovector, int DiffusWindow, float lambdaPar, float sigmaPar, int iterationsNumb, float tau, int penaltytype, float epsil, int dimX, int dimY, int dimZ);
+cdef extern float DiffusionMASK_CPU_main(float *Input, unsigned char *MASK, unsigned char *MASK_upd, unsigned char *SelClassesList, int SelClassesList_length, float *Output, float *infovector, int classesNumb, int DiffusWindow, float lambdaPar, float sigmaPar, int iterationsNumb, float tau, int penaltytype, float epsil, int dimX, int dimY, int dimZ);
 cdef extern float Diffus4th_CPU_main(float *Input, float *Output,  float *infovector, float lambdaPar, float sigmaPar, int iterationsNumb, float tau, float epsil, int dimX, int dimY, int dimZ);
 cdef extern float dTV_FGP_CPU_main(float *Input, float *InputRef, float *Output, float *infovector, float lambdaPar, int iterationsNumb, float epsil, float eta, int methodTV, int nonneg, 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);
@@ -384,14 +384,16 @@ def NDF_3D(np.ndarray[np.float32_t, ndim=3, mode="c"] inputData,
 #****************************************************************#
 #********Constrained Nonlinear(Isotropic) Diffusion**************#
 #****************************************************************#
-def NDF_MASK_CPU(inputData, maskData, diffuswindow, regularisation_parameter, edge_parameter, iterationsNumb,time_marching_parameter, penalty_type, tolerance_param):
+def NDF_MASK_CPU(inputData, maskData, select_classes, total_classesNum, diffuswindow, regularisation_parameter, edge_parameter, iterationsNumb,time_marching_parameter, penalty_type, tolerance_param):
     if inputData.ndim == 2:
-        return NDF_MASK_2D(inputData, maskData, diffuswindow, regularisation_parameter, edge_parameter, iterationsNumb, time_marching_parameter, penalty_type, tolerance_param)
+        return NDF_MASK_2D(inputData, maskData, select_classes, total_classesNum, diffuswindow, regularisation_parameter, edge_parameter, iterationsNumb, time_marching_parameter, penalty_type, tolerance_param)
     elif inputData.ndim == 3:
         return 0
 
 def NDF_MASK_2D(np.ndarray[np.float32_t, ndim=2, mode="c"] inputData,
                     np.ndarray[np.uint8_t, ndim=2, mode="c"] maskData,
+                    np.ndarray[np.uint8_t, ndim=1, mode="c"] select_classes,
+                     int total_classesNum,
                      int diffuswindow,
                      float regularisation_parameter,
                      float edge_parameter,
@@ -403,7 +405,8 @@ def NDF_MASK_2D(np.ndarray[np.float32_t, ndim=2, mode="c"] inputData,
     dims[0] = inputData.shape[0]
     dims[1] = inputData.shape[1]
 
-
+    select_classes_length = select_classes.shape[0]
+    
     cdef np.ndarray[np.uint8_t, ndim=2, mode="c"] mask_upd = \
             np.zeros([dims[0],dims[1]], dtype='uint8')
     cdef np.ndarray[np.float32_t, ndim=2, mode="c"] outputData = \
@@ -413,8 +416,8 @@ def NDF_MASK_2D(np.ndarray[np.float32_t, ndim=2, mode="c"] inputData,
 
 
     # Run constrained nonlinear diffusion iterations for 2D data
-    DiffusionMASK_CPU_main(&inputData[0,0], &maskData[0,0], &mask_upd[0,0], &outputData[0,0], &infovec[0],
-    diffuswindow, regularisation_parameter, edge_parameter, iterationsNumb,
+    DiffusionMASK_CPU_main(&inputData[0,0], &maskData[0,0], &mask_upd[0,0], &select_classes[0], select_classes_length, &outputData[0,0], &infovec[0],
+    total_classesNum, diffuswindow, regularisation_parameter, edge_parameter, iterationsNumb,
     time_marching_parameter, penalty_type,
     tolerance_param,
     dims[1], dims[0], 1)