112 lines
3.7 KiB
Python
112 lines
3.7 KiB
Python
"""点云加载模块"""
|
||
import numpy as np
|
||
from pathlib import Path
|
||
import struct
|
||
|
||
|
||
class PointCloudLoader:
|
||
"""加载点云文件"""
|
||
|
||
@staticmethod
|
||
def load_ply(file_path: str) -> np.ndarray:
|
||
"""
|
||
加载PLY格式点云
|
||
|
||
Args:
|
||
file_path: PLY文件路径
|
||
|
||
Returns:
|
||
np.ndarray: 点云数组 (N, 3)
|
||
"""
|
||
path = Path(file_path)
|
||
if not path.exists():
|
||
raise FileNotFoundError(f"点云文件不存在: {file_path}")
|
||
|
||
print(f"正在加载点云: {file_path}")
|
||
|
||
with open(path, 'rb') as f:
|
||
# 读取头部
|
||
header_lines = []
|
||
while True:
|
||
line = f.readline().decode('ascii').strip()
|
||
header_lines.append(line)
|
||
if line == 'end_header':
|
||
break
|
||
|
||
# 解析头部
|
||
vertex_count = 0
|
||
format_type = None
|
||
properties = []
|
||
|
||
for line in header_lines:
|
||
if line.startswith('element vertex'):
|
||
vertex_count = int(line.split()[-1])
|
||
elif line.startswith('format'):
|
||
format_type = line.split()[1]
|
||
elif line.startswith('property'):
|
||
parts = line.split()
|
||
properties.append({
|
||
'type': parts[1],
|
||
'name': parts[2]
|
||
})
|
||
|
||
if vertex_count == 0:
|
||
raise ValueError("PLY文件中没有顶点数据")
|
||
|
||
# 读取顶点数据
|
||
if format_type == 'binary_little_endian':
|
||
points = PointCloudLoader._read_binary_ply(f, vertex_count, properties)
|
||
elif format_type == 'ascii':
|
||
points = PointCloudLoader._read_ascii_ply(f, vertex_count)
|
||
else:
|
||
raise ValueError(f"不支持的PLY格式: {format_type}")
|
||
|
||
print(f" 加载完成: {len(points)} 个点")
|
||
return points
|
||
|
||
@staticmethod
|
||
def _read_binary_ply(f, vertex_count: int, properties: list) -> np.ndarray:
|
||
"""读取二进制PLY数据（优化版：批量读取）"""
|
||
# 构建numpy dtype
|
||
type_map = {
|
||
'float': np.float32,
|
||
'double': np.float64,
|
||
'uchar': np.uint8,
|
||
'uint': np.uint32,
|
||
'int': np.int32
|
||
}
|
||
|
||
# 构建结构化数组的dtype
|
||
dtype_list = []
|
||
for prop in properties:
|
||
dtype_list.append((prop['name'], type_map.get(prop['type'], np.float32)))
|
||
|
||
dtype = np.dtype(dtype_list)
|
||
|
||
# 一次性读取所有数据
|
||
raw_data = np.fromfile(f, dtype=dtype, count=vertex_count)
|
||
|
||
# 提取xyz坐标（假设前3个属性是x, y, z）
|
||
# 如果属性名不是x,y,z，可以通过索引访问
|
||
if 'x' in raw_data.dtype.names and 'y' in raw_data.dtype.names and 'z' in raw_data.dtype.names:
|
||
points = np.column_stack([raw_data['x'], raw_data['y'], raw_data['z']])
|
||
else:
|
||
# 假设前3个字段是xyz
|
||
points = np.column_stack([raw_data[raw_data.dtype.names[0]],
|
||
raw_data[raw_data.dtype.names[1]],
|
||
raw_data[raw_data.dtype.names[2]]])
|
||
|
||
return points.astype(np.float64)
|
||
|
||
@staticmethod
|
||
def _read_ascii_ply(f, vertex_count: int) -> np.ndarray:
|
||
"""读取ASCII PLY数据（优化版：批量读取）"""
|
||
# 读取剩余所有行
|
||
remaining = f.read().decode('ascii')
|
||
|
||
# 使用numpy的快速文本解析
|
||
from io import StringIO
|
||
points = np.loadtxt(StringIO(remaining), dtype=np.float64, max_rows=vertex_count, usecols=(0, 1, 2))
|
||
|
||
return points
|