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,7 +18,9 @@
18 #-------------------------------------------------------------------------- 18 #--------------------------------------------------------------------------
19 19
20 import os 20 import os
  21 +import multiprocessing
21 import tempfile 22 import tempfile
  23 +import time
22 24
23 import vtk 25 import vtk
24 import wx 26 import wx
@@ -738,6 +740,43 @@ class EditorInteractorStyle(DefaultInteractorStyle): @@ -738,6 +740,43 @@ class EditorInteractorStyle(DefaultInteractorStyle):
738 return x, y, z 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 class WatershedConfig(object): 780 class WatershedConfig(object):
742 __metaclass__= utils.Singleton 781 __metaclass__= utils.Singleton
743 def __init__(self): 782 def __init__(self):
@@ -1201,28 +1240,57 @@ class WaterShedInteractorStyle(DefaultInteractorStyle): @@ -1201,28 +1240,57 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
1201 if BRUSH_BACKGROUND in markers and BRUSH_FOREGROUND in markers: 1240 if BRUSH_BACKGROUND in markers and BRUSH_FOREGROUND in markers:
1202 #w_algorithm = WALGORITHM[self.config.algorithm] 1241 #w_algorithm = WALGORITHM[self.config.algorithm]
1203 bstruct = generate_binary_structure(3, CON3D[self.config.con_3d]) 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 #tmp_image = ndimage.morphological_gradient( 1274 #tmp_image = ndimage.morphological_gradient(
1214 #get_LUT_value(image, ww, wl).astype('uint16'), 1275 #get_LUT_value(image, ww, wl).astype('uint16'),
1215 #self.config.mg_size) 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 #tmp_image = ndimage.morphological_gradient((image - image.min()).astype('uint16'), self.config.mg_size) 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 if self.viewer.overwrite_mask: 1295 if self.viewer.overwrite_mask:
1228 mask[:] = 0 1296 mask[:] = 0
@@ -1241,6 +1309,34 @@ class WaterShedInteractorStyle(DefaultInteractorStyle): @@ -1241,6 +1309,34 @@ class WaterShedInteractorStyle(DefaultInteractorStyle):
1241 Publisher.sendMessage('Reload actual slice') 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 def get_style(style): 1340 def get_style(style):
1245 STYLES = { 1341 STYLES = {
1246 const.STATE_DEFAULT: DefaultInteractorStyle, 1342 const.STATE_DEFAULT: DefaultInteractorStyle,