提交scripts

libs/MVS/utils/get_world_point.py

@@ -0,0 +1,117 @@
+'''
+获取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)
+