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 import wx
10 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 12 # Linux if installed plaidml with pip3 install --user
16 13 if sys.platform.startswith("linux"):
17 14 local_user_plaidml = pathlib.Path("~/.local/share/plaidml/").expanduser().absolute()
... ... @@ -25,11 +22,29 @@ elif sys.platform == "darwin":
25 22 os.environ["RUNFILES_DIR"] = str(local_user_plaidml)
26 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 29 import invesalius.data.slice_ as slc
30 30 from invesalius.segmentation.brain import segment
31 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 50 class BrainSegmenterDialog(wx.Dialog):
... ... @@ -42,13 +57,14 @@ class BrainSegmenterDialog(wx.Dialog):
42 57 backends.append("Theano")
43 58 self.segmenter = segment.BrainSegmenter()
44 59 # self.pg_dialog = None
45   - self.plaidml_devices = utils.get_plaidml_devices()
  60 + self.plaidml_devices = PLAIDML_DEVICES
46 61 self.cb_backends = wx.ComboBox(self, wx.ID_ANY, choices=backends, value=backends[0], style=wx.CB_DROPDOWN | wx.CB_READONLY)
47 62 w, h = self.CalcSizeFromTextSize("MM" * (1 + max(len(i) for i in backends)))
48 63 self.cb_backends.SetMinClientSize((w, -1))
49 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 68 self.sld_threshold = wx.Slider(self, wx.ID_ANY, 75, 0, 100)
53 69 w, h = self.CalcSizeFromTextSize("M" * 20)
54 70 self.sld_threshold.SetMinClientSize((w, -1))
... ... @@ -76,8 +92,9 @@ class BrainSegmenterDialog(wx.Dialog):
76 92 main_sizer.Add(sizer_backends, 0, wx.ALL | wx.EXPAND, 5)
77 93 main_sizer.Add(self.chk_use_gpu, 0, wx.ALL, 5)
78 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 98 main_sizer.Add(sizer_devices, 0, wx.ALL | wx.EXPAND, 5)
82 99 label_5 = wx.StaticText(self, wx.ID_ANY, _("Level of certainty"))
83 100 main_sizer.Add(label_5, 0, wx.ALL, 5)
... ... @@ -118,12 +135,14 @@ class BrainSegmenterDialog(wx.Dialog):
118 135  
119 136 def OnSetBackend(self, evt=None):
120 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 141 self.chk_use_gpu.Hide()
124 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 146 self.chk_use_gpu.Show()
128 147  
129 148 self.main_sizer.Fit(self)
... ... @@ -159,7 +178,7 @@ class BrainSegmenterDialog(wx.Dialog):
159 178 backend = self.cb_backends.GetValue()
160 179 try:
161 180 device_id = self.plaidml_devices[self.cb_devices.GetValue()]
162   - except KeyError:
  181 + except (KeyError, AttributeError):
163 182 device_id = "llvm_cpu.0"
164 183 use_gpu = self.chk_use_gpu.GetValue()
165 184 prob_threshold = self.sld_threshold.GetValue() / 100.0
... ...