Commit 9a0892f4162350e15e09ef5729ade3fb3c8c915c
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): | ... | ... |