在三维模型和点云分割的实验中,通常都需要对分割结果进行可视化,不同的区域用不同的颜色,每个点的颜色要根据分割结果来确定,例如下面这张图就是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)