使用RANSAC算法拟合空间圆

import open3d as o3d
import numpy as np
import glob
import os


def load_point_cloud_from_txt(file_path):
    """
    从TXT文件中加载点云数据
    """
    print(f"Loading point cloud from {file_path}")
    try:
        points = np.genfromtxt(file_path, delimiter=' ')
        print(f"Loaded {points.shape[0]} points")
    except Exception as e:
        print(f"Error loading {file_path}: {e}")
        points = np.array([])  # 返回空数组以继续处理其他文件
    return points


def fit_circle_ransac(points, distance_threshold=2, ransac_n=5, num_iterations=2000):
    """
    使用RANSAC算法拟合平面，并在平面上拟合圆
    """
    point_cloud = o3d.geometry.PointCloud()
    point_cloud.points = o3d.utility.Vector3dVector(points)

    # 使用RANSAC拟合平面
    plane_model, inliers = point_cloud.segment_plane(distance_threshold=distance_threshold,
                                                     ransac_n=ransac_n,
                                                     num_iterations=num_iterations)

    inlier_cloud = point_cloud.select_by_index(inliers)
    outlier_cloud = point_cloud.select_by_index(inliers, invert=True)

    [a, b, c, d] = plane_model
    plane_normal = np.array([a, b, c])

    # 投影点到平面上
    def project_point_to_plane(point, plane_normal, d):
        return point - (np.dot(point, plane_normal) + d) * plane_normal

    points_projected = np.array(
        [project_point_to_plane(point, plane_normal, d) for point in np.asarray(inlier_cloud.points)])

    # 计算圆心和半径
    centroid = np.mean(points_projected, axis=0)
    radii = [np.linalg.norm(point - centroid) for point in points_projected]
    radius = np.mean(radii)

    return centroid, radius, inlier_cloud, outlier_cloud, plane_model


def save_circle_to_txt(filename, centroid, radius, plane_model):
    """
    将拟合的圆保存到TXT文件
    """
    with open(filename, 'w') as f:
        f.write(f"圆心: {centroid[0]}, {centroid[1]}, {centroid[2]}\n")
        f.write(f"半径: {radius}\n")
        f.write(f"平面模型: {plane_model[0]}, {plane_model[1]}, {plane_model[2]}, {plane_model[3]}\n")


def process_point_cloud_files(file_paths, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for file_path in file_paths:
        print(f"Processing {file_path}...")
        points = load_point_cloud_from_txt(file_path)

        if points.shape[0] == 0:
            print(f"No points found in {file_path}")
            continue

        # 拟合圆
        centroid, radius, inlier_cloud, outlier_cloud, plane_model = fit_circle_ransac(points)

        print(f"拟合圆心: {centroid}")
        print(f"拟合半径: {radius}")

        # 保存拟合结果到TXT文件
        output_file = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(file_path))[0]}_circle.txt")
        save_circle_to_txt(output_file, centroid, radius, plane_model)

        # 可视化点云和拟合结果
        inlier_cloud.paint_uniform_color([1, 0, 0])  # 红色表示拟合内点
        outlier_cloud.paint_uniform_color([0, 1, 0])  # 绿色表示拟合外点

        # 生成拟合圆的点云用于可视化
        angles = np.linspace(0, 2 * np.pi, 100)
        circle_points = np.array([[centroid[0] + radius * np.cos(angle),
                                   centroid[1] + radius * np.sin(angle),
                                   centroid[2]] for angle in angles])

        # 投影圆点到拟合平面上
        def project_circle_to_plane(circle_points, plane_model):
            [a, b, c, d] = plane_model
            plane_normal = np.array([a, b, c])
            return np.array([point - (np.dot(point, plane_normal) + d) * plane_normal for point in circle_points])

        circle_points_projected = project_circle_to_plane(circle_points, plane_model)

        # 创建拟合圆的点云
        circle_cloud = o3d.geometry.PointCloud()
        circle_cloud.points = o3d.utility.Vector3dVector(circle_points_projected)
        circle_cloud.paint_uniform_color([0, 0, 1])  # 蓝色表示拟合圆

        # 可视化
        o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud, circle_cloud])

def main():
    # 获取所有TXT点云文件的路径
    file_paths = glob.glob('1.txt')  # 替换为你的TXT文件路径
    output_dir = 'data'

    if not file_paths:
        print("No point cloud files found.")
        return

    process_point_cloud_files(file_paths, output_dir)

    # 示例使用
if __name__ == "__main__":
    main()