如何将vtk文件转换为numpy.array



我用下面的程序对STL文件进行体素化,并将其保存为vtk文件。

代码本身很好,因为下面的程序运行得很好。但是,我想将其转换为numpy.array格式,并将其保存为npz格式。使用vtk_to_numy不起作用,我该怎么办?

最终目标是将其转换为numpy.array格式,并用于三维深度学习。

例如。以333体素数据为例,

我想生成形状为(3,3,3(的np.array,如下

[
[[0,0,0],
[0,1,0],
[0,0,0]]
[[1,1,1],
[1,1,1],
[0,0,0]]
[[1,1,1],
[0,0,0],
[0,0,0]]
]
import vtk
import time
import numpy as np
import vtk.util.numpy_support as vnp
######## data ########
filename_in = "<your STL>.stl"
filename_out = "out.vtk"
mesh_size = 100
tol = 1e-7
cubicORrect = "rect"
##################################
start = time.time()
reader = vtk.vtkSTLReader()
reader.SetFileName(filename_in)
reader.Update()
closed_poly = reader.GetOutput()
# x_min:0 x_max:1, y_min:2,y_max:3,z_min:4,z_max:5
bounds = closed_poly.GetBounds()
max_size = max([bounds[1] - bounds[0], bounds[3] -
bounds[2], bounds[5] - bounds[4]])
cell_dims = [mesh_size, mesh_size, mesh_size]  # x, y, z
if cubicORrect == "cubic":
mesh_pitch = [max_size/cell_dims[0],
max_size/cell_dims[1],
max_size/cell_dims[2]]
else:
mesh_pitch = [(bounds[1] - bounds[0])/cell_dims[0],
(bounds[3] - bounds[2])/cell_dims[1],
(bounds[5] - bounds[4])/cell_dims[2]]
mins = [bounds[0], bounds[2], bounds[4]]

px, py, pz = mesh_pitch
mx, my, mz = (cell_dims+np.array([1, 1, 1])) * mesh_pitch  # max
points = vtk.vtkPoints()
coords = np.stack(np.mgrid[:mx:px, :my:py, :mz:pz], -1).reshape(-1, 3) + mins
points.SetData(vnp.numpy_to_vtk(coords))
structured_base_mesh = vtk.vtkStructuredGrid()
structured_base_mesh.SetExtent(
0, cell_dims[0], 0, cell_dims[1], 0, cell_dims[2])
structured_base_mesh.SetPoints(points)
append = vtk.vtkAppendFilter()
append.AddInputData(structured_base_mesh)
append.Update()
base_mesh = append.GetOutput()
cell_centers = vtk.vtkCellCenters()
cell_centers.SetInputData(base_mesh)
cell_centers.Update()
poly_points = cell_centers.GetOutput()
select_enclosed = vtk.vtkSelectEnclosedPoints()
select_enclosed.SetInputData(poly_points)
select_enclosed.SetSurfaceData(closed_poly)
select_enclosed.SetTolerance(tol)
select_enclosed.Update()
isInsideOrOutside = select_enclosed.GetOutput(
).GetPointData().GetArray("SelectedPoints")
structured_base_mesh.GetCellData().AddArray(isInsideOrOutside)
threshold = vtk.vtkThreshold()
threshold.SetInputArrayToProcess(
0, 0, 0, vtk.vtkDataObject.FIELD_ASSOCIATION_CELLS, "SelectedPoints")
threshold.SetInputData(structured_base_mesh)
threshold.ThresholdBetween(0, 1)
threshold.Update()
writer = vtk.vtkDataSetWriter()
writer.SetFileName(filename_out)
writer.SetInputData(threshold.GetOutput())
writer.Update()
nparray = vnp.vtk_to_numpy(threshold.GetOutput().GetCellData().GetArray("SelectedPoints"))

创建3d形状numpy数组的步骤:

  • 获取所需的vtkDataArray。类似threshold.GetOutput().GetPointData().GetArary(<arrayname>)
  • 使用CCD_ 2将其转换为numpy数组。这将返回一个"平面"数组:所有值的1D列表
  • 根据网格尺寸重塑(threshold.GetOutput().GetDimensions()(

从这个答案中得到启发https://discourse.vtk.org/t/convert-vtk-array-to-numpy-array/3152/4

import vtk
import time
import numpy as np
import vtk.util.numpy_support as vnp
######## data ########
filename_in = "<your STL>.stl"
filename_out = "out.vtk"
mesh_size = 100
tol = 1e-7
cubicORrect = "rect"
##################################
start = time.time()
reader = vtk.vtkSTLReader()
reader.SetFileName(filename_in)
reader.Update()
closed_poly = reader.GetOutput()
# x_min:0 x_max:1, y_min:2,y_max:3,z_min:4,z_max:5
bounds = closed_poly.GetBounds()
max_size = max([bounds[1] - bounds[0], bounds[3] -
bounds[2], bounds[5] - bounds[4]])
cell_dims = [mesh_size, mesh_size, mesh_size]  # x, y, z
if cubicORrect == "cubic":
mesh_pitch = [max_size/cell_dims[0],
max_size/cell_dims[1],
max_size/cell_dims[2]]
else:
mesh_pitch = [(bounds[1] - bounds[0])/cell_dims[0],
(bounds[3] - bounds[2])/cell_dims[1],
(bounds[5] - bounds[4])/cell_dims[2]]
mins = [bounds[0], bounds[2], bounds[4]]

px, py, pz = mesh_pitch
mx, my, mz = (cell_dims+np.array([1, 1, 1])) * mesh_pitch  # max
points = vtk.vtkPoints()
coords = np.stack(np.mgrid[:mx:px, :my:py, :mz:pz], -1).reshape(-1, 3) + mins
points.SetData(vnp.numpy_to_vtk(coords))
structured_base_mesh = vtk.vtkStructuredGrid()
structured_base_mesh.SetExtent(
0, cell_dims[0], 0, cell_dims[1], 0, cell_dims[2])
structured_base_mesh.SetPoints(points)
append = vtk.vtkAppendFilter()
append.AddInputData(structured_base_mesh)
append.Update()
base_mesh = append.GetOutput()
cell_centers = vtk.vtkCellCenters()
cell_centers.SetInputData(base_mesh)
cell_centers.Update()
poly_points = cell_centers.GetOutput()
select_enclosed = vtk.vtkSelectEnclosedPoints()
select_enclosed.SetInputData(poly_points)
select_enclosed.SetSurfaceData(closed_poly)
select_enclosed.SetTolerance(tol)
select_enclosed.Update()
isInsideOrOutside = select_enclosed.GetOutput(
).GetPointData().GetArray("SelectedPoints")
structured_base_mesh.GetCellData().AddArray(isInsideOrOutside)
threshold = vtk.vtkThreshold()
threshold.SetInputArrayToProcess(
0, 0, 0, vtk.vtkDataObject.FIELD_ASSOCIATION_CELLS, "SelectedPoints")
threshold.SetInputData(structured_base_mesh)
threshold.ThresholdBetween(0, 1)
threshold.Update()
writer = vtk.vtkDataSetWriter()
writer.SetFileName(filename_out)
writer.SetInputData(threshold.GetOutput())
writer.Update()
nparray = vnp.vtk_to_numpy(threshold.GetOutput().GetCellData().GetArray("SelectedPoints"))

最新更新