Commit 13bcfe5a0912d2e7672cd8edee4380ed63260281

Authored by Thiago Franco de Moraes
1 parent b331b9f2
Exists in master

There was a error if only theano is disponible

Showing 1 changed file with 32 additions and 13 deletions   Show diff stats
invesalius/gui/brain_seg_dialog.py
@@ -9,9 +9,6 @@ import sys @@ -9,9 +9,6 @@ import sys
9 import wx 9 import wx
10 from wx.lib.pubsub import pub as Publisher 10 from wx.lib.pubsub import pub as Publisher
11 11
12 -HAS_THEANO = bool(importlib.util.find_spec("theano"))  
13 -HAS_PLAIDML = bool(importlib.util.find_spec("plaidml"))  
14 -  
15 # Linux if installed plaidml with pip3 install --user 12 # Linux if installed plaidml with pip3 install --user
16 if sys.platform.startswith("linux"): 13 if sys.platform.startswith("linux"):
17 local_user_plaidml = pathlib.Path("~/.local/share/plaidml/").expanduser().absolute() 14 local_user_plaidml = pathlib.Path("~/.local/share/plaidml/").expanduser().absolute()
@@ -25,11 +22,29 @@ elif sys.platform == "darwin": @@ -25,11 +22,29 @@ elif sys.platform == "darwin":
25 os.environ["RUNFILES_DIR"] = str(local_user_plaidml) 22 os.environ["RUNFILES_DIR"] = str(local_user_plaidml)
26 os.environ["PLAIDML_NATIVE_PATH"] = str(pathlib.Path("/usr/local/lib/libplaidml.dylib").expanduser().absolute()) 23 os.environ["PLAIDML_NATIVE_PATH"] = str(pathlib.Path("/usr/local/lib/libplaidml.dylib").expanduser().absolute())
27 24
  25 +HAS_THEANO = bool(importlib.util.find_spec("theano"))
  26 +HAS_PLAIDML = bool(importlib.util.find_spec("plaidml"))
  27 +PLAIDML_DEVICES = {}
28 28
29 import invesalius.data.slice_ as slc 29 import invesalius.data.slice_ as slc
30 from invesalius.segmentation.brain import segment 30 from invesalius.segmentation.brain import segment
31 from invesalius.segmentation.brain import utils 31 from invesalius.segmentation.brain import utils
32 32
  33 +try:
  34 + import theano
  35 +except ImportError:
  36 + HAS_THEANO = False
  37 +
  38 +try:
  39 + import plaidml
  40 +except ImportError:
  41 + HAS_PLAIDML = False
  42 +
  43 +if HAS_PLAIDML:
  44 + try:
  45 + PLAIDML_DEVICES = utils.get_plaidml_devices()
  46 + except OSError:
  47 + HAS_PLAIDML = False
33 48
34 49
35 class BrainSegmenterDialog(wx.Dialog): 50 class BrainSegmenterDialog(wx.Dialog):
@@ -42,13 +57,14 @@ class BrainSegmenterDialog(wx.Dialog): @@ -42,13 +57,14 @@ class BrainSegmenterDialog(wx.Dialog):
42 backends.append("Theano") 57 backends.append("Theano")
43 self.segmenter = segment.BrainSegmenter() 58 self.segmenter = segment.BrainSegmenter()
44 # self.pg_dialog = None 59 # self.pg_dialog = None
45 - self.plaidml_devices = utils.get_plaidml_devices() 60 + self.plaidml_devices = PLAIDML_DEVICES
46 self.cb_backends = wx.ComboBox(self, wx.ID_ANY, choices=backends, value=backends[0], style=wx.CB_DROPDOWN | wx.CB_READONLY) 61 self.cb_backends = wx.ComboBox(self, wx.ID_ANY, choices=backends, value=backends[0], style=wx.CB_DROPDOWN | wx.CB_READONLY)
47 w, h = self.CalcSizeFromTextSize("MM" * (1 + max(len(i) for i in backends))) 62 w, h = self.CalcSizeFromTextSize("MM" * (1 + max(len(i) for i in backends)))
48 self.cb_backends.SetMinClientSize((w, -1)) 63 self.cb_backends.SetMinClientSize((w, -1))
49 self.chk_use_gpu = wx.CheckBox(self, wx.ID_ANY, _("Use GPU")) 64 self.chk_use_gpu = wx.CheckBox(self, wx.ID_ANY, _("Use GPU"))
50 - self.lbl_device = wx.StaticText(self, -1, _("Device"))  
51 - self.cb_devices = wx.ComboBox(self, wx.ID_ANY, choices=list(self.plaidml_devices.keys()), value=list(self.plaidml_devices.keys())[0],style=wx.CB_DROPDOWN | wx.CB_READONLY) 65 + if HAS_PLAIDML:
  66 + self.lbl_device = wx.StaticText(self, -1, _("Device"))
  67 + self.cb_devices = wx.ComboBox(self, wx.ID_ANY, choices=list(self.plaidml_devices.keys()), value=list(self.plaidml_devices.keys())[0],style=wx.CB_DROPDOWN | wx.CB_READONLY)
52 self.sld_threshold = wx.Slider(self, wx.ID_ANY, 75, 0, 100) 68 self.sld_threshold = wx.Slider(self, wx.ID_ANY, 75, 0, 100)
53 w, h = self.CalcSizeFromTextSize("M" * 20) 69 w, h = self.CalcSizeFromTextSize("M" * 20)
54 self.sld_threshold.SetMinClientSize((w, -1)) 70 self.sld_threshold.SetMinClientSize((w, -1))
@@ -76,8 +92,9 @@ class BrainSegmenterDialog(wx.Dialog): @@ -76,8 +92,9 @@ class BrainSegmenterDialog(wx.Dialog):
76 main_sizer.Add(sizer_backends, 0, wx.ALL | wx.EXPAND, 5) 92 main_sizer.Add(sizer_backends, 0, wx.ALL | wx.EXPAND, 5)
77 main_sizer.Add(self.chk_use_gpu, 0, wx.ALL, 5) 93 main_sizer.Add(self.chk_use_gpu, 0, wx.ALL, 5)
78 sizer_devices = wx.BoxSizer(wx.HORIZONTAL) 94 sizer_devices = wx.BoxSizer(wx.HORIZONTAL)
79 - sizer_devices.Add(self.lbl_device, 0, wx.ALIGN_CENTER, 0)  
80 - sizer_devices.Add(self.cb_devices, 1, wx.LEFT, 5) 95 + if HAS_PLAIDML:
  96 + sizer_devices.Add(self.lbl_device, 0, wx.ALIGN_CENTER, 0)
  97 + sizer_devices.Add(self.cb_devices, 1, wx.LEFT, 5)
81 main_sizer.Add(sizer_devices, 0, wx.ALL | wx.EXPAND, 5) 98 main_sizer.Add(sizer_devices, 0, wx.ALL | wx.EXPAND, 5)
82 label_5 = wx.StaticText(self, wx.ID_ANY, _("Level of certainty")) 99 label_5 = wx.StaticText(self, wx.ID_ANY, _("Level of certainty"))
83 main_sizer.Add(label_5, 0, wx.ALL, 5) 100 main_sizer.Add(label_5, 0, wx.ALL, 5)
@@ -118,12 +135,14 @@ class BrainSegmenterDialog(wx.Dialog): @@ -118,12 +135,14 @@ class BrainSegmenterDialog(wx.Dialog):
118 135
119 def OnSetBackend(self, evt=None): 136 def OnSetBackend(self, evt=None):
120 if self.cb_backends.GetValue().lower() == "plaidml": 137 if self.cb_backends.GetValue().lower() == "plaidml":
121 - self.lbl_device.Show()  
122 - self.cb_devices.Show() 138 + if HAS_PLAIDML:
  139 + self.lbl_device.Show()
  140 + self.cb_devices.Show()
123 self.chk_use_gpu.Hide() 141 self.chk_use_gpu.Hide()
124 else: 142 else:
125 - self.lbl_device.Hide()  
126 - self.cb_devices.Hide() 143 + if HAS_PLAIDML:
  144 + self.lbl_device.Hide()
  145 + self.cb_devices.Hide()
127 self.chk_use_gpu.Show() 146 self.chk_use_gpu.Show()
128 147
129 self.main_sizer.Fit(self) 148 self.main_sizer.Fit(self)
@@ -159,7 +178,7 @@ class BrainSegmenterDialog(wx.Dialog): @@ -159,7 +178,7 @@ class BrainSegmenterDialog(wx.Dialog):
159 backend = self.cb_backends.GetValue() 178 backend = self.cb_backends.GetValue()
160 try: 179 try:
161 device_id = self.plaidml_devices[self.cb_devices.GetValue()] 180 device_id = self.plaidml_devices[self.cb_devices.GetValue()]
162 - except KeyError: 181 + except (KeyError, AttributeError):
163 device_id = "llvm_cpu.0" 182 device_id = "llvm_cpu.0"
164 use_gpu = self.chk_use_gpu.GetValue() 183 use_gpu = self.chk_use_gpu.GetValue()
165 prob_threshold = self.sld_threshold.GetValue() / 100.0 184 prob_threshold = self.sld_threshold.GetValue() / 100.0