created separate module for mask processing

5 files changed, 368 insertions, 79 deletions

diff --git a/src/Core/CMakeLists.txt b/src/Core/CMakeLists.txt
index 6975a89..8aec749 100644
--- a/src/Core/CMakeLists.txt
+++ b/src/Core/CMakeLists.txt
@@ -68,12 +68,12 @@ add_library(cilreg SHARED
 	    ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/SB_TV_core.c
 	    ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/TGV_core.c
 	    ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/Diffusion_core.c
-    	    ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/DiffusionMASK_core.c
+    	    ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/MASK_merge_core.c
 	    ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/Diffus4th_order_core.c
 	    ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/LLT_ROF_core.c
-        ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/ROF_TV_core.c
-        ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/FGP_dTV_core.c
-        ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/TNV_core.c
+            ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/ROF_TV_core.c
+            ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/FGP_dTV_core.c
+            ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/TNV_core.c
             ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/Nonlocal_TV_core.c
             ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/PatchSelect_core.c
 	    ${CMAKE_CURRENT_SOURCE_DIR}/regularisers_CPU/utils.c
diff --git a/src/Core/regularisers_CPU/MASK_merge_core.c b/src/Core/regularisers_CPU/MASK_merge_core.c
new file mode 100644
index 0000000..d77c812
--- /dev/null
+++ b/src/Core/regularisers_CPU/MASK_merge_core.c
@@ -0,0 +1,262 @@
+/*
+ * This work is part of the Core Imaging Library developed by
+ * Visual Analytics and Imaging System Group of the Science Technology
+ * Facilities Council, STFC
+ *
+ * Copyright 2019 Daniil Kazantsev
+ * Copyright 2019 Srikanth Nagella, Edoardo Pasca
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "MASK_merge_core.h"
+#include "utils.h"
+
+/* A method to ensure connectivity within regions of the segmented image/volume. Here we assume 
+ * that the MASK has been obtained using some classification/segmentation method such as k-means or gaussian
+ * mixture. Some pixels/voxels have been misclassified and we check the spatial dependences
+ * and correct the mask. We check the connectivity using the bresenham line algorithm within the non-local window
+ * surrounding the pixel of interest. 
+ *
+ * Input Parameters:
+ * 1. MASK [0:255], the result of some classification algorithm (information-based preferably)
+ * 2. The list of classes (e.g. [3,4]) to apply the method. The given order matters. 
+ * 3. The total number of classes in the MASK. 
+ * 4. The size of the Correction Window inside which the method works. 
+
+ * Output:
+ * 1. MASK_upd - the UPDATED MASK where some regions have been corrected (merged) or removed
+ * 2. CORRECTEDRegions - The array of the same size as MASK where all regions which were 
+ * changed are highlighted and the changes have been counted
+ */
+
+float Mask_merge_main(unsigned char *MASK, unsigned char *MASK_upd, unsigned char *CORRECTEDRegions, unsigned char *SelClassesList, int SelClassesList_length, int classesNumb, int CorrectionWindow, int dimX, int dimY, int dimZ)
+{
+    long i,j,l;
+    int counterG, switcher;
+    long DimTotal;
+    unsigned char *MASK_temp, *ClassesList, CurrClass, temp;
+    DimTotal = (long)(dimX*dimY*dimZ);
+
+    /* defines the list for all classes in the mask */
+    ClassesList = (unsigned char*) calloc (classesNumb,sizeof(unsigned char));
+
+     /* find which classes (values) are present in the segmented data */
+     CurrClass =  MASK[0]; ClassesList[0]= MASK[0]; counterG = 1;
+     for(i=0; i<DimTotal; i++) {
+       if (MASK[i] != CurrClass) {
+          switcher = 1;
+          for(j=0; j<counterG; j++) {
+            if (ClassesList[j] == MASK[i]) {
+              switcher = 0;
+              break;
+            }}
+            if (switcher == 1) {
+                CurrClass = MASK[i];
+                ClassesList[counterG] = MASK[i];
+                /*printf("[%u]\n", ClassesList[counterG]);*/
+                counterG++;
+              }
+        }
+        if (counterG == classesNumb) break;
+      }
+      /* sort from LOW->HIGH the obtained values (classes) */
+      for(i=0; i<classesNumb; i++)	{
+                  for(j=0; j<classesNumb-1; j++) {
+                      if(ClassesList[j] > ClassesList[j+1]) {
+                          temp = ClassesList[j+1];
+                          ClassesList[j+1] = ClassesList[j];
+                          ClassesList[j] = temp;
+                      }}}
+
+    MASK_temp = (unsigned char*) calloc (DimTotal,sizeof(unsigned char));
+
+    /* copy given MASK to MASK_upd*/
+    copyIm_unchar(MASK, MASK_upd, (long)(dimX), (long)(dimY), (long)(dimZ));
+
+    if (dimZ == 1) {
+    /********************** PERFORM 2D MASK PROCESSING ************************/
+    #pragma omp parallel for shared(MASK,MASK_upd) private(i,j)
+    for(i=0; i<dimX; i++) {
+        for(j=0; j<dimY; j++) {
+    /* STEP1: in a smaller neighbourhood check that the current pixel is NOT an outlier */
+    OutiersRemoval2D(MASK, MASK_upd, i, j, (long)(dimX), (long)(dimY));
+    }}
+    /* copy the updated MASK (clean of outliers) */
+    copyIm_unchar(MASK_upd, MASK_temp, (long)(dimX), (long)(dimY), (long)(dimZ));
+
+    for(l=0; l<SelClassesList_length; l++) {
+    /*printf("[%u]\n", ClassesList[SelClassesList[l]]);*/
+    #pragma omp parallel for shared(MASK_temp,MASK_upd,l) private(i,j)
+    for(i=0; i<dimX; i++) {
+        for(j=0; j<dimY; j++) {
+      /* The class of the central pixel has not changed, i.e. the central pixel is not an outlier -> continue */
+      if (MASK_temp[j*dimX+i] == MASK[j*dimX+i]) {
+	    /* !One needs to work with a specific class to avoid overlaps! It is
+        crucial to establish relevant classes first (given as an input in SelClassesList) */
+       if (MASK_temp[j*dimX+i] == ClassesList[SelClassesList[l]]) {
+        /* i = 258; j = 165; */
+        Mask_update2D(MASK_temp, MASK_upd, CORRECTEDRegions, i, j, CorrectionWindow, (long)(dimX), (long)(dimY));
+        }}
+      }}
+      /* copy the updated mask */
+      copyIm_unchar(MASK_upd, MASK_temp, (long)(dimX), (long)(dimY), (long)(dimZ));
+      }
+    }
+    else {
+    /********************** PERFORM 3D MASK PROCESSING ************************/
+    
+    
+    }
+
+    free(MASK_temp);   
+    return 0;
+}
+
+
+/********************************************************************/
+/***************************2D Functions*****************************/
+/********************************************************************/
+float OutiersRemoval2D(unsigned char *MASK, unsigned char *MASK_upd, long i, long j, long dimX, long dimY)
+{
+  /*if the ROI pixel does not belong to the surrondings, turn it into the surronding*/
+  long i_m, j_m, i1, j1, counter;
+    counter = 0;
+    for(i_m=-1; i_m<=1; i_m++) {
+      for(j_m=-1; j_m<=1; j_m++) {
+        i1 = i+i_m;
+        j1 = j+j_m;
+        if (((i1 >= 0) && (i1 < dimX)) && ((j1 >= 0) && (j1 < dimY))) {
+          if (MASK[j*dimX+i] != MASK[j1*dimX+i1]) counter++;
+        }
+      }}
+      if (counter >= 8) MASK_upd[j*dimX+i] = MASK[j1*dimX+i1];
+      return *MASK_upd;
+}
+
+float Mask_update2D(unsigned char *MASK_temp, unsigned char *MASK_upd, unsigned char *CORRECTEDRegions, long i, long j, int CorrectionWindow, long dimX, long dimY)
+{
+  long i_m, j_m, i1, j1, CounterOtherClass;
+
+  /* STEP2: in a larger neighbourhood check that the other class is present  */
+  CounterOtherClass = 0;
+  for(i_m=-CorrectionWindow; i_m<=CorrectionWindow; i_m++) {
+      for(j_m=-CorrectionWindow; j_m<=CorrectionWindow; j_m++) {
+        i1 = i+i_m;
+        j1 = j+j_m;
+        if (((i1 >= 0) && (i1 < dimX)) && ((j1 >= 0) && (j1 < dimY))) {
+          if (MASK_temp[j*dimX+i] != MASK_temp[j1*dimX+i1]) CounterOtherClass++;
+        }
+      }}
+      if (CounterOtherClass > 0) {
+      /* the other class is present in the vicinity of CorrectionWindow, continue to STEP 3 */
+      /*
+      STEP 3: Loop through all neighbours in CorrectionWindow and check the spatial connection.
+      Meaning that we're instrested if there are any classes between points A and B that
+      does not belong to A and B (A,B \in C)
+      */
+      for(i_m=-CorrectionWindow; i_m<=CorrectionWindow; i_m++) {
+          for(j_m=-CorrectionWindow; j_m<=CorrectionWindow; j_m++) {
+            i1 = i+i_m;
+            j1 = j+j_m;
+            if (((i1 >= 0) && (i1 < dimX)) && ((j1 >= 0) && (j1 < dimY))) {
+              if (MASK_temp[j*dimX+i] == MASK_temp[j1*dimX+i1]) {
+                /* A and B points belong to the same class, do STEP 4*/
+                /* STEP 4: Run Bresenham line algorithm between A and B points
+                and convert all points on the way to the class of A/B.  */
+               bresenham2D(i, j, i1, j1, MASK_temp, MASK_upd, CORRECTEDRegions, (long)(dimX), (long)(dimY));
+              }
+            }
+          }}
+      }
+  return 0;
+}
+int bresenham2D(int i, int j, int i1, int j1, unsigned char *MASK, unsigned char *MASK_upd, unsigned char *CORRECTEDRegions, long dimX, long dimY)
+{
+                   int n;
+                   int x[] = {i, i1};
+                   int y[] = {j, j1};
+                   int steep = (fabs(y[1]-y[0]) > fabs(x[1]-x[0]));
+                   int ystep = 0;
+
+                   //printf("[%i][%i][%i][%i]\n", x[1], y[1], steep, kk) ;
+                   //if (steep == 1) {swap(x[0],y[0]); swap(x[1],y[1]);}
+
+                   if (steep == 1) {
+                   // swaping
+                   int a, b;
+
+                   a = x[0];
+                   b = y[0];
+                   x[0] = b;
+                   y[0] = a;
+
+                   a = x[1];
+                   b = y[1];
+                   x[1] = b;
+                   y[1] = a;
+                   }
+
+                   if (x[0] > x[1]) {
+                   int a, b;
+                   a = x[0];
+                   b = x[1];
+                   x[0] = b;
+                   x[1] = a;
+
+                   a = y[0];
+                   b = y[1];
+                   y[0] = b;
+                   y[1] = a;
+                   } //(x[0] > x[1])
+
+                  int delx = x[1]-x[0];
+                  int dely = fabs(y[1]-y[0]);
+                  int error = 0;
+                  int x_n = x[0];
+                  int y_n = y[0];
+                  if (y[0] < y[1]) {ystep = 1;}
+                  else {ystep = -1;}
+
+                  for(n = 0; n<delx+1; n++) {
+                       if (steep == 1) {
+                        /*printf("[%i][%i][%u]\n", x_n, y_n, MASK[y_n*dimX+x_n]);*/
+                        // MASK_upd[x_n*dimX+y_n] = 10;
+                        if (MASK[j*dimX+i] != MASK[x_n*dimX+y_n]) {
+                        	MASK_upd[x_n*dimX+y_n] = MASK[j*dimX+i];
+                        	CORRECTEDRegions[x_n*dimX+y_n] += 1;
+                        }
+                       }
+                       else {
+                        // printf("[%i][%i][%u]\n", y_n, x_n, MASK[x_n*dimX+y_n]);
+                        // MASK_upd[y_n*dimX+x_n] = 20;
+                        if (MASK[j*dimX+i] != MASK[y_n*dimX+x_n]) {
+	                        MASK_upd[y_n*dimX+x_n] = MASK[j*dimX+i];
+                              	CORRECTEDRegions[y_n*dimX+x_n] += 1;
+                        }
+                       }
+                       x_n = x_n + 1;
+                       error = error + dely;
+
+                       if (2*error >= delx) {
+                          y_n = y_n + ystep;
+                         error = error - delx;
+                       } // (2*error >= delx)
+                       //printf("[%i][%i][%i]\n", X_new[n], Y_new[n], n) ;
+                  } // for(int n = 0; n<delx+1; n++)
+                  return 0;
+}
+/********************************************************************/
+/***************************3D Functions*****************************/
+/********************************************************************/
+
+
+
diff --git a/src/Core/regularisers_CPU/MASK_merge_core.h b/src/Core/regularisers_CPU/MASK_merge_core.h
new file mode 100644
index 0000000..287f875
--- /dev/null
+++ b/src/Core/regularisers_CPU/MASK_merge_core.h
@@ -0,0 +1,56 @@
+/*
+This work is part of the Core Imaging Library developed by
+Visual Analytics and Imaging System Group of the Science Technology
+Facilities Council, STFC
+
+Copyright 2019 Daniil Kazantsev
+Copyright 2019 Srikanth Nagella, Edoardo Pasca
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include <math.h>
+#include <stdlib.h>
+#include <memory.h>
+#include <stdio.h>
+#include "omp.h"
+#include "utils.h"
+#include "CCPiDefines.h"
+
+/* A method to ensure connectivity within regions of the segmented image/volume. Here we assume 
+ * that the MASK has been obtained using some classification/segmentation method such as k-means or gaussian
+ * mixture. Some pixels/voxels have been misclassified and we check the spatial dependences
+ * and correct the mask. We check the connectivity using the bresenham line algorithm within the non-local window
+ * surrounding the pixel of interest. 
+ *
+ * Input Parameters:
+ * 1. MASK [0:255], the result of some classification algorithm (information-based preferably)
+ * 2. The list of classes (e.g. [3,4]) to apply the method. The given order matters. 
+ * 3. The total number of classes in the MASK. 
+ * 4. The size of the Correction Window inside which the method works. 
+
+ * Output:
+ * 1. MASK_upd - the UPDATED MASK where some regions have been corrected (merged) or removed
+ * 2. CORRECTEDRegions - The array of the same size as MASK where all regions which were 
+ * changed are highlighted and the changes have been counted
+ */
+
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+CCPI_EXPORT float Mask_merge_main(unsigned char *MASK, unsigned char *MASK_upd, unsigned char *CORRECTEDRegions, unsigned char *SelClassesList, int SelClassesList_length, int classesNumb, int CorrectionWindow, int dimX, int dimY, int dimZ);
+CCPI_EXPORT float OutiersRemoval2D(unsigned char *MASK, unsigned char *MASK_upd, long i, long j, long dimX, long dimY);
+CCPI_EXPORT float Mask_update2D(unsigned char *MASK_temp, unsigned char *MASK_upd, unsigned char *CORRECTEDRegions, long i, long j, int CorrectionWindow, long dimX, long dimY);
+CCPI_EXPORT int bresenham2D(int i, int j, int i1, int j1, unsigned char *MASK, unsigned char *MASK_upd, unsigned char *CORRECTEDRegions, long dimX, long dimY);
+#ifdef __cplusplus
+}
+#endif
diff --git a/src/Python/ccpi/filters/regularisers.py b/src/Python/ccpi/filters/regularisers.py
index 610907d..2fee8b3 100644
--- a/src/Python/ccpi/filters/regularisers.py
+++ b/src/Python/ccpi/filters/regularisers.py
@@ -2,7 +2,7 @@
 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, NDF_MASK_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_SB_CPU, dTV_FGP_CPU, TNV_CPU, NDF_CPU, Diff4th_CPU, TGV_CPU, LLT_ROF_CPU, PATCHSEL_CPU, NLTV_CPU, MASK_CORR_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
     gpu_enabled = True
@@ -127,37 +127,6 @@ 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, 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,
-                     iterations,
-                     time_marching_parameter,
-                     penalty_type,
-                     tolerance_param)
-    elif device == 'gpu' and gpu_enabled:
-        return NDF_MASK_CPU(inputData,
-                     maskdata,
-                     select_classes,
-                     total_classesNum,
-                     diffuswindow,
-                     regularisation_parameter,
-                     edge_parameter,
-                     iterations,
-                     time_marching_parameter,
-                     penalty_type,
-                     tolerance_param)
-    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 Diff4th(inputData, regularisation_parameter, edge_parameter, iterations,
                      time_marching_parameter, tolerance_param, device='cpu'):
     if device == 'cpu':
@@ -243,3 +212,15 @@ 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)
+    
+def MASK_CORR(maskdata, select_classes, total_classesNum, CorrectionWindow, device='cpu'):
+    if device == 'cpu':
+        return MASK_CORR_CPU(maskdata, select_classes, total_classesNum, CorrectionWindow)
+    elif device == 'gpu' and gpu_enabled:
+        return MASK_CORR_CPU(maskdata, select_classes, total_classesNum, CorrectionWindow)
+    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))
+
diff --git a/src/Python/src/cpu_regularisers.pyx b/src/Python/src/cpu_regularisers.pyx
index 78c46e2..a63ecfa 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, 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 Mask_merge_main(unsigned char *MASK, unsigned char *MASK_upd, unsigned char *CORRECTEDRegions, unsigned char *SelClassesList, int SelClassesList_length, int classesNumb, int CorrectionWindow, 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);
@@ -382,48 +382,6 @@ def NDF_3D(np.ndarray[np.float32_t, ndim=3, mode="c"] inputData,
     return (outputData,infovec)
 
 #****************************************************************#
-#********Constrained Nonlinear(Isotropic) Diffusion**************#
-#****************************************************************#
-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, 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,
-                     int iterationsNumb,
-                     float time_marching_parameter,
-                     int penalty_type,
-                     float tolerance_param):
-    cdef long dims[2]
-    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 = \
-            np.zeros([dims[0],dims[1]], dtype='float32')
-    cdef np.ndarray[np.float32_t, ndim=1, mode="c"] infovec = \
-                np.zeros([2], dtype='float32')
-
-
-    # Run constrained nonlinear diffusion iterations for 2D data
-    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)
-    return (mask_upd,outputData,infovec)
-
-#****************************************************************#
 #*************Anisotropic Fourth-Order diffusion*****************#
 #****************************************************************#
 def Diff4th_CPU(inputData, regularisation_parameter, edge_parameter, iterationsNumb, time_marching_parameter,tolerance_param):
@@ -736,6 +694,38 @@ def NVM_INP_2D(np.ndarray[np.float32_t, ndim=2, mode="c"] inputData,
 
     return (outputData, maskData_upd)
 
+##############################################################################
+#****************************************************************#
+#********Mask (segmented image) correction module **************#
+#****************************************************************#
+def MASK_CORR_CPU(maskData, select_classes, total_classesNum, CorrectionWindow):
+    if maskData.ndim == 2:
+        return MASK_CORR_CPU_2D(maskData, select_classes, total_classesNum, CorrectionWindow)
+    elif maskData.ndim == 3:
+        return 0
+
+def MASK_CORR_CPU_2D(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 CorrectionWindow):
+    
+    cdef long dims[2]
+    dims[0] = maskData.shape[0]
+    dims[1] = maskData.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.uint8_t, ndim=2, mode="c"] corr_regions = \
+            np.zeros([dims[0],dims[1]], dtype='uint8')
+
+    # Run the function to process given MASK
+    Mask_merge_main(&maskData[0,0], &mask_upd[0,0], &corr_regions[0,0], &select_classes[0], select_classes_length, 
+    total_classesNum, CorrectionWindow, dims[1], dims[0], 1)
+    return (mask_upd,corr_regions)
+
+##############################################################################
 
 #****************************************************************#
 #***************Calculation of TV-energy functional**************#