超好用！！！VTK几种拾取交互方式的封装（python版），包含Cell、Point、Actor的拾取高亮，可选择单选or多选模式

最近把VTK的拾取功能封装到了InteractorStyle的继承类里面，可以很方便的进行拾取，如下

点选取

Cell选取

源码附上，纯个人手敲和测试，觉得有用的uu别忘点赞收藏一下哈 ，有问题随时评论或私信沟通~（博主python版本为3.9，在函数定义的时候加了一些对参数类型的限制，低于3.9的版本可能会报错，把那些冒号后面的参数类型要求删掉就好）

import vtkmodules.all as vtk
import numpy as np
from typing import Union
from vtkmodules.util.numpy_support import numpy_to_vtk


def genHighlightActor(polydata: vtk.vtkPolyData, cell_id: int, color: tuple[float, float, float],
                      mark_size: int = 4) -> vtk.vtkActor:
    """根据polydata和cell_id生成一个对应的actor"""
    cell = polydata.GetCell(cell_id)
    ptn = cell.GetNumberOfPoints()
    pts = np.zeros((ptn, 3), dtype=np.float64)
    for i in range(ptn):
        point_id = cell.GetPointId(i)
        polydata.GetPoint(point_id, pts[i])
    if ptn == 1:  # 如果是点
        actor = Geometry.genPixelActor(pts[0], mark_size, color)
    elif ptn == 3:  # 如果是三角形
        actor = Geometry.genTriangleFrameActor(pts, mark_size, color)
    else:  # 如果是线段或折线
        actor = Geometry.genWholeBreakLineActor(pts, mark_size, color)
    return actor


class Geometry:
    """几何元素生成函数封装类"""
    @staticmethod
    def genActorFromPolyData(polydata: vtk.vtkPolyData) -> vtk.vtkActor:
        """由polydata生成actor"""
        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputData(polydata)
        mapper.Update()
        actor = vtk.vtkActor()
        actor.SetMapper(mapper)
        actor.Modified()
        return actor

    @staticmethod
    def genSphereActor(position: Union[tuple[float, float, float], list[float], np.ndarray], radius: float = 2.0,
                       color: tuple[float, float, float] = (1.0, 0.0, 0.0)) -> vtk.vtkActor:
        """以球体形式生成点actor，输入位置、半径、颜色"""
        pt = vtk.vtkSphereSource()
        pt.SetRadius(radius)
        pt.SetCenter(position[0], position[1], position[2])
        pt.SetPhiResolution(20)  # 球体显示的分辨率参数
        pt.SetThetaResolution(20)  # 球体显示的分辨率参数
        pt.Update()
        actor = Geometry.genActorFromPolyData(pt.GetOutput())
        actor.GetProperty().SetColor(color)
        actor.Modified()
        return actor

    @staticmethod
    def genPixelActor(position: Union[tuple[float, float, float], list[float], np.ndarray],
                      size: int = 5, color: tuple[float, float, float] = (1.0, 0.0, 0.0)) -> vtk.vtkActor:
        """以像素点形式生成点actor，输入位置、像素大小、颜色"""
        pts = np.array([position])
        actor = Geometry.genPointCloudActor(pts, size, color)
        return actor

    @staticmethod
    def genLineActor(start_pos: Union[tuple[float, float, float], list[float], np.ndarray],
                     end_pos: Union[tuple[float, float, float], list[float], np.ndarray],
                     line_width: int = 2, color: tuple[float, float, float] = (1.0, 0.0, 0.0)) -> vtk.vtkActor:
        """以像素线段形式生成线actor，可自定义线宽颜色"""
        pts = np.array([start_pos, end_pos])
        actor = Geometry.genWholeBreakLineActor(pts, line_width, color)
        return actor

    @staticmethod
    def genTriangleFrameActor(pts: np.ndarray, line_width: int = 2,
                              color: tuple[float, float, float] = (1.0, 0.0, 0.0)) -> vtk.vtkActor:
        """以线框形式生成三角形actor，可自定义线宽颜色"""
        points = vtk.vtkPoints()
        points.SetData(numpy_to_vtk(np.copy(pts)))
        cells = vtk.vtkCellArray()
        polydata = vtk.vtkPolyData()
        cells.InsertNextCell(4, [0, 1, 2, 0])
        polydata.SetPoints(points)
        polydata.SetLines(cells)
        actor = Geometry.genActorFromPolyData(polydata)
        actor.GetProperty().SetLineWidth(line_width)
        actor.GetProperty().SetColor(color)
        return actor

    @staticmethod
    def genWholeBreakLineActor(pts: np.ndarray, line_width: int = 2,
                               color: tuple[float, float, float] = (1.0, 0.0, 0.0)) -> vtk.vtkActor:
        """生成整段的折线actor，其中只有一个Cell，可自定义线宽颜色"""
        nL = pts.shape[0]
        points = vtk.vtkPoints()
        points.SetData(numpy_to_vtk(np.copy(pts)))
        cells = vtk.vtkCellArray()
        polydata = vtk.vtkPolyData()
        cells.InsertNextCell(nL, [i for i in range(nL)])
        polydata.SetPoints(points)
        polydata.SetLines(cells)
        actor = Geometry.genActorFromPolyData(polydata)
        actor.GetProperty().SetLineWidth(line_width)
        actor.GetProperty().SetColor(color)
        return actor

    @staticmethod
    def genPolyLineActor(pts: np.ndarray, line_width: int = 2,
                         color: tuple[float, float, float] = (1.0, 0.0, 0.0)) -> vtk.vtkActor:
        """生成分段的折线actor，其中每一条线段是一个Cell，可自定义线宽颜色"""
        nL = pts.shape[0]
        points = vtk.vtkPoints()
        points.SetData(numpy_to_vtk(np.copy(pts)))
        cells = vtk.vtkCellArray()
        polydata = vtk.vtkPolyData()
        for i in range(nL - 1):
            cells.InsertNextCell(2, [i, i + 1])
        polydata.SetPoints(points)
        polydata.SetLines(cells)
        actor = Geometry.genActorFromPolyData(polydata)
        actor.GetProperty().SetLineWidth(line_width)
        actor.GetProperty().SetColor(color)
        return actor

    @staticmethod
    def genAxesActor(axes_length: int = 600) -> vtk.vtkAxesActor:
        """坐标轴创建，输入三轴长度"""
        axes = vtk.vtkAxesActor()
        axes.SetTotalLength(axes_length, axes_length, axes_length)  # 设置轴长度
        axes.SetShaftType(1)  # 设置轴类型，1为线0为圆柱

        # axes.SetCylinderRadius(1)  # 设置圆柱半径
        axes.SetConeRadius(0.05)  # 设置箭头大小
        axes.GetXAxisCaptionActor2D().SetWidth(0.03)  # 设置坐标轴标签大小
        axes.GetYAxisCaptionActor2D().SetWidth(0.03)
        axes.GetZAxisCaptionActor2D().SetWidth(0.03)
        return axes

    @staticmethod
    def genPointCloudActor(pts: np.ndarray, size: int = 5,
                           color: tuple[float, float, float] = (1.0, 0.0, 0.0)) -> vtk.vtkActor:
        """生成点云actor，用固定像素大小的点表示，可定义大小与颜色"""
        npt = pts.shape[0]
        points = vtk.vtkPoints()
        points.SetData(numpy_to_vtk(pts))
        cells = vtk.vtkCellArray()
        polydata = vtk.vtkPolyData()
        for i in range(npt):
            cells.InsertNextCell(1, [i])
        polydata.SetPoints(points)
        polydata.SetVerts(cells)
        actor = Geometry.genActorFromPolyData(polydata)
        actor.GetProperty().SetPointSize(size)
        actor.GetProperty().SetColor(color)
        return actor

    @staticmethod
    def genPointCloudSpheres(pts: np.ndarray, radius: float,
                             color: tuple[float, float, float] = (1.0, 0.0, 0.0)) -> list[vtk.vtkActor]:
        """用球体对象组表示点云，可定义半径与颜色"""
        npt = pts.shape[0]
        actor_list = []
        for i in range(npt):
            actor = Geometry.genSphereActor(tuple(pts[i]), radius, color)
            actor_list.append(actor)
        return actor_list


class ActorPickInteractorStyle(vtk.vtkInteractorStyleTrackballCamera):
    """Actor选取交互器，color表示选中高亮actor的颜色，select_mode为0表示单选，select_mode为1表示多选"""
    def __init__(self, parent=None, color: tuple[float, float, float] = (1.0, 1.0, 0.0), select_mode: int = 0):  # parent不能删掉
        super().__init__()
        self.AddObserver("LeftButtonPressEvent", self.__leftButtonPressEvent)
        self.picked_actor_list = []  # 已选取的actor列表
        self.original_color_list = []  # 已选取的actor原始颜色列表，与picked_actor_list对应
        self.color = color
        if select_mode != 0 and select_mode != 1:
            raise ValueError
        self.select_mode = select_mode

    def __leftButtonPressEvent(self, obj, event):  # obj和event参数不能删掉
        """左键单击事件"""
        click_pos = self.GetInteractor().GetEventPosition()  # 获取鼠标左键点击时的屏幕坐标
        picker = vtk.vtkPropPicker()
        picker.Pick(click_pos[0], click_pos[1], 0, self.GetDefaultRenderer())
        picked_actor = picker.GetActor()
        if picked_actor is not None:
            # 如果选择的actor已被选，则复原
            if picked_actor in self.picked_actor_list:
                index = self.picked_actor_list.index(picked_actor)
                self.picked_actor_list[index].GetProperty().SetColor(self.original_color_list[index])
                self.picked_actor_list[index].Modified()
                self.picked_actor_list.pop(index)
                self.original_color_list.pop(index)
            # 如果选择的是新actor
            else:
                # 如果是单选模式，需要先把之前选择的actor复原
                if self.select_mode == 0 and len(self.picked_actor_list) > 0:
                    self.picked_actor_list[0].GetProperty().SetColor(self.original_color_list[0])
                    self.picked_actor_list[0].Modified()
                    self.picked_actor_list.pop(0)
                    self.original_color_list.pop(0)
                # 储存选中actor的原始颜色，并把选中的actor高亮
                picked_color = picked_actor.GetProperty().GetColor()
                self.original_color_list.append(picked_color)
                picked_actor.GetProperty().SetColor(self.color)
                picked_actor.Modified()
                self.picked_actor_list.append(picked_actor)
        self.OnLeftButtonDown()


class CellPickInteractorStyle(vtk.vtkInteractorStyleTrackballCamera):
    """
    Cell选取交互器，mark_size表示选中高亮元素的线宽或点大小，color表示选中高亮颜色，select_mode为0表示单选，select_mode为1表示多选
    easy_to_pick若设置为True，则可以降低选取精度以更方便地选取由线和点组成的单元
    """
    def __init__(self, parent=None, mark_size: int = 4, color: tuple[float, float, float] = (1.0, 0.0, 0.0),
                 select_mode: int = 0, easy_to_pick: bool = False):  # parent不能删掉
        super().__init__()
        self.AddObserver("LeftButtonPressEvent", self.__leftButtonPressEvent)
        self.picked_cellUnit_list = []  # 已选取的[actor, cell_id]列表
        self.highlight_cellActor_list = []  # 已生成的高亮actor列表，与picked_cellUnit_list对应
        self.mark_size = mark_size
        self.color = color
        if select_mode != 0 and select_mode != 1:
            raise ValueError
        self.select_mode = select_mode
        self.easy_to_pick = easy_to_pick

    def __get_selected_cellUnit_index(self, actor: vtk.vtkActor, cell_id: int) -> int:
        """看选的cell是否在已选列表里，是就返回index，不是就返回-1"""
        selected_num = len(self.highlight_cellActor_list)
        if selected_num < 1:
            return -1
        for i in range(selected_num):
            if actor == self.picked_cellUnit_list[i][0] and cell_id == self.picked_cellUnit_list[i][1]:
                return i
        else:
            return -1

    def __leftButtonPressEvent(self, obj, event):  # obj和event参数不能删掉
        """左键单击事件"""
        click_pos = self.GetInteractor().GetEventPosition()  # 获取鼠标左键点击时的屏幕坐标
        picker = vtk.vtkCellPicker()
        if self.easy_to_pick:
            picker.SetTolerance(0.001)  # 设置误差限
        picker.Pick(click_pos[0], click_pos[1], 0, self.GetDefaultRenderer())
        picked_cell_id = picker.GetCellId()
        if picked_cell_id == -1:
            picked_cell_id = None
        if picked_cell_id is not None:
            picked_actor = picker.GetActor()
            index = self.__get_selected_cellUnit_index(picked_actor, picked_cell_id)
            # 如果选择的cell未被选
            if index == -1:
                # 如果是单选模式，需要先把之前选择的cell复原
                if self.select_mode == 0 and len(self.highlight_cellActor_list) > 0:
                    self.GetDefaultRenderer().RemoveActor(self.highlight_cellActor_list[0])
                    self.highlight_cellActor_list.pop(0)
                    self.picked_cellUnit_list.pop(0)
                # 生成高亮actor
                picked_mapper = picked_actor.GetMapper()
                picked_polydata = picked_mapper.GetInput()
                newHighlightActor = genHighlightActor(picked_polydata, picked_cell_id, self.color, self.mark_size)
                newHighlightActor.SetPickable(False)
                newHighlightActor.Modified()
                self.picked_cellUnit_list.append([picked_actor, picked_cell_id])
                self.highlight_cellActor_list.append(newHighlightActor)
                self.GetDefaultRenderer().AddActor(newHighlightActor)
            # 如果选择的cell已被选，则直接复原
            else:
                self.GetDefaultRenderer().RemoveActor(self.highlight_cellActor_list[index])
                self.highlight_cellActor_list.pop(index)
                self.picked_cellUnit_list.pop(index)
        self.GetDefaultRenderer().Modified()
        self.OnLeftButtonDown()


class PointPickInteractorStyle(vtk.vtkInteractorStyleTrackballCamera):
    """Point选取交互器，mark_size表示选中高亮点大小，mode为0表示以像素显示，mode为1表示以球体显示，color表示选中高亮颜色，select_mode为0表示单选，select_mode为1表示多选"""
    def __init__(self, parent=None, pt_size: int = 4, color: tuple[float, float, float] = (1.0, 0.0, 0.0),
                 mode: int = 0, select_mode: int = 0):  # parent不能删掉
        super().__init__()
        self.AddObserver("LeftButtonPressEvent", self.__leftButtonPressEvent)
        self.picked_ptUnit_list = []  # 已选取的[actor, pt_id]列表
        self.highlight_ptActor_list = []  # 已生成的高亮actor列表，与picked_cellUnit_list对应
        self.pt_size = pt_size
        self.color = color
        if select_mode != 0 and select_mode != 1:
            raise ValueError
        if mode != 0 and mode != 1:
            raise ValueError
        self.mode = mode
        self.select_mode = select_mode

    def __get_selected_ptUnit_index(self, actor: vtk.vtkActor, pt_id: int) -> int:
        """看选的cell是否在已选列表里，是就返回index，不是就返回-1"""
        selected_num = len(self.highlight_ptActor_list)
        if selected_num < 1:
            return -1
        for i in range(selected_num):
            if actor == self.picked_ptUnit_list[i][0] and pt_id == self.picked_ptUnit_list[i][1]:
                return i
        else:
            return -1

    def __leftButtonPressEvent(self, obj, event):  # obj和event参数不能删掉
        """左键单击事件"""
        click_pos = self.GetInteractor().GetEventPosition()  # 获取鼠标左键点击时的屏幕坐标
        picker = vtk.vtkPointPicker()
        # picker.SetTolerance(0.001)  # 设置误差限
        picker.Pick(click_pos[0], click_pos[1], 0, self.GetDefaultRenderer())
        picked_pt_id = picker.GetPointId()
        if picked_pt_id == -1:
            picked_pt_id = None
        if picked_pt_id is not None:
            picked_actor = picker.GetActor()
            index = self.__get_selected_ptUnit_index(picked_actor, picked_pt_id)
            # 如果选择的point未被选
            if index == -1:
                # 如果是单选模式，需要先把之前选择的point复原
                if self.select_mode == 0 and len(self.highlight_ptActor_list) > 0:
                    self.GetDefaultRenderer().RemoveActor(self.highlight_ptActor_list[0])
                    self.highlight_ptActor_list.pop(0)
                    self.picked_ptUnit_list.pop(0)
                # 生成高亮actor
                picked_mapper = picked_actor.GetMapper()
                picked_polydata = picked_mapper.GetInput()
                pos = picked_polydata.GetPoint(picked_pt_id)
                if self.mode == 0:
                    newHighlightActor = Geometry.genPixelActor(pos, self.pt_size, self.color)
                else:
                    newHighlightActor = Geometry.genSphereActor(pos, self.pt_size, self.color)
                newHighlightActor.SetPickable(False)
                newHighlightActor.Modified()
                self.picked_ptUnit_list.append([picked_actor, picked_pt_id])
                self.highlight_ptActor_list.append(newHighlightActor)
                self.GetDefaultRenderer().AddActor(newHighlightActor)
            # 如果选择的point已被选，则直接复原
            else:
                self.GetDefaultRenderer().RemoveActor(self.highlight_ptActor_list[index])
                self.highlight_ptActor_list.pop(index)
                self.picked_ptUnit_list.pop(index)
        self.GetDefaultRenderer().Modified()
        self.OnLeftButtonDown()

使用方法：

# 生成两个球体来测试
actor1 = Geometry.genSphereActor(pos=(0., 0., 0.), radius=30, color=(1., 1., 1.))
actor2 = Geometry.genSphereActor(pos=(100., 0., 0.), radius=20, color=(0.7, 0.7, 0.7))

renWin = vtk.vtkRenderWindow()  # 渲染窗口
renWin.SetSize(800, 800)  # 设置窗口大小
ren = vtk.vtkRenderer()  # 渲染器
ren.SetBackground(0.1, 0.1, 0.1)  # 场景背景色
renWin.AddRenderer(ren)

# 添加actor到场景
ren.AddActor(actor1）
ren.AddActor(actor2)
ren.Modified()

style = CellPickInteractorStyle(mark_size=4, color=(1., 0., 0.), select_mode=0)  # 换成上述代码中的任一种style,参数自己按需求定义即可
style.SetDefaultRenderer(ren)
style.Modified()
iren = vtk.vtkRenderWindowInteractor()
iren.SetRenderWindow(renWin)
iren.SetInteractorStyle(style)
renWin.Render()
iren.Initialize()
iren.Start()