refactored vtk/tvtk use and ETSConfigToolkit

Author: oesteban

nipype/algorithms/mesh.py

@@ -14,7 +14,6 @@
 from __future__ import division
 from builtins import zip
 
-import os
 import os.path as op
 from warnings import warn
 
@@ -25,52 +24,10 @@
 from ..external.six import string_types
 from ..interfaces.base import (BaseInterface, traits, TraitedSpec, File,
                                BaseInterfaceInputSpec)
-
+from ..interfaces.vtkbase import tvtk
+from ..interfaces import vtkbase as VTKInfo
 iflogger = logging.getLogger('interface')
 
-# Ensure that tvtk is loaded with the appropriate ETS_TOOLKIT env var
-old_ets = os.getenv('ETS_TOOLKIT')
-os.environ['ETS_TOOLKIT'] = 'null'
-have_tvtk = False
-try:
-    from tvtk.api import tvtk
-    have_tvtk = True
-except ImportError:
-    iflogger.warning('tvtk wasn\'t found')
-finally:
-    if old_ets is not None:
-        os.environ['ETS_TOOLKIT'] = old_ets
-    else:
-        del os.environ['ETS_TOOLKIT']
-
-
-class Info(object):
-    """ Handle VTK version information """
-    _vtk_version = None
-
-    @staticmethod
-    def vtk_version():
-        """ Get VTK version """
-        if not Info.no_tvtk():
-            return None
-
-        if Info._vtk_version is None:
-            try:
-                from tvtk.tvtk_classes.vtk_version import vtk_build_version
-                vsplits = vtk_build_version.split('.')
-                Info._vtk_version = tuple([int(vsplits[0]), int(vsplits[1])] + vsplits[2:])
-            except ImportError:
-                iflogger.warning(
-                    'VTK version-major inspection using tvtk failed, assuming VTK == 4.0.')
-                Info._vtk_version = (4, 0)
-
-        return Info._vtk_version
-
-    @staticmethod
-    def no_tvtk():
-        global have_tvtk
-        return not have_tvtk
-
 
 class TVTKBaseInterface(BaseInterface):
 
@@ -79,13 +36,10 @@ class TVTKBaseInterface(BaseInterface):
     _redirect_x = True
 
     def __init__(self, **inputs):
-        if Info.no_tvtk():
+        if VTKInfo.no_tvtk():
             raise ImportError('This interface requires tvtk to run.')
         super(TVTKBaseInterface, self).__init__(**inputs)
 
-    def vtk_version(self):
-        return Info.vtk_version()
-
 
 class WarpPointsInputSpec(BaseInterfaceInputSpec):
     points = File(exists=True, mandatory=True,
@@ -144,17 +98,15 @@ def _run_interface(self, runtime):
         import nibabel as nb
         import numpy as np
         from scipy import ndimage
-        from tvtk.common import configure_input_data
-        from tvtk.common import is_old_pipeline as vtk_old
 
         r = tvtk.PolyDataReader(file_name=self.inputs.points)
         r.update()
-        mesh = r.output if vtk_old() else r.get_output()
+        mesh = VTKInfo.vtk_output(r)
         points = np.array(mesh.points)
         warp_dims = nb.funcs.four_to_three(nb.load(self.inputs.warp))
 
         affine = warp_dims[0].affine
-        voxsize = warp_dims[0].header.get_zooms()
+        # voxsize = warp_dims[0].header.get_zooms()
         vox2ras = affine[0:3, 0:3]
         ras2vox = np.linalg.inv(vox2ras)
         origin = affine[0:3, 3]
@@ -176,7 +128,7 @@ def _run_interface(self, runtime):
         newpoints = [p + d for p, d in zip(points, disps)]
         mesh.points = newpoints
         w = tvtk.PolyDataWriter()
-        configure_input_data(w, mesh)
+        VTKInfo.configure_input_data(w, mesh)
         w.file_name = self._gen_fname(self.inputs.points, suffix='warped', ext='.vtk')
         w.write()
         return runtime
@@ -253,13 +205,10 @@ def _triangle_area(self, A, B, C):
         return area
 
     def _run_interface(self, runtime):
-        from tvtk.common import configure_input_data
-        from tvtk.common import is_old_pipeline as vtk_old
-
         r1 = tvtk.PolyDataReader(file_name=self.inputs.surface1)
         r2 = tvtk.PolyDataReader(file_name=self.inputs.surface2)
-        vtk1 = r1.output if vtk_old() else r1.get_output()
-        vtk2 = r2.output if vtk_old() else r2.get_output()
+        vtk1 = VTKInfo.vtk_output(r1)
+        vtk2 = VTKInfo.vtk_output(r2)
         r1.update()
         r2.update()
         assert(len(vtk1.points) == len(vtk2.points))
@@ -303,7 +252,7 @@ def _run_interface(self, runtime):
         out_mesh.point_data.vectors.name = 'warpings'
         writer = tvtk.PolyDataWriter(
             file_name=op.abspath(self.inputs.out_warp))
-        configure_input_data(writer, out_mesh)
+        VTKInfo.configure_input_data(writer, out_mesh)
         writer.write()
 
         self._distance = np.average(errvector, weights=weights)
@@ -372,11 +321,8 @@ class MeshWarpMaths(TVTKBaseInterface):
     output_spec = MeshWarpMathsOutputSpec
 
     def _run_interface(self, runtime):
-        from tvtk.common import configure_input_data
-        from tvtk.common import is_old_pipeline as vtk_old
-
         r1 = tvtk.PolyDataReader(file_name=self.inputs.in_surf)
-        vtk1 = r1.output if vtk_old() else r1.get_output()
+        vtk1 = VTKInfo.vtk_output(r1)
         r1.update()
         points1 = np.array(vtk1.points)
 
@@ -388,7 +334,7 @@ def _run_interface(self, runtime):
 
         if isinstance(operator, string_types):
             r2 = tvtk.PolyDataReader(file_name=self.inputs.surface2)
-            vtk2 = r2.output if vtk_old() else r2.get_output()
+            vtk2 = VTKInfo.vtk_output(r2)
             r2.update()
             assert(len(points1) == len(vtk2.points))
 
@@ -399,7 +345,7 @@ def _run_interface(self, runtime):
 
             if opfield is None:
                 raise RuntimeError(
-                    ('No operator values found in operator file'))
+                    'No operator values found in operator file')
 
             opfield = np.array(opfield)
 
@@ -422,13 +368,13 @@ def _run_interface(self, runtime):
 
         vtk1.point_data.vectors = warping
         writer = tvtk.PolyDataWriter(file_name=op.abspath(self.inputs.out_warp))
-        configure_input_data(writer, vtk1)
+        VTKInfo.configure_input_data(writer, vtk1)
         writer.write()
 
         vtk1.point_data.vectors = None
         vtk1.points = points1 + warping
         writer = tvtk.PolyDataWriter(file_name=op.abspath(self.inputs.out_file))
-        configure_input_data(writer, vtk1)
+        VTKInfo.configure_input_data(writer, vtk1)
         writer.write()
         return runtime
 

nipype/algorithms/tests/test_mesh_ops.py

@@ -8,20 +8,17 @@
 
 from nipype.testing import (assert_equal, assert_raises, skipif,
                             assert_almost_equal, example_data)
-
 import numpy as np
-
 from nipype.algorithms import mesh as m
-
-import platform
+from ...interfaces import vtkbase as VTKInfo
 
 
 def test_ident_distances():
     tempdir = mkdtemp()
     curdir = os.getcwd()
     os.chdir(tempdir)
 
-    if m.Info.no_tvtk():
+    if VTKInfo.no_tvtk():
         yield assert_raises, ImportError, m.ComputeMeshWarp
     else:
         in_surf = example_data('surf01.vtk')
@@ -45,24 +42,23 @@ def test_trans_distances():
     curdir = os.getcwd()
     os.chdir(tempdir)
 
-    if m.Info.no_tvtk():
+    if VTKInfo.no_tvtk():
         yield assert_raises, ImportError, m.ComputeMeshWarp
     else:
-        from nipype.algorithms.mesh import tvtk
-        from tvtk.common import is_old_pipeline as vtk_old
-        from tvtk.common import configure_input_data
+        from ...interfaces.vtkbase import tvtk
+
         in_surf = example_data('surf01.vtk')
         warped_surf = os.path.join(tempdir, 'warped.vtk')
 
         inc = np.array([0.7, 0.3, -0.2])
 
         r1 = tvtk.PolyDataReader(file_name=in_surf)
-        vtk1 = r1.output if vtk_old() else r1.get_output()
+        vtk1 = VTKInfo.vtk_output(r1)
         r1.update()
         vtk1.points = np.array(vtk1.points) + inc
 
         writer = tvtk.PolyDataWriter(file_name=warped_surf)
-        configure_input_data(writer, vtk1)
+        VTKInfo.configure_input_data(writer, vtk1)
         writer.write()
 
         dist = m.ComputeMeshWarp()
@@ -84,7 +80,7 @@ def test_warppoints():
     curdir = os.getcwd()
     os.chdir(tempdir)
 
-    if m.Info.no_tvtk():
+    if VTKInfo.no_tvtk():
         yield assert_raises, ImportError, m.WarpPoints
 
     # TODO: include regression tests for when tvtk is installed
@@ -98,7 +94,7 @@ def test_meshwarpmaths():
     curdir = os.getcwd()
     os.chdir(tempdir)
 
-    if m.Info.no_tvtk():
+    if VTKInfo.no_tvtk():
         yield assert_raises, ImportError, m.MeshWarpMaths
 
     # TODO: include regression tests for when tvtk is installed

nipype/interfaces/fsl/utils.py

@@ -34,8 +34,6 @@
 from ...utils.filemanip import (load_json, save_json, split_filename,
                                 fname_presuffix, copyfile)
 
-from ...algorithms.mesh import Info as VTKInfo
-
 warn = warnings.warn
 
 
@@ -1895,17 +1893,15 @@ def _parse_inputs(self, skip=None):
         return first_args + [second_args]
 
     def _vtk_to_coords(self, in_file, out_file=None):
-        # Ensure that tvtk is loaded with the appropriate ETS_TOOLKIT env var
+        from ..vtkbase import tvtk
+        from ...interfaces import vtkbase as VTKInfo
+
         if VTKInfo.no_tvtk():
             raise ImportError('TVTK is required and tvtk package was not found')
 
-        from ...algorithms.mesh import tvtk
-        from tvtk.common import is_old_pipeline as vtk_old
-
         reader = tvtk.PolyDataReader(file_name=in_file + '.vtk')
         reader.update()
-
-        mesh = reader.output if vtk_old() else reader.get_output()
+        mesh = VTKInfo.vtk_output(reader)
         points = mesh.points
 
         if out_file is None:
@@ -1915,23 +1911,20 @@ def _vtk_to_coords(self, in_file, out_file=None):
         return out_file
 
     def _coords_to_vtk(self, points, out_file):
-        import os
-        # Ensure that tvtk is loaded with the appropriate ETS_TOOLKIT env var
+        from ..vtkbase import tvtk
+        from ...interfaces import vtkbase as VTKInfo
+
         if VTKInfo.no_tvtk():
             raise ImportError('TVTK is required and tvtk package was not found')
 
-        from ...algorithms.mesh import tvtk
-        from tvtk.common import is_old_pipeline as vtk_old
-        from tvtk.common import configure_input_data
-
         reader = tvtk.PolyDataReader(file_name=self.inputs.in_file)
         reader.update()
 
-        mesh = reader.output if vtk_old() else reader.get_output()
+        mesh = VTKInfo.vtk_output(reader)
         mesh.points = points
 
         writer = tvtk.PolyDataWriter(file_name=out_file)
-        configure_input_data(writer, mesh)
+        VTKInfo.configure_input_data(writer, mesh)
         writer.write()
 
     def _trk_to_coords(self, in_file, out_file=None):

nipype/interfaces/vtkbase.py

@@ -0,0 +1,81 @@
+# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
+# vi: set ft=python sts=4 ts=4 sw=4 et:
+"""
+vtkbase provides some helpers to use VTK through the tvtk package (mayavi)
+
+Code using tvtk should import it through this module
+"""
+
+import os
+from .. import logging
+
+iflogger = logging.getLogger('interface')
+
+# Check that VTK can be imported and get version
+_vtk_version = None
+try:
+    import vtk
+    _vtk_version = (vtk.vtkVersion.GetVTKMajorVersion(),
+                    vtk.vtkVersion.GetVTKMinorVersion())
+except ImportError:
+    iflogger.warning('VTK was not found')
+
+# Ensure that tvtk is loaded with the appropriate ETS_TOOLKIT env var
+old_ets = os.getenv('ETS_TOOLKIT')
+os.environ['ETS_TOOLKIT'] = 'null'
+_have_tvtk = False
+try:
+    from tvtk.api import tvtk
+    _have_tvtk = True
+except ImportError:
+    iflogger.warning('tvtk wasn\'t found')
+    tvtk = None
+finally:
+    if old_ets is not None:
+        os.environ['ETS_TOOLKIT'] = old_ets
+    else:
+        del os.environ['ETS_TOOLKIT']
+
+
+def vtk_version():
+    """ Get VTK version """
+    global _vtk_version
+    return _vtk_version
+
+
+def no_vtk():
+    """ Checks if VTK is installed and the python wrapper is functional """
+    global _vtk_version
+    return _vtk_version is None
+
+
+def no_tvtk():
+    """ Checks if tvtk was found """
+    global _have_tvtk
+    return not _have_tvtk
+
+
+def vtk_old():
+    """ Checks if VTK uses the old-style pipeline (VTK<6.0) """
+    global _vtk_version
+    if _vtk_version is None:
+        raise RuntimeException('VTK is not correctly installed.')
+    return _vtk_version[0] < 6
+
+
+def configure_input_data(obj, data):
+    """
+    Configure the input data for vtk pipeline object obj.
+    Copied from latest version of mayavi
+    """
+    if vtk_old():
+        obj.input = data
+    else:
+        obj.set_input_data(data)
+
+
+def vtk_output(obj):
+    """ Configure the input data for vtk pipeline object obj."""
+    if vtk_old():
+        return obj.output
+    return obj.get_output()