fixes all matlab issues

8 files changed, 150 insertions, 335 deletions

diff --git a/src/Core/regularisers_CPU/MASK_merge_core.c b/src/Core/regularisers_CPU/MASK_merge_core.c
deleted file mode 100644
index d77c812..0000000
--- a/src/Core/regularisers_CPU/MASK_merge_core.c
+++ /dev/null
@@ -1,262 +0,0 @@
-/*
- * 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
deleted file mode 100644
index 287f875..0000000
--- a/src/Core/regularisers_CPU/MASK_merge_core.h
+++ /dev/null
@@ -1,56 +0,0 @@
-/*
-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/Matlab/mex_compile/compileCPU_mex_Linux.m b/src/Matlab/mex_compile/compileCPU_mex_Linux.m
index 19fb1a5..2ed7ea6 100644
--- a/src/Matlab/mex_compile/compileCPU_mex_Linux.m
+++ b/src/Matlab/mex_compile/compileCPU_mex_Linux.m
@@ -77,5 +77,5 @@ delete PatchSelect_core* Nonlocal_TV_core*
 delete Diffusion_Inpaint_core* NonlocalMarching_Inpaint_core*
 fprintf('%s \n', '<<<<<<< Regularisers successfully compiled! >>>>>>>');
 
-pathA2 = sprintf(['..' fsep '..' fsep '..' fsep '..' fsep 'demos'], 1i);
+pathA2 = sprintf(['..' fsep '..' fsep '..' fsep '..' fsep 'demos' fsep 'Matlab_demos'], 1i);
 cd(pathA2);
\ No newline at end of file
diff --git a/src/Matlab/mex_compile/compileCPU_mex_WINDOWS.m b/src/Matlab/mex_compile/compileCPU_mex_WINDOWS.m
index 3a9e2af..7d0917f 100644
--- a/src/Matlab/mex_compile/compileCPU_mex_WINDOWS.m
+++ b/src/Matlab/mex_compile/compileCPU_mex_WINDOWS.m
@@ -7,7 +7,11 @@
 % Here I present two ways how software can be compiled, if you have some
 % other suggestions/remarks please contact me at dkazanc@hotmail.com 
 % >>>>>>>>>>>>>>>>>>>>>>>>>>>>>
-
+% cuda related messgaes can be  solved in Linux:
+%export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/lib64"
+%export CUDA_HOME=/usr/local/cuda
+%export LD_PRELOAD=/home/algol/anaconda3/lib/libstdc++.so.6.0.24
+% >>
 fsep = '/';
 
 pathcopyFrom = sprintf(['..' fsep '..' fsep 'Core' fsep 'regularisers_CPU'], 1i);
diff --git a/src/Matlab/mex_compile/compileGPU_mex.m b/src/Matlab/mex_compile/compileGPU_mex.m
index 7e15233..56fcd38 100644
--- a/src/Matlab/mex_compile/compileGPU_mex.m
+++ b/src/Matlab/mex_compile/compileGPU_mex.m
@@ -66,10 +66,14 @@ fprintf('%s \n', 'Compiling ROF-LLT...');
 mex -g -I/usr/local/cuda-10.0/include -L/usr/local/cuda-10.0/lib64 -lcudart -lcufft -lmwgpu LLT_ROF_GPU.cpp LLT_ROF_GPU_core.o
 movefile('LLT_ROF_GPU.mex*',Pathmove);
 
+fprintf('%s \n', 'Compiling PatchSelect...');
+!/usr/local/cuda/bin/nvcc -O0 -c PatchSelect_GPU_core.cu -Xcompiler -fPIC -I~/SOFT/MATLAB9/extern/include/
+mex -g -I/usr/local/cuda-10.0/include -L/usr/local/cuda-10.0/lib64 -lcudart -lcufft -lmwgpu PatchSelect_GPU.cpp PatchSelect_GPU_core.o
+movefile('PatchSelect_GPU.mex*',Pathmove);
 
 delete TV_ROF_GPU_core* TV_FGP_GPU_core* TV_SB_GPU_core* dTV_FGP_GPU_core* NonlDiff_GPU_core* Diffus_4thO_GPU_core* TGV_GPU_core* LLT_ROF_GPU_core* CCPiDefines.h
-delete PatchSelect_core* Nonlocal_TV_core* shared.h
+delete PatchSelect_GPU_core* Nonlocal_TV_core* shared.h
 fprintf('%s \n', 'All successfully compiled!');
 
-pathA2 = sprintf(['..' fsep '..' fsep '..' fsep '..' fsep 'demos'], 1i);
+pathA2 = sprintf(['..' fsep '..' fsep '..' fsep '..' fsep 'demos' fsep 'Matlab_demos'], 1i);
 cd(pathA2);
diff --git a/src/Matlab/mex_compile/regularisers_CPU/ROF_TV.c b/src/Matlab/mex_compile/regularisers_CPU/ROF_TV.c
index ffe7b91..d6bdaa6 100644
--- a/src/Matlab/mex_compile/regularisers_CPU/ROF_TV.c
+++ b/src/Matlab/mex_compile/regularisers_CPU/ROF_TV.c
@@ -26,7 +26,7 @@
  * 
  * Input Parameters:
  * 1. Noisy image/volume [REQUIRED]
- * 2. lambda - regularization parameter [REQUIRED]
+ * 2. lambda - regularization parameter [REQUIRED] scalar or the same size as 1
  * 3. Number of iterations, for explicit scheme >= 150 is recommended  [REQUIRED]
  * 4. tau - marching step for explicit scheme, ~1 is recommended [REQUIRED]
  * 5. eplsilon: tolerance constant [REQUIRED]
@@ -38,7 +38,7 @@
  * This function is based on the paper by
  * [1] Rudin, Osher, Fatemi, "Nonlinear Total Variation based noise removal algorithms"
  *
- * D. Kazantsev, 2016-18
+ * D. Kazantsev, 2016-19
  */
 
 void mexFunction(
@@ -49,18 +49,29 @@ void mexFunction(
     int number_of_dims, iter_numb;
     mwSize dimX, dimY, dimZ;
     const mwSize *dim_array_i;
-    float *Input, *Output=NULL, lambda, tau, epsil;    
+    float *Input, *Output=NULL, lambda_scalar, tau, epsil;        
     float *infovec=NULL;
+    float *lambda;      
     
     dim_array_i = mxGetDimensions(prhs[0]);
     number_of_dims = mxGetNumberOfDimensions(prhs[0]);
     
     /*Handling Matlab input data*/
     Input  = (float *) mxGetData(prhs[0]);
-    lambda =  (float) mxGetScalar(prhs[1]); /* regularization parameter */
+    /*special check to find the input scalar or an array*/
+    int mrows = mxGetM(prhs[1]);
+    int ncols = mxGetN(prhs[1]); 
+    if (mrows==1 && ncols==1) {        
+        lambda = (float*) calloc (1 ,sizeof(float));
+        lambda_scalar =  (float) mxGetScalar(prhs[1]); /* regularization parameter */        
+        lambda[0] = lambda_scalar;
+    }
+    else {
+        lambda =  (float *) mxGetData(prhs[1]); /* regularization parameter */    
+    }
     iter_numb =  (int) mxGetScalar(prhs[2]); /* iterations number */
     tau =  (float) mxGetScalar(prhs[3]); /* marching step parameter */  
-    epsil = (float) mxGetScalar(prhs[4]); /* tolerance */  
+    epsil = (float) mxGetScalar(prhs[4]); /* tolerance */    
     
     if (mxGetClassID(prhs[0]) != mxSINGLE_CLASS) {mexErrMsgTxt("The input image must be in a single precision"); }
     if(nrhs != 5) mexErrMsgTxt("Four inputs reqired: Image(2D,3D), regularization parameter, iterations number,  marching step constant, tolerance");
@@ -80,6 +91,11 @@ void mexFunction(
     mwSize vecdim[1];
     vecdim[0] = 2;
     infovec = (float*)mxGetPr(plhs[1] = mxCreateNumericArray(1, vecdim, mxSINGLE_CLASS, mxREAL));
-     
-    TV_ROF_CPU_main(Input, Output, infovec, lambda, iter_numb, tau, epsil, dimX, dimY, dimZ);    
+    
+    if (mrows==1 && ncols==1) {
+    TV_ROF_CPU_main(Input, Output, infovec, lambda, 0, iter_numb, tau, epsil, dimX, dimY, dimZ);    
+    free(lambda);
+    }
+    else TV_ROF_CPU_main(Input, Output, infovec, lambda, 1, iter_numb, tau, epsil, dimX, dimY, dimZ);     
+        
 }
diff --git a/src/Matlab/mex_compile/regularisers_GPU/PatchSelect_GPU.cpp b/src/Matlab/mex_compile/regularisers_GPU/PatchSelect_GPU.cpp
new file mode 100644
index 0000000..13319fa
--- /dev/null
+++ b/src/Matlab/mex_compile/regularisers_GPU/PatchSelect_GPU.cpp
@@ -0,0 +1,94 @@
+/*
+ * This work is part of the Core Imaging Library developed by
+ * Visual Analytics and Imaging System Group of the Science Technology
+ * Facilities Council, STFC and Diamond Light Source Ltd. 
+ *
+ * Copyright 2017 Daniil Kazantsev
+ * Copyright 2017 Srikanth Nagella, Edoardo Pasca
+ * Copyright 2018 Diamond Light Source Ltd. 
+ *
+ * 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 "matrix.h"
+#include "mex.h"
+#include "PatchSelect_GPU_core.h"
+
+/* CUDA implementation of non-local weight pre-calculation for non-local priors
+ * Weights and associated indices are stored into pre-allocated arrays and passed
+ * to the regulariser
+ *
+ *
+ * Input Parameters:
+ * 1. 2D/3D grayscale image/volume
+ * 2. Searching window (half-size of the main bigger searching window, e.g. 11)
+ * 3. Similarity window (half-size of the patch window, e.g. 2)
+ * 4. The number of neighbours to take (the most prominent after sorting neighbours will be taken)
+ * 5. noise-related parameter to calculate non-local weights
+ *
+ * Output [2D]:
+ * 1. AR_i - indeces of i neighbours
+ * 2. AR_j - indeces of j neighbours
+ * 3. Weights_ij - associated weights
+ *
+ * Output [3D]:
+ * 1. AR_i - indeces of i neighbours
+ * 2. AR_j - indeces of j neighbours
+ * 3. AR_k - indeces of j neighbours
+ * 4. Weights_ijk - associated weights
+ */
+/**************************************************/
+void mexFunction(
+        int nlhs, mxArray *plhs[],
+        int nrhs, const mxArray *prhs[])
+{
+    int number_of_dims,  SearchWindow, SimilarWin, NumNeighb;
+    mwSize dimX, dimY, dimZ;
+    const mwSize *dim_array;
+    unsigned short *H_i=NULL, *H_j=NULL, *H_k=NULL;    
+    float *A, *Weights = NULL, h;
+    mwSize dim_array2[3]; /* for 2D data */
+    mwSize dim_array3[4]; /* for 3D data */
+    
+    dim_array = mxGetDimensions(prhs[0]);
+    number_of_dims = mxGetNumberOfDimensions(prhs[0]);
+    
+    /*Handling Matlab input data*/
+    A  = (float *) mxGetData(prhs[0]); /* a 2D or 3D image/volume */
+    SearchWindow = (int) mxGetScalar(prhs[1]);    /* Large Searching window */
+    SimilarWin = (int) mxGetScalar(prhs[2]);    /* Similarity window (patch-search)*/
+    NumNeighb = (int) mxGetScalar(prhs[3]); /* the total number of neighbours to take */
+    h = (float) mxGetScalar(prhs[4]); /* NLM parameter */
+
+    dimX = dim_array[0]; dimY = dim_array[1]; dimZ = dim_array[2];
+    dim_array2[0] = dimX; dim_array2[1] = dimY; dim_array2[2] = NumNeighb;  /* 2D case */
+    dim_array3[0] = dimX; dim_array3[1] = dimY; dim_array3[2] = dimZ; dim_array3[3] = NumNeighb;  /* 3D case */
+    
+    /****************2D INPUT ***************/
+    if (number_of_dims == 2) {
+        dimZ = 0;               
+        H_i = (unsigned short*)mxGetPr(plhs[0] = mxCreateNumericArray(3, dim_array2, mxUINT16_CLASS, mxREAL));
+        H_j = (unsigned short*)mxGetPr(plhs[1] = mxCreateNumericArray(3, dim_array2, mxUINT16_CLASS, mxREAL));
+        Weights = (float*)mxGetPr(plhs[2] = mxCreateNumericArray(3, dim_array2, mxSINGLE_CLASS, mxREAL));
+        }
+    /****************3D INPUT ***************/
+    if (number_of_dims == 3) { 
+        /*
+        H_i = (unsigned short*)mxGetPr(plhs[0] = mxCreateNumericArray(4, dim_array3, mxUINT16_CLASS, mxREAL));
+        H_j = (unsigned short*)mxGetPr(plhs[1] = mxCreateNumericArray(4, dim_array3, mxUINT16_CLASS, mxREAL));
+        H_k = (unsigned short*)mxGetPr(plhs[2] = mxCreateNumericArray(4, dim_array3, mxUINT16_CLASS, mxREAL));
+        Weights = (float*)mxGetPr(plhs[3] = mxCreateNumericArray(4, dim_array3, mxSINGLE_CLASS, mxREAL));        
+        */
+        mexErrMsgTxt("3D version is not yet supported");
+    }
+    
+    PatchSelect_GPU_main(A, H_i, H_j, Weights, (long)(dimX), (long)(dimY), SearchWindow, SimilarWin, NumNeighb, h);         
+ }
diff --git a/src/Matlab/mex_compile/regularisers_GPU/ROF_TV_GPU.cpp b/src/Matlab/mex_compile/regularisers_GPU/ROF_TV_GPU.cpp
index d9b7e83..2cd4469 100644
--- a/src/Matlab/mex_compile/regularisers_GPU/ROF_TV_GPU.cpp
+++ b/src/Matlab/mex_compile/regularisers_GPU/ROF_TV_GPU.cpp
@@ -25,7 +25,7 @@
  * 
  * Input Parameters:
  * 1. Noisy image/volume [REQUIRED]
- * 2. lambda - regularization parameter [REQUIRED]
+ * 2. lambda - regularization parameter [REQUIRED], scalar or the same size as 1
  * 3. Number of iterations, for explicit scheme >= 150 is recommended  [REQUIRED]
  * 4. tau - marching step for explicit scheme, ~1 is recommended [REQUIRED]
  * 5. eplsilon: tolerance constant [REQUIRED]
@@ -37,7 +37,7 @@
  * This function is based on the paper by
  * [1] Rudin, Osher, Fatemi, "Nonlinear Total Variation based noise removal algorithms"
  *
- * D. Kazantsev, 2016-18
+ * D. Kazantsev, 2016-19
  */
 void mexFunction(
         int nlhs, mxArray *plhs[],
@@ -47,15 +47,26 @@ void mexFunction(
      int number_of_dims, iter_numb;
     mwSize dimX, dimY, dimZ;
     const mwSize *dim_array_i;
-    float *Input, *Output=NULL, lambda, tau, epsil;    
+    float *Input, *Output=NULL, lambda_scalar, tau, epsil;    
     float *infovec=NULL;
-    
+    float *lambda;  
+        
     dim_array_i = mxGetDimensions(prhs[0]);
     number_of_dims = mxGetNumberOfDimensions(prhs[0]);
     
     /*Handling Matlab input data*/
     Input  = (float *) mxGetData(prhs[0]);
-    lambda =  (float) mxGetScalar(prhs[1]); /* regularization parameter */
+    /*special check to find the input scalar or an array*/
+    int mrows = mxGetM(prhs[1]);
+    int ncols = mxGetN(prhs[1]); 
+    if (mrows==1 && ncols==1) {        
+        lambda = (float*) calloc (1 ,sizeof(float));
+        lambda_scalar =  (float) mxGetScalar(prhs[1]); /* regularization parameter */        
+        lambda[0] = lambda_scalar;
+    }
+    else {
+        lambda =  (float *) mxGetData(prhs[1]); /* regularization parameter */    
+    }
     iter_numb =  (int) mxGetScalar(prhs[2]); /* iterations number */
     tau =  (float) mxGetScalar(prhs[3]); /* marching step parameter */  
     epsil = (float) mxGetScalar(prhs[4]); /* tolerance */  
@@ -79,5 +90,9 @@ void mexFunction(
     vecdim[0] = 2;
     infovec = (float*)mxGetPr(plhs[1] = mxCreateNumericArray(1, vecdim, mxSINGLE_CLASS, mxREAL));
     
-    TV_ROF_GPU_main(Input, Output, infovec, lambda, iter_numb, tau, epsil, dimX, dimY, dimZ);    
+    if (mrows==1 && ncols==1) {
+    TV_ROF_GPU_main(Input, Output, infovec, lambda, 0, iter_numb, tau, epsil, dimX, dimY, dimZ);    
+    free(lambda);
+    }
+    else TV_ROF_GPU_main(Input, Output, infovec, lambda, 1, iter_numb, tau, epsil, dimX, dimY, dimZ);       
 }