Download as
Browse Code »

Commit ec28565e626970c6b63950d27106acbf71c69a44

Authored by Thiago Franco de Moraes
1 parent 07491a64
Exists in watershed_improvements_bkp

runing 3d watershed using multiprocessing

Showing 1 changed file with 114 additions and 18 deletions   Show diff stats
invesalius/data/styles.py
  Diff comments   View file @ec28565
... ... @@ -18,7 +18,9 @@
18 18 #--------------------------------------------------------------------------
19 19  
20 20 import os
  21 +import multiprocessing
21 22 import tempfile
  23 +import time
22 24  
23 25 import vtk
24 26 import wx
... ... @@ -738,6 +740,43 @@ class EditorInteractorStyle(DefaultInteractorStyle):
738 740 return x, y, z
739 741  
740 742  
  743 +class WatershedProgressWindow(wx.Frame):
  744 + def __init__(self, process, parent=None):
  745 + wx.Frame.__init__(self, parent, -1)
  746 + self.process = process
  747 + self._build_gui()
  748 + self._bind_wx_events()
  749 + self.timer = wx.Timer(self)
  750 + self.timer.Start(1000)
  751 +
  752 + def _build_gui(self):
  753 + self.gauge = wx.Gauge(self, -1, 100)
  754 + self.btn_cancel = wx.Button(self, -1, _("&Cancel"))
  755 +
  756 + sizer = wx.BoxSizer(wx.VERTICAL)
  757 + sizer.Add(wx.StaticText(self, -1, _("Applying watershed")))
  758 + sizer.Add(self.gauge, 0, wx.EXPAND)
  759 + sizer.Add(self.btn_cancel, 0, wx.ALIGN_LEFT)
  760 +
  761 + self.SetSizer(sizer)
  762 + sizer.Fit(self)
  763 + self.Layout()
  764 +
  765 + def __del__(self):
  766 + self.timer.Stop()
  767 +
  768 + def _bind_wx_events(self):
  769 + self.Bind(wx.EVT_TIMER, self.TimeHandler)
  770 + self.btn_cancel.Bind(wx.EVT_BUTTON, self.on_cancel)
  771 +
  772 + def on_cancel(self, evt):
  773 + self.timer.Stop()
  774 + self.process.terminate()
  775 +
  776 + def TimeHandler(self, evt):
  777 + self.gauge.Pulse()
  778 +
  779 +
741 780 class WatershedConfig(object):
742 781 __metaclass__= utils.Singleton
743 782 def __init__(self):
... ... @@ -1201,28 +1240,57 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
1201 1240 if BRUSH_BACKGROUND in markers and BRUSH_FOREGROUND in markers:
1202 1241 #w_algorithm = WALGORITHM[self.config.algorithm]
1203 1242 bstruct = generate_binary_structure(3, CON3D[self.config.con_3d])
1204   - if self.config.use_ww_wl:
1205   - if self.config.algorithm == 'Watershed':
1206   - tmp_image = ndimage.morphological_gradient(
1207   - get_LUT_value(image, ww, wl).astype('uint16'),
1208   - self.config.mg_size)
1209   - tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
1210   - else:
1211   - tmp_image = get_LUT_value(image, ww, wl).astype('uint16')
1212   - #tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
  1243 + tfile = tempfile.mktemp()
  1244 + tmp_mask = np.memmap(tfile, shape=mask.shape, dtype=mask.dtype,
  1245 + mode='w+')
  1246 + q = multiprocessing.Queue()
  1247 + p = multiprocessing.Process(target=do_watershed, args=(image,
  1248 + markers, tmp_mask, bstruct,
  1249 + self.config.algorithm,
  1250 + self.config.mg_size,
  1251 + self.config.use_ww_wl, wl, ww, q))
  1252 +
  1253 + wp = WatershedProgressWindow(p)
  1254 + wp.Center(wx.BOTH)
  1255 + wp.Show()
  1256 + wp.MakeModal()
  1257 +
  1258 + p.start()
  1259 +
  1260 + while q.empty() and p.is_alive():
  1261 + time.sleep(0.5)
  1262 + wx.Yield()
  1263 +
  1264 + wp.MakeModal(False)
  1265 + wp.Destroy()
  1266 + del wp
  1267 +
  1268 + if q.empty():
  1269 + return
  1270 + #do_watershed(image, markers, tmp_mask, bstruct, self.config.algorithm,
  1271 + #self.config.mg_size, self.config.use_ww_wl, wl, ww)
  1272 + #if self.config.use_ww_wl:
  1273 + #if self.config.algorithm == 'Watershed':
1213 1274 #tmp_image = ndimage.morphological_gradient(
1214 1275 #get_LUT_value(image, ww, wl).astype('uint16'),
1215 1276 #self.config.mg_size)
1216   - tmp_mask = watershed_ift(tmp_image, markers.astype('int16'), bstruct)
1217   - else:
1218   - if self.config.algorithm == 'Watershed':
1219   - tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), self.config.mg_size)
1220   - tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
1221   - else:
1222   - tmp_image = (image - image.min()).astype('uint16')
1223   - #tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
  1277 + #tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
  1278 + #else:
  1279 + #tmp_image = get_LUT_value(image, ww, wl).astype('uint16')
  1280 + ##tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
  1281 + ##tmp_image = ndimage.morphological_gradient(
  1282 + ##get_LUT_value(image, ww, wl).astype('uint16'),
  1283 + ##self.config.mg_size)
  1284 + #tmp_mask = watershed_ift(tmp_image, markers.astype('int16'), bstruct)
  1285 + #else:
  1286 + #if self.config.algorithm == 'Watershed':
1224 1287 #tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), self.config.mg_size)
1225   - tmp_mask = watershed_ift(tmp_image, markers.astype('int8'), bstruct)
  1288 + #tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
  1289 + #else:
  1290 + #tmp_image = (image - image.min()).astype('uint16')
  1291 + ##tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
  1292 + ##tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), self.config.mg_size)
  1293 + #tmp_mask = watershed_ift(tmp_image, markers.astype('int8'), bstruct)
1226 1294  
1227 1295 if self.viewer.overwrite_mask:
1228 1296 mask[:] = 0
... ... @@ -1241,6 +1309,34 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
1241 1309 Publisher.sendMessage('Reload actual slice')
1242 1310  
1243 1311  
  1312 +def do_watershed(image, markers, mask, bstruct, algorithm, mg_size, use_ww_wl, wl, ww, q):
  1313 + if use_ww_wl:
  1314 + if algorithm == 'Watershed':
  1315 + tmp_image = ndimage.morphological_gradient(
  1316 + get_LUT_value(image, ww, wl).astype('uint16'),
  1317 + mg_size)
  1318 + tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
  1319 + else:
  1320 + tmp_image = get_LUT_value(image, ww, wl).astype('uint16')
  1321 + #tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
  1322 + #tmp_image = ndimage.morphological_gradient(
  1323 + #get_LUT_value(image, ww, wl).astype('uint16'),
  1324 + #self.config.mg_size)
  1325 + tmp_mask = watershed_ift(tmp_image, markers.astype('int16'), bstruct)
  1326 + else:
  1327 + if algorithm == 'Watershed':
  1328 + tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), mg_size)
  1329 + tmp_mask = watershed(tmp_image, markers.astype('int16'), bstruct)
  1330 + else:
  1331 + tmp_image = (image - image.min()).astype('uint16')
  1332 + #tmp_image = ndimage.gaussian_filter(tmp_image, self.config.mg_size)
  1333 + #tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), self.config.mg_size)
  1334 + tmp_mask = watershed_ift(tmp_image, markers.astype('int8'), bstruct)
  1335 + mask[:] = tmp_mask
  1336 + q.put(1)
  1337 +
  1338 +
  1339 +
1244 1340 def get_style(style):
1245 1341 STYLES = {
1246 1342 const.STATE_DEFAULT: DefaultInteractorStyle,
... ...