numpy转vtk并切片

使用的是
python3.8.10
vtk9.2.6

#!/usr/bin/env python
# _*_ coding:utf-8 _*_

# noinspection PyUnresolvedReferences
import vtkmodules.vtkInteractionStyle
# noinspection PyUnresolvedReferences
import vtkmodules.vtkRenderingOpenGL2

from vtkmodules.vtkRenderingCore import (
    vtkImageActor,
    vtkRenderer,
    vtkRenderWindow,
    vtkRenderWindowInteractor
)
from vtkmodules.vtkIOImage import (
    vtkImageImport
)
from vtkmodules.vtkCommonMath import (
    vtkMatrix4x4
)
from vtkmodules.vtkCommonCore import (
    vtkLookupTable
)
from vtkmodules.vtkImagingCore import (
    vtkImageReslice,
    vtkImageMapToColors
)
from vtkmodules.vtkInteractionStyle import (
    vtkInteractorStyleImage
)
import SimpleITK as sitk
import numpy as np


def numpy2vtkImageImport(source_numpy, spacing=None, origin=None, direction=None, as_uint8=False):
    """
    numpy转成vtkImageImport

    :param source_numpy: numpy格式的图片（z,y,x）(A,C,S)
    :param spacing: 像素的间隔
    :param origin: 图片的原点
    :param direction: 方向
    :param as_uint8: 转成uint8
    :return: vtkImageImport
    """
    importer = vtkImageImport()
    origin_type = source_numpy.dtype
    if as_uint8:
        origin_type = np.uint8
        source_numpy = source_numpy.astype('uint8')
    # else:
    #     source_numpy = source_numpy.astype('int32')
    img_string = source_numpy.tobytes()
    dim = source_numpy.shape  # (z,y,x),(A,C,S)

    importer.CopyImportVoidPointer(img_string, len(img_string))
    if as_uint8:
        importer.SetDataScalarTypeToUnsignedChar()
    elif np.int16 == origin_type:
        importer.SetDataScalarTypeToShort()
    elif np.float32 == origin_type:
        importer.SetDataScalarTypeToDouble()
    else:
        importer.SetDataScalarTypeToInt()
    importer.SetNumberOfScalarComponents(1)

    extent = importer.GetDataExtent()  # 0,0,0,0,0,0
    # 图像的维度=(extent[1]-extent[0]+1) * (extent[3]-extent[2]+1) * (extent[5]-DataExtent[4]+1)=(x,y,z)=(S,C,A)
    # importer.SetDataExtent(extent[0], extent[0] + dim[2] - 1,
    #                        extent[2], extent[2] + dim[1] - 1,
    #                        extent[4], extent[4] + dim[0] - 1)
    importer.SetWholeExtent(extent[0], extent[0] + dim[2] - 1,
                            extent[2], extent[2] + dim[1] - 1,
                            extent[4], extent[4] + dim[0] - 1)
    importer.SetDataExtentToWholeExtent()
    if spacing is not None:
        importer.SetDataSpacing(*spacing)
    if origin is not None:
        importer.SetDataOrigin(*origin)
    if direction is not None:
        importer.SetDataDirection(direction)
    importer.Update()
    return importer


def resliceByOrient(vtk_data, center, ori):
    """
    重采样

    :param vtk_data: vtk的数据
    :param center: 采样中心(S,C,A)(x,y,z)
    :param ori: 1:S(x),2:C(y),0:A(z)
    :return: 重采样的数据
    """
    if 1 == ori:
        # 矢状面（sagittal plane）S(x)
        elements = [
            0, 0, -1, 0,
            1, 0, 0, 0,
            0, -1, 0, 0,
            0, 0, 0, 1
        ]
    elif 2 == ori:
        # 冠状面（coronal plane）C(y)
        elements = [
            1, 0, 0, 0,
            0, 0, 1, 0,
            0, -1, 0, 0,
            0, 0, 0, 1
        ]
    elif 0 == ori:
        # 横断面（transverse plane）A(z)
        elements = [
            1, 0, 0, 0,
            0, 1, 0, 0,
            0, 0, 1, 0,
            0, 0, 0, 1
        ]
    else:
        # 斜切切片
        elements = [
            1, 0, 0, 0,
            0, 0, 8.66025, -0.5, 0,
            0, 0.5, 0.866025, 0,
            0, 0, 0, 1
        ]
    resliceAxes = vtkMatrix4x4()
    resliceAxes.DeepCopy(elements)
    for i in range(3):
        resliceAxes.SetElement(i, 3, center[i])
        # resliceAxes.SetElement(i, 3, int(center[i]))
    reslice = vtkImageReslice()
    reslice.SetInputConnection(vtk_data.GetOutputPort())
    reslice.SetOutputDimensionality(2)
    reslice.SetResliceAxes(resliceAxes)
    reslice.SetInterpolationModeToLinear()
    reslice.Update()
    return reslice


def get_color_table(is_image, minimum=None, maximum=None, window=None, level=None):
    """
    获取颜色映射表

    :param is_image: 是否为图像
    :param minimum: 最小值
    :param maximum: 最大值
    :param window: 窗宽
    :param level: 窗位
    :return: 颜色映射表
    """
    colorTable = vtkLookupTable()
    if is_image:
        if minimum is None and maximum is None and window is not None and level is not None:
            minimum = (2 * level - window) / 2.0 + 0.5
            maximum = (2 * level + window) / 2.0 + 0.5
        elif minimum is None or maximum is None:
            raise Exception("can't get right min max")

        # 像素最小值到最大值，（可以根据窗宽窗位计算出像素范围）
        colorTable.SetRange(minimum, maximum)
        colorTable.SetValueRange(0.0, 1.0)
        colorTable.SetSaturationRange(0.0, 0.0)
        colorTable.SetRampToLinear()
    else:
        # TODO 改成正常的颜色
        colorTable.SetNumberOfTableValues(5)
        colorTable.SetRange(0.0, 4.0)
        colorTable.SetTableValue(0, 0, 0.0, 0.0, 0.0)
        colorTable.SetTableValue(1, 0.0, 1.0, 0.0, 0.4)
        colorTable.SetTableValue(2, 1.0, 0.0, 0.0, 0.4)
        colorTable.SetTableValue(3, 0.0, 0.0, 1.0, 0.4)
        colorTable.SetTableValue(4, 0.0, 0.0, 1.0, 0.4)
    colorTable.Build()
    return colorTable


if __name__ == '__main__':
    src = sitk.ReadImage('RibFrac101-image.nii.gz')
    spacing = src.GetSpacing()
    origin = src.GetOrigin()
    direction = src.GetDirection()
    source_numpy = sitk.GetArrayFromImage(src)
    vtk_data = numpy2vtkImageImport(source_numpy, spacing, origin, direction)
    print(vtk_data.GetDataExtent())
    print(vtk_data.GetDataOrigin())
    print(vtk_data.GetDataSpacing())
    print(vtk_data.GetDataDirection())
    # 257,414,177表示对应面的切片
    center = [
        origin[0] + spacing[0] * 257,  # 矢状面（sagittal plane）S
        origin[1] + spacing[1] * 414,  # 冠状面（coronal plane）C
        origin[2] + spacing[2] * 177,  # 横断面（transverse plane）A
    ]
    print(center)
    reslice = resliceByOrient(vtk_data, center, 1)
    colorTable = get_color_table(True, *(vtk_data.GetOutput().GetScalarRange()))

    colorMap = vtkImageMapToColors()
    colorMap.SetLookupTable(colorTable)
    colorMap.SetInputConnection(reslice.GetOutputPort())
    colorMap.Update()

    imgActor = vtkImageActor()
    imgActor.SetInputData(colorMap.GetOutput())

    renderer = vtkRenderer()
    renderer.AddActor(imgActor)
    renderer.SetBackground(0.0, 0.0, 0.0)

    renderWindow = vtkRenderWindow()
    renderWindow.AddRenderer(renderer)
    renderWindow.SetSize(640, 400)

    rwi = vtkRenderWindowInteractor()
    imageStyle = vtkInteractorStyleImage()
    rwi.SetInteractorStyle(imageStyle)
    rwi.SetRenderWindow(renderWindow)
    rwi.Initialize()
    renderWindow.Render()
    rwi.Start()

参考

《VTK图形图像开发进阶》P97 三维图像切面提取