Commit 9a0892f4162350e15e09ef5729ade3fb3c8c915c

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

Added suport to erase a mark

Showing 1 changed file with 68 additions and 15 deletions   Show diff stats
invesalius/data/styles.py
... ... @@ -30,6 +30,7 @@ import converters
30 30 import numpy as np
31 31  
32 32 from scipy import ndimage
  33 +from skimage.morphology import watershed
33 34  
34 35 ORIENTATIONS = {
35 36 "AXIAL": const.AXIAL,
... ... @@ -37,6 +38,16 @@ ORIENTATIONS = {
37 38 "SAGITAL": const.SAGITAL,
38 39 }
39 40  
  41 +BRUSH_FOREGROUND=1
  42 +BRUSH_BACKGROUND=2
  43 +BRUSH_ERASE=3
  44 +
  45 +def get_LUT_value(data, window, level):
  46 + return np.piecewise(data,
  47 + [data <= (level - 0.5 - (window-1)/2),
  48 + data > (level - 0.5 + (window-1)/2)],
  49 + [0, 255, lambda data: ((data - (level - 0.5))/(window-1) + 0.5)*(255-0)])
  50 +
40 51 class BaseImageInteractorStyle(vtk.vtkInteractorStyleImage):
41 52 def __init__(self, viewer):
42 53 self.right_pressed = False
... ... @@ -668,6 +679,10 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
668 679 self.AddObserver("MouseMoveEvent", self.OnBrushMove)
669 680  
670 681 def SetUp(self):
  682 + mask = self.viewer.slice_.current_mask.matrix
  683 + mask[0] = 1
  684 + mask[:, 0, :] = 1
  685 + mask[:, :, 0] = 1
671 686 self._create_mask()
672 687  
673 688 def CleanUp(self):
... ... @@ -723,17 +738,40 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
723 738 coord = slice_data.actor.GetInput().GetPoint(position)
724 739  
725 740 slice_data.cursor.SetPosition(coord)
  741 +
726 742 cursor = slice_data.cursor
  743 + position = slice_data.actor.GetInput().FindPoint(coord)
727 744 radius = cursor.radius
728 745  
729 746 if position < 0:
730 747 position = viewer.calculate_matrix_position(coord)
731 748  
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
  749 + if iren.GetControlKey():
  750 + operation = BRUSH_BACKGROUND
  751 + elif iren.GetShiftKey():
  752 + operation = BRUSH_ERASE
  753 + else:
  754 + operation = BRUSH_FOREGROUND
736 755  
  756 + if operation == const.BRUSH_DRAW:
  757 + self.foreground = True
  758 +
  759 + elif operation == const.BRUSH_ERASE:
  760 + self.foreground = True
  761 +
  762 + n = self.viewer.slice_data.number
  763 + self.edit_mask_pixel(operation, n, cursor.GetPixels(),
  764 + position, radius, self.orientation)
  765 + if self.orientation == 'AXIAL':
  766 + mask = self.matrix[n, :, :]
  767 + elif self.orientation == 'CORONAL':
  768 + mask = self.matrix[:, n, :]
  769 + elif self.orientation == 'SAGITAL':
  770 + mask = self.matrix[:, :, n]
  771 + spacing = self.viewer.slice_.spacing
  772 + vmask = converters.to_vtk(mask, spacing, n, self.orientation)
  773 + cvmask = do_colour_mask(vmask)
  774 + self.viewer.slice_.qblend[self.orientation][n] = cvmask
737 775 # TODO: To create a new function to reload images to viewer.
738 776 viewer.OnScrollBar()
739 777  
... ... @@ -763,14 +801,6 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
763 801  
764 802 coord = self.get_coordinate_cursor()
765 803 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 804  
775 805 # when position == -1 the cursos is not over the image, so is not
776 806 # necessary to set the cursor position to world coordinate center of
... ... @@ -787,6 +817,19 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
787 817  
788 818 if position < 0:
789 819 position = viewer.calculate_matrix_position(coord)
  820 +
  821 + if iren.GetControlKey():
  822 + operation = BRUSH_BACKGROUND
  823 + elif iren.GetShiftKey():
  824 + operation = BRUSH_ERASE
  825 + else:
  826 + operation = BRUSH_FOREGROUND
  827 +
  828 + if operation == const.BRUSH_DRAW:
  829 + self.foreground = True
  830 +
  831 + elif operation == const.BRUSH_ERASE:
  832 + self.foreground = True
790 833  
791 834 n = self.viewer.slice_data.number
792 835 self.edit_mask_pixel(operation, n, cursor.GetPixels(),
... ... @@ -825,9 +868,17 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
825 868 mask = self.viewer.slice_.current_mask.matrix[1: , 1:, n+1]
826 869 markers = self.matrix[:, :, n]
827 870  
828   - tmp_mask = ndimage.watershed_ift((image - image.min()).astype('uint16'), markers)
  871 +
  872 + ww = self.viewer.slice_.window_width
  873 + wl = self.viewer.slice_.window_level
  874 +
  875 + #tmp_image = get_LUT_value(image, ww, wl).astype('uint16')
  876 + tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), 5)
  877 + print tmp_image.dtype, tmp_image.min(), tmp_image.max()
  878 + tmp_mask = watershed(tmp_image, markers)
829 879 mask[:] = 0
830 880 mask[tmp_mask == 1] = 255
  881 + self.viewer._flush_buffer = True
831 882 self.viewer.OnScrollBar(update3D=False)
832 883  
833 884 def get_coordinate_cursor(self):
... ... @@ -910,10 +961,12 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
910 961  
911 962 # Checking if roi_i has at least one element.
912 963 if roi_m.size:
913   - if operation == const.BRUSH_DRAW:
  964 + if operation == BRUSH_FOREGROUND:
914 965 roi_m[index] = 1
915   - elif operation == const.BRUSH_ERASE:
  966 + elif operation == BRUSH_BACKGROUND:
916 967 roi_m[index] = 2
  968 + elif operation == BRUSH_ERASE:
  969 + roi_m[index] = 0
917 970  
918 971  
919 972 def do_colour_mask(imagedata):
... ...