ENH: Generating surface using multiprocessing and with memmap

tfmoraes
1 parent a3614eda
Showing 4 changed files with 156 additions and 147 deletions Show diff stats
invesalius/data/imagedata_utils.py
invesalius/data/slice_.py
invesalius/data/surface.py
invesalius/data/surface_process.py
@@ -497,20 +497,23 @@ def dcm2memmap(files, slice_size, orientation):
     return matrix, temp_file
 def to_vtk(n_array, spacing, slice_number, orientation):
-    dy, dx = n_array.shape
+    try:
+        dz, dy, dx = n_array.shape
+    except ValueError:
+        dy, dx = n_array.shape
+        dz = 1
     v_image = numpy_support.numpy_to_vtk(n_array.flat)
     print orientation
     if orientation == 'AXIAL':
-        dz = 1
-        extent = (0, dx -1, 0, dy -1, slice_number, slice_number)
+        extent = (0, dx -1, 0, dy -1, slice_number, slice_number + dz - 1)
     elif orientation == 'SAGITAL':
-        dx, dy, dz = 1, dx, dy
-        extent = (slice_number, slice_number, 0, dy - 1, 0, dz - 1)
+        dx, dy, dz = dz, dx, dy
+        extent = (slice_number, slice_number + dx - 1, 0, dy - 1, 0, dz - 1)
     elif orientation == 'CORONAL':
-        dx, dy, dz = dx, 1, dy
-        extent = (0, dx - 1, slice_number, slice_number, 0, dz - 1)
+        dx, dy, dz = dx, dz, dy
+        extent = (0, dx - 1, slice_number, slice_number + dy - 1, 0, dz - 1)
     # Generating the vtkImageData
     image = vtk.vtkImageData()
@@ -260,8 +260,8 @@ class Slice(object):
     #---------------------------------------------------------------------------
     def GetSlices(self, orientation, slice_number):
-        print "Getting slice", self.buffer_slices[orientation][0], slice_number
         if self.buffer_slices[orientation][0] == slice_number:
+            print "From buffer"
             image = self.buffer_slices[orientation][1]
             if self.current_mask and self.current_mask.is_shown:
                 n_mask = self.buffer_slices[orientation][2]
@@ -275,7 +275,6 @@ class Slice(object):
             image = self.do_ww_wl(image)
             if self.current_mask and self.current_mask.is_shown:
-                print "Mask"
                 n_mask = self.GetMaskSlice(orientation, slice_number)
                 mask = iu.to_vtk(n_mask, self.spacing, slice_number, orientation)
                 final_image = self.do_blend(image, self.do_colour_mask(mask))
@@ -411,12 +410,12 @@ class Slice(object):
                     m[slice_ < thresh_min] = 0
                     m[slice_ > thresh_max] = 0
                     m[m == 1] = 255
-                    self.current_mask.matrix[n] = m
+                    self.current_mask.matrix[n+1, 1:, 1:] = m
             else:
                 print "Only one slice"
                 slice_ = self.buffer_slices[orientation][3]
-                self.buffer_slices[orientation][2][:] = 0
-                self.buffer_slices[orientation][2][numpy.logical_and(slice_ >= thresh_min,slice_ <= thresh_max)] = 255
+                self.buffer_slices[orientation][2] = 255 * ((slice_ >= thresh_min) & (slice_ <= thresh_max))
+                print self.buffer_slices[orientation][2].dtype
             # Update viewer
             #ps.Publisher().sendMessage('Update slice viewer')
@@ -521,15 +520,15 @@ class Slice(object):
         # This is very important. Do not use masks' imagedata. It would mess up
         # surface quality event when using contour
-        imagedata = self.imagedata
-
         colour = mask.colour
         threshold = mask.threshold_range
         edited_points = mask.edited_points
-        ps.Publisher().sendMessage('Create surface',
-                                   (imagedata,colour,threshold,
-                                    edited_points, overwrite_surface))
+        self.SetMaskThreshold(mask.index, threshold)
+
+        mask.matrix.flush()
+
+        ps.Publisher().sendMessage('Create surface', (mask, self.spacing))
     def GetOutput(self):
         return self.blend_filter.GetOutput()
@@ -376,29 +376,36 @@ class SurfaceManager():
         Create surface actor, save into project and send it to viewer.
         """
         surface_data = pubsub_evt.data
-
-        if len(surface_data) == 5:
-            imagedata, colour, [min_value, max_value], \
-            edited_points, overwrite = pubsub_evt.data
-            quality=_('Optimal *')
-            surface_name = ""
-            fill_holes = True
-            keep_largest = False
-        else:
-            imagedata, colour, [min_value, max_value],\
-            edited_points, overwrite, surface_name,\
-            quality, fill_holes, keep_largest =\
-            pubsub_evt.data
+        mask, spacing = pubsub_evt.data
+        min_value, max_value = mask.threshold_range
+        fill_holes = True
+
+        #if len(surface_data) == 5:
+            #imagedata, colour, [min_value, max_value], \
+            #edited_points, overwrite = pubsub_evt.data
+            #quality=_('Optimal *')
+            #surface_name = ""
+            #fill_holes = True
+            #keep_largest = False
+        #else:
+            #imagedata, colour, [min_value, max_value],\
+            #edited_points, overwrite, surface_name,\
+            #quality, fill_holes, keep_largest =\
+            #pubsub_evt.data
         mode = 'CONTOUR' # 'GRAYSCALE'
-
-        ps.Publisher().sendMessage('Begin busy cursor')
-        imagedata_tmp = None
-        if (edited_points):
-            imagedata_tmp = vtk.vtkImageData()
-            imagedata_tmp.DeepCopy(imagedata)
-            imagedata_tmp.Update()
-            imagedata = iu.BuildEditedImage(imagedata_tmp, edited_points)
+        quality=_('Optimal *')
+        keep_largest = True
+        surface_name = ""
+        colour = mask.colour
+
+        #ps.Publisher().sendMessage('Begin busy cursor')
+        #imagedata_tmp = None
+        #if (edited_points):
+            #imagedata_tmp = vtk.vtkImageData()
+            #imagedata_tmp.DeepCopy(imagedata)
+            #imagedata_tmp.Update()
+            #imagedata = iu.BuildEditedImage(imagedata_tmp, edited_points)
         if quality in const.SURFACE_QUALITY.keys():
             imagedata_resolution = const.SURFACE_QUALITY[quality][0]
@@ -406,61 +413,89 @@ class SurfaceManager():
             smooth_relaxation_factor = const.SURFACE_QUALITY[quality][2]
             decimate_reduction = const.SURFACE_QUALITY[quality][3]
-        if imagedata_resolution:
-            imagedata = iu.ResampleImage3D(imagedata, imagedata_resolution)
+        #if imagedata_resolution:
+            #imagedata = iu.ResampleImage3D(imagedata, imagedata_resolution)
-        pipeline_size = 4
-        if decimate_reduction:
-            pipeline_size += 1
-        if (smooth_iterations and smooth_relaxation_factor):
-            pipeline_size += 1
-        if fill_holes:
-            pipeline_size += 1
-        if keep_largest:
-            pipeline_size += 1
+        #pipeline_size = 4
+        #if decimate_reduction:
+            #pipeline_size += 1
+        #if (smooth_iterations and smooth_relaxation_factor):
+            #pipeline_size += 1
+        #if fill_holes:
+            #pipeline_size += 1
+        #if keep_largest:
+            #pipeline_size += 1
-        # Update progress value in GUI
-        UpdateProgress = vu.ShowProgress(pipeline_size)
-        UpdateProgress(0, _("Generating 3D surface..."))
+        ## Update progress value in GUI
+        #UpdateProgress = vu.ShowProgress(pipeline_size)
+        #UpdateProgress(0, _("Generating 3D surface..."))
-        filename_img = tempfile.mktemp()
+        #filename_img = tempfile.mktemp()
-        writer = vtk.vtkXMLImageDataWriter()
-        writer.SetFileName(filename_img)
-        writer.SetInput(imagedata)
-        writer.Write()
+        #writer = vtk.vtkXMLImageDataWriter()
+        #writer.SetFileName(filename_img)
+        #writer.SetInput(imagedata)
+        #writer.Write()
         language = ses.Session().language
+
+        filename_img = mask.temp_file
+        overwrite = 0
         if (prj.Project().original_orientation == const.CORONAL):
             flip_image = False
         else:
             flip_image = True
+
+        n_processors = 1 # multiprocessing.cpu_count()
         pipe_in, pipe_out = multiprocessing.Pipe()
-        sp = surface_process.SurfaceProcess(pipe_in, filename_img, mode, min_value, max_value,
-                 decimate_reduction, smooth_relaxation_factor,
-                 smooth_iterations, language, fill_holes, keep_largest, flip_image)
-        sp.start()
-
-        while 1:
-            msg = pipe_out.recv()
-            if(msg is None):
-                break
-            UpdateProgress(msg[0],msg[1])
-
-        filename_polydata = pipe_out.recv()
-
-        reader = vtk.vtkXMLPolyDataReader()
-        reader.SetFileName(filename_polydata)
-        reader.Update()
-
-        polydata = reader.GetOutput()
+        o_piece = 2
+        piece_size = 40
+
+        n_pieces = round(mask.matrix.shape[0] / piece_size + 0.5, 0)
+        print "n_pieces", n_pieces, mask.matrix.shape
+
+        q_in = multiprocessing.Queue()
+        q_out = multiprocessing.Queue()
+
+        p = []
+        for i in xrange(n_processors):
+            sp = surface_process.SurfaceProcess(pipe_in, filename_img,
+                    mask.matrix.shape, mask.matrix.dtype, spacing, 
+                    mode, min_value, max_value,
+                    decimate_reduction, smooth_relaxation_factor,
+                    smooth_iterations, language, fill_holes, keep_largest,
+                                                flip_image, q_in, q_out)
+            p.append(sp)
+            sp.start()
+
+        for i in xrange(n_pieces):
+            init = i * piece_size
+            end = init + piece_size + o_piece
+            roi = slice(init, end)
+            q_in.put(roi)
+            print "new_piece", roi
+
+        for i in p:
+            q_in.put(None)
+
+        polydata_append = vtk.vtkAppendPolyData()
+        t = n_pieces
+        while t:
+            filename_polydata = q_out.get()
+
+            reader = vtk.vtkXMLPolyDataReader()
+            reader.SetFileName(filename_polydata)
+            reader.Update()
+            polydata_append.AddInput(reader.GetOutput())
+
+            t -= 1
+
+        polydata = polydata_append.GetOutput()
         # Orient normals from inside to outside
         normals = vtk.vtkPolyDataNormals()
-        normals.AddObserver("ProgressEvent", lambda obj,evt:
-                            UpdateProgress(obj, _("Generating 3D surface...")))
         normals.SetInput(polydata)
         normals.SetFeatureAngle(80)
         normals.AutoOrientNormalsOn()
@@ -468,8 +503,6 @@ class SurfaceManager():
 	# Improve performance
         stripper = vtk.vtkStripper()
-        stripper.AddObserver("ProgressEvent", lambda obj,evt:
-                            UpdateProgress(obj, _("Generating 3D surface...")))
         stripper.SetInput(normals.GetOutput())
         stripper.PassThroughCellIdsOn()
         stripper.PassThroughPointIdsOn()
@@ -526,8 +559,6 @@ class SurfaceManager():
         # The following lines have to be here, otherwise all volumes disappear
         measured_polydata = vtk.vtkMassProperties()
-        measured_polydata.AddObserver("ProgressEvent", lambda obj,evt:
-                            UpdateProgress(obj, _("Generating 3D surface...")))
         measured_polydata.SetInput(polydata)
         volume =  measured_polydata.GetVolume()
         surface.volume = volume
@@ -552,10 +583,6 @@ class SurfaceManager():
                                     surface.colour, surface.volume,
                                     surface.transparency))
-        #Destroy Copy original imagedata
-        if(imagedata_tmp):
-            del imagedata_tmp
-
         ps.Publisher().sendMessage('End busy cursor')
     def RemoveActor(self, index):
-import vtk
 import multiprocessing
 import tempfile
+import time
+
+import numpy
+import vtk
 import i18n
+import imagedata_utils
+
+from scipy import ndimage
 class SurfaceProcess(multiprocessing.Process):
-    def __init__(self, pipe, filename, mode, min_value, max_value,
+    def __init__(self, pipe, filename, shape, dtype, spacing, mode, min_value, max_value,
                  decimate_reduction, smooth_relaxation_factor,
                  smooth_iterations, language,  fill_holes, keep_largest, 
-                 flip_image):
+                 flip_image, q_in, q_out):
         multiprocessing.Process.__init__(self)
         self.pipe = pipe
+        self.spacing = spacing
         self.filename = filename
         self.mode = mode
         self.min_value = min_value
@@ -24,64 +31,59 @@ class SurfaceProcess(multiprocessing.Process):
         self.fill_holes = fill_holes
         self.keep_largest = keep_largest
         self.flip_image = flip_image
+        self.q_in = q_in
+        self.q_out = q_out
+        self.mask = numpy.memmap(filename, mode='r', dtype=dtype,
+                                 shape=shape)
     def run(self):
-        self.CreateSurface()
+        while 1:
+            roi = self.q_in.get()
+            print roi
+            if roi is None:
+                break
+            self.CreateSurface(roi)
     def SendProgress(self, obj, msg):
         prog = obj.GetProgress()
         self.pipe.send([prog, msg])
-    def CreateSurface(self):
-        _ = i18n.InstallLanguage(self.language)
-
-        reader = vtk.vtkXMLImageDataReader()
-        reader.SetFileName(self.filename)
-        reader.Update()
-        
-        image = reader.GetOutput()
-        
-        if (self.flip_image):
-            # Flip original vtkImageData
-            flip = vtk.vtkImageFlip()
-            flip.SetInput(reader.GetOutput())
-            flip.SetFilteredAxis(1)
-            flip.FlipAboutOriginOn()
-            image = flip.GetOutput()
-            
+    def CreateSurface(self, roi):
+        smoothed = ndimage.gaussian_filter(self.mask[roi], (1, 1, 1))
+        image = imagedata_utils.to_vtk(smoothed, self.spacing, roi.start,
+                                       "AXIAL")
+
         # Create vtkPolyData from vtkImageData
+        print "Generating Polydata"
         if self.mode == "CONTOUR":
+            print "Contour"
             contour = vtk.vtkContourFilter()
             contour.SetInput(image)
-            contour.SetValue(0, self.min_value) # initial threshold
-            contour.SetValue(1, self.max_value) # final threshold
-            contour.GetOutput().ReleaseDataFlagOn()
-            contour.AddObserver("ProgressEvent", lambda obj,evt:
-                    self.SendProgress(obj, _("Generating 3D surface...")))
+            contour.SetValue(0, 127.5) # initial threshold
+            contour.ComputeScalarsOn()
+            contour.ComputeGradientsOn()
+            contour.ComputeNormalsOn()
             polydata = contour.GetOutput()
         else: #mode == "GRAYSCALE":
             mcubes = vtk.vtkMarchingCubes()
             mcubes.SetInput(image)
-            mcubes.SetValue(0, 255)
+            mcubes.SetValue(0, 127.5)
             mcubes.ComputeScalarsOn()
             mcubes.ComputeGradientsOn()
             mcubes.ComputeNormalsOn()
-            mcubes.ThresholdBetween(self.min_value, self.max_value)
-            mcubes.GetOutput().ReleaseDataFlagOn()
-            mcubes.AddObserver("ProgressEvent", lambda obj,evt:
-                    self.SendProgress(obj, _("Generating 3D surface...")))
             polydata = mcubes.GetOutput()
+        print "Decimating"
         if self.decimate_reduction:
-            decimation = vtk.vtkQuadricDecimation()
+            decimation = vtk.vtkDecimatePro()
             decimation.SetInput(polydata)
             decimation.SetTargetReduction(self.decimate_reduction)
-            decimation.GetOutput().ReleaseDataFlagOn()
-            decimation.AddObserver("ProgressEvent", lambda obj,evt:
-                    self.SendProgress(obj, _("Generating 3D surface...")))
+            decimation.PreserveTopologyOn()
+            decimation.SplittingOff()
             polydata = decimation.GetOutput()
+        print "Smoothing"
         if self.smooth_iterations and self.smooth_relaxation_factor:
             smoother = vtk.vtkSmoothPolyDataFilter()
             smoother.SetInput(polydata)
@@ -90,38 +92,16 @@ class SurfaceProcess(multiprocessing.Process):
             smoother.SetRelaxationFactor(self.smooth_relaxation_factor)
             smoother.FeatureEdgeSmoothingOn()
             smoother.BoundarySmoothingOn()
-            smoother.GetOutput().ReleaseDataFlagOn()
-            smoother.AddObserver("ProgressEvent", lambda obj,evt:
-                    self.SendProgress(obj, _("Generating 3D surface...")))
             polydata = smoother.GetOutput()
-
-        if self.keep_largest:
-            conn = vtk.vtkPolyDataConnectivityFilter()
-            conn.SetInput(polydata)
-            conn.SetExtractionModeToLargestRegion()
-            conn.AddObserver("ProgressEvent", lambda obj,evt:
-                    self.SendProgress(obj, _("Generating 3D surface...")))
-            polydata = conn.GetOutput()
-
-        # Filter used to detect and fill holes. Only fill boundary edges holes.
-        #TODO: Hey! This piece of code is the same from
-        # polydata_utils.FillSurfaceHole, we need to review this.
-        if self.fill_holes:
-            filled_polydata = vtk.vtkFillHolesFilter()
-            filled_polydata.SetInput(polydata)
-            filled_polydata.SetHoleSize(300)
-            filled_polydata.AddObserver("ProgressEvent", lambda obj,evt:
-                    self.SendProgress(obj, _("Generating 3D surface...")))
-            polydata = filled_polydata.GetOutput()
-
-
-
+        print "Saving"
         filename = tempfile.mktemp()
         writer = vtk.vtkXMLPolyDataWriter()
         writer.SetInput(polydata)
         writer.SetFileName(filename)
         writer.Write()
+        print filename
+
+        time.sleep(1)
-        self.pipe.send(None)
-        self.pipe.send(filename)
+        self.q_out.put(filename)