118 lines
3.3 KiB
Python
118 lines
3.3 KiB
Python
'''
|
||
获取Mesh的深度图
|
||
|
||
required:
|
||
- open3d
|
||
- numpy
|
||
- numba
|
||
'''
|
||
|
||
import open3d as o3d
|
||
import numpy as np
|
||
from numba import njit
|
||
|
||
@njit
|
||
def _process_all_points(Z, width, height, cx, cy, fx, fy, c2w):
|
||
"""JIT优化的所有点处理函数"""
|
||
# 手动创建像素网格（替代np.meshgrid）
|
||
u = np.zeros((height, width), dtype=np.float64)
|
||
v = np.zeros((height, width), dtype=np.float64)
|
||
|
||
for y in range(height):
|
||
for x in range(width):
|
||
u[y, x] = x
|
||
v[y, x] = y
|
||
|
||
# 创建输出数组
|
||
Z_all = np.zeros_like(Z)
|
||
|
||
# 手动处理有效点
|
||
for y in range(height):
|
||
for x in range(width):
|
||
if Z[y, x] > 0 and Z[y, x] < 1e9 and not np.isnan(Z[y, x]):
|
||
Z_all[y, x] = Z[y, x]
|
||
else:
|
||
Z_all[y, x] = -10000
|
||
|
||
# 计算所有点的3D坐标
|
||
X_all = (u - cx) * Z_all / fx
|
||
Y_all = (v - cy) * Z_all / fy
|
||
|
||
# 创建齐次坐标
|
||
points_3d = np.zeros((height * width, 4))
|
||
points_3d[:, 0] = X_all.reshape(-1)
|
||
points_3d[:, 1] = Y_all.reshape(-1)
|
||
points_3d[:, 2] = Z_all.reshape(-1)
|
||
points_3d[:, 3] = 1.0
|
||
|
||
# 坐标变换
|
||
return points_3d @ c2w.T
|
||
|
||
@njit
|
||
def _process_valid_points(Z, width, height, cx, cy, fx, fy, c2w):
|
||
"""JIT优化的有效点处理函数"""
|
||
# 首先计算有效点的数量
|
||
valid_count = 0
|
||
for y in range(height):
|
||
for x in range(width):
|
||
if Z[y, x] > 0 and Z[y, x] < 1e9 and not np.isnan(Z[y, x]):
|
||
valid_count += 1
|
||
|
||
# 创建输出数组
|
||
X_valid = np.zeros(valid_count)
|
||
Y_valid = np.zeros(valid_count)
|
||
Z_valid = np.zeros(valid_count)
|
||
|
||
# 填充有效点数据
|
||
idx = 0
|
||
for y in range(height):
|
||
for x in range(width):
|
||
if Z[y, x] > 0 and Z[y, x] < 1e9 and not np.isnan(Z[y, x]):
|
||
Z_valid[idx] = Z[y, x]
|
||
X_valid[idx] = (x - cx) * Z_valid[idx] / fx
|
||
Y_valid[idx] = (y - cy) * Z_valid[idx] / fy
|
||
idx += 1
|
||
|
||
# 创建齐次坐标
|
||
points_3d = np.zeros((valid_count, 4))
|
||
points_3d[:, 0] = X_valid
|
||
points_3d[:, 1] = Y_valid
|
||
points_3d[:, 2] = Z_valid
|
||
points_3d[:, 3] = 1.0
|
||
|
||
# 坐标变换
|
||
return points_3d @ c2w.T
|
||
|
||
def get_world_points(depth_map: o3d.geometry.Image, K: np.ndarray, c2w: np.ndarray, all_point: bool = False) -> np.ndarray:
|
||
'''
|
||
获取深度图对应的世界坐标点云
|
||
|
||
param:
|
||
depth_map: 深度图
|
||
K: 相机内参矩阵
|
||
c2w: 相机坐标系到世界坐标系的变换矩阵
|
||
all_point: 是否返回所有点（包括无效点）
|
||
|
||
return:
|
||
np.ndarray: 点云坐标（世界坐标系）
|
||
- 当all_point=False时：只包含有效深度值的点云
|
||
- 当all_point=True时：包含所有点的点云，无效点深度值设为10000
|
||
'''
|
||
# 转换为numpy数组
|
||
Z = np.asarray(depth_map)
|
||
height, width = Z.shape
|
||
|
||
# 获取内参
|
||
fx, fy = K[0, 0], K[1, 1]
|
||
cx, cy = K[0, 2], K[1, 2]
|
||
|
||
# 确保数据类型兼容性（Numba要求）
|
||
Z = Z.astype(np.float64)
|
||
c2w = c2w.astype(np.float64)
|
||
|
||
if all_point:
|
||
return _process_all_points(Z, width, height, cx, cy, fx, fy, c2w)
|
||
else:
|
||
return _process_valid_points(Z, width, height, cx, cy, fx, fy, c2w)
|
||
|