Commit f7f83def96c311a84e7b30a2d572892dfe3a3a1a

Authored by Thiago Franco de Moraes
1 parent a4f6c33f
Exists in watershed

Starting to using watershed

invesalius/constants.py
... ... @@ -485,6 +485,7 @@ STATE_MEASURE_ANGLE = 1008
485 485 SLICE_STATE_CROSS = 3006
486 486 SLICE_STATE_SCROLL = 3007
487 487 SLICE_STATE_EDITOR = 3008
  488 +SLICE_STATE_WATERSHED = 3009
488 489  
489 490 VOLUME_STATE_SEED = 2001
490 491 #STATE_LINEAR_MEASURE = 3001
... ... @@ -500,6 +501,7 @@ TOOL_SLICE_STATES = [SLICE_STATE_CROSS, SLICE_STATE_SCROLL]
500 501 SLICE_STYLES = TOOL_STATES + TOOL_SLICE_STATES
501 502 SLICE_STYLES.append(STATE_DEFAULT)
502 503 SLICE_STYLES.append(SLICE_STATE_EDITOR)
  504 +SLICE_STYLES.append(SLICE_STATE_WATERSHED)
503 505  
504 506 VOLUME_STYLES = TOOL_STATES + [VOLUME_STATE_SEED, STATE_MEASURE_DISTANCE,
505 507 STATE_MEASURE_ANGLE]
... ... @@ -507,6 +509,7 @@ VOLUME_STYLES.append(STATE_DEFAULT)
507 509  
508 510  
509 511 STYLE_LEVEL = {SLICE_STATE_EDITOR: 1,
  512 + SLICE_STATE_WATERSHED: 1,
510 513 SLICE_STATE_CROSS: 2,
511 514 SLICE_STATE_SCROLL: 2,
512 515 STATE_ANNOTATE: 2,
... ...
invesalius/data/slice_.py
... ... @@ -107,6 +107,9 @@ class Slice(object):
107 107  
108 108 self.from_ = OTHER
109 109 self.__bind_events()
  110 + self.qblend = {'AXIAL': {},
  111 + 'CORONAL': {},
  112 + 'SAGITAL': {}}
110 113  
111 114 @property
112 115 def matrix(self):
... ... @@ -187,6 +190,11 @@ class Slice(object):
187 190 elif orientation == 'SAGITAL':
188 191 return shape[2] - 1
189 192  
  193 + def discard_all_buffers(self):
  194 + for buffer_ in self.buffer_slices.values():
  195 + buffer_.discard_vtk_mask()
  196 + buffer_.discard_mask()
  197 +
190 198 def OnRemoveMasks(self, pubsub_evt):
191 199 selected_items = pubsub_evt.data
192 200 proj = Project()
... ... @@ -380,6 +388,12 @@ class Slice(object):
380 388 value = False
381 389 Publisher.sendMessage('Show mask', (index, value))
382 390  
  391 + def create_temp_mask(self):
  392 + temp_file = tempfile.mktemp()
  393 + shape = self.matrix.shape
  394 + matrix = numpy.memmap(temp_file, mode='w+', dtype='int8', shape=shape)
  395 + return temp_file, matrix
  396 +
383 397 def edit_mask_pixel(self, operation, index, position, radius, orientation):
384 398 mask = self.buffer_slices[orientation].mask
385 399 image = self.buffer_slices[orientation].image
... ... @@ -509,6 +523,11 @@ class Slice(object):
509 523 self.buffer_slices[orientation].vtk_image = image
510 524 self.buffer_slices[orientation].vtk_mask = mask
511 525  
  526 + print self.qblend
  527 + if self.qblend[orientation].get(slice_number, None) is not None:
  528 + print "BLENDING"
  529 + final_image = self.do_blend(final_image,
  530 + self.qblend[orientation][slice_number])
512 531 return final_image
513 532  
514 533 def get_image_slice(self, orientation, slice_number, number_slices=1,
... ...
invesalius/data/styles.py
... ... @@ -17,12 +17,19 @@
17 17 # detalhes.
18 18 #--------------------------------------------------------------------------
19 19  
  20 +import os
  21 +import tempfile
  22 +
20 23 import vtk
21 24 import wx
22 25  
23 26 from wx.lib.pubsub import pub as Publisher
24 27  
25 28 import constants as const
  29 +import converters
  30 +import numpy as np
  31 +
  32 +from scipy import ndimage
26 33  
27 34 ORIENTATIONS = {
28 35 "AXIAL": const.AXIAL,
... ... @@ -640,6 +647,306 @@ class EditorInteractorStyle(DefaultInteractorStyle):
640 647 return x, y, z
641 648  
642 649  
  650 +class WaterShedInteractorStyle(DefaultInteractorStyle):
  651 + matrix = None
  652 + def __init__(self, viewer):
  653 + DefaultInteractorStyle.__init__(self, viewer)
  654 +
  655 + self.viewer = viewer
  656 + self.orientation = self.viewer.orientation
  657 +
  658 + self.foreground = False
  659 + self.background = False
  660 +
  661 + self.picker = vtk.vtkWorldPointPicker()
  662 +
  663 + self.AddObserver("EnterEvent", self.OnEnterInteractor)
  664 + self.AddObserver("LeaveEvent", self.OnLeaveInteractor)
  665 +
  666 + self.AddObserver("LeftButtonPressEvent", self.OnBrushClick)
  667 + self.AddObserver("LeftButtonReleaseEvent", self.OnBrushRelease)
  668 + self.AddObserver("MouseMoveEvent", self.OnBrushMove)
  669 +
  670 + def SetUp(self):
  671 + self._create_mask()
  672 +
  673 + def CleanUp(self):
  674 + self._remove_mask()
  675 +
  676 + def _create_mask(self):
  677 + if self.matrix is None:
  678 + self.temp_file, self.matrix = self.viewer.slice_.create_temp_mask()
  679 + print "created", self.temp_file
  680 +
  681 + def _remove_mask(self):
  682 + if self.matrix is not None:
  683 + self.matrix = None
  684 + os.remove(self.temp_file)
  685 + print "deleting", self.temp_file
  686 +
  687 + def OnEnterInteractor(self, obj, evt):
  688 + if (self.viewer.slice_.buffer_slices[self.orientation].mask is None):
  689 + return
  690 + self.viewer.slice_data.cursor.Show()
  691 + self.viewer.interactor.SetCursor(wx.StockCursor(wx.CURSOR_BLANK))
  692 + self.viewer.interactor.Render()
  693 +
  694 + def OnLeaveInteractor(self, obj, evt):
  695 + self.viewer.slice_data.cursor.Show(0)
  696 + self.viewer.interactor.SetCursor(wx.StockCursor(wx.CURSOR_DEFAULT))
  697 + self.viewer.interactor.Render()
  698 +
  699 + def OnBrushClick(self, obj, evt):
  700 + if (self.viewer.slice_.buffer_slices[self.orientation].mask is None):
  701 + return
  702 +
  703 + viewer = self.viewer
  704 + iren = viewer.interactor
  705 +
  706 + viewer._set_editor_cursor_visibility(1)
  707 +
  708 + mouse_x, mouse_y = iren.GetEventPosition()
  709 + render = iren.FindPokedRenderer(mouse_x, mouse_y)
  710 + slice_data = viewer.get_slice_data(render)
  711 +
  712 + # TODO: Improve!
  713 + #for i in self.slice_data_list:
  714 + #i.cursor.Show(0)
  715 + slice_data.cursor.Show()
  716 +
  717 + self.picker.Pick(mouse_x, mouse_y, 0, render)
  718 +
  719 + coord = self.get_coordinate_cursor()
  720 + position = slice_data.actor.GetInput().FindPoint(coord)
  721 +
  722 + if position != -1:
  723 + coord = slice_data.actor.GetInput().GetPoint(position)
  724 +
  725 + slice_data.cursor.SetPosition(coord)
  726 + cursor = slice_data.cursor
  727 + radius = cursor.radius
  728 +
  729 + if position < 0:
  730 + position = viewer.calculate_matrix_position(coord)
  731 +
  732 + n = self.viewer.slice_data.number
  733 + self.edit_mask_pixel(viewer._brush_cursor_op, n, cursor.GetPixels(),
  734 + position, radius, viewer.orientation)
  735 + viewer._flush_buffer = True
  736 +
  737 + # TODO: To create a new function to reload images to viewer.
  738 + viewer.OnScrollBar()
  739 +
  740 + def OnBrushMove(self, obj, evt):
  741 + if (self.viewer.slice_.buffer_slices[self.orientation].mask is None):
  742 + return
  743 +
  744 + viewer = self.viewer
  745 + iren = viewer.interactor
  746 +
  747 + viewer._set_editor_cursor_visibility(1)
  748 +
  749 + mouse_x, mouse_y = iren.GetEventPosition()
  750 + render = iren.FindPokedRenderer(mouse_x, mouse_y)
  751 + slice_data = viewer.get_slice_data(render)
  752 +
  753 + # TODO: Improve!
  754 + #for i in self.slice_data_list:
  755 + #i.cursor.Show(0)
  756 +
  757 + self.picker.Pick(mouse_x, mouse_y, 0, render)
  758 +
  759 + #if (self.pick.GetViewProp()):
  760 + #self.interactor.SetCursor(wx.StockCursor(wx.CURSOR_BLANK))
  761 + #else:
  762 + #self.interactor.SetCursor(wx.StockCursor(wx.CURSOR_DEFAULT))
  763 +
  764 + coord = self.get_coordinate_cursor()
  765 + position = viewer.slice_data.actor.GetInput().FindPoint(coord)
  766 + operations = [const.BRUSH_DRAW, const.BRUSH_ERASE]
  767 + operation = operations[iren.GetControlKey()]
  768 +
  769 + if operation == const.BRUSH_DRAW:
  770 + self.foreground = True
  771 +
  772 + elif operation == const.BRUSH_ERASE:
  773 + self.foreground = True
  774 +
  775 + # when position == -1 the cursos is not over the image, so is not
  776 + # necessary to set the cursor position to world coordinate center of
  777 + # pixel from slice image.
  778 + if position != -1:
  779 + coord = slice_data.actor.GetInput().GetPoint(position)
  780 + slice_data.cursor.SetPosition(coord)
  781 + #self.__update_cursor_position(slice_data, coord)
  782 +
  783 + if (self.left_pressed):
  784 + cursor = slice_data.cursor
  785 + position = slice_data.actor.GetInput().FindPoint(coord)
  786 + radius = cursor.radius
  787 +
  788 + if position < 0:
  789 + position = viewer.calculate_matrix_position(coord)
  790 +
  791 + n = self.viewer.slice_data.number
  792 + self.edit_mask_pixel(operation, n, cursor.GetPixels(),
  793 + position, radius, self.orientation)
  794 + if self.orientation == 'AXIAL':
  795 + mask = self.matrix[n, :, :]
  796 + elif self.orientation == 'CORONAL':
  797 + mask = self.matrix[:, n, :]
  798 + elif self.orientation == 'SAGITAL':
  799 + mask = self.matrix[:, :, n]
  800 + spacing = self.viewer.slice_.spacing
  801 + vmask = converters.to_vtk(mask, spacing, n, self.orientation)
  802 + cvmask = do_colour_mask(vmask)
  803 + self.viewer.slice_.qblend[self.orientation][n] = cvmask
  804 + # TODO: To create a new function to reload images to viewer.
  805 + viewer.OnScrollBar(update3D=False)
  806 +
  807 + else:
  808 + viewer.interactor.Render()
  809 +
  810 + def OnBrushRelease(self, evt, obj):
  811 + n = self.viewer.slice_data.number
  812 + self.viewer.slice_.discard_all_buffers()
  813 + if self.orientation == 'AXIAL':
  814 + image = self.viewer.slice_.matrix[n]
  815 + mask = self.viewer.slice_.current_mask.matrix[n+1, 1:, 1:]
  816 + markers = self.matrix[n]
  817 +
  818 + elif self.orientation == 'CORONAL':
  819 + image = self.viewer.slice_.matrix[:, n, :]
  820 + mask = self.viewer.slice_.current_mask.matrix[1:, n+1, 1:]
  821 + markers = self.matrix[:, n, :]
  822 +
  823 + elif self.orientation == 'SAGITAL':
  824 + image = self.viewer.slice_.matrix[:, :, n]
  825 + mask = self.viewer.slice_.current_mask.matrix[1: , 1:, n+1]
  826 + markers = self.matrix[:, :, n]
  827 +
  828 + tmp_mask = ndimage.watershed_ift((image - image.min()).astype('uint16'), markers)
  829 + mask[:] = 0
  830 + mask[tmp_mask == 1] = 255
  831 + self.viewer.OnScrollBar(update3D=False)
  832 +
  833 + def get_coordinate_cursor(self):
  834 + # Find position
  835 + x, y, z = self.picker.GetPickPosition()
  836 + bounds = self.viewer.slice_data.actor.GetBounds()
  837 + if bounds[0] == bounds[1]:
  838 + x = bounds[0]
  839 + elif bounds[2] == bounds[3]:
  840 + y = bounds[2]
  841 + elif bounds[4] == bounds[5]:
  842 + z = bounds[4]
  843 + return x, y, z
  844 +
  845 + def edit_mask_pixel(self, operation, n, index, position, radius, orientation):
  846 + if orientation == 'AXIAL':
  847 + mask = self.matrix[n, :, :]
  848 + elif orientation == 'CORONAL':
  849 + mask = self.matrix[:, n, :]
  850 + elif orientation == 'SAGITAL':
  851 + mask = self.matrix[:, :, n]
  852 +
  853 + spacing = self.viewer.slice_.spacing
  854 + if hasattr(position, '__iter__'):
  855 + py, px = position
  856 + if orientation == 'AXIAL':
  857 + sx = spacing[0]
  858 + sy = spacing[1]
  859 + elif orientation == 'CORONAL':
  860 + sx = spacing[0]
  861 + sy = spacing[2]
  862 + elif orientation == 'SAGITAL':
  863 + sx = spacing[2]
  864 + sy = spacing[1]
  865 +
  866 + else:
  867 + if orientation == 'AXIAL':
  868 + sx = spacing[0]
  869 + sy = spacing[1]
  870 + py = position / mask.shape[1]
  871 + px = position % mask.shape[1]
  872 + elif orientation == 'CORONAL':
  873 + sx = spacing[0]
  874 + sy = spacing[2]
  875 + py = position / mask.shape[1]
  876 + px = position % mask.shape[1]
  877 + elif orientation == 'SAGITAL':
  878 + sx = spacing[2]
  879 + sy = spacing[1]
  880 + py = position / mask.shape[1]
  881 + px = position % mask.shape[1]
  882 +
  883 + cx = index.shape[1] / 2 + 1
  884 + cy = index.shape[0] / 2 + 1
  885 + xi = px - index.shape[1] + cx
  886 + xf = xi + index.shape[1]
  887 + yi = py - index.shape[0] + cy
  888 + yf = yi + index.shape[0]
  889 +
  890 + if yi < 0:
  891 + index = index[abs(yi):,:]
  892 + yi = 0
  893 + if yf > mask.shape[0]:
  894 + index = index[:index.shape[0]-(yf-image.shape[0]), :]
  895 + yf = mask.shape[0]
  896 +
  897 + if xi < 0:
  898 + index = index[:,abs(xi):]
  899 + xi = 0
  900 + if xf > mask.shape[1]:
  901 + index = index[:,:index.shape[1]-(xf-image.shape[1])]
  902 + xf = mask.shape[1]
  903 +
  904 + # Verifying if the points is over the image array.
  905 + if (not 0 <= xi <= mask.shape[1] and not 0 <= xf <= mask.shape[1]) or \
  906 + (not 0 <= yi <= mask.shape[0] and not 0 <= yf <= mask.shape[0]):
  907 + return
  908 +
  909 + roi_m = mask[yi:yf,xi:xf]
  910 +
  911 + # Checking if roi_i has at least one element.
  912 + if roi_m.size:
  913 + if operation == const.BRUSH_DRAW:
  914 + roi_m[index] = 1
  915 + elif operation == const.BRUSH_ERASE:
  916 + roi_m[index] = 2
  917 +
  918 +
  919 +def do_colour_mask(imagedata):
  920 + scalar_range = int(imagedata.GetScalarRange()[1])
  921 + r,g,b = 0, 1, 0
  922 +
  923 + # map scalar values into colors
  924 + lut_mask = vtk.vtkLookupTable()
  925 + lut_mask.SetNumberOfColors(3)
  926 + lut_mask.SetHueRange(const.THRESHOLD_HUE_RANGE)
  927 + lut_mask.SetSaturationRange(1, 1)
  928 + lut_mask.SetValueRange(0, 2)
  929 + lut_mask.SetRange(0, 2)
  930 + lut_mask.SetNumberOfTableValues(3)
  931 + lut_mask.SetTableValue(0, 0, 0, 0, 0.0)
  932 + lut_mask.SetTableValue(1, 0, 1, 0, 1.0)
  933 + lut_mask.SetTableValue(2, 1, 0, 0, 1.0)
  934 + lut_mask.SetRampToLinear()
  935 + lut_mask.Build()
  936 + # self.lut_mask = lut_mask
  937 +
  938 + # map the input image through a lookup table
  939 + img_colours_mask = vtk.vtkImageMapToColors()
  940 + img_colours_mask.SetLookupTable(lut_mask)
  941 + img_colours_mask.SetOutputFormatToRGBA()
  942 + img_colours_mask.SetInput(imagedata)
  943 + img_colours_mask.Update()
  944 + # self.img_colours_mask = img_colours_mask
  945 +
  946 + return img_colours_mask.GetOutput()
  947 +
  948 +
  949 +
643 950 def get_style(style):
644 951 STYLES = {
645 952 const.STATE_DEFAULT: DefaultInteractorStyle,
... ... @@ -653,5 +960,6 @@ def get_style(style):
653 960 const.STATE_ZOOM_SL: ZoomSLInteractorStyle,
654 961 const.SLICE_STATE_SCROLL: ChangeSliceInteractorStyle,
655 962 const.SLICE_STATE_EDITOR: EditorInteractorStyle,
  963 + const.SLICE_STATE_WATERSHED: WaterShedInteractorStyle,
656 964 }
657 965 return STYLES[style]
... ...
invesalius/gui/frame.py
... ... @@ -26,6 +26,8 @@ import webbrowser
26 26 import wx
27 27 import wx.aui
28 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 32 import constants as const
31 33 import default_tasks as tasks
... ... @@ -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 79 self.Center(wx.BOTH)
65 80 icon_path = os.path.join(const.ICON_DIR, "invesalius.ico")
66 81 self.SetIcon(wx.Icon(icon_path, wx.BITMAP_TYPE_ICO))
  82 +
  83 + self.mw = None
67 84  
68 85 if sys.platform != 'darwin':
69 86 self.Maximize()
... ... @@ -104,6 +121,7 @@ class Frame(wx.Frame):
104 121 sub(self._SetProjectName, 'Set project name')
105 122 sub(self._ShowContentPanel, 'Show content panel')
106 123 sub(self._ShowImportPanel, 'Show import panel in frame')
  124 + #sub(self._ShowHelpMessage, 'Show help message')
107 125 sub(self._ShowImportNetwork, 'Show retrieve dicom panel')
108 126 sub(self._ShowTask, 'Show task panel')
109 127 sub(self._UpdateAUI, 'Update AUI')
... ... @@ -116,6 +134,7 @@ class Frame(wx.Frame):
116 134 self.Bind(wx.EVT_SIZE, self.OnSize)
117 135 self.Bind(wx.EVT_MENU, self.OnMenuClick)
118 136 self.Bind(wx.EVT_CLOSE, self.OnClose)
  137 + #self.Bind(wx.EVT_MOVE, self.OnMove)
119 138  
120 139 def __init_aui(self):
121 140 """
... ... @@ -289,6 +308,14 @@ class Frame(wx.Frame):
289 308 aui_manager.GetPane("Import").Show(0)
290 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 319 def _ShowImportPanel(self, evt_pubsub):
293 320 """
294 321 Show only DICOM import panel.
... ... @@ -378,6 +405,12 @@ class Frame(wx.Frame):
378 405 Publisher.sendMessage(('ProgressBar Reposition'))
379 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 414 def ShowPreferences(self):
382 415  
383 416 if self.preferences.ShowModal() == wx.ID_OK:
... ...
invesalius/gui/task_slice.py
... ... @@ -252,6 +252,13 @@ class InnerFoldPanel(wx.Panel):
252 252 self.__id_editor = item.GetId()
253 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 262 #fold_panel.Expand(fold_panel.GetFoldPanel(1))
256 263  
257 264 # Panel sizer to expand fold panel
... ... @@ -284,8 +291,18 @@ class InnerFoldPanel(wx.Panel):
284 291 Publisher.sendMessage('Disable style', const.SLICE_STATE_EDITOR)
285 292 self.last_style = None
286 293 else:
287   - Publisher.sendMessage('Enable style', const.SLICE_STATE_EDITOR)
  294 + Publisher.sendMessage('Enable style',
  295 + const.SLICE_STATE_EDITOR)
288 296 self.last_style = const.SLICE_STATE_EDITOR
  297 + elif self.__id_watershed == id:
  298 + if closed:
  299 + Publisher.sendMessage('Disable style',
  300 + const.SLICE_STATE_WATERSHED)
  301 + self.last_style = None
  302 + else:
  303 + Publisher.sendMessage('Enable style', const.SLICE_STATE_WATERSHED)
  304 + Publisher.sendMessage('Show help message', 'Mark the object and the background')
  305 + self.last_style = const.SLICE_STATE_WATERSHED
289 306 else:
290 307 Publisher.sendMessage('Disable style', const.SLICE_STATE_EDITOR)
291 308 self.last_style = None
... ... @@ -688,3 +705,7 @@ class EditionTools(wx.Panel):
688 705 Publisher.sendMessage('Set edition operation', brush_op_id)
689 706  
690 707  
  708 +class WatershedTool(EditionTools):
  709 + pass
  710 +
  711 +
... ...