三维网格模型和点云的分割结果可视化

羡长江

已于 2023-08-16 16:37:03 修改

阅读量747

点赞数 3

分类专栏： VTK 文章标签： python 3d 数据可视化

于 2023-08-02 23:39:16 首次发布

本文链接：https://blog.csdn.net/weixin_43850909/article/details/132072933

版权

VTK 专栏收录该内容

8 篇文章 2 订阅

订阅专栏

在三维模型和点云分割的实验中，通常都需要对分割结果进行可视化，不同的区域用不同的颜色，每个点的颜色要根据分割结果来确定，例如下面这张图就是PointNet中部件分割的结果可视化。这里用一个三维网格模型作为例子在实现这一过程

一、读取三维网格模型

对三维网格模型进行分割，第一当然是读取数据，这里我读一个obj格式的模型

import vtk

def polydata_visualize(poly):
    colors = vtk.vtkNamedColors()
    mapper = vtk.vtkPolyDataMapper()
    mapper.SetInputData(poly)
    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
    ren = vtk.vtkRenderer()
    ren.AddActor(actor)
    ren.SetBackground(colors.GetColor3d('White'))
    renWin = vtk.vtkRenderWindow()
    renWin.AddRenderer(ren)
    renWin.SetWindowName('ReadSTL')
    iren = vtk.vtkRenderWindowInteractor()
    iren.SetRenderWindow(renWin)
    iren.Initialize()
    renWin.Render()
    iren.Start()

reader = vtk.vtkOBJReader()
reader.SetFileName('XXX.obj') #读取一个obj模型
reader.Update()
mesh = reader.GetOutput()

polydata_visualize(mesh) #网格模型可视化

二、将三维网格模型转换成点云

用点云分割网络来进行三维模型分割，输入数据一定是点云，所以要在三维模型上获得点云，其实就是网格顶点

def point_visualize(point_cloud):
    numPts = point_cloud.shape[0]
    x = point_cloud[:, 0]
    y = point_cloud[:, 1]
    z = point_cloud[:, 2]
    # labels = point_cloud[:, -1]

    # setup points and vertices
    Points = vtk.vtkPoints()
    Vertices = vtk.vtkCellArray()

    # setup colors
    colors = vtk.vtkNamedColors()
    Colors = vtk.vtkUnsignedCharArray()
    Colors.SetNumberOfComponents(3)
    Colors.SetName("Colors")

    for i in range(numPts):
        id = Points.InsertNextPoint(x[i], y[i], z[i])
        Vertices.InsertNextCell(1)
        Vertices.InsertCellPoint(id)
        Colors.InsertNextTuple3(*colors.GetColor3ub("Red"))

    poly = vtk.vtkPolyData()
    poly.SetPoints(Points)
    poly.SetVerts(Vertices)

    poly.GetPointData().SetScalars(Colors)
    poly.Modified()

    polydata_visualize(poly)
points = vtk_to_numpy(mesh.GetPoints().GetData()) # 从三维网格模型上获取顶点
point_visualize(points)  # 对点云进行可视化

这个点云数据被存储在一个维度为N×3的numpy矩阵中，N表示点数，3为三维坐标。

三、根据分割结果对点云进行上色

我们将从三维模型上得到的点云数据传入某个分割网络进行分割后，会得到每个点的一个预测值，我称它为pred，接下来每个点将根据对应的pred进行上色。这里我把pred保存在一个txt文件当中。

# 设置一个颜色表,每个标签对应一种颜色，颜色通过RGB值表示
dt_Color = {11: [190, 190, 190], 12: [0, 0, 255], 13: [0, 255, 0], 14: [255, 255, 0], 15: [255, 0, 0],
            16: [0, 0, 0], 17: [255, 181, 197], 18: [255, 165, 0], 21: [83, 134, 139], 22: [155, 48, 255],
            23: [255, 250, 205], 24: [135, 206, 255], 25: [105, 105, 105], 26: [139, 0, 0], 27: [0, 139, 0],
            28: [255, 255, 224],
            31: [190, 190, 190], 32: [0, 0, 255], 33: [0, 255, 0], 34: [255, 255, 0], 35: [255, 0, 0],
            36: [0, 0, 0], 37: [255, 181, 197], 38: [255, 165, 0], 41: [83, 134, 139], 42: [155, 48, 255],
            43: [255, 250, 205], 44: [135, 206, 255], 45: [105, 105, 105], 46: [139, 0, 0], 47: [0, 139, 0],
            48: [255, 255, 224], 49: [135, 206, 255]}
def point_visualize2(point_cloud, pred):
    numPts = point_cloud.shape[0]
    x = point_cloud[:, 0]
    y = point_cloud[:, 1]
    z = point_cloud[:, 2]


    # setup points and vertices
    Points = vtk.vtkPoints()
    Vertices = vtk.vtkCellArray()

    # setup colors 创建一个名为Colors的标量,用这标量来存储每个点的颜色
    colors = vtk.vtkNamedColors()
    Colors = vtk.vtkUnsignedCharArray()
    Colors.SetNumberOfComponents(3)
    Colors.SetName("Colors")

    for i in range(numPts):
        id = Points.InsertNextPoint(x[i], y[i], z[i])
        Vertices.InsertNextCell(1)
        Vertices.InsertCellPoint(id)
        color = dt_Color[int(pred[i])]
        Colors.InsertNextTuple3(color[0], color[1], color[2])

    poly = vtk.vtkPolyData()
    poly.SetPoints(Points)
    poly.SetVerts(Vertices)

    poly.GetPointData().SetScalars(Colors)
    poly.Modified()

    polydata_visualize(poly)

pred = np.loadtxt('lables.txt') # 读取分割结果

point_visualize2(points, pred) # 进行可视化

四、对三维模型进行上色

Colors = vtk.vtkUnsignedCharArray()
Colors.SetNumberOfComponents(3)
Colors.SetName("Colors")

for p in pred:
    color = dt_Color[int(p)]
    Colors.InsertNextTuple3(color[0], color[1], color[2])

mesh.GetPointData().SetScalars(Colors)
mesh.Modified()
polydata_visualize(mesh)

完整代码：

import vtk
from vtk.util.numpy_support import vtk_to_numpy
import numpy as np


def polydata_visualize(poly):
    colors = vtk.vtkNamedColors()
    mapper = vtk.vtkPolyDataMapper()
    mapper.SetInputData(poly)
    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
    ren = vtk.vtkRenderer()
    ren.AddActor(actor)
    ren.SetBackground(colors.GetColor3d('White'))
    renWin = vtk.vtkRenderWindow()
    renWin.AddRenderer(ren)
    renWin.SetWindowName('ReadSTL')
    iren = vtk.vtkRenderWindowInteractor()
    iren.SetRenderWindow(renWin)
    iren.Initialize()
    renWin.Render()
    iren.Start()

def point_visualize(point_cloud):
    numPts = point_cloud.shape[0]
    x = point_cloud[:, 0]
    y = point_cloud[:, 1]
    z = point_cloud[:, 2]
    # labels = point_cloud[:, -1]

    # setup points and vertices
    Points = vtk.vtkPoints()
    Vertices = vtk.vtkCellArray()

    # setup colors
    colors = vtk.vtkNamedColors()
    Colors = vtk.vtkUnsignedCharArray()
    Colors.SetNumberOfComponents(3)
    Colors.SetName("Colors")

    for i in range(numPts):
        id = Points.InsertNextPoint(x[i], y[i], z[i])
        Vertices.InsertNextCell(1)
        Vertices.InsertCellPoint(id)
        Colors.InsertNextTuple3(*colors.GetColor3ub("Red"))

    poly = vtk.vtkPolyData()
    poly.SetPoints(Points)
    poly.SetVerts(Vertices)

    poly.GetPointData().SetScalars(Colors)
    poly.Modified()

    polydata_visualize(poly)

def point_visualize2(point_cloud, pred):
    numPts = point_cloud.shape[0]
    x = point_cloud[:, 0]
    y = point_cloud[:, 1]
    z = point_cloud[:, 2]


    # setup points and vertices
    Points = vtk.vtkPoints()
    Vertices = vtk.vtkCellArray()

    # setup colors 创建一个名为Colors的标量,用这标量来存储每个点的颜色
    colors = vtk.vtkNamedColors()
    Colors = vtk.vtkUnsignedCharArray()
    Colors.SetNumberOfComponents(3)
    Colors.SetName("Colors")

    for i in range(numPts):
        id = Points.InsertNextPoint(x[i], y[i], z[i])
        Vertices.InsertNextCell(1)
        Vertices.InsertCellPoint(id)
        color = dt_Color[int(pred[i])]
        Colors.InsertNextTuple3(color[0], color[1], color[2])

    poly = vtk.vtkPolyData()
    poly.SetPoints(Points)
    poly.SetVerts(Vertices)

    poly.GetPointData().SetScalars(Colors)
    poly.Modified()

    polydata_visualize(poly)


# 一、读取三维网格模型
reader = vtk.vtkOBJReader()
reader.SetFileName('xxx.obj')
reader.Update()
mesh = reader.GetOutput()
polydata_visualize(mesh) # 可视化

# 二、将三维网格模型转换成点云
points = vtk_to_numpy(mesh.GetPoints().GetData())
point_visualize(points)  # 对点云进行可视化



# 三、根据分割结果对点云进行上色
# 设置一个颜色表,每个标签对应一种颜色，颜色通过RGB值表示
dt_Color = {11: [190, 190, 190], 12: [0, 0, 255], 13: [0, 255, 0], 14: [255, 255, 0], 15: [255, 0, 0],
            16: [0, 0, 0], 17: [255, 181, 197], 18: [255, 165, 0], 21: [83, 134, 139], 22: [155, 48, 255],
            23: [255, 250, 205], 24: [135, 206, 255], 25: [105, 105, 105], 26: [139, 0, 0], 27: [0, 139, 0],
            28: [255, 255, 224],
            31: [190, 190, 190], 32: [0, 0, 255], 33: [0, 255, 0], 34: [255, 255, 0], 35: [255, 0, 0],
            36: [0, 0, 0], 37: [255, 181, 197], 38: [255, 165, 0], 41: [83, 134, 139], 42: [155, 48, 255],
            43: [255, 250, 205], 44: [135, 206, 255], 45: [105, 105, 105], 46: [139, 0, 0], 47: [0, 139, 0],
            48: [255, 255, 224], 49: [135, 206, 255]}


pred = np.loadtxt('pred.txt') # 读取分割结果,每个点对应一个预测值

point_visualize2(points, pred) # 进行可视化


# 四、对三维模型进行上色
Colors = vtk.vtkUnsignedCharArray()
Colors.SetNumberOfComponents(3)
Colors.SetName("Colors")

for p in pred:
    color = dt_Color[int(p)]
    Colors.InsertNextTuple3(color[0], color[1], color[2])

mesh.GetPointData().SetScalars(Colors)
mesh.Modified()
polydata_visualize(mesh)