PolyDataToImageData的Python实现注意事项

为将STL模型转化为具有空间局部性的图像（体素）格式，模仿官方例子用Python实现了算法。

def stl_to_voxel(stl, padding=0.0, spacing=None, scale=None):
    """ Transform STL file to voxelbased volumes
    Args:
        stl (vtkPolyData): vtk 3D poly data in the form of points and lines
        padding (float): adds padding space in calculating the bounds of stl 
            file, in order to preserve the surface attached to orignal bounds.
        spacing (float or tuple): spacing in three directions.
        scale (int or tuple): the resolution of output 3D data. Not used
            if spacing is specified.
    Return:
        voxel (vtkImageData): vtk 3D image data which contain voxel representation
        of input stl.
    """
    assert isinstance(stl, vtk.vtkPolyData)
    assert isinstance(padding, float)
    assert isinstance(spacing, (int, tuple)) or isinstance(scale, (int, tuple))

    bounds = np.array(stl.GetBounds())

    # Add small padding to bounds
    if padding != 0:
        padding = (bounds[1::2] - bounds[::2]) * padding
        bounds[::2] = bounds[::2] - padding
        bounds[1::2] = bounds[1::2] + padding

    if spacing:
        if isinstance(spacing, float):
            spacing = [spacing for i in range(3)]
    else:
        if isinstance(scale, int):
            scale = [scale for i in range(3)]
        spacing = [(bounds[i * 2 + 1] - bounds[i * 2]) / scale[i]
                   for i in range(3)]

    whiteImage = vtk.vtkImageData()
    whiteImage.SetSpacing(spacing)

    # Set dim
    dim = [
        math.ceil((bounds[i * 2 + 1] - bounds[i * 2]) / spacing[i])
        for i in range(3)
    ]

    whiteImage.SetDimensions(dim)
    whiteImage.SetExtent(0, dim[0] - 1, 0, dim[1] - 1, 0, dim[2] - 1)

    # Set origin
    origin = [bounds[i * 2] + spacing[i] / 2 for i in range(3)]
    whiteImage.SetOrigin(origin)

    whiteImage.AllocateScalars(vtk.VTK_UNSIGNED_CHAR, 1)

    # Fill the image with foreground voxels
    inval = 1
    outval = 0
    for i in range(whiteImage.GetNumberOfPoints()):
        whiteImage.GetPointData().GetScalars().SetTuple1(i, inval)

    # Polygonal data --> image stencil
    pol2stenc = vtk.vtkPolyDataToImageStencil()
    pol2stenc.SetInputData(stl)
    pol2stenc.SetOutputOrigin(origin)
    pol2stenc.SetOutputSpacing(spacing)
    pol2stenc.SetOutputWholeExtent(whiteImage.GetExtent())
    pol2stenc.Update()

    # Cut the corresponding white image and set the background
    imgstenc = vtk.vtkImageStencil()
    imgstenc.SetInputData(whiteImage)
    imgstenc.SetStencilConnection(pol2stenc.GetOutputPort())
    imgstenc.ReverseStencilOff()
    imgstenc.SetBackgroundValue(outval)
    imgstenc.Update()

    return imgstenc.GetOutput()

但是在使用过程中，发现这个转换特别的慢，经过运行时间分析发现，上述代码对于输入的STL模型的点数与面数不敏感，但对模型输出大小特别敏感。这说明，代码在建立输出时产生了特别耗时的操作，尤其有可能是初始化操作。经过研究发现，符合了“Python请不要使用for循环定律”，问题出现在下述代码中。

for i in range(whiteImage.GetNumberOfPoints()):
    whiteImage.GetPointData().GetScalars().SetTuple1(i, inval)

采用如下代码替换后，效率明显改善。

whiteImage.GetPointData().GetScalars().Fill(inval)

这次事件的启示是：python用for循环请三思。转换C代码时千万不要无脑照抄for循环。