Commit 536f91e77b5dc857e3b9187ff61381031203f0c7

Authored by Thiago Franco de Moraes
2 parents dbecc12e 6b7a5567

Merge pull request #28 from tfmoraes/watershed_improvements

Improvements to watershed segmentation tool (algorithm, connectivity, gui to config).
icons/configuration.png 0 → 100644

614 Bytes

invesalius/data/styles.py
@@ -18,7 +18,9 @@ @@ -18,7 +18,9 @@
18 #-------------------------------------------------------------------------- 18 #--------------------------------------------------------------------------
19 19
20 import os 20 import os
  21 +import multiprocessing
21 import tempfile 22 import tempfile
  23 +import time
22 24
23 import vtk 25 import vtk
24 import wx 26 import wx
@@ -27,13 +29,17 @@ from wx.lib.pubsub import pub as Publisher @@ -27,13 +29,17 @@ from wx.lib.pubsub import pub as Publisher
27 29
28 import constants as const 30 import constants as const
29 import converters 31 import converters
  32 +import cursor_actors as ca
30 import numpy as np 33 import numpy as np
31 34
32 from scipy import ndimage 35 from scipy import ndimage
33 from scipy.misc import imsave 36 from scipy.misc import imsave
  37 +from scipy.ndimage import watershed_ift, generate_binary_structure
34 from skimage.morphology import watershed 38 from skimage.morphology import watershed
35 from skimage import filter 39 from skimage import filter
36 40
  41 +import utils
  42 +
37 ORIENTATIONS = { 43 ORIENTATIONS = {
38 "AXIAL": const.AXIAL, 44 "AXIAL": const.AXIAL,
39 "CORONAL": const.CORONAL, 45 "CORONAL": const.CORONAL,
@@ -734,6 +740,86 @@ class EditorInteractorStyle(DefaultInteractorStyle): @@ -734,6 +740,86 @@ class EditorInteractorStyle(DefaultInteractorStyle):
734 return x, y, z 740 return x, y, z
735 741
736 742
  743 +class WatershedProgressWindow(wx.Frame):
  744 + def __init__(self, process, parent=None):
  745 + wx.Frame.__init__(self, parent, -1)
  746 + self.process = process
  747 + self._build_gui()
  748 + self._bind_wx_events()
  749 + self.timer = wx.Timer(self)
  750 + self.timer.Start(1000)
  751 +
  752 + def _build_gui(self):
  753 + self.gauge = wx.Gauge(self, -1, 100)
  754 + self.btn_cancel = wx.Button(self, wx.ID_CANCEL)
  755 +
  756 + sizer = wx.BoxSizer(wx.VERTICAL)
  757 + sizer.Add(wx.StaticText(self, -1, _("Applying watershed")))
  758 + sizer.Add(self.gauge, 0, wx.EXPAND)
  759 + sizer.Add(self.btn_cancel, 0, wx.ALIGN_LEFT)
  760 +
  761 + self.SetSizer(sizer)
  762 + sizer.Fit(self)
  763 + self.Layout()
  764 +
  765 + def __del__(self):
  766 + self.timer.Stop()
  767 +
  768 + def _bind_wx_events(self):
  769 + self.Bind(wx.EVT_TIMER, self.TimeHandler)
  770 + self.btn_cancel.Bind(wx.EVT_BUTTON, self.on_cancel)
  771 +
  772 + def on_cancel(self, evt):
  773 + self.timer.Stop()
  774 + self.process.terminate()
  775 +
  776 + def TimeHandler(self, evt):
  777 + self.gauge.Pulse()
  778 +
  779 +
  780 +class WatershedConfig(object):
  781 + __metaclass__= utils.Singleton
  782 + def __init__(self):
  783 + self.algorithm = "Watershed"
  784 + self.con_2d = 4
  785 + self.con_3d = 6
  786 + self.mg_size = 3
  787 + self.use_ww_wl = True
  788 + self.operation = BRUSH_FOREGROUND
  789 + self.cursor_type = const.BRUSH_CIRCLE
  790 + self.cursor_size = const.BRUSH_SIZE
  791 +
  792 + Publisher.subscribe(self.set_operation, 'Set watershed operation')
  793 + Publisher.subscribe(self.set_use_ww_wl, 'Set use ww wl')
  794 +
  795 + Publisher.subscribe(self.set_algorithm, "Set watershed algorithm")
  796 + Publisher.subscribe(self.set_2dcon, "Set watershed 2d con")
  797 + Publisher.subscribe(self.set_3dcon, "Set watershed 3d con")
  798 + Publisher.subscribe(self.set_gaussian_size, "Set watershed gaussian size")
  799 +
  800 + def set_operation(self, pubsub_evt):
  801 + self.operation = WATERSHED_OPERATIONS[pubsub_evt.data]
  802 +
  803 + def set_use_ww_wl(self, pubsub_evt):
  804 + self.use_ww_wl = pubsub_evt.data
  805 +
  806 + def set_algorithm(self, pubsub_evt):
  807 + self.algorithm = pubsub_evt.data
  808 +
  809 + def set_2dcon(self, pubsub_evt):
  810 + self.con_2d = pubsub_evt.data
  811 +
  812 + def set_3dcon(self, pubsub_evt):
  813 + self.con_3d = pubsub_evt.data
  814 +
  815 + def set_gaussian_size(self, pubsub_evt):
  816 + self.mg_size = pubsub_evt.data
  817 +
  818 +WALGORITHM = {"Watershed": watershed,
  819 + "Watershed IFT": watershed_ift}
  820 +CON2D = {4: 1, 8: 2}
  821 +CON3D = {6: 1, 18: 2, 26: 3}
  822 +
737 class WaterShedInteractorStyle(DefaultInteractorStyle): 823 class WaterShedInteractorStyle(DefaultInteractorStyle):
738 def __init__(self, viewer): 824 def __init__(self, viewer):
739 DefaultInteractorStyle.__init__(self, viewer) 825 DefaultInteractorStyle.__init__(self, viewer)
@@ -742,9 +828,7 @@ class WaterShedInteractorStyle(DefaultInteractorStyle): @@ -742,9 +828,7 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
742 self.orientation = self.viewer.orientation 828 self.orientation = self.viewer.orientation
743 self.matrix = None 829 self.matrix = None
744 830
745 - self.operation = BRUSH_FOREGROUND  
746 -  
747 - self.mg_size = 3 831 + self.config = WatershedConfig()
748 832
749 self.picker = vtk.vtkWorldPointPicker() 833 self.picker = vtk.vtkWorldPointPicker()
750 834
@@ -761,7 +845,10 @@ class WaterShedInteractorStyle(DefaultInteractorStyle): @@ -761,7 +845,10 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
761 self.AddObserver("MouseMoveEvent", self.OnBrushMove) 845 self.AddObserver("MouseMoveEvent", self.OnBrushMove)
762 846
763 Publisher.subscribe(self.expand_watershed, 'Expand watershed to 3D ' + self.orientation) 847 Publisher.subscribe(self.expand_watershed, 'Expand watershed to 3D ' + self.orientation)
764 - Publisher.subscribe(self.set_operation, 'Set watershed operation') 848 + Publisher.subscribe(self.set_bsize, 'Set watershed brush size')
  849 + Publisher.subscribe(self.set_bformat, 'Set watershed brush format')
  850 +
  851 + self._set_cursor()
765 852
766 def SetUp(self): 853 def SetUp(self):
767 mask = self.viewer.slice_.current_mask.matrix 854 mask = self.viewer.slice_.current_mask.matrix
@@ -772,7 +859,8 @@ class WaterShedInteractorStyle(DefaultInteractorStyle): @@ -772,7 +859,8 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
772 def CleanUp(self): 859 def CleanUp(self):
773 #self._remove_mask() 860 #self._remove_mask()
774 Publisher.unsubscribe(self.expand_watershed, 'Expand watershed to 3D ' + self.orientation) 861 Publisher.unsubscribe(self.expand_watershed, 'Expand watershed to 3D ' + self.orientation)
775 - Publisher.unsubscribe(self.set_operation, 'Set watershed operation') 862 + Publisher.unsubscribe(self.set_bformat, 'Set watershed brush format')
  863 + Publisher.unsubscribe(self.set_bsize, 'Set watershed brush size')
776 self.RemoveAllObservers() 864 self.RemoveAllObservers()
777 self.viewer.slice_.to_show_aux = '' 865 self.viewer.slice_.to_show_aux = ''
778 self.viewer.OnScrollBar() 866 self.viewer.OnScrollBar()
@@ -791,8 +879,33 @@ class WaterShedInteractorStyle(DefaultInteractorStyle): @@ -791,8 +879,33 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
791 os.remove(self.temp_file) 879 os.remove(self.temp_file)
792 print "deleting", self.temp_file 880 print "deleting", self.temp_file
793 881
794 - def set_operation(self, pubsub_evt):  
795 - self.operation = WATERSHED_OPERATIONS[pubsub_evt.data] 882 + def _set_cursor(self):
  883 + if self.config.cursor_type == const.BRUSH_SQUARE:
  884 + cursor = ca.CursorRectangle()
  885 + elif self.config.cursor_type == const.BRUSH_CIRCLE:
  886 + cursor = ca.CursorCircle()
  887 +
  888 + cursor.SetOrientation(self.orientation)
  889 + n = self.viewer.slice_data.number
  890 + coordinates = {"SAGITAL": [n, 0, 0],
  891 + "CORONAL": [0, n, 0],
  892 + "AXIAL": [0, 0, n]}
  893 + cursor.SetPosition(coordinates[self.orientation])
  894 + spacing = self.viewer.slice_.spacing
  895 + cursor.SetSpacing(spacing)
  896 + cursor.SetColour(self.viewer._brush_cursor_colour)
  897 + cursor.SetSize(self.config.cursor_size)
  898 + self.viewer.slice_data.SetCursor(cursor)
  899 + self.viewer.interactor.Render()
  900 +
  901 + def set_bsize(self, pubsub_evt):
  902 + size = pubsub_evt.data
  903 + self.config.cursor_size = size
  904 + self.viewer.slice_data.cursor.SetSize(size)
  905 +
  906 + def set_bformat(self, pubsub_evt):
  907 + self.config.cursor_type = pubsub_evt.data
  908 + self._set_cursor()
796 909
797 def OnEnterInteractor(self, obj, evt): 910 def OnEnterInteractor(self, obj, evt):
798 if (self.viewer.slice_.buffer_slices[self.orientation].mask is None): 911 if (self.viewer.slice_.buffer_slices[self.orientation].mask is None):
@@ -869,7 +982,7 @@ class WaterShedInteractorStyle(DefaultInteractorStyle): @@ -869,7 +982,7 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
869 if position < 0: 982 if position < 0:
870 position = viewer.calculate_matrix_position(coord) 983 position = viewer.calculate_matrix_position(coord)
871 984
872 - operation = self.operation 985 + operation = self.config.operation
873 986
874 if operation == BRUSH_FOREGROUND: 987 if operation == BRUSH_FOREGROUND:
875 if iren.GetControlKey(): 988 if iren.GetControlKey():
@@ -937,7 +1050,7 @@ class WaterShedInteractorStyle(DefaultInteractorStyle): @@ -937,7 +1050,7 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
937 if position < 0: 1050 if position < 0:
938 position = viewer.calculate_matrix_position(coord) 1051 position = viewer.calculate_matrix_position(coord)
939 1052
940 - operation = self.operation 1053 + operation = self.config.operation
941 1054
942 if operation == BRUSH_FOREGROUND: 1055 if operation == BRUSH_FOREGROUND:
943 if iren.GetControlKey(): 1056 if iren.GetControlKey():
@@ -991,8 +1104,35 @@ class WaterShedInteractorStyle(DefaultInteractorStyle): @@ -991,8 +1104,35 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
991 wl = self.viewer.slice_.window_level 1104 wl = self.viewer.slice_.window_level
992 1105
993 if BRUSH_BACKGROUND in markers and BRUSH_FOREGROUND in markers: 1106 if BRUSH_BACKGROUND in markers and BRUSH_FOREGROUND in markers:
994 - tmp_image = ndimage.morphological_gradient(get_LUT_value(image, ww, wl).astype('uint16'), self.mg_size)  
995 - tmp_mask = watershed(tmp_image, markers) 1107 + #w_algorithm = WALGORITHM[self.config.algorithm]
  1108 + bstruct = generate_binary_structure(2, CON2D[self.config.con_2d])
  1109 + if self.config.use_ww_wl:
  1110 + if self.config.algorithm == 'Watershed':
  1111 + tmp_image = ndimage.morphological_gradient(
  1112 + get_LUT_value(image, ww, wl).astype('uint16'),
  1113 + self.config.mg_size)
  1114 + tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
  1115 + else:
  1116 + #tmp_image = ndimage.gaussian_filter(get_LUT_value(image, ww, wl).astype('uint16'), self.config.mg_size)
  1117 + #tmp_image = ndimage.morphological_gradient(
  1118 + #get_LUT_value(image, ww, wl).astype('uint16'),
  1119 + #self.config.mg_size)
  1120 + tmp_image = get_LUT_value(image, ww, wl).astype('uint16')
  1121 + #markers[markers == 2] = -1
  1122 + tmp_mask = watershed_ift(tmp_image, markers.astype('int16'), bstruct)
  1123 + #markers[markers == -1] = 2
  1124 + #tmp_mask[tmp_mask == -1] = 2
  1125 +
  1126 + else:
  1127 + if self.config.algorithm == 'Watershed':
  1128 + tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), self.config.mg_size)
  1129 + tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
  1130 + else:
  1131 + #tmp_image = (image - image.min()).astype('uint16')
  1132 + #tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
  1133 + #tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), self.config.mg_size)
  1134 + tmp_image = image - image.min().astype('uint16')
  1135 + tmp_mask = watershed_ift(tmp_image, markers.astype('int16'), bstruct)
996 1136
997 if self.viewer.overwrite_mask: 1137 if self.viewer.overwrite_mask:
998 mask[:] = 0 1138 mask[:] = 0
@@ -1098,8 +1238,59 @@ class WaterShedInteractorStyle(DefaultInteractorStyle): @@ -1098,8 +1238,59 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
1098 ww = self.viewer.slice_.window_width 1238 ww = self.viewer.slice_.window_width
1099 wl = self.viewer.slice_.window_level 1239 wl = self.viewer.slice_.window_level
1100 if BRUSH_BACKGROUND in markers and BRUSH_FOREGROUND in markers: 1240 if BRUSH_BACKGROUND in markers and BRUSH_FOREGROUND in markers:
1101 - tmp_image = ndimage.morphological_gradient(get_LUT_value(image, ww, wl).astype('uint16'), self.mg_size)  
1102 - tmp_mask = watershed(tmp_image, markers) 1241 + #w_algorithm = WALGORITHM[self.config.algorithm]
  1242 + bstruct = generate_binary_structure(3, CON3D[self.config.con_3d])
  1243 + tfile = tempfile.mktemp()
  1244 + tmp_mask = np.memmap(tfile, shape=mask.shape, dtype=mask.dtype,
  1245 + mode='w+')
  1246 + q = multiprocessing.Queue()
  1247 + p = multiprocessing.Process(target=do_watershed, args=(image,
  1248 + markers, tmp_mask, bstruct,
  1249 + self.config.algorithm,
  1250 + self.config.mg_size,
  1251 + self.config.use_ww_wl, wl, ww, q))
  1252 +
  1253 + wp = WatershedProgressWindow(p)
  1254 + wp.Center(wx.BOTH)
  1255 + wp.Show()
  1256 + wp.MakeModal()
  1257 +
  1258 + p.start()
  1259 +
  1260 + while q.empty() and p.is_alive():
  1261 + time.sleep(0.5)
  1262 + wx.Yield()
  1263 +
  1264 + wp.MakeModal(False)
  1265 + wp.Destroy()
  1266 + del wp
  1267 +
  1268 + if q.empty():
  1269 + return
  1270 + #do_watershed(image, markers, tmp_mask, bstruct, self.config.algorithm,
  1271 + #self.config.mg_size, self.config.use_ww_wl, wl, ww)
  1272 + #if self.config.use_ww_wl:
  1273 + #if self.config.algorithm == 'Watershed':
  1274 + #tmp_image = ndimage.morphological_gradient(
  1275 + #get_LUT_value(image, ww, wl).astype('uint16'),
  1276 + #self.config.mg_size)
  1277 + #tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
  1278 + #else:
  1279 + #tmp_image = get_LUT_value(image, ww, wl).astype('uint16')
  1280 + ##tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
  1281 + ##tmp_image = ndimage.morphological_gradient(
  1282 + ##get_LUT_value(image, ww, wl).astype('uint16'),
  1283 + ##self.config.mg_size)
  1284 + #tmp_mask = watershed_ift(tmp_image, markers.astype('int16'), bstruct)
  1285 + #else:
  1286 + #if self.config.algorithm == 'Watershed':
  1287 + #tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), self.config.mg_size)
  1288 + #tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
  1289 + #else:
  1290 + #tmp_image = (image - image.min()).astype('uint16')
  1291 + ##tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
  1292 + ##tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), self.config.mg_size)
  1293 + #tmp_mask = watershed_ift(tmp_image, markers.astype('int8'), bstruct)
1103 1294
1104 if self.viewer.overwrite_mask: 1295 if self.viewer.overwrite_mask:
1105 mask[:] = 0 1296 mask[:] = 0
@@ -1118,6 +1309,34 @@ class WaterShedInteractorStyle(DefaultInteractorStyle): @@ -1118,6 +1309,34 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
1118 Publisher.sendMessage('Reload actual slice') 1309 Publisher.sendMessage('Reload actual slice')
1119 1310
1120 1311
  1312 +def do_watershed(image, markers, mask, bstruct, algorithm, mg_size, use_ww_wl, wl, ww, q):
  1313 + if use_ww_wl:
  1314 + if algorithm == 'Watershed':
  1315 + tmp_image = ndimage.morphological_gradient(
  1316 + get_LUT_value(image, ww, wl).astype('uint16'),
  1317 + mg_size)
  1318 + tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
  1319 + else:
  1320 + tmp_image = get_LUT_value(image, ww, wl).astype('uint16')
  1321 + #tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
  1322 + #tmp_image = ndimage.morphological_gradient(
  1323 + #get_LUT_value(image, ww, wl).astype('uint16'),
  1324 + #self.config.mg_size)
  1325 + tmp_mask = watershed_ift(tmp_image, markers.astype('int16'), bstruct)
  1326 + else:
  1327 + if algorithm == 'Watershed':
  1328 + tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), mg_size)
  1329 + tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
  1330 + else:
  1331 + tmp_image = (image - image.min()).astype('uint16')
  1332 + #tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
  1333 + #tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), self.config.mg_size)
  1334 + tmp_mask = watershed_ift(tmp_image, markers.astype('int8'), bstruct)
  1335 + mask[:] = tmp_mask
  1336 + q.put(1)
  1337 +
  1338 +
  1339 +
1121 def get_style(style): 1340 def get_style(style):
1122 STYLES = { 1341 STYLES = {
1123 const.STATE_DEFAULT: DefaultInteractorStyle, 1342 const.STATE_DEFAULT: DefaultInteractorStyle,
invesalius/gui/dialogs.py
@@ -1384,3 +1384,88 @@ class ClutImagedataDialog(wx.Dialog): @@ -1384,3 +1384,88 @@ class ClutImagedataDialog(wx.Dialog):
1384 super(wx.Dialog, self).Show(show) 1384 super(wx.Dialog, self).Show(show)
1385 if gen_evt: 1385 if gen_evt:
1386 self.clut_widget._generate_event() 1386 self.clut_widget._generate_event()
  1387 +
  1388 +
  1389 +class WatershedOptions(wx.Panel):
  1390 + def __init__(self, parent):
  1391 + wx.Panel.__init__(self, parent)
  1392 +
  1393 + self.algorithms = ("Watershed", "Watershed IFT")
  1394 + self.con2d_choices = (4, 8)
  1395 + self.con3d_choices = (6, 18, 26)
  1396 +
  1397 + self._init_gui()
  1398 + self._bind_events()
  1399 +
  1400 + def _init_gui(self):
  1401 + self.choice_algorithm = wx.RadioBox(self, -1, "Algorithm",
  1402 + choices=("Watershed", "Watershed IFT"),
  1403 + style=wx.NO_BORDER | wx.HORIZONTAL)
  1404 +
  1405 + self.choice_2dcon = wx.RadioBox(self, -1, "2D",
  1406 + choices=[str(i) for i in self.con2d_choices],
  1407 + style=wx.NO_BORDER | wx.HORIZONTAL)
  1408 +
  1409 + self.choice_3dcon = wx.RadioBox(self, -1, "3D",
  1410 + choices=[str(i) for i in self.con3d_choices],
  1411 + style=wx.NO_BORDER | wx.HORIZONTAL)
  1412 +
  1413 + self.gaussian_size = wx.SpinCtrl(self, -1, "", min=1, max=10)
  1414 +
  1415 + box_sizer = wx.StaticBoxSizer(wx.StaticBox(self, -1, "Conectivity"), wx.VERTICAL)
  1416 + box_sizer.Add(self.choice_2dcon, 0, wx.ALIGN_CENTER_VERTICAL,2)
  1417 + box_sizer.Add(self.choice_3dcon, 0, wx.ALIGN_CENTER_VERTICAL,2)
  1418 +
  1419 + g_sizer = wx.BoxSizer(wx.HORIZONTAL)
  1420 + g_sizer.Add(wx.StaticText(self, -1, "Gaussian size"), 0, wx.ALIGN_RIGHT | wx.ALL, 5)
  1421 + g_sizer.Add(self.gaussian_size, 0, wx.ALIGN_LEFT | wx.ALL, 5)
  1422 +
  1423 + sizer = wx.BoxSizer(wx.VERTICAL)
  1424 + sizer.Add(self.choice_algorithm, 0, wx.ALIGN_CENTER_VERTICAL,2)
  1425 + sizer.Add(box_sizer, 1, wx.EXPAND,2)
  1426 + sizer.Add(g_sizer, 0, wx.ALIGN_LEFT, 2)
  1427 +
  1428 + self.SetSizer(sizer)
  1429 + sizer.Fit(self)
  1430 + self.Layout()
  1431 +
  1432 + def _bind_events(self):
  1433 + self.choice_algorithm.Bind(wx.EVT_RADIOBOX, self.OnSetAlgorithm)
  1434 + self.gaussian_size.Bind(wx.EVT_SPINCTRL, self.OnSetGaussianSize)
  1435 + self.choice_2dcon.Bind(wx.EVT_RADIOBOX, self.OnSetCon2D)
  1436 + self.choice_3dcon.Bind(wx.EVT_RADIOBOX, self.OnSetCon3D)
  1437 +
  1438 + def OnSetAlgorithm(self, evt):
  1439 + v = self.algorithms[evt.GetInt()]
  1440 + Publisher.sendMessage("Set watershed algorithm", v)
  1441 +
  1442 + def OnSetGaussianSize(self, evt):
  1443 + v = self.gaussian_size.GetValue()
  1444 + Publisher.sendMessage("Set watershed gaussian size", v)
  1445 +
  1446 + def OnSetCon2D(self, evt):
  1447 + v = self.con2d_choices[evt.GetInt()]
  1448 + Publisher.sendMessage("Set watershed 2d con", v)
  1449 +
  1450 + def OnSetCon3D(self, evt):
  1451 + v = self.con3d_choices[evt.GetInt()]
  1452 + Publisher.sendMessage("Set watershed 3d con", v)
  1453 +
  1454 +
  1455 +class WatershedOptionsDialog(wx.Dialog):
  1456 + def __init__(self):
  1457 + pre = wx.PreDialog()
  1458 + pre.Create(wx.GetApp().GetTopWindow(), -1, style=wx.DEFAULT_DIALOG_STYLE|wx.FRAME_FLOAT_ON_PARENT)
  1459 + self.PostCreate(pre)
  1460 +
  1461 + self._init_gui()
  1462 +
  1463 + def _init_gui(self):
  1464 + wop = WatershedOptions(self)
  1465 +
  1466 + sizer = wx.BoxSizer(wx.VERTICAL)
  1467 + sizer.Add(wop, 0, wx.EXPAND)
  1468 +
  1469 + self.SetSizer(sizer)
  1470 + sizer.Fit(self)
  1471 + self.Layout()
invesalius/gui/task_slice.py
@@ -720,7 +720,7 @@ class WatershedTool(EditionTools): @@ -720,7 +720,7 @@ class WatershedTool(EditionTools):
720 self.SetBackgroundColour(default_colour) 720 self.SetBackgroundColour(default_colour)
721 721
722 ## LINE 1 722 ## LINE 1
723 - text1 = wx.StaticText(self, -1, _("Choose brush type and size:")) 723 + text1 = wx.StaticText(self, -1, _("Choose brush type, size or operation:"))
724 724
725 ## LINE 2 725 ## LINE 2
726 menu = wx.Menu() 726 menu = wx.Menu()
@@ -773,17 +773,28 @@ class WatershedTool(EditionTools): @@ -773,17 +773,28 @@ class WatershedTool(EditionTools):
773 773
774 # LINE 5 774 # LINE 5
775 check_box = wx.CheckBox(self, -1, _("Overwrite mask")) 775 check_box = wx.CheckBox(self, -1, _("Overwrite mask"))
  776 + ww_wl_cbox = wx.CheckBox(self, -1, _("Use WW&WL"))
  777 + ww_wl_cbox.SetValue(True)
776 self.check_box = check_box 778 self.check_box = check_box
  779 + self.ww_wl_cbox = ww_wl_cbox
777 780
778 # Line 6 781 # Line 6
  782 + bmp = wx.Bitmap("../icons/configuration.png", wx.BITMAP_TYPE_PNG)
  783 + self.btn_wconfig = wx.BitmapButton(self, -1, bitmap=bmp,
  784 + size=(bmp.GetWidth()+10, bmp.GetHeight()+10))
779 self.btn_exp_watershed = wx.Button(self, -1, _('Expand watershed to 3D')) 785 self.btn_exp_watershed = wx.Button(self, -1, _('Expand watershed to 3D'))
780 786
  787 + sizer_btns = wx.BoxSizer(wx.HORIZONTAL)
  788 + sizer_btns.Add(self.btn_wconfig, 0, wx.ALIGN_LEFT | wx.LEFT | wx.TOP | wx.DOWN, 5)
  789 + sizer_btns.Add(self.btn_exp_watershed, 0, wx.GROW|wx.EXPAND| wx.ALL, 5)
  790 +
781 # Add lines into main sizer 791 # Add lines into main sizer
782 sizer = wx.BoxSizer(wx.VERTICAL) 792 sizer = wx.BoxSizer(wx.VERTICAL)
783 sizer.Add(text1, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5) 793 sizer.Add(text1, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
784 sizer.Add(line2, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5) 794 sizer.Add(line2, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
785 sizer.Add(check_box, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5) 795 sizer.Add(check_box, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
786 - sizer.Add(self.btn_exp_watershed, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5) 796 + sizer.Add(ww_wl_cbox, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
  797 + sizer.Add(sizer_btns, 0, wx.EXPAND)
787 sizer.Fit(self) 798 sizer.Fit(self)
788 799
789 self.SetSizer(sizer) 800 self.SetSizer(sizer)
@@ -797,7 +808,9 @@ class WatershedTool(EditionTools): @@ -797,7 +808,9 @@ class WatershedTool(EditionTools):
797 self.Bind(wx.EVT_MENU, self.OnMenu) 808 self.Bind(wx.EVT_MENU, self.OnMenu)
798 self.combo_brush_op.Bind(wx.EVT_COMBOBOX, self.OnComboBrushOp) 809 self.combo_brush_op.Bind(wx.EVT_COMBOBOX, self.OnComboBrushOp)
799 self.check_box.Bind(wx.EVT_CHECKBOX, self.OnCheckOverwriteMask) 810 self.check_box.Bind(wx.EVT_CHECKBOX, self.OnCheckOverwriteMask)
  811 + self.ww_wl_cbox.Bind(wx.EVT_CHECKBOX, self.OnCheckWWWL)
800 self.btn_exp_watershed.Bind(wx.EVT_BUTTON, self.OnExpandWatershed) 812 self.btn_exp_watershed.Bind(wx.EVT_BUTTON, self.OnExpandWatershed)
  813 + self.btn_wconfig.Bind(wx.EVT_BUTTON, self.OnConfig)
801 814
802 def ChangeMaskColour(self, pubsub_evt): 815 def ChangeMaskColour(self, pubsub_evt):
803 colour = pubsub_evt.data 816 colour = pubsub_evt.data
@@ -834,14 +847,14 @@ class WatershedTool(EditionTools): @@ -834,14 +847,14 @@ class WatershedTool(EditionTools):
834 847
835 self.btn_brush_format.SetBitmap(bitmap[evt.GetId()]) 848 self.btn_brush_format.SetBitmap(bitmap[evt.GetId()])
836 849
837 - Publisher.sendMessage('Set brush format', brush[evt.GetId()]) 850 + Publisher.sendMessage('Set watershed brush format', brush[evt.GetId()])
838 851
839 def OnBrushSize(self, evt): 852 def OnBrushSize(self, evt):
840 """ """ 853 """ """
841 # FIXME: Using wx.EVT_SPINCTRL in MacOS it doesnt capture changes only 854 # FIXME: Using wx.EVT_SPINCTRL in MacOS it doesnt capture changes only
842 # in the text ctrl - so we are capturing only changes on text 855 # in the text ctrl - so we are capturing only changes on text
843 # Strangelly this is being called twice 856 # Strangelly this is being called twice
844 - Publisher.sendMessage('Set edition brush size',self.spin.GetValue()) 857 + Publisher.sendMessage('Set watershed brush size',self.spin.GetValue())
845 858
846 def OnComboBrushOp(self, evt): 859 def OnComboBrushOp(self, evt):
847 brush_op = self.combo_brush_op.GetValue() 860 brush_op = self.combo_brush_op.GetValue()
@@ -851,5 +864,12 @@ class WatershedTool(EditionTools): @@ -851,5 +864,12 @@ class WatershedTool(EditionTools):
851 value = self.check_box.GetValue() 864 value = self.check_box.GetValue()
852 Publisher.sendMessage('Set overwrite mask', value) 865 Publisher.sendMessage('Set overwrite mask', value)
853 866
  867 + def OnCheckWWWL(self, evt):
  868 + value = self.ww_wl_cbox.GetValue()
  869 + Publisher.sendMessage('Set use ww wl', value)
  870 +
  871 + def OnConfig(self, evt):
  872 + dlg.WatershedOptionsDialog().Show()
  873 +
854 def OnExpandWatershed(self, evt): 874 def OnExpandWatershed(self, evt):
855 Publisher.sendMessage('Expand watershed to 3D AXIAL') 875 Publisher.sendMessage('Expand watershed to 3D AXIAL')