406 lines
16 KiB
Python
406 lines
16 KiB
Python
import numpy as np
|
||
import open3d as o3d
|
||
import cv2
|
||
import time
|
||
|
||
import time
|
||
from functools import wraps
|
||
|
||
def print_runtime(func):
|
||
@wraps(func)
|
||
def wrapper(*args, **kwargs):
|
||
start = time.time() # 记录开始时间
|
||
result = func(*args, **kwargs) # 执行函数
|
||
end = time.time() # 记录结束时间
|
||
print(f"{func.__name__} 运行耗时: {end - start:.3f}秒") # 打印耗时
|
||
return result
|
||
return wrapper
|
||
|
||
class SimpleNormalVector:
|
||
def __init__(self, depth_path, rgb_path, intrinsics):
|
||
"""
|
||
初始化类，设置深度图像文件路径和相机内参。
|
||
将耗时的点云生成和法向量计算提前到初始化阶段完成。
|
||
"""
|
||
self.depth_path = depth_path
|
||
self.rgb_path = rgb_path
|
||
self.intrinsics = intrinsics
|
||
|
||
# 预先读取深度图像
|
||
self.depth_image = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
|
||
|
||
# 在初始化时就创建RGBD图像和点云
|
||
self.rgbd_image = self.create_rgbd_image()
|
||
self.pcd, self.points, self.normals = self.generate_point_cloud(self.rgbd_image)
|
||
|
||
# 创建KD树用于快速最近邻搜索
|
||
self.pcd_tree = o3d.geometry.KDTreeFlann(self.pcd)
|
||
|
||
def create_rgbd_image(self):
|
||
"""
|
||
创建RGBD图像，使用深度图像和RGB图像
|
||
"""
|
||
color_raw = o3d.io.read_image(self.rgb_path)
|
||
depth_raw = o3d.io.read_image(self.depth_path)
|
||
|
||
# 将深度图像转换为 NumPy 数组
|
||
depth_image = np.asarray(depth_raw)
|
||
|
||
# 创建掩膜：深度值为0的区域
|
||
mask = (depth_image == 0).astype(np.uint8)
|
||
|
||
# 使用 OpenCV 的 inpaint 函数进行深度修复
|
||
depth_image_inpainted = cv2.inpaint(depth_image, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
|
||
|
||
# 将修复后的深度图像转换回 Open3D 图像
|
||
depth_fixed = o3d.geometry.Image(depth_image_inpainted)
|
||
|
||
# 创建RGBD图像
|
||
rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
|
||
color_raw, depth_fixed, convert_rgb_to_intensity=False)
|
||
|
||
return rgbd_image
|
||
|
||
def generate_point_cloud(self, rgbd_image):
|
||
"""
|
||
从RGBD图像生成点云
|
||
返回点云对象、点坐标数组和法向量数组
|
||
"""
|
||
# 使用相机内参生成点云
|
||
fx = self.intrinsics['fx']
|
||
fy = self.intrinsics['fy']
|
||
cx = self.intrinsics['cx']
|
||
cy = self.intrinsics['cy']
|
||
intrinsic = o3d.camera.PinholeCameraIntrinsic()
|
||
intrinsic.set_intrinsics(640, 480, fx, fy, cx, cy)
|
||
|
||
pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, intrinsic)
|
||
pcd = pcd.voxel_down_sample(voxel_size=0.01) # 降采样
|
||
|
||
# 估算点云的法向量
|
||
pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
|
||
|
||
# 获取点云的点和法向量
|
||
points = np.asarray(pcd.points)
|
||
normals = np.asarray(pcd.normals)
|
||
|
||
return pcd, points, normals
|
||
|
||
def get_normal_at_point(self, target_point):
|
||
"""
|
||
使用KD树快速查找最近点，并获取其法向量
|
||
"""
|
||
# 使用KD树查找最近的点
|
||
_, idx, _ = self.pcd_tree.search_knn_vector_3d(target_point, 1)
|
||
|
||
# 获取最近点的法向量
|
||
normal = self.normals[idx[0]]
|
||
|
||
return {'coordinates': target_point, 'normal': normal}
|
||
|
||
def calculate_normals_batch(self, target_points):
|
||
"""
|
||
批量计算多个目标点的法向量，使用向量化操作减少循环次数
|
||
"""
|
||
n_points = len(target_points)
|
||
# 预分配结果数组
|
||
normals_out = np.zeros((n_points, 3))
|
||
indices = np.zeros(n_points, dtype=np.int32)
|
||
|
||
# 批量查找最近点的索引
|
||
for i, point in enumerate(target_points):
|
||
_, idx, _ = self.pcd_tree.search_knn_vector_3d(point, 1)
|
||
indices[i] = idx[0]
|
||
|
||
# 使用索引批量获取法向量
|
||
normals_out = self.normals[indices]
|
||
|
||
# 构建结果字典列表
|
||
results = []
|
||
for i, point in enumerate(target_points):
|
||
results.append({
|
||
'coordinates': point,
|
||
'normal': normals_out[i]
|
||
})
|
||
|
||
return results
|
||
|
||
def calculate_normals_vectorized(self, target_points):
|
||
"""
|
||
使用向量化操作计算法向量，适用于大量点的情况
|
||
"""
|
||
# 转换为numpy数组确保向量化操作
|
||
points_array = np.array(target_points)
|
||
|
||
# 为每个目标点找到点云中最近的点
|
||
# 计算所有目标点到所有点云点的距离矩阵
|
||
# 这会消耗大量内存，但速度很快
|
||
# target_points: (n_targets, 3), self.points: (n_cloud, 3)
|
||
# 使用广播计算距离矩阵: (n_targets, n_cloud)
|
||
n_targets = points_array.shape[0]
|
||
n_cloud = self.points.shape[0]
|
||
|
||
# 对于大型点云，分批处理以避免内存溢出
|
||
batch_size = 1000 # 每批处理的目标点数
|
||
closest_indices = np.zeros(n_targets, dtype=np.int32)
|
||
|
||
for batch_start in range(0, n_targets, batch_size):
|
||
batch_end = min(batch_start + batch_size, n_targets)
|
||
batch_points = points_array[batch_start:batch_end]
|
||
|
||
# 计算这批点到所有点云点的距离
|
||
# 展开 (a-b)^2 = a^2 - 2ab + b^2 以避免大矩阵
|
||
batch_sq_dist = np.sum(batch_points**2, axis=1, keepdims=True)
|
||
cloud_sq_dist = np.sum(self.points**2, axis=1)
|
||
dot_product = np.dot(batch_points, self.points.T)
|
||
|
||
# 计算欧氏距离的平方
|
||
distances = batch_sq_dist - 2 * dot_product + cloud_sq_dist
|
||
|
||
# 找到每行（每个目标点）的最小距离索引
|
||
min_indices = np.argmin(distances, axis=1)
|
||
closest_indices[batch_start:batch_end] = min_indices
|
||
|
||
# 使用最近点的索引获取对应的法向量
|
||
normals_out = self.normals[closest_indices]
|
||
|
||
# 构建结果
|
||
results = []
|
||
for i in range(n_targets):
|
||
results.append({
|
||
'coordinates': target_points[i],
|
||
'normal': normals_out[i]
|
||
})
|
||
|
||
return results
|
||
|
||
def calculate_normals(self, target_points):
|
||
"""
|
||
计算多个目标点的法向量，根据点数量自动选择最佳方法
|
||
"""
|
||
# if len(target_points) < 10:
|
||
# # 少量点时用循环，避免向量化的开销
|
||
# normals_data = []
|
||
# for point in target_points:
|
||
# normal_data = self.get_normal_at_point(point)
|
||
# normals_data.append(normal_data)
|
||
# return normals_data
|
||
# elif len(target_points) < 500:
|
||
# # 中等数量点时，使用批量查询但保留KD树
|
||
# return self.calculate_normals_batch(target_points)
|
||
# else:
|
||
# # 大量点时，使用完全向量化的方法
|
||
# return self.calculate_normals_vectorized(target_points)
|
||
return self.calculate_normals_batch(target_points)
|
||
|
||
def pixel_to_camera_point(self, pixel_x, pixel_y, depth_value=None):
|
||
"""
|
||
将像素坐标转换为相机坐标系中的3D点
|
||
pixel_x, pixel_y: 像素坐标
|
||
depth_value: 可选，指定的深度值。如果为None，则从深度图中获取
|
||
"""
|
||
if depth_value is None:
|
||
# 确保像素坐标在图像范围内
|
||
if pixel_x < 0 or pixel_x >= self.depth_image.shape[1] or pixel_y < 0 or pixel_y >= self.depth_image.shape[0]:
|
||
return None
|
||
# 获取该像素的深度值（单位：毫米）
|
||
depth_value = self.depth_image[pixel_y, pixel_x]
|
||
|
||
# 相机内参
|
||
fx = self.intrinsics['fx']
|
||
fy = self.intrinsics['fy']
|
||
cx = self.intrinsics['cx']
|
||
cy = self.intrinsics['cy']
|
||
|
||
# 计算相机坐标系中的3D点（单位：米）
|
||
x = (pixel_x - cx) * depth_value / fx / 1000.0
|
||
y = (pixel_y - cy) * depth_value / fy / 1000.0
|
||
z = depth_value / 1000.0
|
||
|
||
return np.array([x, y, z])
|
||
|
||
def get_normal_at_pixel(self, pixel_x, pixel_y, depth_value=None):
|
||
"""
|
||
获取特定像素位置的法向量
|
||
"""
|
||
# 将像素坐标转换为相机坐标系下的3D点
|
||
camera_point = self.pixel_to_camera_point(pixel_x, pixel_y, depth_value)
|
||
if camera_point is None:
|
||
return None
|
||
|
||
# 计算该点的法向量
|
||
return self.get_normal_at_point(camera_point)
|
||
|
||
def pixels_to_camera_points(self, pixel_coords):
|
||
"""
|
||
批量将像素坐标转换为相机坐标系中的3D点
|
||
pixel_coords: 像素坐标数组，形状为(n, 2)，每行为(x, y)
|
||
返回：相机坐标系中的3D点数组，形状为(n, 3)
|
||
"""
|
||
# 预分配结果数组
|
||
n_pixels = len(pixel_coords)
|
||
camera_points = np.zeros((n_pixels, 3))
|
||
|
||
# 获取所有像素点的深度值
|
||
depth_values = np.zeros(n_pixels)
|
||
for i, (px, py) in enumerate(pixel_coords):
|
||
if 0 <= px < self.depth_image.shape[1] and 0 <= py < self.depth_image.shape[0]:
|
||
depth_values[i] = self.depth_image[py, px]
|
||
|
||
# 相机内参
|
||
fx = self.intrinsics['fx']
|
||
fy = self.intrinsics['fy']
|
||
cx = self.intrinsics['cx']
|
||
cy = self.intrinsics['cy']
|
||
|
||
# 批量计算相机坐标
|
||
camera_points[:, 0] = (pixel_coords[:, 0] - cx) * depth_values / fx / 1000.0
|
||
camera_points[:, 1] = (pixel_coords[:, 1] - cy) * depth_values / fy / 1000.0
|
||
camera_points[:, 2] = depth_values / 1000.0
|
||
|
||
return camera_points
|
||
|
||
def get_normals_at_pixels(self, pixel_coords):
|
||
"""
|
||
批量获取多个像素位置的法向量
|
||
pixel_coords: 像素坐标数组，形状为(n, 2)，每行为(x, y)
|
||
"""
|
||
# 将像素坐标批量转换为相机坐标系下的3D点
|
||
camera_points = self.pixels_to_camera_points(pixel_coords)
|
||
|
||
# 计算法向量
|
||
return self.calculate_normals(camera_points)
|
||
|
||
@print_runtime
|
||
def run(self, points_3d_camera_with_labels, visualize=False):
|
||
"""
|
||
与原始vector.py中的run函数兼容的方法
|
||
points_3d_camera_with_labels: 包含label和coordinates的点列表
|
||
visualize: 是否可视化结果（此参数仅为保持接口兼容，在此实现中不起作用）
|
||
返回: 包含label、coordinates和normal的字典列表
|
||
"""
|
||
# 提取目标点坐标列表
|
||
target_points = [point_data['coordinates'] for point_data in points_3d_camera_with_labels]
|
||
|
||
# 使用优化的批量计算法向量
|
||
normals_data = self.calculate_normals(target_points)
|
||
|
||
# 构建与原始输出格式一致的结果
|
||
result = []
|
||
for i, point_data in enumerate(points_3d_camera_with_labels):
|
||
result.append({
|
||
'label': point_data.get('label', ''),
|
||
'coordinates': normals_data[i]['coordinates'],
|
||
'normal': normals_data[i]['normal']
|
||
})
|
||
|
||
return result
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 示例用法
|
||
# 设置相机内参
|
||
intrinsics = {'fx': 453.17746, 'fy': 453.17746, 'cx': 325.943024, 'cy': 243.559982}
|
||
|
||
# 深度图和RGB图路径
|
||
depth_image_path = 'aucpuncture2point/configs/using_img/depth.png'
|
||
rgb_image_path = 'aucpuncture2point/configs/using_img/color.png'
|
||
|
||
print("开始初始化...")
|
||
start_time = time.time()
|
||
normal_vector = SimpleNormalVector(depth_image_path, rgb_image_path, intrinsics)
|
||
init_time = time.time() - start_time
|
||
print(f"初始化耗时: {init_time:.4f} 秒")
|
||
|
||
# 方法1：直接传入图像像素坐标，从深度图获取深度值
|
||
pixel_x, pixel_y = 320, 240 # 图像中心点
|
||
|
||
# 测试多次调用性能
|
||
num_trials = 100
|
||
print(f"\n测试方法1（像素坐标法向量）{num_trials}次调用性能...")
|
||
start_time = time.time()
|
||
for _ in range(num_trials):
|
||
result = normal_vector.get_normal_at_pixel(pixel_x, pixel_y)
|
||
method1_time = time.time() - start_time
|
||
print(f"方法1平均耗时: {method1_time/num_trials:.6f} 秒/次")
|
||
if result:
|
||
print(f"像素坐标 ({pixel_x}, {pixel_y}) 对应的3D点: {result['coordinates']}")
|
||
print(f"该点的法向量: {result['normal']}")
|
||
|
||
# 方法2：已知3D点坐标，直接计算法向量
|
||
# 给定一组3D点（假设已经转换到相机坐标系，单位：米）
|
||
camera_points = [
|
||
np.array([-0.08326775896, -0.15066889276, 0.771]),
|
||
np.array([0.15773010017, -0.15340477408, 0.785])
|
||
]
|
||
|
||
# 测试多次调用性能
|
||
print(f"\n测试方法2（已知3D点法向量）{num_trials}次调用性能...")
|
||
start_time = time.time()
|
||
for _ in range(num_trials):
|
||
normals_data = normal_vector.calculate_normals(camera_points)
|
||
method2_time = time.time() - start_time
|
||
print(f"方法2平均耗时: {method2_time/num_trials:.6f} 秒/次")
|
||
|
||
# 输出结果
|
||
for i, data in enumerate(normals_data):
|
||
print(f"点 {i+1}: 坐标: {data['coordinates']}, 法向量: {data['normal']}")
|
||
|
||
# 测试单点性能
|
||
print(f"\n测试单点法向量查询{num_trials}次调用性能...")
|
||
single_point = camera_points[0]
|
||
start_time = time.time()
|
||
for _ in range(num_trials):
|
||
normal_data = normal_vector.get_normal_at_point(single_point)
|
||
single_point_time = time.time() - start_time
|
||
print(f"单点法向量查询平均耗时: {single_point_time/num_trials:.6f} 秒/次")
|
||
|
||
# 测试与原vector.py兼容的run方法
|
||
print("\n测试run方法（兼容原vector.py）...")
|
||
points_3d_camera_with_labels = [
|
||
{"label": "point1", "coordinates": camera_points[0]},
|
||
{"label": "point2", "coordinates": camera_points[1]}
|
||
]
|
||
start_time = time.time()
|
||
run_results = normal_vector.run(points_3d_camera_with_labels)
|
||
run_time = time.time() - start_time
|
||
print(f"run方法耗时: {run_time:.6f} 秒")
|
||
|
||
# 输出结果
|
||
for data in run_results:
|
||
print(f"Label: {data['label']}, Coordinates: {data['coordinates']}, Normal: {data['normal']}")
|
||
|
||
# 测试批量像素处理性能
|
||
print("\n测试批量像素处理性能...")
|
||
# 创建多个像素点
|
||
n_pixels = 1000
|
||
pixel_coords = np.random.randint(0, 480, size=(n_pixels, 2))
|
||
start_time = time.time()
|
||
pixel_results = normal_vector.get_normals_at_pixels(pixel_coords)
|
||
pixel_batch_time = time.time() - start_time
|
||
print(f"处理 {n_pixels} 个像素点耗时: {pixel_batch_time:.6f} 秒")
|
||
print(f"平均每点耗时: {pixel_batch_time/n_pixels:.6f} 秒/点")
|
||
|
||
# 测试不同规模下的向量化处理性能
|
||
print("\n测试不同规模下的向量化处理性能...")
|
||
|
||
# 小规模测试 (10个点)
|
||
small_points = [np.random.rand(3) for _ in range(10)]
|
||
start_time = time.time()
|
||
small_results = normal_vector.calculate_normals(small_points)
|
||
small_time = time.time() - start_time
|
||
print(f"处理 10 个点耗时: {small_time:.6f} 秒 (每点 {small_time/10:.6f} 秒)")
|
||
|
||
# 中规模测试 (100个点)
|
||
medium_points = [np.random.rand(3) for _ in range(100)]
|
||
start_time = time.time()
|
||
medium_results = normal_vector.calculate_normals(medium_points)
|
||
medium_time = time.time() - start_time
|
||
print(f"处理 100 个点耗时: {medium_time:.6f} 秒 (每点 {medium_time/100:.6f} 秒)")
|
||
|
||
# 大规模测试 (1000个点)
|
||
large_points = [np.random.rand(3) for _ in range(1000)]
|
||
start_time = time.time()
|
||
large_results = normal_vector.calculate_normals(large_points)
|
||
large_time = time.time() - start_time
|
||
print(f"处理 1000 个点耗时: {large_time:.6f} 秒 (每点 {large_time/1000:.6f} 秒)") |