Created TVTKBaseInterface. Should fix #1218

Author: oesteban

nipype/algorithms/mesh.py

@@ -22,9 +22,40 @@
 from ..interfaces.base import (BaseInterface, traits, TraitedSpec, File,
                                BaseInterfaceInputSpec)
 from warnings import warn
+
+
+have_tvtk = False
+
+try:
+    from tvtk.api import tvtk
+    have_tvtk = True
+except ImportError:
+    raise ImportError('Interface requires tvtk')
+
+try:
+    from enthought.etsconfig.api import ETSConfig
+    ETSConfig.toolkit = 'null'
+except ImportError:
+    iflogger.warn(('ETS toolkit could not be imported'))
+except ValueError:
+    iflogger.warn(('ETS toolkit could not be set to null'))
+
+if have_tvtk:
+    tvtk.tvtk_classes.vtk_version import vtk_build_version
+    vtk_major = int(vtk_build_version[0])
+
 iflogger = logging.getLogger('interface')
 
 
+class TVTKBaseInterface(BaseInterface):
+    _redirect_x = True
+
+    def __init__(self, **inputs):
+        if not have_tvtk:
+            raise RuntimeError('Interface requires tvtk')
+        super(TVTKBaseInterface, self).__init__(**inputs)
+
+
 class WarpPointsInputSpec(BaseInterfaceInputSpec):
     points = File(exists=True, mandatory=True,
                   desc=('file containing the point set'))
@@ -41,7 +72,7 @@ class WarpPointsOutputSpec(TraitedSpec):
     out_points = File(desc='the warped point set')
 
 
-class WarpPoints(BaseInterface):
+class WarpPoints(TVTKBaseInterface):
 
     """
     Applies a displacement field to a point set given in vtk format.
@@ -61,7 +92,6 @@ class WarpPoints(BaseInterface):
     """
     input_spec = WarpPointsInputSpec
     output_spec = WarpPointsOutputSpec
-    _redirect_x = True
 
     def _gen_fname(self, in_file, suffix='generated', ext=None):
         import os.path as op
@@ -80,26 +110,6 @@ def _gen_fname(self, in_file, suffix='generated', ext=None):
         return op.abspath('%s_%s.%s' % (fname, suffix, ext))
 
     def _run_interface(self, runtime):
-        vtk_major = 6
-        try:
-            import vtk
-            vtk_major = vtk.VTK_MAJOR_VERSION
-        except ImportError:
-            iflogger.warn(('python-vtk could not be imported'))
-
-        try:
-            from tvtk.api import tvtk
-        except ImportError:
-            raise ImportError('Interface requires tvtk')
-
-        try:
-            from enthought.etsconfig.api import ETSConfig
-            ETSConfig.toolkit = 'null'
-        except ImportError:
-            iflogger.warn(('ETS toolkit could not be imported'))
-        except ValueError:
-            iflogger.warn(('ETS toolkit could not be set to null'))
-
         import nibabel as nb
         import numpy as np
         from scipy import ndimage
@@ -182,7 +192,7 @@ class ComputeMeshWarpOutputSpec(TraitedSpec):
                     desc='numpy file keeping computed distances and weights')
 
 
-class ComputeMeshWarp(BaseInterface):
+class ComputeMeshWarp(TVTKBaseInterface):
 
     """
     Calculates a the vertex-wise warping to get surface2 from surface1.
@@ -207,7 +217,6 @@ class ComputeMeshWarp(BaseInterface):
 
     input_spec = ComputeMeshWarpInputSpec
     output_spec = ComputeMeshWarpOutputSpec
-    _redirect_x = True
 
     def _triangle_area(self, A, B, C):
         A = np.array(A)
@@ -220,19 +229,6 @@ def _triangle_area(self, A, B, C):
         return area
 
     def _run_interface(self, runtime):
-        try:
-            from tvtk.api import tvtk
-        except ImportError:
-            raise ImportError('Interface ComputeMeshWarp requires tvtk')
-
-        try:
-            from enthought.etsconfig.api import ETSConfig
-            ETSConfig.toolkit = 'null'
-        except ImportError:
-            iflogger.warn(('ETS toolkit could not be imported'))
-        except ValueError:
-            iflogger.warn(('ETS toolkit is already set'))
-
         r1 = tvtk.PolyDataReader(file_name=self.inputs.surface1)
         r2 = tvtk.PolyDataReader(file_name=self.inputs.surface2)
         vtk1 = r1.output
@@ -280,7 +276,12 @@ def _run_interface(self, runtime):
         out_mesh.point_data.vectors.name = 'warpings'
         writer = tvtk.PolyDataWriter(
             file_name=op.abspath(self.inputs.out_warp))
-        writer.set_input_data(out_mesh)
+
+        if vtk_major <= 5:
+            writer.input = mesh
+        else:
+            writer.set_input_data_object(mesh)
+
         writer.write()
 
         self._distance = np.average(errvector, weights=weights)
@@ -322,7 +323,7 @@ class MeshWarpMathsOutputSpec(TraitedSpec):
                     desc='vtk with surface warped')
 
 
-class MeshWarpMaths(BaseInterface):
+class MeshWarpMaths(TVTKBaseInterface):
 
     """
     Performs the most basic mathematical operations on the warping field
@@ -348,22 +349,8 @@ class MeshWarpMaths(BaseInterface):
 
     input_spec = MeshWarpMathsInputSpec
     output_spec = MeshWarpMathsOutputSpec
-    _redirect_x = True
 
     def _run_interface(self, runtime):
-        try:
-            from tvtk.api import tvtk
-        except ImportError:
-            raise ImportError('Interface ComputeMeshWarp requires tvtk')
-
-        try:
-            from enthought.etsconfig.api import ETSConfig
-            ETSConfig.toolkit = 'null'
-        except ImportError:
-            iflogger.warn(('ETS toolkit could not be imported'))
-        except ValueError:
-            iflogger.warn(('ETS toolkit is already set'))
-
         r1 = tvtk.PolyDataReader(file_name=self.inputs.in_surf)
         vtk1 = r1.output
         r1.update()
@@ -412,14 +399,21 @@ def _run_interface(self, runtime):
         vtk1.point_data.vectors = warping
         writer = tvtk.PolyDataWriter(
             file_name=op.abspath(self.inputs.out_warp))
-        writer.set_input_data(vtk1)
+        if vtk_major <= 5:
+            writer.input = vtk1
+        else:
+            writer.set_input_data_object(vtk1)
         writer.write()
 
         vtk1.point_data.vectors = None
         vtk1.points = points1 + warping
         writer = tvtk.PolyDataWriter(
             file_name=op.abspath(self.inputs.out_file))
-        writer.set_input_data(vtk1)
+
+        if vtk_major <= 5:
+            writer.input = vtk1
+        else:
+            writer.set_input_data_object(vtk1)
         writer.write()
 
         return runtime