Easy to make the input point cloud be Axis-aligned

[zhirui-gao]

We have implemented an easy-to-use script to get axis-aligned point clouds. script

Hope that helps!

[wild-deer]

base on your code

Supports placing mesh.ply directly under dataset, as well as multiple subdirectories under dataset.

Automatically downsamples: if there are more than 500,000 points, voxel_down_sample(0.02) will be applied.

Empty point clouds are skipped to prevent freezing.

num_iterations changed from 1000 to 500 for faster speed.

All steps have [DEBUG] output so you can see in real time where it’s stuck.

import open3d as o3d
import numpy as np
import os
from sklearn.decomposition import PCA


def align_scene_to_z_up(pcd, save_pcd_path=None, save_transform_path=None, visualize=False, fit_ground=True):
    print("=== 开始对齐 ===")
    points = np.asarray(pcd.points)
    print(f"[DEBUG] 原始点云点数: {points.shape[0]}")
    centroid = np.mean(points, axis=0)
    print(f"[DEBUG] 点云质心: {centroid}")

    distance_threshold = 0.02
    R_ground = np.eye(3)

    if fit_ground:
        print("[DEBUG] 开始检测地面平面")
        plane_model, ground_inliers = pcd.segment_plane(
            distance_threshold=distance_threshold,
            ransac_n=3,
            num_iterations=500
        )
        [a, b, c, d] = plane_model
        ground_normal = np.array([a, b, c])
        print(f"[DEBUG] 地面法向量: {ground_normal}, d={d}, 内点数={len(ground_inliers)}")

        z_axis = np.array([0, 0, 1])
        ground_normal = ground_normal / np.linalg.norm(ground_normal)
        rotation_axis = np.cross(ground_normal, z_axis)

        if np.linalg.norm(rotation_axis) < 1e-6:
            print("[DEBUG] 地面已接近Z轴，不需要旋转")
            rotation_axis = np.array([0, 1, 0])
        else:
            rotation_axis = rotation_axis / np.linalg.norm(rotation_axis)

        cos_theta = np.dot(ground_normal, z_axis)
        angle = np.arccos(np.clip(cos_theta, -1.0, 1.0))
        print(f"[DEBUG] 地面旋转角度: {np.degrees(angle):.4f}°, 旋转轴: {rotation_axis}")

        K = np.array([
            [0, -rotation_axis[2], rotation_axis[1]],
            [rotation_axis[2], 0, -rotation_axis[0]],
            [-rotation_axis[1], rotation_axis[0], 0]
        ])
        R_ground = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K)

        centered_points = points - centroid
        ground_rotated = (R_ground @ centered_points.T).T
        all_indices = np.arange(len(points))
        non_ground_indices = np.setdiff1d(all_indices, ground_inliers)
        non_ground_points = ground_rotated[non_ground_indices]

        temp_pcd = o3d.geometry.PointCloud()
        temp_pcd.points = o3d.utility.Vector3dVector(non_ground_points)
        remaining_pcd = temp_pcd
    else:
        remaining_pcd = pcd

    wall_directions = []
    max_points = 0
    best_direction = None
    R_walls = np.eye(3)

    print("[DEBUG] 开始检测墙面")
    for i in range(6):
        if len(remaining_pcd.points) < 100:
            break

        plane_model, inliers = remaining_pcd.segment_plane(
            distance_threshold=distance_threshold,
            ransac_n=3,
            num_iterations=500
        )
        [a, b, c, d] = plane_model
        normal = np.array([a, b, c])

        if abs(normal[2]) < 0.1 and np.linalg.norm(normal[:2]) > 0.5:
            horizontal_dir = normal[:2] / np.linalg.norm(normal[:2])
            if max_points < len(inliers):
                max_points = len(inliers)
                best_direction = horizontal_dir
                angle = np.arctan2(best_direction[1], best_direction[0])
                R_walls = np.array([
                    [np.cos(-angle), -np.sin(-angle), 0],
                    [np.sin(-angle), np.cos(-angle), 0],
                    [0, 0, 1]
                ])
        remaining_pcd = remaining_pcd.select_by_index(inliers, invert=True)

    if fit_ground:
        combined_rotation = R_walls @ R_ground
        transform_matrix = np.eye(4)
        transform_matrix[:3, :3] = combined_rotation
        transform_matrix[:3, 3] = -combined_rotation @ centroid
    else:
        transform_matrix = np.eye(4)
        transform_matrix[:3, :3] = R_walls
        transform_matrix[:3, 3] = -R_walls @ centroid

    aligned_points = (transform_matrix[:3, :3] @ points.T + transform_matrix[:3, 3:4]).T
    aligned_pcd = o3d.geometry.PointCloud()
    aligned_pcd.points = o3d.utility.Vector3dVector(aligned_points)

    if pcd.colors:
        aligned_pcd.colors = pcd.colors

    if save_pcd_path:
        o3d.io.write_point_cloud(save_pcd_path, aligned_pcd)
        print(f"[DEBUG] 对齐后的点云已保存到 {save_pcd_path}")

    if save_transform_path:
        np.savetxt(save_transform_path, transform_matrix, fmt='%.8f')
        print(f"[DEBUG] 变换矩阵已保存到 {save_transform_path}")

    return aligned_pcd, transform_matrix


def process_point_cloud(input_path, output_path, fit_ground=True):
    print(f"\n=== 处理文件: {input_path} ===")
    pcd = o3d.io.read_point_cloud(input_path)
    print(f"[DEBUG] 读取点云: {len(pcd.points)} points")

    if len(pcd.points) == 0:
        print(f"[ERROR] 点云为空: {input_path}")
        return

    if len(pcd.points) > 500000:
        print("[DEBUG] 点数过多，进行降采样...")
        pcd = pcd.voxel_down_sample(voxel_size=0.02)

    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    transform_path = os.path.splitext(output_path)[0] + "_transform.txt"

    transformed_pcd, transform = align_scene_to_z_up(
        pcd,
        save_transform_path=transform_path,
        save_pcd_path=output_path,
        visualize=False,
        fit_ground=fit_ground
    )

    aligned_points = np.asarray(transformed_pcd.points)
    min_z = np.min(aligned_points[:, 2])
    max_z = np.max(aligned_points[:, 2])
    height = max_z - min_z
    print(f"[DEBUG] 对齐后高度: {height:.4f}")
    estimated_height = 2.5
    scale = estimated_height / height
    print(f"[DEBUG] 缩放系数: {scale:.4f}")


if __name__ == "__main__":
    # 示例: 输入文件 和 输出文件
    input_file = "dataset/mesh2.ply"
    output_file = "dataset/mesh2_align.ply"
    process_point_cloud(input_file, output_file, fit_ground=True)