158 lines
4.6 KiB
Python
158 lines
4.6 KiB
Python
|
import json
|
|||
|
import vtk
|
|||
|
import numpy as np
|
|||
|
import geojson
|
|||
|
from pygltflib import GLTF2, Buffer, Accessor, Mesh, Primitive, BufferView
|
|||
|
# # Create a grid
|
|||
|
# grid_size = 50
|
|||
|
# x = np.linspace(0, 10, grid_size)
|
|||
|
# y = np.linspace(0, 10, grid_size)
|
|||
|
# z = np.linspace(0, 10, grid_size)
|
|||
|
# x, y, z = np.meshgrid(x, y, z)
|
|||
|
|
|||
|
# # Create scalar values based on a Gaussian distribution
|
|||
|
# values = np.exp(-((x - 5)**2 + (y - 5)**2 + (z - 5)**2) / 10)
|
|||
|
|
|||
|
# # Convert grid points and values to VTK format
|
|||
|
# points = vtk.vtkPoints()
|
|||
|
# scalars = vtk.vtkDoubleArray()
|
|||
|
# scalars.SetName("HeatValues")
|
|||
|
|
|||
|
# for i in range(grid_size):
|
|||
|
# for j in range(grid_size):
|
|||
|
# for k in range(grid_size):
|
|||
|
# points.InsertNextPoint(x[i, j, k], y[i, j, k], z[i, j, k])
|
|||
|
# scalars.InsertNextValue(values[i, j, k])
|
|||
|
|
|||
|
# # Create a structured grid
|
|||
|
# structured_grid = vtk.vtkStructuredGrid()
|
|||
|
# structured_grid.SetDimensions(grid_size, grid_size, grid_size)
|
|||
|
# structured_grid.SetPoints(points)
|
|||
|
# structured_grid.GetPointData().SetScalars(scalars)
|
|||
|
|
|||
|
# # Write the grid to a VTK file
|
|||
|
# vtk_writer = vtk.vtkStructuredGridWriter()
|
|||
|
# vtk_writer.SetFileName("heatmap.vtk")
|
|||
|
# vtk_writer.SetInputData(structured_grid)
|
|||
|
# vtk_writer.Write()
|
|||
|
|
|||
|
|
|||
|
|
|||
|
def load_vtk(vtk_file):
|
|||
|
# 使用 vtkStructuredGridReader 来读取 structured_grid 数据类型
|
|||
|
reader = vtk.vtkStructuredGridReader()
|
|||
|
reader.SetFileName(vtk_file)
|
|||
|
reader.Update()
|
|||
|
|
|||
|
# 获取数据集
|
|||
|
dataset = reader.GetOutput()
|
|||
|
|
|||
|
# 提取点数据
|
|||
|
points = dataset.GetPoints()
|
|||
|
point_data = points.GetData()
|
|||
|
|
|||
|
# 由于结构化网格的数据是按规则网格组织的,我们直接获取网格尺寸
|
|||
|
dimensions = dataset.GetDimensions()
|
|||
|
|
|||
|
return point_data, dimensions # 返回维度而不是 cells
|
|||
|
|
|||
|
|
|||
|
def vtk_to_geojson(point_data, dimensions):
|
|||
|
# 转换 VTK 点数据
|
|||
|
points = []
|
|||
|
for i in range(point_data.GetNumberOfTuples()):
|
|||
|
points.append(list(point_data.GetTuple3(i)))
|
|||
|
|
|||
|
# 这里我们不直接使用 cells,而是基于维度构建网格的多边形
|
|||
|
features = []
|
|||
|
x_dim, y_dim, z_dim = dimensions
|
|||
|
for x in range(x_dim - 1):
|
|||
|
for y in range(y_dim - 1):
|
|||
|
for z in range(z_dim - 1):
|
|||
|
# 获取每个方格的四个角
|
|||
|
pts = [
|
|||
|
points[x + y * x_dim + z * x_dim * y_dim],
|
|||
|
points[(x + 1) + y * x_dim + z * x_dim * y_dim],
|
|||
|
points[(x + 1) + (y + 1) * x_dim + z * x_dim * y_dim],
|
|||
|
points[x + (y + 1) * x_dim + z * x_dim * y_dim]
|
|||
|
]
|
|||
|
# 创建多边形面
|
|||
|
features.append({
|
|||
|
"type": "Feature",
|
|||
|
"geometry": {
|
|||
|
"type": "Polygon",
|
|||
|
"coordinates": [pts]
|
|||
|
},
|
|||
|
"properties": {}
|
|||
|
})
|
|||
|
|
|||
|
# 构建 GeoJSON 数据
|
|||
|
geojson_data = geojson.FeatureCollection(features)
|
|||
|
return geojson_data
|
|||
|
|
|||
|
|
|||
|
def vtk_to_gltf(point_data, dimensions):
|
|||
|
# 提取点数据
|
|||
|
vertices = []
|
|||
|
for i in range(point_data.GetNumberOfTuples()):
|
|||
|
vertices.append(list(point_data.GetTuple3(i)))
|
|||
|
|
|||
|
# 创建 GLTF 数据结构
|
|||
|
gltf = GLTF2()
|
|||
|
|
|||
|
# 将点数据转换为 numpy 数组
|
|||
|
vertices = np.array(vertices, dtype=np.float32).flatten()
|
|||
|
|
|||
|
# 创建缓冲区
|
|||
|
buffer = Buffer(uri="data.bin", byteLength=len(vertices) * 4)
|
|||
|
gltf.buffers.append(buffer)
|
|||
|
|
|||
|
# 创建缓冲区视图
|
|||
|
bufferViewVertices = BufferView(
|
|||
|
buffer=0,
|
|||
|
byteOffset=0,
|
|||
|
byteLength=len(vertices) * 4,
|
|||
|
target=34962
|
|||
|
)
|
|||
|
gltf.bufferViews.append(bufferViewVertices)
|
|||
|
|
|||
|
# 创建访问器
|
|||
|
accessorVertices = Accessor(
|
|||
|
bufferView=0,
|
|||
|
byteOffset=0,
|
|||
|
componentType=5126,
|
|||
|
count=len(vertices) // 3,
|
|||
|
type="VEC3"
|
|||
|
)
|
|||
|
gltf.accessors.append(accessorVertices)
|
|||
|
|
|||
|
# 返回 GLTF 数据
|
|||
|
return gltf
|
|||
|
|
|||
|
|
|||
|
def save_geojson(geojson_data, filename):
|
|||
|
with open(filename, 'w') as f:
|
|||
|
json.dump(geojson_data, f)
|
|||
|
|
|||
|
|
|||
|
def save_gltf(gltf, filename):
|
|||
|
gltf.save(filename)
|
|||
|
|
|||
|
|
|||
|
def main():
|
|||
|
# 读取 VTK 文件
|
|||
|
vtk_file = 'heatmap.vtk'
|
|||
|
point_data, dimensions = load_vtk(vtk_file)
|
|||
|
|
|||
|
# 转换为 GeoJSON
|
|||
|
geojson_data = vtk_to_geojson(point_data, dimensions)
|
|||
|
save_geojson(geojson_data, 'output.geojson')
|
|||
|
|
|||
|
# 转换为 GLTF
|
|||
|
gltf = vtk_to_gltf(point_data, dimensions)
|
|||
|
save_gltf(gltf, 'output.gltf')
|
|||
|
|
|||
|
|
|||
|
if __name__ == '__main__':
|
|||
|
main()
|