Download as
Browse Code »

Commit b659fde7ba4869f43bb39cdea417cb666ca5d200

Authored by Thiago Franco de Moraes
Committed by GitHub
1 parent 5c51e20f
Exists in master

Brain Segmentation using Deep Learning (#223) 

Add a function to segment brain in InVesalius. It uses a neural network based on U-Net implemented using Keras. For default it will use Plaidml as backend, but it's possible to set other backend using the GUI.
Showing 11 changed files with 333 additions and 2 deletions   Show diff stats
invesalius/constants.py
  Diff comments   View file @b659fde
... ... @@ -508,6 +508,8 @@ ID_MANUAL_SEGMENTATION = wx.NewId()
508 508 ID_WATERSHED_SEGMENTATION = wx.NewId()
509 509 ID_THRESHOLD_SEGMENTATION = wx.NewId()
510 510 ID_FLOODFILL_SEGMENTATION = wx.NewId()
  511 +ID_FLOODFILL_SEGMENTATION = wx.NewId()
  512 +ID_SEGMENTATION_BRAIN = wx.NewId()
511 513 ID_CROP_MASK = wx.NewId()
512 514 ID_DENSITY_MEASURE = wx.NewId()
513 515 ID_MASK_DENSITY_MEASURE = wx.NewId()
... ...
invesalius/data/imagedata_utils.py
  Diff comments   View file @b659fde
... ... @@ -485,13 +485,13 @@ def img2memmap(group):
485 485  
486 486 data = group.get_data()
487 487 # Normalize image pixel values and convert to int16
488   - data = imgnormalize(data)
  488 + # data = imgnormalize(data)
489 489  
490 490 # Convert RAS+ to default InVesalius orientation ZYX
491 491 data = numpy.swapaxes(data, 0, 2)
492 492 data = numpy.fliplr(data)
493 493  
494   - matrix = numpy.memmap(temp_file, mode='w+', dtype=data.dtype, shape=data.shape)
  494 + matrix = numpy.memmap(temp_file, mode='w+', dtype=np.int16, shape=data.shape)
495 495 matrix[:] = data[:]
496 496 matrix.flush()
497 497  
... ... @@ -537,3 +537,8 @@ def get_LUT_value_255(data, window, level):
537 537 [0, 255, lambda data_: ((data_ - (level - 0.5))/(window-1) + 0.5)*(255)])
538 538 data.shape = shape
539 539 return data
  540 +
  541 +
  542 +def image_normalize(image, min_=0.0, max_=1.0):
  543 + imin, imax = image.min(), image.max()
  544 + return (image - imin) * ((max_ - min_) / (imax - imin)) + min_
... ...
invesalius/data/segmentation.py 0 → 100644
  Diff comments   View file @b659fde
... ... @@ -0,0 +1 @@
  1 +import numpy as np
... ...
invesalius/gui/brain_seg_dialog.py 0 → 100644
  Diff comments   View file @b659fde
... ... @@ -0,0 +1,183 @@
  1 +#!/usr/bin/env python
  2 +# -*- coding: UTF-8 -*-
  3 +
  4 +import os
  5 +import pathlib
  6 +import sys
  7 +
  8 +import wx
  9 +from wx.lib.pubsub import pub as Publisher
  10 +
  11 +HAS_THEANO = True
  12 +HAS_PLAIDML = True
  13 +
  14 +try:
  15 + import theano
  16 +except ImportError:
  17 + HAS_THEANO = False
  18 +
  19 +# Linux if installed plaidml with pip3 install --user
  20 +if sys.platform.startswith("linux"):
  21 + local_user_plaidml = pathlib.Path("~/.local/share/plaidml/").expanduser().absolute()
  22 + if local_user_plaidml.exists():
  23 + os.environ["RUNFILES_DIR"] = str(local_user_plaidml)
  24 + os.environ["PLAIDML_NATIVE_PATH"] = str(pathlib.Path("~/.local/lib/libplaidml.so").expanduser().absolute())
  25 +# Mac if using python3 from homebrew
  26 +elif sys.platform == "darwin":
  27 + local_user_plaidml = pathlib.Path("/usr/local/share/plaidml")
  28 + if local_user_plaidml.exists():
  29 + os.environ["RUNFILES_DIR"] = str(local_user_plaidml)
  30 + os.environ["PLAIDML_NATIVE_PATH"] = str(pathlib.Path("/usr/local/lib/libplaidml.dylib").expanduser().absolute())
  31 +
  32 +try:
  33 + import plaidml
  34 +except ImportError:
  35 + HAS_PLAIDML = False
  36 +
  37 +import invesalius.data.slice_ as slc
  38 +from invesalius.segmentation.brain import segment
  39 +
  40 +
  41 +
  42 +class BrainSegmenterDialog(wx.Dialog):
  43 + def __init__(self, parent):
  44 + wx.Dialog.__init__(self, parent, -1, _(u"Brain segmentation"), style=wx.DEFAULT_DIALOG_STYLE|wx.FRAME_FLOAT_ON_PARENT)
  45 + backends = []
  46 + if HAS_PLAIDML:
  47 + backends.append("PlaidML")
  48 + if HAS_THEANO:
  49 + backends.append("Theano")
  50 + self.segmenter = segment.BrainSegmenter()
  51 + self.pg_dialog = None
  52 +
  53 + self.cb_backends = wx.ComboBox(self, wx.ID_ANY, choices=backends, value=backends[0], style=wx.CB_DROPDOWN | wx.CB_READONLY)
  54 + w, h = self.CalcSizeFromTextSize("MM" * (1 + max(len(i) for i in backends)))
  55 + self.cb_backends.SetMinClientSize((w, -1))
  56 + self.chk_use_gpu = wx.CheckBox(self, wx.ID_ANY, "Use GPU")
  57 + self.sld_threshold = wx.Slider(self, wx.ID_ANY, 75, 0, 100)
  58 + w, h = self.CalcSizeFromTextSize("M" * 20)
  59 + self.sld_threshold.SetMinClientSize((w, -1))
  60 + self.txt_threshold = wx.TextCtrl(self, wx.ID_ANY, "")
  61 + w, h = self.CalcSizeFromTextSize("MMMMM")
  62 + self.txt_threshold.SetMinClientSize((w, -1))
  63 + # self.progress = wx.Gauge(self, -1)
  64 + self.btn_segment = wx.Button(self, wx.ID_ANY, "Segment")
  65 + # self.btn_stop = wx.Button(self, wx.ID_ANY, _("Stop"))
  66 + # self.btn_stop.Disable()
  67 +
  68 + self.txt_threshold.SetValue("{:3d}%".format(self.sld_threshold.GetValue()))
  69 +
  70 + self.__do_layout()
  71 + self.__set_events()
  72 +
  73 + def __do_layout(self):
  74 + main_sizer = wx.BoxSizer(wx.VERTICAL)
  75 + sizer_3 = wx.BoxSizer(wx.HORIZONTAL)
  76 + sizer_backends = wx.BoxSizer(wx.HORIZONTAL)
  77 + label_1 = wx.StaticText(self, wx.ID_ANY, "Backend")
  78 + sizer_backends.Add(label_1, 0, wx.ALIGN_CENTER, 0)
  79 + sizer_backends.Add(self.cb_backends, 1, wx.LEFT, 5)
  80 + main_sizer.Add(sizer_backends, 0, wx.ALL | wx.EXPAND, 5)
  81 + main_sizer.Add(self.chk_use_gpu, 0, wx.ALL, 5)
  82 + label_5 = wx.StaticText(self, wx.ID_ANY, "Level of certainty")
  83 + main_sizer.Add(label_5, 0, wx.ALL, 5)
  84 + sizer_3.Add(self.sld_threshold, 1, wx.ALIGN_CENTER | wx.BOTTOM | wx.EXPAND | wx.LEFT | wx.RIGHT, 5)
  85 + sizer_3.Add(self.txt_threshold, 0, wx.ALL, 5)
  86 + main_sizer.Add(sizer_3, 0, wx.EXPAND, 0)
  87 + # main_sizer.Add(self.progress, 0, wx.EXPAND | wx.ALL, 5)
  88 + sizer_buttons = wx.BoxSizer(wx.HORIZONTAL)
  89 + # sizer_buttons.Add(self.btn_stop, 0, wx.ALIGN_BOTTOM | wx.ALIGN_RIGHT | wx.ALL, 5)
  90 + sizer_buttons.Add(self.btn_segment, 0, wx.ALIGN_BOTTOM | wx.ALIGN_RIGHT | wx.ALL, 5)
  91 + main_sizer.Add(sizer_buttons, 0, wx.ALIGN_BOTTOM | wx.ALIGN_RIGHT | wx.ALL, 0)
  92 + self.SetSizer(main_sizer)
  93 + main_sizer.Fit(self)
  94 + main_sizer.SetSizeHints(self)
  95 + self.Layout()
  96 + self.Centre()
  97 +
  98 + def __set_events(self):
  99 + self.sld_threshold.Bind(wx.EVT_SCROLL, self.OnScrollThreshold)
  100 + self.txt_threshold.Bind(wx.EVT_KILL_FOCUS, self.OnKillFocus)
  101 + self.btn_segment.Bind(wx.EVT_BUTTON, self.OnSegment)
  102 + # self.btn_stop.Bind(wx.EVT_BUTTON, self.OnStop)
  103 + self.Bind(wx.EVT_CLOSE, self.OnClose)
  104 +
  105 + def CalcSizeFromTextSize(self, text):
  106 + dc = wx.WindowDC(self)
  107 + dc.SetFont(self.GetFont())
  108 + width, height = dc.GetTextExtent(text)
  109 + return width, height
  110 +
  111 + def OnScrollThreshold(self, evt):
  112 + value = self.sld_threshold.GetValue()
  113 + self.txt_threshold.SetValue("{:3d}%".format(self.sld_threshold.GetValue()))
  114 + if self.segmenter.segmented:
  115 + threshold = value / 100.0
  116 + self.segmenter.set_threshold(threshold)
  117 + image = slc.Slice().discard_all_buffers()
  118 + Publisher.sendMessage('Reload actual slice')
  119 +
  120 + def OnKillFocus(self, evt):
  121 + value = self.txt_threshold.GetValue()
  122 + value = value.replace('%', '')
  123 + try:
  124 + value = int(value)
  125 + except ValueError:
  126 + value = self.sld_threshold.GetValue()
  127 + self.sld_threshold.SetValue(value)
  128 + self.txt_threshold.SetValue("{:3d}%".format(value))
  129 +
  130 + if self.segmenter.segmented:
  131 + threshold = value / 100.0
  132 + self.segmenter.set_threshold(threshold)
  133 + image = slc.Slice().discard_all_buffers()
  134 + Publisher.sendMessage('Reload actual slice')
  135 +
  136 + def OnSegment(self, evt):
  137 + image = slc.Slice().matrix
  138 + backend = self.cb_backends.GetValue()
  139 + use_gpu = self.chk_use_gpu.GetValue()
  140 + prob_threshold = self.sld_threshold.GetValue() / 100.0
  141 + # self.btn_stop.Enable()
  142 + self.btn_segment.Disable()
  143 + self.pg_dialog = wx.ProgressDialog(_("Brain segmenter"), _("Segmenting brain"), parent=self, style= wx.FRAME_FLOAT_ON_PARENT | wx.PD_CAN_ABORT | wx.PD_AUTO_HIDE | wx.PD_ELAPSED_TIME)
  144 + self.pg_dialog.Bind(wx.EVT_BUTTON, self.OnStop)
  145 + self.pg_dialog.Show()
  146 + self.segmenter.segment(image, prob_threshold, backend, use_gpu, self.SetProgress, self.AfterSegment)
  147 +
  148 + def OnStop(self, evt):
  149 + self.segmenter.stop = True
  150 + # self.btn_stop.Disable()
  151 + self.pg_dialog.Hide()
  152 + self.pg_dialog = None
  153 + self.btn_segment.Enable()
  154 + evt.Skip()
  155 +
  156 + def AfterSegment(self):
  157 + Publisher.sendMessage('Reload actual slice')
  158 +
  159 + def SetProgress(self, progress):
  160 + # self.progress.SetValue(progress * 100)
  161 + self.pg_dialog.Update(progress * 100)
  162 + wx.GetApp().Yield()
  163 +
  164 + def OnClose(self, evt):
  165 + self.segmenter.stop = True
  166 + # self.btn_stop.Disable()
  167 + self.btn_segment.Enable()
  168 + # self.progress.SetValue(0)
  169 + self.pg_dialog.Destroy()
  170 + self.Destroy()
  171 +
  172 +
  173 +class MyApp(wx.App):
  174 + def OnInit(self):
  175 + self.dlg_brain_seg = MyDialog(None, wx.ID_ANY, "")
  176 + self.SetTopWindow(self.dlg_brain_seg)
  177 + self.dlg_brain_seg.ShowModal()
  178 + self.dlg_brain_seg.Destroy()
  179 + return True
  180 +
  181 +if __name__ == "__main__":
  182 + app = MyApp(0)
  183 + app.MainLoop()
... ...
invesalius/gui/frame.py
  Diff comments   View file @b659fde
... ... @@ -508,6 +508,9 @@ class Frame(wx.Frame):
508 508 elif id == const.ID_FLOODFILL_SEGMENTATION:
509 509 self.OnFFillSegmentation()
510 510  
  511 + elif id == const.ID_SEGMENTATION_BRAIN:
  512 + self.OnBrainSegmentation()
  513 +
511 514 elif id == const.ID_VIEW_INTERPOLATED:
512 515 st = self.actived_interpolated_slices.IsChecked(const.ID_VIEW_INTERPOLATED)
513 516 if st:
... ... @@ -740,6 +743,11 @@ class Frame(wx.Frame):
740 743 def OnFFillSegmentation(self):
741 744 Publisher.sendMessage('Enable style', style=const.SLICE_STATE_FFILL_SEGMENTATION)
742 745  
  746 + def OnBrainSegmentation(self):
  747 + from invesalius.gui.brain_seg_dialog import BrainSegmenterDialog
  748 + dlg = BrainSegmenterDialog(self)
  749 + dlg.Show()
  750 +
743 751 def OnInterpolatedSlices(self, status):
744 752 Publisher.sendMessage('Set interpolated slices', flag=status)
745 753  
... ... @@ -797,6 +805,7 @@ class MenuBar(wx.MenuBar):
797 805 const.ID_WATERSHED_SEGMENTATION,
798 806 const.ID_THRESHOLD_SEGMENTATION,
799 807 const.ID_FLOODFILL_SEGMENTATION,
  808 + const.ID_SEGMENTATION_BRAIN,
800 809 const.ID_MASK_DENSITY_MEASURE,
801 810 const.ID_CREATE_SURFACE,
802 811 const.ID_CREATE_MASK,
... ... @@ -936,6 +945,8 @@ class MenuBar(wx.MenuBar):
936 945 self.watershed_segmentation = segmentation_menu.Append(const.ID_WATERSHED_SEGMENTATION, _(u"Watershed\tCtrl+Shift+W"))
937 946 self.ffill_segmentation = segmentation_menu.Append(const.ID_FLOODFILL_SEGMENTATION, _(u"Region growing\tCtrl+Shift+G"))
938 947 self.ffill_segmentation.Enable(False)
  948 + segmentation_menu.AppendSeparator()
  949 + segmentation_menu.Append(const.ID_SEGMENTATION_BRAIN, _("Brain segmentation"))
939 950  
940 951 # Surface Menu
941 952 surface_menu = wx.Menu()
... ...
invesalius/segmentation/__init__.py 0 → 100644
  Diff comments   View file @b659fde
invesalius/segmentation/brain/__init__.py 0 → 100644
  Diff comments   View file @b659fde
invesalius/segmentation/brain/model.h5 0 → 100644
View file @b659fde
No preview for this file type
invesalius/segmentation/brain/model.json 0 → 100644
  Diff comments   View file @b659fde
... ... @@ -0,0 +1 @@
  1 +{"class_name": "Model", "config": {"name": "model_1", "layers": [{"name": "img", "class_name": "InputLayer", "config": {"batch_input_shape": [null, 48, 48, 48, 1], "dtype": "float32", "sparse": false, "name": "img"}, "inbound_nodes": []}, {"name": "conv3d_1", "class_name": "Conv3D", "config": {"name": "conv3d_1", "trainable": true, "filters": 8, "kernel_size": [5, 5, 5], "strides": [1, 1, 1], "padding": "same", "data_format": "channels_last", "dilation_rate": [1, 1, 1], "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["img", 0, 0, {}]]]}, {"name": "conv3d_2", "class_name": "Conv3D", "config": {"name": "conv3d_2", "trainable": true, "filters": 8, "kernel_size": [5, 5, 5], "strides": [1, 1, 1], "padding": "same", "data_format": "channels_last", "dilation_rate": [1, 1, 1], "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["conv3d_1", 0, 0, {}]]]}, {"name": "max_pooling3d_1", "class_name": "MaxPooling3D", "config": {"name": "max_pooling3d_1", "trainable": true, "pool_size": [2, 2, 2], "padding": "valid", "strides": [2, 2, 2], "data_format": "channels_last"}, "inbound_nodes": [[["conv3d_2", 0, 0, {}]]]}, {"name": "dropout_1", "class_name": "Dropout", "config": {"name": "dropout_1", "trainable": true, "rate": 0.3, "noise_shape": null, "seed": null}, "inbound_nodes": [[["max_pooling3d_1", 0, 0, {}]]]}, {"name": "conv3d_3", "class_name": "Conv3D", "config": {"name": "conv3d_3", "trainable": true, "filters": 16, "kernel_size": [5, 5, 5], "strides": [1, 1, 1], "padding": "same", "data_format": "channels_last", "dilation_rate": [1, 1, 1], "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dropout_1", 0, 0, {}]]]}, {"name": "conv3d_4", "class_name": "Conv3D", "config": {"name": "conv3d_4", "trainable": true, "filters": 16, "kernel_size": [5, 5, 5], "strides": [1, 1, 1], "padding": "same", "data_format": "channels_last", "dilation_rate": [1, 1, 1], "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["conv3d_3", 0, 0, {}]]]}, {"name": "max_pooling3d_2", "class_name": "MaxPooling3D", "config": {"name": "max_pooling3d_2", "trainable": true, "pool_size": [2, 2, 2], "padding": "valid", "strides": [2, 2, 2], "data_format": "channels_last"}, "inbound_nodes": [[["conv3d_4", 0, 0, {}]]]}, {"name": "dropout_2", "class_name": "Dropout", "config": {"name": "dropout_2", "trainable": true, "rate": 0.3, "noise_shape": null, "seed": null}, "inbound_nodes": [[["max_pooling3d_2", 0, 0, {}]]]}, {"name": "conv3d_5", "class_name": "Conv3D", "config": {"name": "conv3d_5", "trainable": true, "filters": 32, "kernel_size": [5, 5, 5], "strides": [1, 1, 1], "padding": "same", "data_format": "channels_last", "dilation_rate": [1, 1, 1], "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dropout_2", 0, 0, {}]]]}, {"name": "conv3d_6", "class_name": "Conv3D", "config": {"name": "conv3d_6", "trainable": true, "filters": 32, "kernel_size": [5, 5, 5], "strides": [1, 1, 1], "padding": "same", "data_format": "channels_last", "dilation_rate": [1, 1, 1], "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["conv3d_5", 0, 0, {}]]]}, {"name": "max_pooling3d_3", "class_name": "MaxPooling3D", "config": {"name": "max_pooling3d_3", "trainable": true, "pool_size": [2, 2, 2], "padding": "valid", "strides": [2, 2, 2], "data_format": "channels_last"}, "inbound_nodes": [[["conv3d_6", 0, 0, {}]]]}, {"name": "dropout_3", "class_name": "Dropout", "config": {"name": "dropout_3", "trainable": true, "rate": 0.3, "noise_shape": null, "seed": null}, "inbound_nodes": [[["max_pooling3d_3", 0, 0, {}]]]}, {"name": "conv3d_transpose_1", "class_name": "Conv3DTranspose", "config": {"name": "conv3d_transpose_1", "trainable": true, "filters": 32, "kernel_size": [5, 5, 5], "strides": [2, 2, 2], "padding": "same", "data_format": "channels_last", "activation": "linear", "use_bias": false, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "output_padding": null}, "inbound_nodes": [[["dropout_3", 0, 0, {}]]]}, {"name": "concatenate_1", "class_name": "Concatenate", "config": {"name": "concatenate_1", "trainable": true, "axis": -1}, "inbound_nodes": [[["conv3d_transpose_1", 0, 0, {}], ["conv3d_6", 0, 0, {}]]]}, {"name": "conv3d_7", "class_name": "Conv3D", "config": {"name": "conv3d_7", "trainable": true, "filters": 32, "kernel_size": [5, 5, 5], "strides": [1, 1, 1], "padding": "same", "data_format": "channels_last", "dilation_rate": [1, 1, 1], "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["concatenate_1", 0, 0, {}]]]}, {"name": "dropout_4", "class_name": "Dropout", "config": {"name": "dropout_4", "trainable": true, "rate": 0.3, "noise_shape": null, "seed": null}, "inbound_nodes": [[["conv3d_7", 0, 0, {}]]]}, {"name": "conv3d_transpose_2", "class_name": "Conv3DTranspose", "config": {"name": "conv3d_transpose_2", "trainable": true, "filters": 16, "kernel_size": [5, 5, 5], "strides": [2, 2, 2], "padding": "same", "data_format": "channels_last", "activation": "linear", "use_bias": false, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "output_padding": null}, "inbound_nodes": [[["dropout_4", 0, 0, {}]]]}, {"name": "concatenate_2", "class_name": "Concatenate", "config": {"name": "concatenate_2", "trainable": true, "axis": -1}, "inbound_nodes": [[["conv3d_transpose_2", 0, 0, {}], ["conv3d_4", 0, 0, {}]]]}, {"name": "conv3d_8", "class_name": "Conv3D", "config": {"name": "conv3d_8", "trainable": true, "filters": 16, "kernel_size": [5, 5, 5], "strides": [1, 1, 1], "padding": "same", "data_format": "channels_last", "dilation_rate": [1, 1, 1], "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["concatenate_2", 0, 0, {}]]]}, {"name": "dropout_5", "class_name": "Dropout", "config": {"name": "dropout_5", "trainable": true, "rate": 0.3, "noise_shape": null, "seed": null}, "inbound_nodes": [[["conv3d_8", 0, 0, {}]]]}, {"name": "conv3d_transpose_3", "class_name": "Conv3DTranspose", "config": {"name": "conv3d_transpose_3", "trainable": true, "filters": 8, "kernel_size": [5, 5, 5], "strides": [2, 2, 2], "padding": "same", "data_format": "channels_last", "activation": "linear", "use_bias": false, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null, "output_padding": null}, "inbound_nodes": [[["dropout_5", 0, 0, {}]]]}, {"name": "concatenate_3", "class_name": "Concatenate", "config": {"name": "concatenate_3", "trainable": true, "axis": -1}, "inbound_nodes": [[["conv3d_transpose_3", 0, 0, {}], ["conv3d_2", 0, 0, {}]]]}, {"name": "conv3d_9", "class_name": "Conv3D", "config": {"name": "conv3d_9", "trainable": true, "filters": 8, "kernel_size": [5, 5, 5], "strides": [1, 1, 1], "padding": "same", "data_format": "channels_last", "dilation_rate": [1, 1, 1], "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["concatenate_3", 0, 0, {}]]]}, {"name": "dropout_6", "class_name": "Dropout", "config": {"name": "dropout_6", "trainable": true, "rate": 0.3, "noise_shape": null, "seed": null}, "inbound_nodes": [[["conv3d_9", 0, 0, {}]]]}, {"name": "conv3d_10", "class_name": "Conv3D", "config": {"name": "conv3d_10", "trainable": true, "filters": 1, "kernel_size": [1, 1, 1], "strides": [1, 1, 1], "padding": "same", "data_format": "channels_last", "dilation_rate": [1, 1, 1], "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["dropout_6", 0, 0, {}]]]}, {"name": "dense_1", "class_name": "Dense", "config": {"name": "dense_1", "trainable": true, "units": 1, "activation": "sigmoid", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}, "inbound_nodes": [[["conv3d_10", 0, 0, {}]]]}], "input_layers": [["img", 0, 0]], "output_layers": [["dense_1", 0, 0]]}, "keras_version": "2.2.4", "backend": "plaidml.keras.backend"}
0 2 \ No newline at end of file
... ...
invesalius/segmentation/brain/segment.py 0 → 100644
  Diff comments   View file @b659fde
... ... @@ -0,0 +1,111 @@
  1 +import itertools
  2 +import os
  3 +import pathlib
  4 +import sys
  5 +
  6 +import numpy as np
  7 +from skimage.transform import resize
  8 +
  9 +from invesalius.data import imagedata_utils
  10 +from invesalius.utils import timing
  11 +
  12 +from . import utils
  13 +
  14 +SIZE = 48
  15 +OVERLAP = SIZE // 2 + 1
  16 +
  17 +
  18 +def gen_patches(image, patch_size, overlap):
  19 + sz, sy, sx = image.shape
  20 + i_cuts = list(itertools.product(
  21 + range(0, sz, patch_size - OVERLAP),
  22 + range(0, sy, patch_size - OVERLAP),
  23 + range(0, sx, patch_size - OVERLAP),
  24 + ))
  25 + sub_image = np.empty(shape=(patch_size, patch_size, patch_size), dtype="float32")
  26 + for idx, (iz, iy, ix) in enumerate(i_cuts):
  27 + sub_image[:] = 0
  28 + _sub_image = image[
  29 + iz : iz + patch_size, iy : iy + patch_size, ix : ix + patch_size
  30 + ]
  31 + sz, sy, sx = _sub_image.shape
  32 + sub_image[0:sz, 0:sy, 0:sx] = _sub_image
  33 + ez = iz + sz
  34 + ey = iy + sy
  35 + ex = ix + sx
  36 +
  37 + yield (idx + 1.0)/len(i_cuts), sub_image, ((iz, ez), (iy, ey), (ix, ex))
  38 +
  39 +
  40 +def predict_patch(sub_image, patch, nn_model, patch_size=SIZE):
  41 + (iz, ez), (iy, ey), (ix, ex) = patch
  42 + sub_mask = nn_model.predict(sub_image.reshape(1, patch_size, patch_size, patch_size, 1))
  43 + return sub_mask.reshape(patch_size, patch_size, patch_size)[0:ez-iz, 0:ey-iy, 0:ex-ix]
  44 +
  45 +
  46 +class BrainSegmenter:
  47 + def __init__(self):
  48 + self.mask = None
  49 + self.propability_array = None
  50 + self.stop = False
  51 + self.segmented = False
  52 +
  53 + def segment(self, image, prob_threshold, backend, use_gpu, progress_callback=None, after_segment=None):
  54 + print("backend", backend)
  55 + if backend.lower() == 'plaidml':
  56 + os.environ["KERAS_BACKEND"] = "plaidml.keras.backend"
  57 + device = utils.get_plaidml_devices(use_gpu)
  58 + os.environ["PLAIDML_DEVICE_IDS"] = device.id.decode("utf8")
  59 + elif backend.lower() == 'theano':
  60 + os.environ["KERAS_BACKEND"] = "theano"
  61 + else:
  62 + raise TypeError("Wrong backend")
  63 +
  64 + import keras
  65 + import invesalius.data.slice_ as slc
  66 +
  67 + image = imagedata_utils.image_normalize(image, 0.0, 1.0)
  68 +
  69 + # Loading model
  70 + folder = pathlib.Path(__file__).parent.resolve()
  71 + with open(folder.joinpath("model.json"), "r") as json_file:
  72 + model = keras.models.model_from_json(json_file.read())
  73 + model.load_weights(str(folder.joinpath("model.h5")))
  74 + model.compile("Adam", "binary_crossentropy")
  75 +
  76 + # segmenting by patches
  77 + msk = np.zeros_like(image, dtype="float32")
  78 + sums = np.zeros_like(image)
  79 + for completion, sub_image, patch in gen_patches(image, SIZE, OVERLAP):
  80 + if self.stop:
  81 + self.stop = False
  82 + return
  83 +
  84 + if progress_callback is not None:
  85 + progress_callback(completion)
  86 + print("completion", completion)
  87 + (iz, ez), (iy, ey), (ix, ex) = patch
  88 + sub_mask = predict_patch(sub_image, patch, model, SIZE)
  89 + msk[iz:ez, iy:ey, ix:ex] += sub_mask
  90 + sums[iz:ez, iy:ey, ix:ex] += 1
  91 +
  92 + propability_array = msk / sums
  93 +
  94 + mask = slc.Slice().create_new_mask()
  95 + mask.was_edited = True
  96 + mask.matrix[:] = 1
  97 + mask.matrix[1:, 1:, 1:] = (propability_array >= prob_threshold) * 255
  98 +
  99 + self.mask = mask
  100 + self.propability_array = propability_array
  101 + self.segmented = True
  102 + if after_segment is not None:
  103 + after_segment()
  104 +
  105 + def set_threshold(self, threshold):
  106 + if threshold < 0:
  107 + threshold = 0
  108 + elif threshold > 1:
  109 + threshold = 1
  110 + self.mask.matrix[:] = 1
  111 + self.mask.matrix[1:, 1:, 1:] = (self.propability_array >= threshold) * 255
... ...
invesalius/segmentation/brain/utils.py 0 → 100644
  Diff comments   View file @b659fde
... ... @@ -0,0 +1,17 @@
  1 +def get_plaidml_devices(gpu=False):
  2 + import plaidml
  3 +
  4 + ctx = plaidml.Context()
  5 + plaidml.settings._setup_for_test(plaidml.settings.user_settings)
  6 + plaidml.settings.experimental = True
  7 + devices, _ = plaidml.devices(ctx, limit=100, return_all=True)
  8 + if gpu:
  9 + for device in devices:
  10 + if b"cuda" in device.description.lower():
  11 + return device
  12 + for device in devices:
  13 + if b"opencl" in device.description.lower():
  14 + return device
  15 + for device in devices:
  16 + if b"llvm" in device.description.lower():
  17 + return device
... ...