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()
|