Commit 92f1f1a7a51b402386f88fa4d23a850a5d4ddece

Authored by Thiago Franco de Moraes
1 parent 5926d404

Added a gui to set some configs from watershed

invesalius/data/styles.py
... ... @@ -32,6 +32,7 @@ import numpy as np
32 32  
33 33 from scipy import ndimage
34 34 from scipy.misc import imsave
  35 +from scipy.ndimage import watershed_ift, generate_binary_structure
35 36 from skimage.morphology import watershed
36 37 from skimage import filter
37 38  
... ... @@ -740,20 +741,45 @@ class EditorInteractorStyle(DefaultInteractorStyle):
740 741 class WatershedConfig(object):
741 742 __metaclass__= utils.Singleton
742 743 def __init__(self):
743   - self.operation = BRUSH_FOREGROUND
  744 + self.algorithm = "Watershed"
  745 + self.con_2d = 4
  746 + self.con_3d = 18
  747 + self.mg_size = 3
744 748 self.use_ww_wl = True
  749 + self.operation = BRUSH_FOREGROUND
745 750 self.cursor_type = const.BRUSH_CIRCLE
746 751 self.cursor_size = const.BRUSH_SIZE
747 752  
748 753 Publisher.subscribe(self.set_operation, 'Set watershed operation')
749 754 Publisher.subscribe(self.set_use_ww_wl, 'Set use ww wl')
750 755  
  756 + Publisher.subscribe(self.set_algorithm, "Set watershed algorithm")
  757 + Publisher.subscribe(self.set_2dcon, "Set watershed 2d con")
  758 + Publisher.subscribe(self.set_3dcon, "Set watershed 3d con")
  759 + Publisher.subscribe(self.set_gaussian_size, "Set watershed gaussian size")
  760 +
751 761 def set_operation(self, pubsub_evt):
752 762 self.operation = WATERSHED_OPERATIONS[pubsub_evt.data]
753 763  
754 764 def set_use_ww_wl(self, pubsub_evt):
755 765 self.use_ww_wl = pubsub_evt.data
756 766  
  767 + def set_algorithm(self, pubsub_evt):
  768 + self.algorithm = pubsub_evt.data
  769 +
  770 + def set_2dcon(self, pubsub_evt):
  771 + self.con_2d = pubsub_evt.data
  772 +
  773 + def set_3dcon(self, pubsub_evt):
  774 + self.con_3d = pubsub_evt.data
  775 +
  776 + def set_gaussian_size(self, pubsub_evt):
  777 + self.mg_size = pubsub_evt.data
  778 +
  779 +WALGORITHM = {"Watershed": watershed,
  780 + "Watershed IFT": watershed_ift}
  781 +CON2D = {4: 1, 8: 2}
  782 +CON3D = {6: 1, 18: 2, 26: 3}
757 783  
758 784 class WaterShedInteractorStyle(DefaultInteractorStyle):
759 785 def __init__(self, viewer):
... ... @@ -765,7 +791,6 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
765 791  
766 792 self.operation = BRUSH_FOREGROUND
767 793  
768   - self.mg_size = 3
769 794 self.config = WatershedConfig()
770 795  
771 796 self.picker = vtk.vtkWorldPointPicker()
... ... @@ -1047,12 +1072,17 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
1047 1072 wl = self.viewer.slice_.window_level
1048 1073  
1049 1074 if BRUSH_BACKGROUND in markers and BRUSH_FOREGROUND in markers:
  1075 + w_algorithm = WALGORITHM[self.config.algorithm]
  1076 + bstruct = generate_binary_structure(2, CON2D[self.config.con_2d])
1050 1077 if self.config.use_ww_wl:
1051   - tmp_image = ndimage.morphological_gradient(get_LUT_value(image, ww, wl).astype('uint16'), self.mg_size)
1052   - tmp_mask = watershed(tmp_image, markers)
  1078 + tmp_image = ndimage.morphological_gradient(
  1079 + get_LUT_value(image, ww, wl).astype('uint16'),
  1080 + self.config.mg_size)
  1081 +
  1082 + tmp_mask = w_algorithm(tmp_image, markers.astype('int16'), bstruct)
1053 1083 else:
1054   - tmp_image = ndimage.morphological_gradient(image, self.mg_size)
1055   - tmp_mask = watershed(tmp_image, markers)
  1084 + tmp_image = ndimage.morphological_gradient(image, self.config.mg_size)
  1085 + tmp_mask = w_algorithm(tmp_image, markers.astype('int16'), bstruct)
1056 1086  
1057 1087 if self.viewer.overwrite_mask:
1058 1088 mask[:] = 0
... ... @@ -1158,12 +1188,16 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
1158 1188 ww = self.viewer.slice_.window_width
1159 1189 wl = self.viewer.slice_.window_level
1160 1190 if BRUSH_BACKGROUND in markers and BRUSH_FOREGROUND in markers:
  1191 + w_algorithm = WALGORITHM[self.config.algorithm]
  1192 + bstruct = generate_binary_structure(3, CON3D[self.config.con_3d])
1161 1193 if self.config.use_ww_wl:
1162   - tmp_image = ndimage.morphological_gradient(get_LUT_value(image, ww, wl).astype('uint16'), self.mg_size)
1163   - tmp_mask = watershed(tmp_image, markers)
  1194 + tmp_image = ndimage.morphological_gradient(
  1195 + get_LUT_value(image, ww, wl).astype('uint16'),
  1196 + self.config.mg_size)
  1197 + tmp_mask = w_algorithm(tmp_image, markers.astype('int16'), bstruct)
1164 1198 else:
1165   - tmp_image = ndimage.morphological_gradient(image, self.mg_size)
1166   - tmp_mask = watershed(tmp_image, markers)
  1199 + tmp_image = ndimage.morphological_gradient(image, self.config.mg_size)
  1200 + tmp_mask = w_algorithm(tmp_image, markers.astype('int16'), bstruct)
1167 1201  
1168 1202 if self.viewer.overwrite_mask:
1169 1203 mask[:] = 0
... ...
invesalius/gui/dialogs.py
... ... @@ -1384,3 +1384,88 @@ class ClutImagedataDialog(wx.Dialog):
1384 1384 super(wx.Dialog, self).Show(show)
1385 1385 if gen_evt:
1386 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
... ... @@ -779,15 +779,20 @@ class WatershedTool(EditionTools):
779 779 self.ww_wl_cbox = ww_wl_cbox
780 780  
781 781 # Line 6
  782 + self.btn_wconfig = wx.Button(self, -1, "C", size=(48, -1))
782 783 self.btn_exp_watershed = wx.Button(self, -1, _('Expand watershed to 3D'))
783 784  
  785 + sizer_btns = wx.BoxSizer(wx.HORIZONTAL)
  786 + sizer_btns.Add(self.btn_wconfig, 0, wx.ALIGN_LEFT | wx.LEFT | wx.TOP | wx.DOWN, 5)
  787 + sizer_btns.Add(self.btn_exp_watershed, 0, wx.GROW|wx.EXPAND| wx.ALL, 5)
  788 +
784 789 # Add lines into main sizer
785 790 sizer = wx.BoxSizer(wx.VERTICAL)
786 791 sizer.Add(text1, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
787 792 sizer.Add(line2, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
788 793 sizer.Add(check_box, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
789 794 sizer.Add(ww_wl_cbox, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
790   - sizer.Add(self.btn_exp_watershed, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
  795 + sizer.Add(sizer_btns, 0, wx.EXPAND)
791 796 sizer.Fit(self)
792 797  
793 798 self.SetSizer(sizer)
... ... @@ -803,6 +808,7 @@ class WatershedTool(EditionTools):
803 808 self.check_box.Bind(wx.EVT_CHECKBOX, self.OnCheckOverwriteMask)
804 809 self.ww_wl_cbox.Bind(wx.EVT_CHECKBOX, self.OnCheckWWWL)
805 810 self.btn_exp_watershed.Bind(wx.EVT_BUTTON, self.OnExpandWatershed)
  811 + self.btn_wconfig.Bind(wx.EVT_BUTTON, self.OnConfig)
806 812  
807 813 def ChangeMaskColour(self, pubsub_evt):
808 814 colour = pubsub_evt.data
... ... @@ -860,5 +866,8 @@ class WatershedTool(EditionTools):
860 866 value = self.ww_wl_cbox.GetValue()
861 867 Publisher.sendMessage('Set use ww wl', value)
862 868  
  869 + def OnConfig(self, evt):
  870 + dlg.WatershedOptionsDialog().Show()
  871 +
863 872 def OnExpandWatershed(self, evt):
864 873 Publisher.sendMessage('Expand watershed to 3D AXIAL')
... ...