Vti_generator/vtk_test.py

import os
import vtk
from vtk.util import numpy_support
from vtk.util.numpy_support import numpy_to_vtk, get_vtk_array_type
import shutil


def numpy_to_vti(array, origin, spacing, filename):
    """This function write a VtkImageData vti file from a numpy array.

    :param array: input array
    :type array: :class:`numpy.ndarray`
    :param origin: the origin of the array
    :type origin: array like object of values
    :param spacing: the step in each dimension
    :type spacing: array like object of values
    :param filename: output filename (.vti)
    :type filename: str
    """

    if array.ndim != 3:
        raise ValueError("Only works with 3 dimensional arrays")

    vtkArray = numpy_to_vtk(num_array=array.flatten('F'), deep=True,
                            array_type=get_vtk_array_type(array.dtype))

    imageData = vtk.vtkImageData()
    imageData.SetOrigin(origin)
    imageData.SetSpacing(spacing)
    imageData.SetDimensions(array.shape)
    imageData.GetPointData().SetScalars(vtkArray)

    writer = vtk.vtkXMLImageDataWriter()
    writer.SetFileName(filename)
    writer.SetInputData(imageData)
    writer.Write()

    return ()