Download as
Browse Code »

Commit c60264663d8996611a3f055a5de54c86e2a3ea05

Authored by Thiago Franco de Moraes
Committed by GitHub
1 parent 6dc0fd02
Exists in master

Cancelable surface creation (#147) closes #142 

* Creating the surface using function instead of a class

* Generating and joining surface using multiprocessing

* Stripper and normal calc in surface process

* Calculating volume and area in the surface process

* Better messages in the progress dialog

* Cancel working in windows

* Using callback

* Using callback errors

* python2 doesnt have error_callback
Showing 5 changed files with 603 additions and 489 deletions   Show diff stats
invesalius/data/interpolation.pyx
  Diff comments   View file @c602646
... ... @@ -4,7 +4,7 @@ import numpy as np
4 4 cimport numpy as np
5 5 cimport cython
6 6  
7   -from libc.math cimport floor, ceil, sqrt, fabs, round, sin, M_PI
  7 +from libc.math cimport floor, ceil, sqrt, fabs, sin, M_PI
8 8 from cython.parallel import prange
9 9  
10 10 DEF LANCZOS_A = 4
... ... @@ -81,7 +81,7 @@ cdef double[64][64] temp = [
81 81 @cython.cdivision(True)
82 82 @cython.wraparound(False)
83 83 cdef double nearest_neighbour_interp(image_t[:, :, :] V, double x, double y, double z) nogil:
84   - return V[<int>round(z), <int>round(y), <int>round(x)]
  84 + return V[<int>(z), <int>(y), <int>(x)]
85 85  
86 86 @cython.boundscheck(False) # turn of bounds-checking for entire function
87 87 @cython.cdivision(True)
... ...
invesalius/data/surface.py
  Diff comments   View file @c602646
... ... @@ -17,6 +17,7 @@
17 17 # detalhes.
18 18 #--------------------------------------------------------------------------
19 19  
  20 +import functools
20 21 import multiprocessing
21 22 import os
22 23 import plistlib
... ... @@ -24,10 +25,19 @@ import random
24 25 import shutil
25 26 import sys
26 27 import tempfile
  28 +import time
  29 +import traceback
27 30 import weakref
28 31  
  32 +try:
  33 + import queue
  34 +except ImportError:
  35 + import Queue as queue
  36 +
29 37 import vtk
30 38 import wx
  39 +import wx.lib.agw.genericmessagedialog as GMD
  40 +
31 41 from wx.lib.pubsub import pub as Publisher
32 42  
33 43 if sys.platform == 'win32':
... ... @@ -48,9 +58,12 @@ import invesalius.data.surface_process as surface_process
48 58 import invesalius.utils as utl
49 59 import invesalius.data.vtk_utils as vu
50 60  
  61 +from invesalius.gui import dialogs
  62 +
51 63 from invesalius.data import cy_mesh
52 64 # TODO: Verificar ReleaseDataFlagOn and SetSource
53 65  
  66 +
54 67 class Surface():
55 68 """
56 69 Represent both vtkPolyData and associated properties.
... ... @@ -438,10 +451,81 @@ class SurfaceManager():
438 451 ####
439 452 #(mask_index, surface_name, quality, fill_holes, keep_largest)
440 453  
  454 + def _on_complete_surface_creation(self, args, overwrite, surface_name, colour, dialog):
  455 + surface_filename, surface_measures = args
  456 + print(surface_filename, surface_measures)
  457 + reader = vtk.vtkXMLPolyDataReader()
  458 + reader.SetFileName(surface_filename)
  459 + reader.Update()
  460 + polydata = reader.GetOutput()
  461 +
  462 + # Map polygonal data (vtkPolyData) to graphics primitives.
  463 + mapper = vtk.vtkPolyDataMapper()
  464 + mapper.SetInputData(polydata)
  465 + mapper.ScalarVisibilityOff()
  466 + # mapper.ReleaseDataFlagOn()
  467 + mapper.ImmediateModeRenderingOn() # improve performance
  468 +
  469 + # Represent an object (geometry & properties) in the rendered scene
  470 + actor = vtk.vtkActor()
  471 + actor.SetMapper(mapper)
  472 + del mapper
  473 + #Create Surface instance
  474 + if overwrite:
  475 + surface = Surface(index = self.last_surface_index)
  476 + else:
  477 + surface = Surface(name=surface_name)
  478 + surface.colour = colour
  479 + surface.polydata = polydata
  480 + surface.volume = surface_measures['volume']
  481 + surface.area = surface_measures['area']
  482 + del polydata
  483 +
  484 + # Set actor colour and transparency
  485 + actor.GetProperty().SetColor(colour)
  486 + actor.GetProperty().SetOpacity(1-surface.transparency)
  487 +
  488 + prop = actor.GetProperty()
  489 +
  490 + interpolation = int(ses.Session().surface_interpolation)
  491 +
  492 + prop.SetInterpolation(interpolation)
  493 +
  494 + proj = prj.Project()
  495 + if overwrite:
  496 + proj.ChangeSurface(surface)
  497 + else:
  498 + index = proj.AddSurface(surface)
  499 + surface.index = index
  500 + self.last_surface_index = index
  501 +
  502 + session = ses.Session()
  503 + session.ChangeProject()
  504 +
  505 + Publisher.sendMessage('Load surface actor into viewer', actor=actor)
  506 +
  507 + # Send actor by pubsub to viewer's render
  508 + if overwrite and self.actors_dict.keys():
  509 + old_actor = self.actors_dict[self.last_surface_index]
  510 + Publisher.sendMessage('Remove surface actor from viewer', actor=old_actor)
  511 +
  512 + # Save actor for future management tasks
  513 + self.actors_dict[surface.index] = actor
  514 + Publisher.sendMessage('Update surface info in GUI', surface=surface)
  515 + Publisher.sendMessage('End busy cursor')
  516 +
  517 + dialog.running = False
  518 +
  519 + def _on_callback_error(self, e, dialog=None):
  520 + dialog.running = False
  521 + msg = utl.log_traceback(e)
  522 + dialog.error = msg
  523 +
441 524 def AddNewActor(self, slice_, mask, surface_parameters):
442 525 """
443 526 Create surface actor, save into project and send it to viewer.
444 527 """
  528 + t_init = time.time()
445 529 matrix = slice_.matrix
446 530 filename_img = slice_.matrix_filename
447 531 spacing = slice_.spacing
... ... @@ -470,9 +554,6 @@ class SurfaceManager():
470 554 smooth_relaxation_factor = const.SURFACE_QUALITY[quality][2]
471 555 decimate_reduction = const.SURFACE_QUALITY[quality][3]
472 556  
473   - #if imagedata_resolution:
474   - #imagedata = iu.ResampleImage3D(imagedata, imagedata_resolution)
475   -
476 557 pipeline_size = 4
477 558 if decimate_reduction:
478 559 pipeline_size += 1
... ... @@ -483,10 +564,6 @@ class SurfaceManager():
483 564 if keep_largest:
484 565 pipeline_size += 1
485 566  
486   - ## Update progress value in GUI
487   - UpdateProgress = vu.ShowProgress(pipeline_size)
488   - UpdateProgress(0, _("Creating 3D surface..."))
489   -
490 567 language = ses.Session().language
491 568  
492 569 if (prj.Project().original_orientation == const.CORONAL):
... ... @@ -496,220 +573,59 @@ class SurfaceManager():
496 573  
497 574 n_processors = multiprocessing.cpu_count()
498 575  
499   - pipe_in, pipe_out = multiprocessing.Pipe()
500 576 o_piece = 1
501   - piece_size = 2000
  577 + piece_size = 20
502 578  
503 579 n_pieces = int(round(matrix.shape[0] / piece_size + 0.5, 0))
504 580  
505   - q_in = multiprocessing.Queue()
506   - q_out = multiprocessing.Queue()
507   -
508   - p = []
509   - for i in range(n_processors):
510   - sp = surface_process.SurfaceProcess(pipe_in, filename_img,
511   - matrix.shape, matrix.dtype,
512   - mask.temp_file,
513   - mask.matrix.shape,
514   - mask.matrix.dtype,
515   - spacing,
516   - mode, min_value, max_value,
517   - decimate_reduction,
518   - smooth_relaxation_factor,
519   - smooth_iterations, language,
520   - flip_image, q_in, q_out,
521   - algorithm != 'Default',
522   - algorithm,
523   - imagedata_resolution)
524   - p.append(sp)
525   - sp.start()
526   -
527   - for i in range(n_pieces):
528   - init = i * piece_size
529   - end = init + piece_size + o_piece
530   - roi = slice(init, end)
531   - q_in.put(roi)
532   - print("new_piece", roi)
533   -
534   - for i in p:
535   - q_in.put(None)
536   -
537   - none_count = 1
538   - while 1:
539   - msg = pipe_out.recv()
540   - if(msg is None):
541   - none_count += 1
542   - else:
543   - UpdateProgress(msg[0]/(n_pieces * pipeline_size), msg[1])
544   -
545   - if none_count > n_pieces:
546   - break
  581 + filenames = []
  582 + pool = multiprocessing.Pool(processes=min(n_pieces, n_processors))
  583 + manager = multiprocessing.Manager()
  584 + msg_queue = manager.Queue(1)
547 585  
548   - polydata_append = vtk.vtkAppendPolyData()
549   - # polydata_append.ReleaseDataFlagOn()
550   - t = n_pieces
551   - while t:
552   - filename_polydata = q_out.get()
  586 + # If InVesalius is running without GUI
  587 + if wx.GetApp() is None:
  588 + for i in range(n_pieces):
  589 + init = i * piece_size
  590 + end = init + piece_size + o_piece
  591 + roi = slice(init, end)
  592 + print("new_piece", roi)
  593 + f = pool.apply_async(surface_process.create_surface_piece,
  594 + args = (filename_img, matrix.shape, matrix.dtype,
  595 + mask.temp_file, mask.matrix.shape,
  596 + mask.matrix.dtype, roi, spacing, mode,
  597 + min_value, max_value, decimate_reduction,
  598 + smooth_relaxation_factor,
  599 + smooth_iterations, language, flip_image,
  600 + algorithm != 'Default', algorithm,
  601 + imagedata_resolution),
  602 + callback=lambda x: filenames.append(x))
  603 +
  604 + while len(filenames) != n_pieces:
  605 + time.sleep(0.25)
  606 +
  607 + f = pool.apply_async(surface_process.join_process_surface,
  608 + args=(filenames, algorithm, smooth_iterations,
  609 + smooth_relaxation_factor,
  610 + decimate_reduction, keep_largest,
  611 + fill_holes, options, msg_queue))
  612 +
  613 + while not f.ready():
  614 + time.sleep(0.25)
  615 +
  616 + try:
  617 + surface_filename, surface_measures = f.get()
  618 + except Exception as e:
  619 + print(_("InVesalius was not able to create the surface"))
  620 + print(traceback.print_exc())
  621 + return
553 622  
554 623 reader = vtk.vtkXMLPolyDataReader()
555   - reader.SetFileName(filename_polydata)
556   - # reader.ReleaseDataFlagOn()
  624 + reader.SetFileName(surface_filename)
557 625 reader.Update()
558   - # reader.GetOutput().ReleaseDataFlagOn()
559 626  
560 627 polydata = reader.GetOutput()
561   - # polydata.SetSource(None)
562   -
563   - polydata_append.AddInputData(polydata)
564   - del reader
565   - del polydata
566   - t -= 1
567   -
568   - polydata_append.Update()
569   - # polydata_append.GetOutput().ReleaseDataFlagOn()
570   - polydata = polydata_append.GetOutput()
571   - #polydata.Register(None)
572   - # polydata.SetSource(None)
573   - del polydata_append
574   -
575   - if algorithm == 'ca_smoothing':
576   - normals = vtk.vtkPolyDataNormals()
577   - normals_ref = weakref.ref(normals)
578   - normals_ref().AddObserver("ProgressEvent", lambda obj,evt:
579   - UpdateProgress(normals_ref(), _("Creating 3D surface...")))
580   - normals.SetInputData(polydata)
581   - # normals.ReleaseDataFlagOn()
582   - #normals.SetFeatureAngle(80)
583   - #normals.AutoOrientNormalsOn()
584   - normals.ComputeCellNormalsOn()
585   - # normals.GetOutput().ReleaseDataFlagOn()
586   - normals.Update()
587   - del polydata
588   - polydata = normals.GetOutput()
589   - # polydata.SetSource(None)
590   - del normals
591   -
592   - clean = vtk.vtkCleanPolyData()
593   - # clean.ReleaseDataFlagOn()
594   - # clean.GetOutput().ReleaseDataFlagOn()
595   - clean_ref = weakref.ref(clean)
596   - clean_ref().AddObserver("ProgressEvent", lambda obj,evt:
597   - UpdateProgress(clean_ref(), _("Creating 3D surface...")))
598   - clean.SetInputData(polydata)
599   - clean.PointMergingOn()
600   - clean.Update()
601   -
602   - del polydata
603   - polydata = clean.GetOutput()
604   - # polydata.SetSource(None)
605   - del clean
606   -
607   - # try:
608   - # polydata.BuildLinks()
609   - # except TypeError:
610   - # polydata.BuildLinks(0)
611   - # polydata = ca_smoothing.ca_smoothing(polydata, options['angle'],
612   - # options['max distance'],
613   - # options['min weight'],
614   - # options['steps'])
615   -
616   - mesh = cy_mesh.Mesh(polydata)
617   - cy_mesh.ca_smoothing(mesh, options['angle'],
618   - options['max distance'],
619   - options['min weight'],
620   - options['steps'])
621   - # polydata = mesh.to_vtk()
622   -
623   - # polydata.SetSource(None)
624   - # polydata.DebugOn()
625   - else:
626   - #smoother = vtk.vtkWindowedSincPolyDataFilter()
627   - smoother = vtk.vtkSmoothPolyDataFilter()
628   - smoother_ref = weakref.ref(smoother)
629   - smoother_ref().AddObserver("ProgressEvent", lambda obj,evt:
630   - UpdateProgress(smoother_ref(), _("Creating 3D surface...")))
631   - smoother.SetInputData(polydata)
632   - smoother.SetNumberOfIterations(smooth_iterations)
633   - smoother.SetRelaxationFactor(smooth_relaxation_factor)
634   - smoother.SetFeatureAngle(80)
635   - #smoother.SetEdgeAngle(90.0)
636   - #smoother.SetPassBand(0.1)
637   - smoother.BoundarySmoothingOn()
638   - smoother.FeatureEdgeSmoothingOn()
639   - #smoother.NormalizeCoordinatesOn()
640   - #smoother.NonManifoldSmoothingOn()
641   - # smoother.ReleaseDataFlagOn()
642   - # smoother.GetOutput().ReleaseDataFlagOn()
643   - smoother.Update()
644   - del polydata
645   - polydata = smoother.GetOutput()
646   - #polydata.Register(None)
647   - # polydata.SetSource(None)
648   - del smoother
649   -
650   -
651   - if decimate_reduction:
652   - print("Decimating", decimate_reduction)
653   - decimation = vtk.vtkQuadricDecimation()
654   - # decimation.ReleaseDataFlagOn()
655   - decimation.SetInputData(polydata)
656   - decimation.SetTargetReduction(decimate_reduction)
657   - decimation_ref = weakref.ref(decimation)
658   - decimation_ref().AddObserver("ProgressEvent", lambda obj,evt:
659   - UpdateProgress(decimation_ref(), _("Creating 3D surface...")))
660   - #decimation.PreserveTopologyOn()
661   - #decimation.SplittingOff()
662   - #decimation.BoundaryVertexDeletionOff()
663   - # decimation.GetOutput().ReleaseDataFlagOn()
664   - decimation.Update()
665   - del polydata
666   - polydata = decimation.GetOutput()
667   - #polydata.Register(None)
668   - # polydata.SetSource(None)
669   - del decimation
670   -
671   - #to_measure.Register(None)
672   - # to_measure.SetSource(None)
673 628  
674   - if keep_largest:
675   - conn = vtk.vtkPolyDataConnectivityFilter()
676   - conn.SetInputData(polydata)
677   - conn.SetExtractionModeToLargestRegion()
678   - conn_ref = weakref.ref(conn)
679   - conn_ref().AddObserver("ProgressEvent", lambda obj,evt:
680   - UpdateProgress(conn_ref(), _("Creating 3D surface...")))
681   - conn.Update()
682   - # conn.GetOutput().ReleaseDataFlagOn()
683   - del polydata
684   - polydata = conn.GetOutput()
685   - #polydata.Register(None)
686   - # polydata.SetSource(None)
687   - del conn
688   -
689   - #Filter used to detect and fill holes. Only fill boundary edges holes.
690   - #TODO: Hey! This piece of code is the same from
691   - #polydata_utils.FillSurfaceHole, we need to review this.
692   - if fill_holes:
693   - filled_polydata = vtk.vtkFillHolesFilter()
694   - # filled_polydata.ReleaseDataFlagOn()
695   - filled_polydata.SetInputData(polydata)
696   - filled_polydata.SetHoleSize(300)
697   - filled_polydata_ref = weakref.ref(filled_polydata)
698   - filled_polydata_ref().AddObserver("ProgressEvent", lambda obj,evt:
699   - UpdateProgress(filled_polydata_ref(), _("Creating 3D surface...")))
700   - filled_polydata.Update()
701   - # filled_polydata.GetOutput().ReleaseDataFlagOn()
702   - del polydata
703   - polydata = filled_polydata.GetOutput()
704   - #polydata.Register(None)
705   - # polydata.SetSource(None)
706   - # polydata.DebugOn()
707   - del filled_polydata
708   -
709   - to_measure = polydata
710   -
711   - # If InVesalius is running without GUI
712   - if wx.GetApp() is None:
713 629 proj = prj.Project()
714 630 #Create Surface instance
715 631 if overwrite:
... ... @@ -720,115 +636,108 @@ class SurfaceManager():
720 636 index = proj.AddSurface(surface)
721 637 surface.index = index
722 638 self.last_surface_index = index
  639 +
723 640 surface.colour = colour
724 641 surface.polydata = polydata
  642 + surface.volume = surface_measures['volume']
  643 + surface.area = surface_measures['area']
725 644  
726 645 # With GUI
727 646 else:
728   - normals = vtk.vtkPolyDataNormals()
729   - # normals.ReleaseDataFlagOn()
730   - normals_ref = weakref.ref(normals)
731   - normals_ref().AddObserver("ProgressEvent", lambda obj,evt:
732   - UpdateProgress(normals_ref(), _("Creating 3D surface...")))
733   - normals.SetInputData(polydata)
734   - normals.SetFeatureAngle(80)
735   - normals.AutoOrientNormalsOn()
736   - # normals.GetOutput().ReleaseDataFlagOn()
737   - normals.Update()
738   - del polydata
739   - polydata = normals.GetOutput()
740   - #polydata.Register(None)
741   - # polydata.SetSource(None)
742   - del normals
743   -
744   - # Improve performance
745   - stripper = vtk.vtkStripper()
746   - # stripper.ReleaseDataFlagOn()
747   - stripper_ref = weakref.ref(stripper)
748   - stripper_ref().AddObserver("ProgressEvent", lambda obj,evt:
749   - UpdateProgress(stripper_ref(), _("Creating 3D surface...")))
750   - stripper.SetInputData(polydata)
751   - stripper.PassThroughCellIdsOn()
752   - stripper.PassThroughPointIdsOn()
753   - # stripper.GetOutput().ReleaseDataFlagOn()
754   - stripper.Update()
755   - del polydata
756   - polydata = stripper.GetOutput()
757   - #polydata.Register(None)
758   - # polydata.SetSource(None)
759   - del stripper
760   -
761   - # Map polygonal data (vtkPolyData) to graphics primitives.
762   - mapper = vtk.vtkPolyDataMapper()
763   - mapper.SetInputData(polydata)
764   - mapper.ScalarVisibilityOff()
765   - # mapper.ReleaseDataFlagOn()
766   - mapper.ImmediateModeRenderingOn() # improve performance
767   -
768   - # Represent an object (geometry & properties) in the rendered scene
769   - actor = vtk.vtkActor()
770   - actor.SetMapper(mapper)
771   - del mapper
772   - #Create Surface instance
773   - if overwrite:
774   - surface = Surface(index = self.last_surface_index)
775   - else:
776   - surface = Surface(name=surface_name)
777   - surface.colour = colour
778   - surface.polydata = polydata
779   - del polydata
780   -
781   - # Set actor colour and transparency
782   - actor.GetProperty().SetColor(colour)
783   - actor.GetProperty().SetOpacity(1-surface.transparency)
784   -
785   - prop = actor.GetProperty()
786   -
787   - interpolation = int(ses.Session().surface_interpolation)
788   -
789   - prop.SetInterpolation(interpolation)
790   -
791   - proj = prj.Project()
792   - if overwrite:
793   - proj.ChangeSurface(surface)
794   - else:
795   - index = proj.AddSurface(surface)
796   - surface.index = index
797   - self.last_surface_index = index
798   -
799   - session = ses.Session()
800   - session.ChangeProject()
801   -
802   - measured_polydata = vtk.vtkMassProperties()
803   - # measured_polydata.ReleaseDataFlagOn()
804   - measured_polydata.SetInputData(to_measure)
805   - volume = float(measured_polydata.GetVolume())
806   - area = float(measured_polydata.GetSurfaceArea())
807   - surface.volume = volume
808   - surface.area = area
809   - self.last_surface_index = surface.index
810   - del measured_polydata
811   - del to_measure
812   -
813   - Publisher.sendMessage('Load surface actor into viewer', actor=actor)
814   -
815   - # Send actor by pubsub to viewer's render
816   - if overwrite and self.actors_dict.keys():
817   - old_actor = self.actors_dict[self.last_surface_index]
818   - Publisher.sendMessage('Remove surface actor from viewer', actor=old_actor)
819   -
820   - # Save actor for future management tasks
821   - self.actors_dict[surface.index] = actor
822   -
823   - Publisher.sendMessage('Update surface info in GUI', surface=surface)
824   -
825   - #When you finalize the progress. The bar is cleaned.
826   - UpdateProgress = vu.ShowProgress(1)
827   - UpdateProgress(0, _("Ready"))
828   - Publisher.sendMessage('Update status text in GUI', label=_("Ready"))
829   -
830   - Publisher.sendMessage('End busy cursor')
831   - del actor
  647 + sp = dialogs.SurfaceProgressWindow()
  648 + for i in range(n_pieces):
  649 + init = i * piece_size
  650 + end = init + piece_size + o_piece
  651 + roi = slice(init, end)
  652 + print("new_piece", roi)
  653 + try:
  654 + f = pool.apply_async(surface_process.create_surface_piece,
  655 + args = (filename_img, matrix.shape, matrix.dtype,
  656 + mask.temp_file, mask.matrix.shape,
  657 + mask.matrix.dtype, roi, spacing, mode,
  658 + min_value, max_value, decimate_reduction,
  659 + smooth_relaxation_factor,
  660 + smooth_iterations, language, flip_image,
  661 + algorithm != 'Default', algorithm,
  662 + imagedata_resolution),
  663 + callback=lambda x: filenames.append(x),
  664 + error_callback=functools.partial(self._on_callback_error,
  665 + dialog=sp))
  666 + # python2
  667 + except TypeError:
  668 + f = pool.apply_async(surface_process.create_surface_piece,
  669 + args = (filename_img, matrix.shape, matrix.dtype,
  670 + mask.temp_file, mask.matrix.shape,
  671 + mask.matrix.dtype, roi, spacing, mode,
  672 + min_value, max_value, decimate_reduction,
  673 + smooth_relaxation_factor,
  674 + smooth_iterations, language, flip_image,
  675 + algorithm != 'Default', algorithm,
  676 + imagedata_resolution),
  677 + callback=lambda x: filenames.append(x))
  678 +
  679 + while len(filenames) != n_pieces:
  680 + if sp.WasCancelled() or not sp.running:
  681 + break
  682 + time.sleep(0.25)
  683 + sp.Update(_("Creating 3D surface..."))
  684 + wx.Yield()
  685 +
  686 + if not sp.WasCancelled() or sp.running:
  687 + try:
  688 + f = pool.apply_async(surface_process.join_process_surface,
  689 + args=(filenames, algorithm, smooth_iterations,
  690 + smooth_relaxation_factor,
  691 + decimate_reduction, keep_largest,
  692 + fill_holes, options, msg_queue),
  693 + callback=functools.partial(self._on_complete_surface_creation,
  694 + overwrite=overwrite,
  695 + surface_name=surface_name,
  696 + colour=colour,
  697 + dialog=sp),
  698 + error_callback=functools.partial(self._on_callback_error,
  699 + dialog=sp))
  700 + # python2
  701 + except TypeError:
  702 + f = pool.apply_async(surface_process.join_process_surface,
  703 + args=(filenames, algorithm, smooth_iterations,
  704 + smooth_relaxation_factor,
  705 + decimate_reduction, keep_largest,
  706 + fill_holes, options, msg_queue),
  707 + callback=functools.partial(self._on_complete_surface_creation,
  708 + overwrite=overwrite,
  709 + surface_name=surface_name,
  710 + colour=colour,
  711 + dialog=sp))
  712 +
  713 + while sp.running:
  714 + if sp.WasCancelled():
  715 + break
  716 + time.sleep(0.25)
  717 + try:
  718 + msg = msg_queue.get_nowait()
  719 + sp.Update(msg)
  720 + except:
  721 + sp.Update(None)
  722 + wx.Yield()
  723 +
  724 + t_end = time.time()
  725 + print("Elapsed time - {}".format(t_end-t_init))
  726 + sp.Close()
  727 + if sp.error:
  728 + dlg = GMD.GenericMessageDialog(None, sp.error,
  729 + "Exception!",
  730 + wx.OK|wx.ICON_ERROR)
  731 + dlg.ShowModal()
  732 + del sp
  733 +
  734 + pool.close()
  735 + pool.terminate()
  736 + del pool
  737 + del manager
  738 + del msg_queue
  739 + import gc
  740 + gc.collect()
832 741  
833 742 def UpdateSurfaceInterpolation(self):
834 743 interpolation = int(ses.Session().surface_interpolation)
... ...
invesalius/data/surface_process.py
  Diff comments   View file @c602646
1 1 import multiprocessing
  2 +import os
2 3 import tempfile
3 4 import time
4 5  
  6 +try:
  7 + import queue
  8 +except ImportError:
  9 + import Queue as queue
  10 +
5 11 import numpy
6 12 import vtk
7 13  
8 14 import invesalius.i18n as i18n
9 15 import invesalius.data.converters as converters
  16 +from invesalius.data import cy_mesh
10 17 # import invesalius.data.imagedata_utils as iu
11 18  
  19 +import weakref
12 20 from scipy import ndimage
13 21  
14 22 # TODO: Code duplicated from file {imagedata_utils.py}.
... ... @@ -32,182 +40,329 @@ def ResampleImage3D(imagedata, value):
32 40  
33 41 return resample.GetOutput()
34 42  
35   -class SurfaceProcess(multiprocessing.Process):
36   -
37   - def __init__(self, pipe, filename, shape, dtype, mask_filename,
38   - mask_shape, mask_dtype, spacing, mode, min_value, max_value,
39   - decimate_reduction, smooth_relaxation_factor,
40   - smooth_iterations, language, flip_image, q_in, q_out,
41   - from_binary, algorithm, imagedata_resolution):
42   -
43   - multiprocessing.Process.__init__(self)
44   - self.pipe = pipe
45   - self.spacing = spacing
46   - self.filename = filename
47   - self.mode = mode
48   - self.min_value = min_value
49   - self.max_value = max_value
50   - self.decimate_reduction = decimate_reduction
51   - self.smooth_relaxation_factor = smooth_relaxation_factor
52   - self.smooth_iterations = smooth_iterations
53   - self.language = language
54   - self.flip_image = flip_image
55   - self.q_in = q_in
56   - self.q_out = q_out
57   - self.dtype = dtype
58   - self.shape = shape
59   - self.from_binary = from_binary
60   - self.algorithm = algorithm
61   - self.imagedata_resolution = imagedata_resolution
62   -
63   - self.mask_filename = mask_filename
64   - self.mask_shape = mask_shape
65   - self.mask_dtype = mask_dtype
66   -
67   - def run(self):
68   - if self.from_binary:
69   - self.mask = numpy.memmap(self.mask_filename, mode='r',
70   - dtype=self.mask_dtype,
71   - shape=self.mask_shape)
72   - else:
73   - self.image = numpy.memmap(self.filename, mode='r', dtype=self.dtype,
74   - shape=self.shape)
75   - self.mask = numpy.memmap(self.mask_filename, mode='r',
76   - dtype=self.mask_dtype,
77   - shape=self.mask_shape)
78   -
79   - while 1:
80   - roi = self.q_in.get()
81   - if roi is None:
82   - break
83   - self.CreateSurface(roi)
84   -
85   - def SendProgress(self, obj, msg):
86   - prog = obj.GetProgress()
87   - self.pipe.send([prog, msg])
88   -
89   - def CreateSurface(self, roi):
90   - if self.from_binary:
91   - a_mask = numpy.array(self.mask[roi.start + 1: roi.stop + 1,
  43 +
  44 +def create_surface_piece(filename, shape, dtype, mask_filename, mask_shape,
  45 + mask_dtype, roi, spacing, mode, min_value, max_value,
  46 + decimate_reduction, smooth_relaxation_factor,
  47 + smooth_iterations, language, flip_image,
  48 + from_binary, algorithm, imagedata_resolution):
  49 + if from_binary:
  50 + mask = numpy.memmap(mask_filename, mode='r',
  51 + dtype=mask_dtype,
  52 + shape=mask_shape)
  53 + a_mask = numpy.array(mask[roi.start + 1: roi.stop + 1,
  54 + 1:, 1:])
  55 + image = converters.to_vtk(a_mask, spacing, roi.start,
  56 + "AXIAL")
  57 + del a_mask
  58 + else:
  59 + image = numpy.memmap(filename, mode='r', dtype=dtype,
  60 + shape=shape)
  61 + mask = numpy.memmap(mask_filename, mode='r',
  62 + dtype=mask_dtype,
  63 + shape=mask_shape)
  64 + a_image = numpy.array(image[roi])
  65 +
  66 + if algorithm == u'InVesalius 3.b2':
  67 + a_mask = numpy.array(mask[roi.start + 1: roi.stop + 1,
92 68 1:, 1:])
93   - image = converters.to_vtk(a_mask, self.spacing, roi.start,
  69 + a_image[a_mask == 1] = a_image.min() - 1
  70 + a_image[a_mask == 254] = (min_value + max_value) / 2.0
  71 +
  72 + image = converters.to_vtk(a_image, spacing, roi.start,
94 73 "AXIAL")
  74 +
  75 + gauss = vtk.vtkImageGaussianSmooth()
  76 + gauss.SetInputData(image)
  77 + gauss.SetRadiusFactor(0.3)
  78 + gauss.ReleaseDataFlagOn()
  79 + gauss.Update()
  80 +
  81 + del image
  82 + image = gauss.GetOutput()
  83 + del gauss
95 84 del a_mask
96 85 else:
97   - a_image = numpy.array(self.image[roi])
98   -
99   - if self.algorithm == u'InVesalius 3.b2':
100   - a_mask = numpy.array(self.mask[roi.start + 1: roi.stop + 1,
101   - 1:, 1:])
102   - a_image[a_mask == 1] = a_image.min() - 1
103   - a_image[a_mask == 254] = (self.min_value + self.max_value) / 2.0
104   -
105   - image = converters.to_vtk(a_image, self.spacing, roi.start,
106   - "AXIAL")
107   -
108   - gauss = vtk.vtkImageGaussianSmooth()
109   - gauss.SetInputData(image)
110   - gauss.SetRadiusFactor(0.3)
111   - gauss.ReleaseDataFlagOn()
112   - gauss.Update()
113   -
114   - del image
115   - image = gauss.GetOutput()
116   - del gauss
117   - del a_mask
118   - else:
119   - image = converters.to_vtk(a_image, self.spacing, roi.start,
120   - "AXIAL")
121   - del a_image
122   -
123   - if self.imagedata_resolution:
124   - # image = iu.ResampleImage3D(image, self.imagedata_resolution)
125   - image = ResampleImage3D(image, self.imagedata_resolution)
126   -
127   - flip = vtk.vtkImageFlip()
128   - flip.SetInputData(image)
129   - flip.SetFilteredAxis(1)
130   - flip.FlipAboutOriginOn()
131   - flip.ReleaseDataFlagOn()
132   - flip.Update()
133   -
134   - del image
135   - image = flip.GetOutput()
136   - del flip
137   -
138   - #filename = tempfile.mktemp(suffix='_%s.vti' % (self.pid))
139   - #writer = vtk.vtkXMLImageDataWriter()
140   - #writer.SetInput(mask_vtk)
141   - #writer.SetFileName(filename)
142   - #writer.Write()
143   -
144   - #print "Writing piece", roi, "to", filename
145   -
146   - # Create vtkPolyData from vtkImageData
147   - #print "Generating Polydata"
148   - #if self.mode == "CONTOUR":
149   - #print "Contour"
150   - contour = vtk.vtkContourFilter()
151   - contour.SetInputData(image)
152   - #contour.SetInput(flip.GetOutput())
153   - if self.from_binary:
154   - contour.SetValue(0, 127) # initial threshold
155   - else:
156   - contour.SetValue(0, self.min_value) # initial threshold
157   - contour.SetValue(1, self.max_value) # final threshold
158   - contour.ComputeScalarsOn()
159   - contour.ComputeGradientsOn()
160   - contour.ComputeNormalsOn()
161   - contour.ReleaseDataFlagOn()
162   - contour.Update()
163   - #contour.AddObserver("ProgressEvent", lambda obj,evt:
164   - # self.SendProgress(obj, _("Generating 3D surface...")))
165   - polydata = contour.GetOutput()
166   - del image
167   - del contour
168   -
169   - #else: #mode == "GRAYSCALE":
170   - #mcubes = vtk.vtkMarchingCubes()
171   - #mcubes.SetInput(flip.GetOutput())
172   - #mcubes.SetValue(0, self.min_value)
173   - #mcubes.SetValue(1, self.max_value)
174   - #mcubes.ComputeScalarsOff()
175   - #mcubes.ComputeGradientsOff()
176   - #mcubes.ComputeNormalsOff()
177   - #mcubes.AddObserver("ProgressEvent", lambda obj,evt:
178   - #self.SendProgress(obj, _("Generating 3D surface...")))
179   - #polydata = mcubes.GetOutput()
180   -
181   - #triangle = vtk.vtkTriangleFilter()
182   - #triangle.SetInput(polydata)
183   - #triangle.AddObserver("ProgressEvent", lambda obj,evt:
184   - #self.SendProgress(obj, _("Generating 3D surface...")))
185   - #triangle.Update()
186   - #polydata = triangle.GetOutput()
187   -
188   - #if self.decimate_reduction:
189   -
190   - ##print "Decimating"
191   - #decimation = vtk.vtkDecimatePro()
192   - #decimation.SetInput(polydata)
193   - #decimation.SetTargetReduction(0.3)
194   - #decimation.AddObserver("ProgressEvent", lambda obj,evt:
195   - #self.SendProgress(obj, _("Generating 3D surface...")))
196   - ##decimation.PreserveTopologyOn()
197   - #decimation.SplittingOff()
198   - #decimation.BoundaryVertexDeletionOff()
199   - #polydata = decimation.GetOutput()
200   -
201   - self.pipe.send(None)
202   -
203   - filename = tempfile.mktemp(suffix='_%s.vtp' % (self.pid))
204   - writer = vtk.vtkXMLPolyDataWriter()
205   - writer.SetInputData(polydata)
206   - writer.SetFileName(filename)
207   - writer.Write()
208   -
209   - print("Writing piece", roi, "to", filename)
  86 + image = converters.to_vtk(a_image, spacing, roi.start,
  87 + "AXIAL")
  88 + del a_image
  89 +
  90 + if imagedata_resolution:
  91 + image = ResampleImage3D(image, imagedata_resolution)
  92 +
  93 + flip = vtk.vtkImageFlip()
  94 + flip.SetInputData(image)
  95 + flip.SetFilteredAxis(1)
  96 + flip.FlipAboutOriginOn()
  97 + flip.ReleaseDataFlagOn()
  98 + flip.Update()
  99 +
  100 + del image
  101 + image = flip.GetOutput()
  102 + del flip
  103 +
  104 + contour = vtk.vtkContourFilter()
  105 + contour.SetInputData(image)
  106 + if from_binary:
  107 + contour.SetValue(0, 127) # initial threshold
  108 + else:
  109 + contour.SetValue(0, min_value) # initial threshold
  110 + contour.SetValue(1, max_value) # final threshold
  111 + # contour.ComputeScalarsOn()
  112 + # contour.ComputeGradientsOn()
  113 + # contour.ComputeNormalsOn()
  114 + contour.ReleaseDataFlagOn()
  115 + contour.Update()
  116 +
  117 + polydata = contour.GetOutput()
  118 + del image
  119 + del contour
  120 +
  121 + filename = tempfile.mktemp(suffix='_%d_%d.vtp' % (roi.start, roi.stop))
  122 + writer = vtk.vtkXMLPolyDataWriter()
  123 + writer.SetInputData(polydata)
  124 + writer.SetFileName(filename)
  125 + writer.Write()
  126 +
  127 + print("Writing piece", roi, "to", filename)
  128 + print("MY PID MC", os.getpid())
  129 + return filename
  130 +
  131 +
  132 +def join_process_surface(filenames, algorithm, smooth_iterations, smooth_relaxation_factor, decimate_reduction, keep_largest, fill_holes, options, msg_queue):
  133 + def send_message(msg):
  134 + try:
  135 + msg_queue.put_nowait(msg)
  136 + except queue.Full as e:
  137 + print(e)
  138 +
  139 + send_message('Joining surfaces ...')
  140 + polydata_append = vtk.vtkAppendPolyData()
  141 + for f in filenames:
  142 + reader = vtk.vtkXMLPolyDataReader()
  143 + reader.SetFileName(f)
  144 + reader.Update()
  145 +
  146 + polydata = reader.GetOutput()
  147 +
  148 + polydata_append.AddInputData(polydata)
  149 + del reader
210 150 del polydata
211   - del writer
212 151  
213   - self.q_out.put(filename)
  152 + polydata_append.Update()
  153 + # polydata_append.GetOutput().ReleaseDataFlagOn()
  154 + polydata = polydata_append.GetOutput()
  155 + #polydata.Register(None)
  156 + # polydata.SetSource(None)
  157 + del polydata_append
  158 +
  159 + send_message('Cleaning surface ...')
  160 + clean = vtk.vtkCleanPolyData()
  161 + # clean.ReleaseDataFlagOn()
  162 + # clean.GetOutput().ReleaseDataFlagOn()
  163 + clean_ref = weakref.ref(clean)
  164 + # clean_ref().AddObserver("ProgressEvent", lambda obj,evt:
  165 + # UpdateProgress(clean_ref(), _("Creating 3D surface...")))
  166 + clean.SetInputData(polydata)
  167 + clean.PointMergingOn()
  168 + clean.Update()
  169 +
  170 + del polydata
  171 + polydata = clean.GetOutput()
  172 + # polydata.SetSource(None)
  173 + del clean
  174 +
  175 + if algorithm == 'ca_smoothing':
  176 + send_message('Calculating normals ...')
  177 + normals = vtk.vtkPolyDataNormals()
  178 + normals_ref = weakref.ref(normals)
  179 + # normals_ref().AddObserver("ProgressEvent", lambda obj,evt:
  180 + # UpdateProgress(normals_ref(), _("Creating 3D surface...")))
  181 + normals.SetInputData(polydata)
  182 + # normals.ReleaseDataFlagOn()
  183 + #normals.SetFeatureAngle(80)
  184 + #normals.AutoOrientNormalsOn()
  185 + normals.ComputeCellNormalsOn()
  186 + # normals.GetOutput().ReleaseDataFlagOn()
  187 + normals.Update()
  188 + del polydata
  189 + polydata = normals.GetOutput()
  190 + # polydata.SetSource(None)
  191 + del normals
  192 +
  193 + clean = vtk.vtkCleanPolyData()
  194 + # clean.ReleaseDataFlagOn()
  195 + # clean.GetOutput().ReleaseDataFlagOn()
  196 + clean_ref = weakref.ref(clean)
  197 + # clean_ref().AddObserver("ProgressEvent", lambda obj,evt:
  198 + # UpdateProgress(clean_ref(), _("Creating 3D surface...")))
  199 + clean.SetInputData(polydata)
  200 + clean.PointMergingOn()
  201 + clean.Update()
  202 +
  203 + del polydata
  204 + polydata = clean.GetOutput()
  205 + # polydata.SetSource(None)
  206 + del clean
  207 +
  208 + # try:
  209 + # polydata.BuildLinks()
  210 + # except TypeError:
  211 + # polydata.BuildLinks(0)
  212 + # polydata = ca_smoothing.ca_smoothing(polydata, options['angle'],
  213 + # options['max distance'],
  214 + # options['min weight'],
  215 + # options['steps'])
  216 +
  217 + send_message('Context Aware smoothing ...')
  218 + mesh = cy_mesh.Mesh(polydata)
  219 + cy_mesh.ca_smoothing(mesh, options['angle'],
  220 + options['max distance'],
  221 + options['min weight'],
  222 + options['steps'])
  223 + # polydata = mesh.to_vtk()
  224 +
  225 + # polydata.SetSource(None)
  226 + # polydata.DebugOn()
  227 + else:
  228 + #smoother = vtk.vtkWindowedSincPolyDataFilter()
  229 + send_message('Smoothing ...')
  230 + smoother = vtk.vtkSmoothPolyDataFilter()
  231 + smoother_ref = weakref.ref(smoother)
  232 + # smoother_ref().AddObserver("ProgressEvent", lambda obj,evt:
  233 + # UpdateProgress(smoother_ref(), _("Creating 3D surface...")))
  234 + smoother.SetInputData(polydata)
  235 + smoother.SetNumberOfIterations(smooth_iterations)
  236 + smoother.SetRelaxationFactor(smooth_relaxation_factor)
  237 + smoother.SetFeatureAngle(80)
  238 + #smoother.SetEdgeAngle(90.0)
  239 + #smoother.SetPassBand(0.1)
  240 + smoother.BoundarySmoothingOn()
  241 + smoother.FeatureEdgeSmoothingOn()
  242 + #smoother.NormalizeCoordinatesOn()
  243 + #smoother.NonManifoldSmoothingOn()
  244 + # smoother.ReleaseDataFlagOn()
  245 + # smoother.GetOutput().ReleaseDataFlagOn()
  246 + smoother.Update()
  247 + del polydata
  248 + polydata = smoother.GetOutput()
  249 + #polydata.Register(None)
  250 + # polydata.SetSource(None)
  251 + del smoother
  252 +
  253 +
  254 + if decimate_reduction:
  255 + print("Decimating", decimate_reduction)
  256 + send_message('Decimating ...')
  257 + decimation = vtk.vtkQuadricDecimation()
  258 + # decimation.ReleaseDataFlagOn()
  259 + decimation.SetInputData(polydata)
  260 + decimation.SetTargetReduction(decimate_reduction)
  261 + decimation_ref = weakref.ref(decimation)
  262 + # decimation_ref().AddObserver("ProgressEvent", lambda obj,evt:
  263 + # UpdateProgress(decimation_ref(), _("Creating 3D surface...")))
  264 + #decimation.PreserveTopologyOn()
  265 + #decimation.SplittingOff()
  266 + #decimation.BoundaryVertexDeletionOff()
  267 + # decimation.GetOutput().ReleaseDataFlagOn()
  268 + decimation.Update()
  269 + del polydata
  270 + polydata = decimation.GetOutput()
  271 + #polydata.Register(None)
  272 + # polydata.SetSource(None)
  273 + del decimation
  274 +
  275 + #to_measure.Register(None)
  276 + # to_measure.SetSource(None)
  277 +
  278 + if keep_largest:
  279 + send_message('Finding the largest ...')
  280 + conn = vtk.vtkPolyDataConnectivityFilter()
  281 + conn.SetInputData(polydata)
  282 + conn.SetExtractionModeToLargestRegion()
  283 + conn_ref = weakref.ref(conn)
  284 + # conn_ref().AddObserver("ProgressEvent", lambda obj,evt:
  285 + # UpdateProgress(conn_ref(), _("Creating 3D surface...")))
  286 + conn.Update()
  287 + # conn.GetOutput().ReleaseDataFlagOn()
  288 + del polydata
  289 + polydata = conn.GetOutput()
  290 + #polydata.Register(None)
  291 + # polydata.SetSource(None)
  292 + del conn
  293 +
  294 + #Filter used to detect and fill holes. Only fill boundary edges holes.
  295 + #TODO: Hey! This piece of code is the same from
  296 + #polydata_utils.FillSurfaceHole, we need to review this.
  297 + if fill_holes:
  298 + send_message('Filling holes ...')
  299 + filled_polydata = vtk.vtkFillHolesFilter()
  300 + # filled_polydata.ReleaseDataFlagOn()
  301 + filled_polydata.SetInputData(polydata)
  302 + filled_polydata.SetHoleSize(300)
  303 + filled_polydata_ref = weakref.ref(filled_polydata)
  304 + # filled_polydata_ref().AddObserver("ProgressEvent", lambda obj,evt:
  305 + # UpdateProgress(filled_polydata_ref(), _("Creating 3D surface...")))
  306 + filled_polydata.Update()
  307 + # filled_polydata.GetOutput().ReleaseDataFlagOn()
  308 + del polydata
  309 + polydata = filled_polydata.GetOutput()
  310 + #polydata.Register(None)
  311 + # polydata.SetSource(None)
  312 + # polydata.DebugOn()
  313 + del filled_polydata
  314 +
  315 + to_measure = polydata
  316 +
  317 + normals = vtk.vtkPolyDataNormals()
  318 + # normals.ReleaseDataFlagOn()
  319 + # normals_ref = weakref.ref(normals)
  320 + # normals_ref().AddObserver("ProgressEvent", lambda obj,evt:
  321 + # UpdateProgress(normals_ref(), _("Creating 3D surface...")))
  322 + normals.SetInputData(polydata)
  323 + # normals.SetFeatureAngle(80)
  324 + # normals.SplittingOff()
  325 + # normals.AutoOrientNormalsOn()
  326 + # normals.GetOutput().ReleaseDataFlagOn()
  327 + normals.Update()
  328 + del polydata
  329 + polydata = normals.GetOutput()
  330 + #polydata.Register(None)
  331 + # polydata.SetSource(None)
  332 + del normals
  333 +
  334 +
  335 + # Improve performance
  336 + stripper = vtk.vtkStripper()
  337 + # stripper.ReleaseDataFlagOn()
  338 + # stripper_ref = weakref.ref(stripper)
  339 + # stripper_ref().AddObserver("ProgressEvent", lambda obj,evt:
  340 + # UpdateProgress(stripper_ref(), _("Creating 3D surface...")))
  341 + stripper.SetInputData(polydata)
  342 + stripper.PassThroughCellIdsOn()
  343 + stripper.PassThroughPointIdsOn()
  344 + # stripper.GetOutput().ReleaseDataFlagOn()
  345 + stripper.Update()
  346 + del polydata
  347 + polydata = stripper.GetOutput()
  348 + #polydata.Register(None)
  349 + # polydata.SetSource(None)
  350 + del stripper
  351 +
  352 + send_message('Calculating area and volume ...')
  353 + measured_polydata = vtk.vtkMassProperties()
  354 + measured_polydata.SetInputData(to_measure)
  355 + measured_polydata.Update()
  356 + volume = float(measured_polydata.GetVolume())
  357 + area = float(measured_polydata.GetSurfaceArea())
  358 + del measured_polydata
  359 +
  360 + filename = tempfile.mktemp(suffix='_full.vtp')
  361 + writer = vtk.vtkXMLPolyDataWriter()
  362 + writer.SetInputData(polydata)
  363 + writer.SetFileName(filename)
  364 + writer.Write()
  365 + del writer
  366 +
  367 + print("MY PID", os.getpid())
  368 + return filename, {'volume': volume, 'area': area}
... ...
invesalius/gui/dialogs.py
  Diff comments   View file @c602646
... ... @@ -3523,3 +3523,30 @@ class ObjectCalibrationDialog(wx.Dialog):
3523 3523  
3524 3524 def GetValue(self):
3525 3525 return self.obj_fiducials, self.obj_orients, self.obj_ref_id, self.obj_name
  3526 +
  3527 +
  3528 +class SurfaceProgressWindow(object):
  3529 + def __init__(self):
  3530 + self.title = "InVesalius 3"
  3531 + self.msg = _("Creating 3D surface ...")
  3532 + self.style = wx.PD_APP_MODAL | wx.PD_APP_MODAL | wx.PD_CAN_ABORT | wx.PD_ELAPSED_TIME
  3533 + self.dlg = wx.ProgressDialog(self.title,
  3534 + self.msg,
  3535 + parent=None,
  3536 + style=self.style)
  3537 + self.running = True
  3538 + self.error = None
  3539 + self.dlg.Show()
  3540 +
  3541 + def WasCancelled(self):
  3542 + # print("Cancelled?", self.dlg.WasCancelled())
  3543 + return self.dlg.WasCancelled()
  3544 +
  3545 + def Update(self, msg=None, value=None):
  3546 + if msg is None:
  3547 + self.dlg.Pulse()
  3548 + else:
  3549 + self.dlg.Pulse(msg)
  3550 +
  3551 + def Close(self):
  3552 + self.dlg.Destroy()
... ...
invesalius/utils.py
  Diff comments   View file @c602646
... ... @@ -22,8 +22,10 @@ import sys
22 22 import re
23 23 import locale
24 24 import math
  25 +import traceback
25 26  
26 27 from distutils.version import LooseVersion
  28 +from functools import wraps
27 29  
28 30 import numpy as np
29 31  
... ... @@ -462,3 +464,24 @@ def encode(text, encoding, *args):
462 464 return text.encode(encoding, *args)
463 465 except AttributeError:
464 466 return text
  467 +
  468 +
  469 +def timing(f):
  470 + @wraps(f)
  471 + def wrapper(*args, **kwargs):
  472 + start = time.time()
  473 + result = f(*args, **kwargs)
  474 + end = time.time()
  475 + print('{} elapsed time: {}'.format(f.__name__, end-start))
  476 + return result
  477 + return wrapper
  478 +
  479 +
  480 +def log_traceback(ex):
  481 + if hasattr(ex, '__traceback__'):
  482 + ex_traceback = ex.__traceback__
  483 + else:
  484 + _, _, ex_traceback = sys.exc_info()
  485 + tb_lines = [line.rstrip('\n') for line in
  486 + traceback.format_exception(ex.__class__, ex, ex_traceback)]
  487 + return ''.join(tb_lines)
... ...