Commit 27f9a4ea304dfd30d5d0d21e52a306a453282ffb

Authored by Thiago Franco de Moraes
2 parents 9c8083ac c180da8d

Merge pull request #20 from tfmoraes/watershed_merge

This implements the segmentation watershed algorithm to InVesalius.
invesalius/constants.py
@@ -485,6 +485,7 @@ STATE_MEASURE_ANGLE = 1008 @@ -485,6 +485,7 @@ STATE_MEASURE_ANGLE = 1008
485 SLICE_STATE_CROSS = 3006 485 SLICE_STATE_CROSS = 3006
486 SLICE_STATE_SCROLL = 3007 486 SLICE_STATE_SCROLL = 3007
487 SLICE_STATE_EDITOR = 3008 487 SLICE_STATE_EDITOR = 3008
  488 +SLICE_STATE_WATERSHED = 3009
488 489
489 VOLUME_STATE_SEED = 2001 490 VOLUME_STATE_SEED = 2001
490 #STATE_LINEAR_MEASURE = 3001 491 #STATE_LINEAR_MEASURE = 3001
@@ -500,6 +501,7 @@ TOOL_SLICE_STATES = [SLICE_STATE_CROSS, SLICE_STATE_SCROLL] @@ -500,6 +501,7 @@ TOOL_SLICE_STATES = [SLICE_STATE_CROSS, SLICE_STATE_SCROLL]
500 SLICE_STYLES = TOOL_STATES + TOOL_SLICE_STATES 501 SLICE_STYLES = TOOL_STATES + TOOL_SLICE_STATES
501 SLICE_STYLES.append(STATE_DEFAULT) 502 SLICE_STYLES.append(STATE_DEFAULT)
502 SLICE_STYLES.append(SLICE_STATE_EDITOR) 503 SLICE_STYLES.append(SLICE_STATE_EDITOR)
  504 +SLICE_STYLES.append(SLICE_STATE_WATERSHED)
503 505
504 VOLUME_STYLES = TOOL_STATES + [VOLUME_STATE_SEED, STATE_MEASURE_DISTANCE, 506 VOLUME_STYLES = TOOL_STATES + [VOLUME_STATE_SEED, STATE_MEASURE_DISTANCE,
505 STATE_MEASURE_ANGLE] 507 STATE_MEASURE_ANGLE]
@@ -507,6 +509,7 @@ VOLUME_STYLES.append(STATE_DEFAULT) @@ -507,6 +509,7 @@ VOLUME_STYLES.append(STATE_DEFAULT)
507 509
508 510
509 STYLE_LEVEL = {SLICE_STATE_EDITOR: 1, 511 STYLE_LEVEL = {SLICE_STATE_EDITOR: 1,
  512 + SLICE_STATE_WATERSHED: 1,
510 SLICE_STATE_CROSS: 2, 513 SLICE_STATE_CROSS: 2,
511 SLICE_STATE_SCROLL: 2, 514 SLICE_STATE_SCROLL: 2,
512 STATE_ANNOTATE: 2, 515 STATE_ANNOTATE: 2,
invesalius/control.py
@@ -127,6 +127,11 @@ class Controller(): @@ -127,6 +127,11 @@ class Controller():
127 answer = dialog.SaveChangesDialog2(filename) 127 answer = dialog.SaveChangesDialog2(filename)
128 if answer: 128 if answer:
129 self.ShowDialogSaveProject() 129 self.ShowDialogSaveProject()
  130 + self.CloseProject()
  131 + #Publisher.sendMessage("Enable state project", False)
  132 + Publisher.sendMessage('Set project name')
  133 + Publisher.sendMessage("Stop Config Recording")
  134 + Publisher.sendMessage("Set slice interaction style", const.STATE_DEFAULT)
130 # Import project 135 # Import project
131 dirpath = dialog.ShowImportDirDialog() 136 dirpath = dialog.ShowImportDirDialog()
132 if dirpath and not os.listdir(dirpath): 137 if dirpath and not os.listdir(dirpath):
invesalius/data/slice_.py
@@ -84,6 +84,10 @@ class Slice(object): @@ -84,6 +84,10 @@ class Slice(object):
84 self.blend_filter = None 84 self.blend_filter = None
85 self.histogram = None 85 self.histogram = None
86 self._matrix = None 86 self._matrix = None
  87 + self.aux_matrices = {}
  88 + self.state = const.STATE_DEFAULT
  89 +
  90 + self.to_show_aux = ''
87 91
88 self._type_projection = const.PROJECTION_NORMAL 92 self._type_projection = const.PROJECTION_NORMAL
89 self.n_border = const.PROJECTION_BORDER_SIZE 93 self.n_border = const.PROJECTION_BORDER_SIZE
@@ -107,6 +111,7 @@ class Slice(object): @@ -107,6 +111,7 @@ class Slice(object):
107 111
108 self.from_ = OTHER 112 self.from_ = OTHER
109 self.__bind_events() 113 self.__bind_events()
  114 + self.opacity = 0.8
110 115
111 @property 116 @property
112 def matrix(self): 117 def matrix(self):
@@ -187,6 +192,11 @@ class Slice(object): @@ -187,6 +192,11 @@ class Slice(object):
187 elif orientation == 'SAGITAL': 192 elif orientation == 'SAGITAL':
188 return shape[2] - 1 193 return shape[2] - 1
189 194
  195 + def discard_all_buffers(self):
  196 + for buffer_ in self.buffer_slices.values():
  197 + buffer_.discard_vtk_mask()
  198 + buffer_.discard_mask()
  199 +
190 def OnRemoveMasks(self, pubsub_evt): 200 def OnRemoveMasks(self, pubsub_evt):
191 selected_items = pubsub_evt.data 201 selected_items = pubsub_evt.data
192 proj = Project() 202 proj = Project()
@@ -228,6 +238,7 @@ class Slice(object): @@ -228,6 +238,7 @@ class Slice(object):
228 if (state in const.SLICE_STYLES): 238 if (state in const.SLICE_STYLES):
229 new_state = self.interaction_style.AddState(state) 239 new_state = self.interaction_style.AddState(state)
230 Publisher.sendMessage('Set slice interaction style', new_state) 240 Publisher.sendMessage('Set slice interaction style', new_state)
  241 + self.state = state
231 242
232 def OnDisableStyle(self, pubsub_evt): 243 def OnDisableStyle(self, pubsub_evt):
233 state = pubsub_evt.data 244 state = pubsub_evt.data
@@ -237,6 +248,7 @@ class Slice(object): @@ -237,6 +248,7 @@ class Slice(object):
237 248
238 if (state == const.SLICE_STATE_EDITOR): 249 if (state == const.SLICE_STATE_EDITOR):
239 Publisher.sendMessage('Set interactor default cursor') 250 Publisher.sendMessage('Set interactor default cursor')
  251 + self.state = new_state
240 252
241 def OnCloseProject(self, pubsub_evt): 253 def OnCloseProject(self, pubsub_evt):
242 self.CloseProject() 254 self.CloseProject()
@@ -249,9 +261,18 @@ class Slice(object): @@ -249,9 +261,18 @@ class Slice(object):
249 os.remove(f) 261 os.remove(f)
250 self.current_mask = None 262 self.current_mask = None
251 263
  264 + for name in self.aux_matrices:
  265 + m = self.aux_matrices[name]
  266 + f = m.filename
  267 + m._mmap.close()
  268 + m = None
  269 + os.remove(f)
  270 + self.aux_matrices = {}
  271 +
252 self.values = None 272 self.values = None
253 self.nodes = None 273 self.nodes = None
254 self.from_= OTHER 274 self.from_= OTHER
  275 + self.state = const.STATE_DEFAULT
255 276
256 self.number_of_colours = 256 277 self.number_of_colours = 256
257 self.saturation_range = (0, 0) 278 self.saturation_range = (0, 0)
@@ -380,6 +401,12 @@ class Slice(object): @@ -380,6 +401,12 @@ class Slice(object):
380 value = False 401 value = False
381 Publisher.sendMessage('Show mask', (index, value)) 402 Publisher.sendMessage('Show mask', (index, value))
382 403
  404 + def create_temp_mask(self):
  405 + temp_file = tempfile.mktemp()
  406 + shape = self.matrix.shape
  407 + matrix = numpy.memmap(temp_file, mode='w+', dtype='uint8', shape=shape)
  408 + return temp_file, matrix
  409 +
383 def edit_mask_pixel(self, operation, index, position, radius, orientation): 410 def edit_mask_pixel(self, operation, index, position, radius, orientation):
384 mask = self.buffer_slices[orientation].mask 411 mask = self.buffer_slices[orientation].mask
385 image = self.buffer_slices[orientation].image 412 image = self.buffer_slices[orientation].image
@@ -479,7 +506,7 @@ class Slice(object): @@ -479,7 +506,7 @@ class Slice(object):
479 print "Do not getting from buffer" 506 print "Do not getting from buffer"
480 n_mask = self.get_mask_slice(orientation, slice_number) 507 n_mask = self.get_mask_slice(orientation, slice_number)
481 mask = converters.to_vtk(n_mask, self.spacing, slice_number, orientation) 508 mask = converters.to_vtk(n_mask, self.spacing, slice_number, orientation)
482 - mask = self.do_colour_mask(mask) 509 + mask = self.do_colour_mask(mask, self.opacity)
483 self.buffer_slices[orientation].mask = n_mask 510 self.buffer_slices[orientation].mask = n_mask
484 final_image = self.do_blend(image, mask) 511 final_image = self.do_blend(image, mask)
485 self.buffer_slices[orientation].vtk_mask = mask 512 self.buffer_slices[orientation].vtk_mask = mask
@@ -496,7 +523,7 @@ class Slice(object): @@ -496,7 +523,7 @@ class Slice(object):
496 if self.current_mask and self.current_mask.is_shown: 523 if self.current_mask and self.current_mask.is_shown:
497 n_mask = self.get_mask_slice(orientation, slice_number) 524 n_mask = self.get_mask_slice(orientation, slice_number)
498 mask = converters.to_vtk(n_mask, self.spacing, slice_number, orientation) 525 mask = converters.to_vtk(n_mask, self.spacing, slice_number, orientation)
499 - mask = self.do_colour_mask(mask) 526 + mask = self.do_colour_mask(mask, self.opacity)
500 final_image = self.do_blend(image, mask) 527 final_image = self.do_blend(image, mask)
501 else: 528 else:
502 n_mask = None 529 n_mask = None
@@ -509,6 +536,13 @@ class Slice(object): @@ -509,6 +536,13 @@ class Slice(object):
509 self.buffer_slices[orientation].vtk_image = image 536 self.buffer_slices[orientation].vtk_image = image
510 self.buffer_slices[orientation].vtk_mask = mask 537 self.buffer_slices[orientation].vtk_mask = mask
511 538
  539 + if self.to_show_aux == 'watershed' and self.current_mask.is_shown:
  540 + m = self.get_aux_slice('watershed', orientation, slice_number)
  541 + tmp_vimage = converters.to_vtk(m, self.spacing, slice_number, orientation)
  542 + cimage = self.do_custom_colour(tmp_vimage, {0: (0.0, 0.0, 0.0, 0.0),
  543 + 1: (0.0, 1.0, 0.0, 1.0),
  544 + 2: (1.0, 0.0, 0.0, 1.0)})
  545 + final_image = self.do_blend(final_image, cimage)
512 return final_image 546 return final_image
513 547
514 def get_image_slice(self, orientation, slice_number, number_slices=1, 548 def get_image_slice(self, orientation, slice_number, number_slices=1,
@@ -701,6 +735,15 @@ class Slice(object): @@ -701,6 +735,15 @@ class Slice(object):
701 735
702 return n_mask 736 return n_mask
703 737
  738 + def get_aux_slice(self, name, orientation, n):
  739 + m = self.aux_matrices[name]
  740 + if orientation == 'AXIAL':
  741 + return numpy.array(m[n])
  742 + elif orientation == 'CORONAL':
  743 + return numpy.array(m[:, n, :])
  744 + elif orientation == 'SAGITAL':
  745 + return numpy.array(m[:, :, n])
  746 +
704 def GetNumberOfSlices(self, orientation): 747 def GetNumberOfSlices(self, orientation):
705 if orientation == 'AXIAL': 748 if orientation == 'AXIAL':
706 return self.matrix.shape[0] 749 return self.matrix.shape[0]
@@ -1162,6 +1205,19 @@ class Slice(object): @@ -1162,6 +1205,19 @@ class Slice(object):
1162 m[mask == 254] = 254 1205 m[mask == 254] = 254
1163 return m.astype('uint8') 1206 return m.astype('uint8')
1164 1207
  1208 + def do_threshold_to_all_slices(self):
  1209 + mask = self.current_mask
  1210 +
  1211 + # This is very important. Do not use masks' imagedata. It would mess up
  1212 + # surface quality event when using contour
  1213 + #self.SetMaskThreshold(mask.index, threshold)
  1214 + for n in xrange(1, mask.matrix.shape[0]):
  1215 + if mask.matrix[n, 0, 0] == 0:
  1216 + m = mask.matrix[n, 1:, 1:]
  1217 + mask.matrix[n, 1:, 1:] = self.do_threshold_to_a_slice(self.matrix[n-1], m)
  1218 +
  1219 + mask.matrix.flush()
  1220 +
1165 def do_colour_image(self, imagedata): 1221 def do_colour_image(self, imagedata):
1166 if self.from_ in (PLIST, WIDGET): 1222 if self.from_ in (PLIST, WIDGET):
1167 return imagedata 1223 return imagedata
@@ -1183,7 +1239,7 @@ class Slice(object): @@ -1183,7 +1239,7 @@ class Slice(object):
1183 1239
1184 return img_colours_bg.GetOutput() 1240 return img_colours_bg.GetOutput()
1185 1241
1186 - def do_colour_mask(self, imagedata): 1242 + def do_colour_mask(self, imagedata, opacity):
1187 scalar_range = int(imagedata.GetScalarRange()[1]) 1243 scalar_range = int(imagedata.GetScalarRange()[1])
1188 r, g, b = self.current_mask.colour 1244 r, g, b = self.current_mask.colour
1189 1245
@@ -1197,8 +1253,42 @@ class Slice(object): @@ -1197,8 +1253,42 @@ class Slice(object):
1197 lut_mask.SetNumberOfTableValues(256) 1253 lut_mask.SetNumberOfTableValues(256)
1198 lut_mask.SetTableValue(0, 0, 0, 0, 0.0) 1254 lut_mask.SetTableValue(0, 0, 0, 0, 0.0)
1199 lut_mask.SetTableValue(1, 0, 0, 0, 0.0) 1255 lut_mask.SetTableValue(1, 0, 0, 0, 0.0)
1200 - lut_mask.SetTableValue(254, r, g, b, 1.0)  
1201 - lut_mask.SetTableValue(255, r, g, b, 1.0) 1256 + lut_mask.SetTableValue(2, 0, 0, 0, 0.0)
  1257 + lut_mask.SetTableValue(253, r, g, b, opacity)
  1258 + lut_mask.SetTableValue(254, r, g, b, opacity)
  1259 + lut_mask.SetTableValue(255, r, g, b, opacity)
  1260 + lut_mask.SetRampToLinear()
  1261 + lut_mask.Build()
  1262 + # self.lut_mask = lut_mask
  1263 +
  1264 + # map the input image through a lookup table
  1265 + img_colours_mask = vtk.vtkImageMapToColors()
  1266 + img_colours_mask.SetLookupTable(lut_mask)
  1267 + img_colours_mask.SetOutputFormatToRGBA()
  1268 + img_colours_mask.SetInput(imagedata)
  1269 + img_colours_mask.Update()
  1270 + # self.img_colours_mask = img_colours_mask
  1271 +
  1272 + return img_colours_mask.GetOutput()
  1273 +
  1274 + def do_custom_colour(self, imagedata, map_colours):
  1275 + # map scalar values into colors
  1276 + minv = min(map_colours)
  1277 + maxv = max(map_colours)
  1278 + ncolours = maxv - minv + 1
  1279 +
  1280 + lut_mask = vtk.vtkLookupTable()
  1281 + lut_mask.SetNumberOfColors(ncolours)
  1282 + lut_mask.SetHueRange(const.THRESHOLD_HUE_RANGE)
  1283 + lut_mask.SetSaturationRange(1, 1)
  1284 + lut_mask.SetValueRange(minv, maxv)
  1285 + lut_mask.SetRange(minv, maxv)
  1286 + lut_mask.SetNumberOfTableValues(ncolours)
  1287 +
  1288 + for v in map_colours:
  1289 + r,g, b,a = map_colours[v]
  1290 + lut_mask.SetTableValue(v, r, g, b, a)
  1291 +
1202 lut_mask.SetRampToLinear() 1292 lut_mask.SetRampToLinear()
1203 lut_mask.Build() 1293 lut_mask.Build()
1204 # self.lut_mask = lut_mask 1294 # self.lut_mask = lut_mask
invesalius/data/styles.py
@@ -17,12 +17,22 @@ @@ -17,12 +17,22 @@
17 # detalhes. 17 # detalhes.
18 #-------------------------------------------------------------------------- 18 #--------------------------------------------------------------------------
19 19
  20 +import os
  21 +import tempfile
  22 +
20 import vtk 23 import vtk
21 import wx 24 import wx
22 25
23 from wx.lib.pubsub import pub as Publisher 26 from wx.lib.pubsub import pub as Publisher
24 27
25 import constants as const 28 import constants as const
  29 +import converters
  30 +import numpy as np
  31 +
  32 +from scipy import ndimage
  33 +from scipy.misc import imsave
  34 +from skimage.morphology import watershed
  35 +from skimage import filter
26 36
27 ORIENTATIONS = { 37 ORIENTATIONS = {
28 "AXIAL": const.AXIAL, 38 "AXIAL": const.AXIAL,
@@ -30,6 +40,20 @@ ORIENTATIONS = { @@ -30,6 +40,20 @@ ORIENTATIONS = {
30 "SAGITAL": const.SAGITAL, 40 "SAGITAL": const.SAGITAL,
31 } 41 }
32 42
  43 +BRUSH_FOREGROUND=1
  44 +BRUSH_BACKGROUND=2
  45 +BRUSH_ERASE=0
  46 +
  47 +WATERSHED_OPERATIONS = {_("Erase"): BRUSH_ERASE,
  48 + _("Foreground"): BRUSH_FOREGROUND,
  49 + _("Background"): BRUSH_BACKGROUND,}
  50 +
  51 +def get_LUT_value(data, window, level):
  52 + return np.piecewise(data,
  53 + [data <= (level - 0.5 - (window-1)/2),
  54 + data > (level - 0.5 + (window-1)/2)],
  55 + [0, 255, lambda data: ((data - (level - 0.5))/(window-1) + 0.5)*(255-0)])
  56 +
33 class BaseImageInteractorStyle(vtk.vtkInteractorStyleImage): 57 class BaseImageInteractorStyle(vtk.vtkInteractorStyleImage):
34 def __init__(self, viewer): 58 def __init__(self, viewer):
35 self.right_pressed = False 59 self.right_pressed = False
@@ -640,6 +664,394 @@ class EditorInteractorStyle(DefaultInteractorStyle): @@ -640,6 +664,394 @@ class EditorInteractorStyle(DefaultInteractorStyle):
640 return x, y, z 664 return x, y, z
641 665
642 666
  667 +class WaterShedInteractorStyle(DefaultInteractorStyle):
  668 + matrix = None
  669 + def __init__(self, viewer):
  670 + DefaultInteractorStyle.__init__(self, viewer)
  671 +
  672 + self.viewer = viewer
  673 + self.orientation = self.viewer.orientation
  674 + self.matrix = None
  675 +
  676 + self.operation = BRUSH_FOREGROUND
  677 +
  678 + self.mg_size = 3
  679 +
  680 + self.picker = vtk.vtkWorldPointPicker()
  681 +
  682 + self.AddObserver("EnterEvent", self.OnEnterInteractor)
  683 + self.AddObserver("LeaveEvent", self.OnLeaveInteractor)
  684 +
  685 + self.RemoveObservers("MouseWheelForwardEvent")
  686 + self.RemoveObservers("MouseWheelBackwardEvent")
  687 + self.AddObserver("MouseWheelForwardEvent",self.WOnScrollForward)
  688 + self.AddObserver("MouseWheelBackwardEvent", self.WOnScrollBackward)
  689 +
  690 + self.AddObserver("LeftButtonPressEvent", self.OnBrushClick)
  691 + self.AddObserver("LeftButtonReleaseEvent", self.OnBrushRelease)
  692 + self.AddObserver("MouseMoveEvent", self.OnBrushMove)
  693 +
  694 + Publisher.subscribe(self.expand_watershed, 'Expand watershed to 3D ' + self.orientation)
  695 + Publisher.subscribe(self.set_operation, 'Set watershed operation')
  696 +
  697 + def SetUp(self):
  698 + self.viewer.slice_.do_threshold_to_all_slices()
  699 + mask = self.viewer.slice_.current_mask.matrix
  700 + mask[0] = 1
  701 + mask[:, 0, :] = 1
  702 + mask[:, :, 0] = 1
  703 + self._create_mask()
  704 + self.viewer.slice_.to_show_aux = 'watershed'
  705 + self.viewer.OnScrollBar()
  706 +
  707 + def CleanUp(self):
  708 + #self._remove_mask()
  709 + Publisher.unsubscribe(self.expand_watershed, 'Expand watershed to 3D ' + self.orientation)
  710 + Publisher.unsubscribe(self.set_operation, 'Set watershed operation')
  711 + self.RemoveAllObservers()
  712 + self.viewer.slice_.to_show_aux = ''
  713 + self.viewer.OnScrollBar()
  714 +
  715 + def _create_mask(self):
  716 + if self.matrix is None:
  717 + try:
  718 + self.matrix = self.viewer.slice_.aux_matrices['watershed']
  719 + except KeyError:
  720 + self.temp_file, self.matrix = self.viewer.slice_.create_temp_mask()
  721 + self.viewer.slice_.aux_matrices['watershed'] = self.matrix
  722 +
  723 + def _remove_mask(self):
  724 + if self.matrix is not None:
  725 + self.matrix = None
  726 + os.remove(self.temp_file)
  727 + print "deleting", self.temp_file
  728 +
  729 + def set_operation(self, pubsub_evt):
  730 + self.operation = WATERSHED_OPERATIONS[pubsub_evt.data]
  731 +
  732 + def OnEnterInteractor(self, obj, evt):
  733 + if (self.viewer.slice_.buffer_slices[self.orientation].mask is None):
  734 + return
  735 + self.viewer.slice_data.cursor.Show()
  736 + #self.viewer.interactor.SetCursor(wx.StockCursor(wx.CURSOR_BLANK))
  737 + self.viewer.interactor.Render()
  738 +
  739 + def OnLeaveInteractor(self, obj, evt):
  740 + self.viewer.slice_data.cursor.Show(0)
  741 + #self.viewer.interactor.SetCursor(wx.StockCursor(wx.CURSOR_DEFAULT))
  742 + self.viewer.interactor.Render()
  743 +
  744 + def WOnScrollBackward(self, obj, evt):
  745 + viewer = self.viewer
  746 + iren = viewer.interactor
  747 + if iren.GetControlKey():
  748 + if viewer.slice_.opacity > 0:
  749 + viewer.slice_.opacity -= 0.1
  750 + self.viewer.slice_.buffer_slices['AXIAL'].discard_vtk_mask()
  751 + self.viewer.slice_.buffer_slices['CORONAL'].discard_vtk_mask()
  752 + self.viewer.slice_.buffer_slices['SAGITAL'].discard_vtk_mask()
  753 + viewer.OnScrollBar()
  754 + else:
  755 + self.OnScrollBackward(obj, evt)
  756 +
  757 +
  758 + def WOnScrollForward(self, obj, evt):
  759 + viewer = self.viewer
  760 + iren = viewer.interactor
  761 + if iren.GetControlKey():
  762 + if viewer.slice_.opacity < 1:
  763 + viewer.slice_.opacity += 0.1
  764 + self.viewer.slice_.buffer_slices['AXIAL'].discard_vtk_mask()
  765 + self.viewer.slice_.buffer_slices['CORONAL'].discard_vtk_mask()
  766 + self.viewer.slice_.buffer_slices['SAGITAL'].discard_vtk_mask()
  767 + viewer.OnScrollBar()
  768 + else:
  769 + self.OnScrollForward(obj, evt)
  770 +
  771 +
  772 + def OnBrushClick(self, obj, evt):
  773 + if (self.viewer.slice_.buffer_slices[self.orientation].mask is None):
  774 + return
  775 +
  776 + viewer = self.viewer
  777 + iren = viewer.interactor
  778 +
  779 + viewer._set_editor_cursor_visibility(1)
  780 +
  781 + mouse_x, mouse_y = iren.GetEventPosition()
  782 + render = iren.FindPokedRenderer(mouse_x, mouse_y)
  783 + slice_data = viewer.get_slice_data(render)
  784 +
  785 + # TODO: Improve!
  786 + #for i in self.slice_data_list:
  787 + #i.cursor.Show(0)
  788 + slice_data.cursor.Show()
  789 +
  790 + self.picker.Pick(mouse_x, mouse_y, 0, render)
  791 +
  792 + coord = self.get_coordinate_cursor()
  793 + position = slice_data.actor.GetInput().FindPoint(coord)
  794 +
  795 + if position != -1:
  796 + coord = slice_data.actor.GetInput().GetPoint(position)
  797 +
  798 + slice_data.cursor.SetPosition(coord)
  799 +
  800 + cursor = slice_data.cursor
  801 + position = slice_data.actor.GetInput().FindPoint(coord)
  802 + radius = cursor.radius
  803 +
  804 + if position < 0:
  805 + position = viewer.calculate_matrix_position(coord)
  806 +
  807 + operation = self.operation
  808 +
  809 + if operation == BRUSH_FOREGROUND:
  810 + if iren.GetControlKey():
  811 + operation = BRUSH_BACKGROUND
  812 + elif iren.GetShiftKey():
  813 + operation = BRUSH_ERASE
  814 + elif operation == BRUSH_BACKGROUND:
  815 + if iren.GetControlKey():
  816 + operation = BRUSH_FOREGROUND
  817 + elif iren.GetShiftKey():
  818 + operation = BRUSH_ERASE
  819 +
  820 + n = self.viewer.slice_data.number
  821 + self.edit_mask_pixel(operation, n, cursor.GetPixels(),
  822 + position, radius, self.orientation)
  823 + if self.orientation == 'AXIAL':
  824 + mask = self.matrix[n, :, :]
  825 + elif self.orientation == 'CORONAL':
  826 + mask = self.matrix[:, n, :]
  827 + elif self.orientation == 'SAGITAL':
  828 + mask = self.matrix[:, :, n]
  829 + # TODO: To create a new function to reload images to viewer.
  830 + viewer.OnScrollBar()
  831 +
  832 + def OnBrushMove(self, obj, evt):
  833 + if (self.viewer.slice_.buffer_slices[self.orientation].mask is None):
  834 + return
  835 +
  836 + viewer = self.viewer
  837 + iren = viewer.interactor
  838 +
  839 + viewer._set_editor_cursor_visibility(1)
  840 +
  841 + mouse_x, mouse_y = iren.GetEventPosition()
  842 + render = iren.FindPokedRenderer(mouse_x, mouse_y)
  843 + slice_data = viewer.get_slice_data(render)
  844 +
  845 + # TODO: Improve!
  846 + #for i in self.slice_data_list:
  847 + #i.cursor.Show(0)
  848 +
  849 + self.picker.Pick(mouse_x, mouse_y, 0, render)
  850 +
  851 + #if (self.pick.GetViewProp()):
  852 + #self.interactor.SetCursor(wx.StockCursor(wx.CURSOR_BLANK))
  853 + #else:
  854 + #self.interactor.SetCursor(wx.StockCursor(wx.CURSOR_DEFAULT))
  855 +
  856 + coord = self.get_coordinate_cursor()
  857 + position = viewer.slice_data.actor.GetInput().FindPoint(coord)
  858 +
  859 + # when position == -1 the cursos is not over the image, so is not
  860 + # necessary to set the cursor position to world coordinate center of
  861 + # pixel from slice image.
  862 + if position != -1:
  863 + coord = slice_data.actor.GetInput().GetPoint(position)
  864 + slice_data.cursor.SetPosition(coord)
  865 + #self.__update_cursor_position(slice_data, coord)
  866 +
  867 + if (self.left_pressed):
  868 + cursor = slice_data.cursor
  869 + position = slice_data.actor.GetInput().FindPoint(coord)
  870 + radius = cursor.radius
  871 +
  872 + if position < 0:
  873 + position = viewer.calculate_matrix_position(coord)
  874 +
  875 + operation = self.operation
  876 +
  877 + if operation == BRUSH_FOREGROUND:
  878 + if iren.GetControlKey():
  879 + operation = BRUSH_BACKGROUND
  880 + elif iren.GetShiftKey():
  881 + operation = BRUSH_ERASE
  882 + elif operation == BRUSH_BACKGROUND:
  883 + if iren.GetControlKey():
  884 + operation = BRUSH_FOREGROUND
  885 + elif iren.GetShiftKey():
  886 + operation = BRUSH_ERASE
  887 +
  888 + n = self.viewer.slice_data.number
  889 + self.edit_mask_pixel(operation, n, cursor.GetPixels(),
  890 + position, radius, self.orientation)
  891 + if self.orientation == 'AXIAL':
  892 + mask = self.matrix[n, :, :]
  893 + elif self.orientation == 'CORONAL':
  894 + mask = self.matrix[:, n, :]
  895 + elif self.orientation == 'SAGITAL':
  896 + mask = self.matrix[:, :, n]
  897 + # TODO: To create a new function to reload images to viewer.
  898 + viewer.OnScrollBar(update3D=False)
  899 +
  900 + else:
  901 + viewer.interactor.Render()
  902 +
  903 + def OnBrushRelease(self, evt, obj):
  904 + n = self.viewer.slice_data.number
  905 + self.viewer.slice_.discard_all_buffers()
  906 + if self.orientation == 'AXIAL':
  907 + image = self.viewer.slice_.matrix[n]
  908 + mask = self.viewer.slice_.current_mask.matrix[n+1, 1:, 1:]
  909 + self.viewer.slice_.current_mask.matrix[n+1, 0, 0] = 1
  910 + markers = self.matrix[n]
  911 +
  912 + elif self.orientation == 'CORONAL':
  913 + image = self.viewer.slice_.matrix[:, n, :]
  914 + mask = self.viewer.slice_.current_mask.matrix[1:, n+1, 1:]
  915 + self.viewer.slice_.current_mask.matrix[0, n+1, 0]
  916 + markers = self.matrix[:, n, :]
  917 +
  918 + elif self.orientation == 'SAGITAL':
  919 + image = self.viewer.slice_.matrix[:, :, n]
  920 + mask = self.viewer.slice_.current_mask.matrix[1: , 1:, n+1]
  921 + self.viewer.slice_.current_mask.matrix[0 , 0, n+1]
  922 + markers = self.matrix[:, :, n]
  923 +
  924 +
  925 + ww = self.viewer.slice_.window_width
  926 + wl = self.viewer.slice_.window_level
  927 +
  928 + if BRUSH_BACKGROUND in markers and BRUSH_FOREGROUND in markers:
  929 + tmp_image = ndimage.morphological_gradient(get_LUT_value(image, ww, wl).astype('uint16'), self.mg_size)
  930 + tmp_mask = watershed(tmp_image, markers)
  931 +
  932 + if self.viewer.overwrite_mask:
  933 + mask[:] = 0
  934 + mask[tmp_mask == 1] = 253
  935 + else:
  936 + mask[(tmp_mask==2) & ((mask == 0) | (mask == 2) | (mask == 253))] = 2
  937 + mask[(tmp_mask==1) & ((mask == 0) | (mask == 2) | (mask == 253))] = 253
  938 +
  939 +
  940 + self.viewer.slice_.current_mask.was_edited = True
  941 + self.viewer.slice_.current_mask.clear_history()
  942 + Publisher.sendMessage('Reload actual slice')
  943 + else:
  944 + self.viewer.OnScrollBar(update3D=False)
  945 +
  946 + def get_coordinate_cursor(self):
  947 + # Find position
  948 + x, y, z = self.picker.GetPickPosition()
  949 + bounds = self.viewer.slice_data.actor.GetBounds()
  950 + if bounds[0] == bounds[1]:
  951 + x = bounds[0]
  952 + elif bounds[2] == bounds[3]:
  953 + y = bounds[2]
  954 + elif bounds[4] == bounds[5]:
  955 + z = bounds[4]
  956 + return x, y, z
  957 +
  958 + def edit_mask_pixel(self, operation, n, index, position, radius, orientation):
  959 + if orientation == 'AXIAL':
  960 + mask = self.matrix[n, :, :]
  961 + elif orientation == 'CORONAL':
  962 + mask = self.matrix[:, n, :]
  963 + elif orientation == 'SAGITAL':
  964 + mask = self.matrix[:, :, n]
  965 +
  966 + spacing = self.viewer.slice_.spacing
  967 + if hasattr(position, '__iter__'):
  968 + py, px = position
  969 + if orientation == 'AXIAL':
  970 + sx = spacing[0]
  971 + sy = spacing[1]
  972 + elif orientation == 'CORONAL':
  973 + sx = spacing[0]
  974 + sy = spacing[2]
  975 + elif orientation == 'SAGITAL':
  976 + sx = spacing[2]
  977 + sy = spacing[1]
  978 +
  979 + else:
  980 + if orientation == 'AXIAL':
  981 + sx = spacing[0]
  982 + sy = spacing[1]
  983 + py = position / mask.shape[1]
  984 + px = position % mask.shape[1]
  985 + elif orientation == 'CORONAL':
  986 + sx = spacing[0]
  987 + sy = spacing[2]
  988 + py = position / mask.shape[1]
  989 + px = position % mask.shape[1]
  990 + elif orientation == 'SAGITAL':
  991 + sx = spacing[2]
  992 + sy = spacing[1]
  993 + py = position / mask.shape[1]
  994 + px = position % mask.shape[1]
  995 +
  996 + cx = index.shape[1] / 2 + 1
  997 + cy = index.shape[0] / 2 + 1
  998 + xi = px - index.shape[1] + cx
  999 + xf = xi + index.shape[1]
  1000 + yi = py - index.shape[0] + cy
  1001 + yf = yi + index.shape[0]
  1002 +
  1003 + if yi < 0:
  1004 + index = index[abs(yi):,:]
  1005 + yi = 0
  1006 + if yf > mask.shape[0]:
  1007 + index = index[:index.shape[0]-(yf-mask.shape[0]), :]
  1008 + yf = mask.shape[0]
  1009 +
  1010 + if xi < 0:
  1011 + index = index[:,abs(xi):]
  1012 + xi = 0
  1013 + if xf > mask.shape[1]:
  1014 + index = index[:,:index.shape[1]-(xf-mask.shape[1])]
  1015 + xf = mask.shape[1]
  1016 +
  1017 + # Verifying if the points is over the image array.
  1018 + if (not 0 <= xi <= mask.shape[1] and not 0 <= xf <= mask.shape[1]) or \
  1019 + (not 0 <= yi <= mask.shape[0] and not 0 <= yf <= mask.shape[0]):
  1020 + return
  1021 +
  1022 + roi_m = mask[yi:yf,xi:xf]
  1023 +
  1024 + # Checking if roi_i has at least one element.
  1025 + if roi_m.size:
  1026 + roi_m[index] = operation
  1027 +
  1028 + def expand_watershed(self, pubsub_evt):
  1029 + markers = self.matrix
  1030 + image = self.viewer.slice_.matrix
  1031 + mask = self.viewer.slice_.current_mask.matrix[1:, 1:, 1:]
  1032 + ww = self.viewer.slice_.window_width
  1033 + wl = self.viewer.slice_.window_level
  1034 + if BRUSH_BACKGROUND in markers and BRUSH_FOREGROUND in markers:
  1035 + tmp_image = ndimage.morphological_gradient(get_LUT_value(image, ww, wl).astype('uint16'), self.mg_size)
  1036 + tmp_mask = watershed(tmp_image, markers)
  1037 +
  1038 + if self.viewer.overwrite_mask:
  1039 + mask[:] = 0
  1040 + mask[tmp_mask == 1] = 253
  1041 + else:
  1042 + mask[(tmp_mask==2) & ((mask == 0) | (mask == 2) | (mask == 253))] = 2
  1043 + mask[(tmp_mask==1) & ((mask == 0) | (mask == 2) | (mask == 253))] = 253
  1044 +
  1045 + #mask[:] = tmp_mask
  1046 + self.viewer.slice_.current_mask.matrix[0] = 1
  1047 + self.viewer.slice_.current_mask.matrix[:, 0, :] = 1
  1048 + self.viewer.slice_.current_mask.matrix[:, :, 0] = 1
  1049 +
  1050 + self.viewer.slice_.discard_all_buffers()
  1051 + self.viewer.slice_.current_mask.clear_history()
  1052 + Publisher.sendMessage('Reload actual slice')
  1053 +
  1054 +
643 def get_style(style): 1055 def get_style(style):
644 STYLES = { 1056 STYLES = {
645 const.STATE_DEFAULT: DefaultInteractorStyle, 1057 const.STATE_DEFAULT: DefaultInteractorStyle,
@@ -653,5 +1065,6 @@ def get_style(style): @@ -653,5 +1065,6 @@ def get_style(style):
653 const.STATE_ZOOM_SL: ZoomSLInteractorStyle, 1065 const.STATE_ZOOM_SL: ZoomSLInteractorStyle,
654 const.SLICE_STATE_SCROLL: ChangeSliceInteractorStyle, 1066 const.SLICE_STATE_SCROLL: ChangeSliceInteractorStyle,
655 const.SLICE_STATE_EDITOR: EditorInteractorStyle, 1067 const.SLICE_STATE_EDITOR: EditorInteractorStyle,
  1068 + const.SLICE_STATE_WATERSHED: WaterShedInteractorStyle,
656 } 1069 }
657 return STYLES[style] 1070 return STYLES[style]
invesalius/data/viewer_slice.py
@@ -166,6 +166,8 @@ class Viewer(wx.Panel): @@ -166,6 +166,8 @@ class Viewer(wx.Panel):
166 self.last_position_mouse_move = () 166 self.last_position_mouse_move = ()
167 self.state = const.STATE_DEFAULT 167 self.state = const.STATE_DEFAULT
168 168
  169 + self.overwrite_mask = False
  170 +
169 # All renderers and image actors in this viewer 171 # All renderers and image actors in this viewer
170 self.slice_data_list = [] 172 self.slice_data_list = []
171 self.slice_data = None 173 self.slice_data = None
@@ -281,6 +283,8 @@ class Viewer(wx.Panel): @@ -281,6 +283,8 @@ class Viewer(wx.Panel):
281 if cleanup: 283 if cleanup:
282 self.style.CleanUp() 284 self.style.CleanUp()
283 285
  286 + del self.style
  287 +
284 style = styles.get_style(state)(self) 288 style = styles.get_style(state)(self)
285 289
286 setup = getattr(style, 'SetUp', None) 290 setup = getattr(style, 'SetUp', None)
@@ -751,6 +755,8 @@ class Viewer(wx.Panel): @@ -751,6 +755,8 @@ class Viewer(wx.Panel):
751 Publisher.subscribe(self.OnSetMIPInvert, 'Set MIP Invert %s' % self.orientation) 755 Publisher.subscribe(self.OnSetMIPInvert, 'Set MIP Invert %s' % self.orientation)
752 Publisher.subscribe(self.OnShowMIPInterface, 'Show MIP interface') 756 Publisher.subscribe(self.OnShowMIPInterface, 'Show MIP interface')
753 757
  758 + Publisher.subscribe(self.OnSetOverwriteMask, "Set overwrite mask")
  759 +
754 def SetDefaultCursor(self, pusub_evt): 760 def SetDefaultCursor(self, pusub_evt):
755 self.interactor.SetCursor(wx.StockCursor(wx.CURSOR_DEFAULT)) 761 self.interactor.SetCursor(wx.StockCursor(wx.CURSOR_DEFAULT))
756 762
@@ -1247,7 +1253,10 @@ class Viewer(wx.Panel): @@ -1247,7 +1253,10 @@ class Viewer(wx.Panel):
1247 self.mip_ctrls.Hide() 1253 self.mip_ctrls.Hide()
1248 self.GetSizer().Remove(self.mip_ctrls) 1254 self.GetSizer().Remove(self.mip_ctrls)
1249 self.Layout() 1255 self.Layout()
1250 - 1256 +
  1257 + def OnSetOverwriteMask(self, pubsub_evt):
  1258 + value = pubsub_evt.data
  1259 + self.overwrite_mask = value
1251 1260
1252 def set_slice_number(self, index): 1261 def set_slice_number(self, index):
1253 inverted = self.mip_ctrls.inverted.GetValue() 1262 inverted = self.mip_ctrls.inverted.GetValue()
invesalius/gui/frame.py
@@ -26,6 +26,8 @@ import webbrowser @@ -26,6 +26,8 @@ import webbrowser
26 import wx 26 import wx
27 import wx.aui 27 import wx.aui
28 from wx.lib.pubsub import pub as Publisher 28 from wx.lib.pubsub import pub as Publisher
  29 +import wx.lib.agw.toasterbox as TB
  30 +import wx.lib.popupctl as pc
29 31
30 import constants as const 32 import constants as const
31 import default_tasks as tasks 33 import default_tasks as tasks
@@ -45,6 +47,19 @@ VIEW_TOOLS = [ID_LAYOUT, ID_TEXT] =\ @@ -45,6 +47,19 @@ VIEW_TOOLS = [ID_LAYOUT, ID_TEXT] =\
45 47
46 48
47 49
  50 +class MessageWatershed(wx.PopupWindow):
  51 + def __init__(self, prnt, msg):
  52 + wx.PopupWindow.__init__(self, prnt, -1)
  53 + self.txt = wx.StaticText(self, -1, msg)
  54 +
  55 + self.sizer = wx.BoxSizer(wx.HORIZONTAL)
  56 + self.sizer.Add(self.txt, 1, wx.EXPAND)
  57 + self.SetSizer(self.sizer)
  58 +
  59 + self.sizer.Fit(self)
  60 + self.Layout()
  61 + self.Update()
  62 + self.SetAutoLayout(1)
48 63
49 64
50 65
@@ -64,6 +79,8 @@ class Frame(wx.Frame): @@ -64,6 +79,8 @@ class Frame(wx.Frame):
64 self.Center(wx.BOTH) 79 self.Center(wx.BOTH)
65 icon_path = os.path.join(const.ICON_DIR, "invesalius.ico") 80 icon_path = os.path.join(const.ICON_DIR, "invesalius.ico")
66 self.SetIcon(wx.Icon(icon_path, wx.BITMAP_TYPE_ICO)) 81 self.SetIcon(wx.Icon(icon_path, wx.BITMAP_TYPE_ICO))
  82 +
  83 + self.mw = None
67 84
68 if sys.platform != 'darwin': 85 if sys.platform != 'darwin':
69 self.Maximize() 86 self.Maximize()
@@ -104,6 +121,7 @@ class Frame(wx.Frame): @@ -104,6 +121,7 @@ class Frame(wx.Frame):
104 sub(self._SetProjectName, 'Set project name') 121 sub(self._SetProjectName, 'Set project name')
105 sub(self._ShowContentPanel, 'Show content panel') 122 sub(self._ShowContentPanel, 'Show content panel')
106 sub(self._ShowImportPanel, 'Show import panel in frame') 123 sub(self._ShowImportPanel, 'Show import panel in frame')
  124 + #sub(self._ShowHelpMessage, 'Show help message')
107 sub(self._ShowImportNetwork, 'Show retrieve dicom panel') 125 sub(self._ShowImportNetwork, 'Show retrieve dicom panel')
108 sub(self._ShowTask, 'Show task panel') 126 sub(self._ShowTask, 'Show task panel')
109 sub(self._UpdateAUI, 'Update AUI') 127 sub(self._UpdateAUI, 'Update AUI')
@@ -116,6 +134,7 @@ class Frame(wx.Frame): @@ -116,6 +134,7 @@ class Frame(wx.Frame):
116 self.Bind(wx.EVT_SIZE, self.OnSize) 134 self.Bind(wx.EVT_SIZE, self.OnSize)
117 self.Bind(wx.EVT_MENU, self.OnMenuClick) 135 self.Bind(wx.EVT_MENU, self.OnMenuClick)
118 self.Bind(wx.EVT_CLOSE, self.OnClose) 136 self.Bind(wx.EVT_CLOSE, self.OnClose)
  137 + #self.Bind(wx.EVT_MOVE, self.OnMove)
119 138
120 def __init_aui(self): 139 def __init_aui(self):
121 """ 140 """
@@ -289,6 +308,14 @@ class Frame(wx.Frame): @@ -289,6 +308,14 @@ class Frame(wx.Frame):
289 aui_manager.GetPane("Import").Show(0) 308 aui_manager.GetPane("Import").Show(0)
290 aui_manager.Update() 309 aui_manager.Update()
291 310
  311 + def _ShowHelpMessage(self, evt_pubsub):
  312 + aui_manager = self.aui_manager
  313 + pos = aui_manager.GetPane("Data").window.GetScreenPosition()
  314 + msg = evt_pubsub.data
  315 + self.mw = MessageWatershed(self, msg)
  316 + self.mw.SetPosition(pos)
  317 + self.mw.Show()
  318 +
292 def _ShowImportPanel(self, evt_pubsub): 319 def _ShowImportPanel(self, evt_pubsub):
293 """ 320 """
294 Show only DICOM import panel. 321 Show only DICOM import panel.
@@ -378,6 +405,12 @@ class Frame(wx.Frame): @@ -378,6 +405,12 @@ class Frame(wx.Frame):
378 Publisher.sendMessage(('ProgressBar Reposition')) 405 Publisher.sendMessage(('ProgressBar Reposition'))
379 evt.Skip() 406 evt.Skip()
380 407
  408 +
  409 + def OnMove(self, evt):
  410 + aui_manager = self.aui_manager
  411 + pos = aui_manager.GetPane("Data").window.GetScreenPosition()
  412 + self.mw.SetPosition(pos)
  413 +
381 def ShowPreferences(self): 414 def ShowPreferences(self):
382 415
383 if self.preferences.ShowModal() == wx.ID_OK: 416 if self.preferences.ShowModal() == wx.ID_OK:
invesalius/gui/task_slice.py
@@ -228,7 +228,7 @@ class InnerFoldPanel(wx.Panel): @@ -228,7 +228,7 @@ class InnerFoldPanel(wx.Panel):
228 # parent panel. Perhaps we need to insert the item into the sizer also... 228 # parent panel. Perhaps we need to insert the item into the sizer also...
229 # Study this. 229 # Study this.
230 fold_panel = fpb.FoldPanelBar(self, -1, wx.DefaultPosition, 230 fold_panel = fpb.FoldPanelBar(self, -1, wx.DefaultPosition,
231 - (10, 190), 0,fpb.FPB_SINGLE_FOLD) 231 + (10, 220), 0,fpb.FPB_SINGLE_FOLD)
232 232
233 # Fold panel style 233 # Fold panel style
234 style = fpb.CaptionBarStyle() 234 style = fpb.CaptionBarStyle()
@@ -252,6 +252,13 @@ class InnerFoldPanel(wx.Panel): @@ -252,6 +252,13 @@ class InnerFoldPanel(wx.Panel):
252 self.__id_editor = item.GetId() 252 self.__id_editor = item.GetId()
253 self.last_panel_opened = None 253 self.last_panel_opened = None
254 254
  255 + # Fold 3 - Watershed
  256 + item = fold_panel.AddFoldPanel(_("Watershed"), collapsed=True)
  257 + fold_panel.ApplyCaptionStyle(item, style)
  258 + fold_panel.AddFoldPanelWindow(item, WatershedTool(item), Spacing= 0,
  259 + leftSpacing=0, rightSpacing=0)
  260 + self.__id_watershed = item.GetId()
  261 +
255 #fold_panel.Expand(fold_panel.GetFoldPanel(1)) 262 #fold_panel.Expand(fold_panel.GetFoldPanel(1))
256 263
257 # Panel sizer to expand fold panel 264 # Panel sizer to expand fold panel
@@ -274,6 +281,7 @@ class InnerFoldPanel(wx.Panel): @@ -274,6 +281,7 @@ class InnerFoldPanel(wx.Panel):
274 def __bind_pubsub_evt(self): 281 def __bind_pubsub_evt(self):
275 Publisher.subscribe(self.OnRetrieveStyle, 'Retrieve task slice style') 282 Publisher.subscribe(self.OnRetrieveStyle, 'Retrieve task slice style')
276 Publisher.subscribe(self.OnDisableStyle, 'Disable task slice style') 283 Publisher.subscribe(self.OnDisableStyle, 'Disable task slice style')
  284 + Publisher.subscribe(self.OnCloseProject, 'Close project data')
277 285
278 def OnFoldPressCaption(self, evt): 286 def OnFoldPressCaption(self, evt):
279 id = evt.GetTag().GetId() 287 id = evt.GetTag().GetId()
@@ -284,8 +292,18 @@ class InnerFoldPanel(wx.Panel): @@ -284,8 +292,18 @@ class InnerFoldPanel(wx.Panel):
284 Publisher.sendMessage('Disable style', const.SLICE_STATE_EDITOR) 292 Publisher.sendMessage('Disable style', const.SLICE_STATE_EDITOR)
285 self.last_style = None 293 self.last_style = None
286 else: 294 else:
287 - Publisher.sendMessage('Enable style', const.SLICE_STATE_EDITOR) 295 + Publisher.sendMessage('Enable style',
  296 + const.SLICE_STATE_EDITOR)
288 self.last_style = const.SLICE_STATE_EDITOR 297 self.last_style = const.SLICE_STATE_EDITOR
  298 + elif self.__id_watershed == id:
  299 + if closed:
  300 + Publisher.sendMessage('Disable style',
  301 + const.SLICE_STATE_WATERSHED)
  302 + self.last_style = None
  303 + else:
  304 + Publisher.sendMessage('Enable style', const.SLICE_STATE_WATERSHED)
  305 + Publisher.sendMessage('Show help message', 'Mark the object and the background')
  306 + self.last_style = const.SLICE_STATE_WATERSHED
289 else: 307 else:
290 Publisher.sendMessage('Disable style', const.SLICE_STATE_EDITOR) 308 Publisher.sendMessage('Disable style', const.SLICE_STATE_EDITOR)
291 self.last_style = None 309 self.last_style = None
@@ -300,6 +318,9 @@ class InnerFoldPanel(wx.Panel): @@ -300,6 +318,9 @@ class InnerFoldPanel(wx.Panel):
300 if (self.last_style == const.SLICE_STATE_EDITOR): 318 if (self.last_style == const.SLICE_STATE_EDITOR):
301 Publisher.sendMessage('Disable style', const.SLICE_STATE_EDITOR) 319 Publisher.sendMessage('Disable style', const.SLICE_STATE_EDITOR)
302 320
  321 + def OnCloseProject(self, pubsub_evt):
  322 + self.fold_panel.Expand(self.fold_panel.GetFoldPanel(0))
  323 +
303 def GetMaskSelected(self): 324 def GetMaskSelected(self):
304 x= self.mask_prop_panel.GetMaskSelected() 325 x= self.mask_prop_panel.GetMaskSelected()
305 return self.mask_prop_panel.GetMaskSelected() 326 return self.mask_prop_panel.GetMaskSelected()
@@ -688,3 +709,143 @@ class EditionTools(wx.Panel): @@ -688,3 +709,143 @@ class EditionTools(wx.Panel):
688 Publisher.sendMessage('Set edition operation', brush_op_id) 709 Publisher.sendMessage('Set edition operation', brush_op_id)
689 710
690 711
  712 +class WatershedTool(EditionTools):
  713 + def __init__(self, parent):
  714 + wx.Panel.__init__(self, parent, size=(50,150))
  715 + default_colour = wx.SystemSettings_GetColour(wx.SYS_COLOUR_MENUBAR)
  716 + self.SetBackgroundColour(default_colour)
  717 +
  718 + ## LINE 1
  719 + text1 = wx.StaticText(self, -1, _("Choose brush type and size:"))
  720 +
  721 + ## LINE 2
  722 + menu = wx.Menu()
  723 +
  724 + CIRCLE_BMP = wx.Bitmap("../icons/brush_circle.jpg", wx.BITMAP_TYPE_JPEG)
  725 + item = wx.MenuItem(menu, MENU_BRUSH_CIRCLE, _("Circle"))
  726 + item.SetBitmap(CIRCLE_BMP)
  727 +
  728 + SQUARE_BMP = wx.Bitmap("../icons/brush_square.jpg", wx.BITMAP_TYPE_JPEG)
  729 + item2 = wx.MenuItem(menu, MENU_BRUSH_SQUARE, _("Square"))
  730 + item2.SetBitmap(SQUARE_BMP)
  731 +
  732 + menu.AppendItem(item)
  733 + menu.AppendItem(item2)
  734 +
  735 + bmp_brush_format = {const.BRUSH_CIRCLE: CIRCLE_BMP,
  736 + const.BRUSH_SQUARE: SQUARE_BMP}
  737 + selected_bmp = bmp_brush_format[const.DEFAULT_BRUSH_FORMAT]
  738 +
  739 + btn_brush_format = pbtn.PlateButton(self, wx.ID_ANY,"", selected_bmp,
  740 + style=pbtn.PB_STYLE_SQUARE)
  741 + btn_brush_format.SetMenu(menu)
  742 + self.btn_brush_format = btn_brush_format
  743 +
  744 + spin_brush_size = wx.SpinCtrl(self, -1, "", (20, 50))
  745 + spin_brush_size.SetRange(1,100)
  746 + spin_brush_size.SetValue(const.BRUSH_SIZE)
  747 + spin_brush_size.Bind(wx.EVT_TEXT, self.OnBrushSize)
  748 + self.spin = spin_brush_size
  749 +
  750 + combo_brush_op = wx.ComboBox(self, -1, "", size=(15,-1),
  751 + choices = (_("Foreground"),
  752 + _("Background"),
  753 + _("Erase")),
  754 + style = wx.CB_DROPDOWN|wx.CB_READONLY)
  755 + combo_brush_op.SetSelection(0)
  756 + if sys.platform != 'win32':
  757 + combo_brush_op.SetWindowVariant(wx.WINDOW_VARIANT_SMALL)
  758 + self.combo_brush_op = combo_brush_op
  759 +
  760 + # Sizer which represents the second line
  761 + line2 = wx.BoxSizer(wx.HORIZONTAL)
  762 + line2.Add(btn_brush_format, 0, wx.EXPAND|wx.GROW|wx.TOP|wx.RIGHT, 0)
  763 + line2.Add(spin_brush_size, 0, wx.RIGHT, 5)
  764 + line2.Add(combo_brush_op, 1, wx.EXPAND|wx.TOP|wx.RIGHT|wx.LEFT, 5)
  765 +
  766 + ## LINE 3
  767 +
  768 + ## LINE 4
  769 +
  770 + # LINE 5
  771 + check_box = wx.CheckBox(self, -1, _("Overwrite mask"))
  772 + self.check_box = check_box
  773 +
  774 + # Line 6
  775 + self.btn_exp_watershed = wx.Button(self, -1, _('Expand watershed to 3D'))
  776 +
  777 + # Add lines into main sizer
  778 + sizer = wx.BoxSizer(wx.VERTICAL)
  779 + sizer.Add(text1, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
  780 + sizer.Add(line2, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
  781 + sizer.Add(check_box, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
  782 + sizer.Add(self.btn_exp_watershed, 0, wx.GROW|wx.EXPAND|wx.LEFT|wx.RIGHT|wx.TOP, 5)
  783 + sizer.Fit(self)
  784 +
  785 + self.SetSizer(sizer)
  786 + self.Update()
  787 + self.SetAutoLayout(1)
  788 +
  789 + self.__bind_events_wx()
  790 +
  791 +
  792 + def __bind_events_wx(self):
  793 + self.Bind(wx.EVT_MENU, self.OnMenu)
  794 + self.combo_brush_op.Bind(wx.EVT_COMBOBOX, self.OnComboBrushOp)
  795 + self.check_box.Bind(wx.EVT_CHECKBOX, self.OnCheckOverwriteMask)
  796 + self.btn_exp_watershed.Bind(wx.EVT_BUTTON, self.OnExpandWatershed)
  797 +
  798 + def ChangeMaskColour(self, pubsub_evt):
  799 + colour = pubsub_evt.data
  800 + self.gradient_thresh.SetColour(colour)
  801 +
  802 + def SetGradientColour(self, pubsub_evt):
  803 + vtk_colour = pubsub_evt.data[3]
  804 + wx_colour = [c*255 for c in vtk_colour]
  805 + self.gradient_thresh.SetColour(wx_colour)
  806 +
  807 + def SetThresholdValues(self, pubsub_evt):
  808 + thresh_min, thresh_max = pubsub_evt.data
  809 + self.bind_evt_gradient = False
  810 + self.gradient_thresh.SetMinValue(thresh_min)
  811 + self.gradient_thresh.SetMaxValue(thresh_max)
  812 + self.bind_evt_gradient = True
  813 +
  814 + def SetThresholdBounds(self, pubsub_evt):
  815 + thresh_min = pubsub_evt.data[0]
  816 + thresh_max = pubsub_evt.data[1]
  817 + self.gradient_thresh.SetMinRange(thresh_min)
  818 + self.gradient_thresh.SetMaxRange(thresh_max)
  819 + self.gradient_thresh.SetMinValue(thresh_min)
  820 + self.gradient_thresh.SetMaxValue(thresh_max)
  821 +
  822 + def OnMenu(self, evt):
  823 + SQUARE_BMP = wx.Bitmap("../icons/brush_square.jpg", wx.BITMAP_TYPE_JPEG)
  824 + CIRCLE_BMP = wx.Bitmap("../icons/brush_circle.jpg", wx.BITMAP_TYPE_JPEG)
  825 +
  826 + brush = {MENU_BRUSH_CIRCLE: const.BRUSH_CIRCLE,
  827 + MENU_BRUSH_SQUARE: const.BRUSH_SQUARE}
  828 + bitmap = {MENU_BRUSH_CIRCLE: CIRCLE_BMP,
  829 + MENU_BRUSH_SQUARE: SQUARE_BMP}
  830 +
  831 + self.btn_brush_format.SetBitmap(bitmap[evt.GetId()])
  832 +
  833 + Publisher.sendMessage('Set brush format', brush[evt.GetId()])
  834 +
  835 + def OnBrushSize(self, evt):
  836 + """ """
  837 + # FIXME: Using wx.EVT_SPINCTRL in MacOS it doesnt capture changes only
  838 + # in the text ctrl - so we are capturing only changes on text
  839 + # Strangelly this is being called twice
  840 + Publisher.sendMessage('Set edition brush size',self.spin.GetValue())
  841 +
  842 + def OnComboBrushOp(self, evt):
  843 + brush_op = self.combo_brush_op.GetValue()
  844 + Publisher.sendMessage('Set watershed operation', brush_op)
  845 +
  846 + def OnCheckOverwriteMask(self, evt):
  847 + value = self.check_box.GetValue()
  848 + Publisher.sendMessage('Set overwrite mask', value)
  849 +
  850 + def OnExpandWatershed(self, evt):
  851 + Publisher.sendMessage('Expand watershed to 3D AXIAL')