diff --git a/invesalius/data/styles.py b/invesalius/data/styles.py
index 4c7c13a..b258afc 100644
--- a/invesalius/data/styles.py
+++ b/invesalius/data/styles.py
@@ -59,10 +59,14 @@ WATERSHED_OPERATIONS = {_("Erase"): BRUSH_ERASE,
                         _("Background"): BRUSH_BACKGROUND,}
 
 def get_LUT_value(data, window, level):
-    return np.piecewise(data,
-                        [data <= (level - 0.5 - (window-1)/2),
-                         data > (level - 0.5 + (window-1)/2)],
-                        [0, window, lambda data: ((data - (level - 0.5))/(window-1) + 0.5)*(window)])
+    shape = data.shape
+    data_ = data.ravel()
+    data = np.piecewise(data_,
+                        [data_ <= (level - 0.5 - (window-1)/2),
+                         data_ > (level - 0.5 + (window-1)/2)],
+                        [0, window, lambda data_: ((data_ - (level - 0.5))/(window-1) + 0.5)*(window)])
+    data.shape = shape
+    return data
 
 class BaseImageInteractorStyle(vtk.vtkInteractorStyleImage):
     def __init__(self, viewer):
diff --git a/invesalius/data/watershed_process.py b/invesalius/data/watershed_process.py
index 41ff61b..af3de65 100644
--- a/invesalius/data/watershed_process.py
+++ b/invesalius/data/watershed_process.py
@@ -1,43 +1,46 @@
-import numpy as np
-from scipy import ndimage
-from scipy.ndimage import watershed_ift, generate_binary_structure
-from skimage.morphology import watershed
-from skimage import filter
-
-
-def get_LUT_value(data, window, level):
-    return np.piecewise(data, 
-                        [data <= (level - 0.5 - (window-1)/2),
-                         data > (level - 0.5 + (window-1)/2)],
-                        [0, 255, lambda data: ((data - (level - 0.5))/(window-1) + 0.5)*(255-0)])
-
-
-
-def do_watershed(image, markers,  tfile, shape, bstruct, algorithm, mg_size, use_ww_wl, wl, ww, q):
-    mask = np.memmap(tfile, shape=shape, dtype='uint8', mode='r+')
-                
-    if use_ww_wl:
-        if algorithm == 'Watershed':
-            tmp_image = ndimage.morphological_gradient(
-                           get_LUT_value(image, ww, wl).astype('uint16'),
-                           mg_size)
-            tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
-        else:
-            tmp_image = get_LUT_value(image, ww, wl).astype('uint16')
-            #tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
-            #tmp_image = ndimage.morphological_gradient(
-                           #get_LUT_value(image, ww, wl).astype('uint16'),
-                           #self.config.mg_size)
-            tmp_mask = watershed_ift(tmp_image, markers.astype('int16'), bstruct)
-    else:
-        if algorithm == 'Watershed':
-            tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), mg_size)
-            tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
-        else:
-            tmp_image = (image - image.min()).astype('uint16')
-            #tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
-            #tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), self.config.mg_size)
-            tmp_mask = watershed_ift(tmp_image, markers.astype('int8'), bstruct)
-    mask[:] = tmp_mask
-    mask.flush()
-    q.put(1)
+import numpy as np
+from scipy import ndimage
+from scipy.ndimage import watershed_ift, generate_binary_structure
+from skimage.morphology import watershed
+from skimage import filter
+
+
+def get_LUT_value(data, window, level):
+    shape = data.shape
+    data_ = data.ravel()
+    data = np.piecewise(data_,
+                        [data_ <= (level - 0.5 - (window-1)/2),
+                         data_ > (level - 0.5 + (window-1)/2)],
+                        [0, window, lambda data_: ((data_ - (level - 0.5))/(window-1) + 0.5)*(window)])
+    data.shape = shape
+    return data
+
+
+def do_watershed(image, markers,  tfile, shape, bstruct, algorithm, mg_size, use_ww_wl, wl, ww, q):
+    mask = np.memmap(tfile, shape=shape, dtype='uint8', mode='r+')
+                
+    if use_ww_wl:
+        if algorithm == 'Watershed':
+            tmp_image = ndimage.morphological_gradient(
+                           get_LUT_value(image, ww, wl).astype('uint16'),
+                           mg_size)
+            tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
+        else:
+            tmp_image = get_LUT_value(image, ww, wl).astype('uint16')
+            #tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
+            #tmp_image = ndimage.morphological_gradient(
+                           #get_LUT_value(image, ww, wl).astype('uint16'),
+                           #self.config.mg_size)
+            tmp_mask = watershed_ift(tmp_image, markers.astype('int16'), bstruct)
+    else:
+        if algorithm == 'Watershed':
+            tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), mg_size)
+            tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
+        else:
+            tmp_image = (image - image.min()).astype('uint16')
+            #tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
+            #tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), self.config.mg_size)
+            tmp_mask = watershed_ift(tmp_image, markers.astype('int8'), bstruct)
+    mask[:] = tmp_mask
+    mask.flush()
+    q.put(1)
--
libgit2 0.21.2