import open3d as o3d
import numpy as np


def dbscan_open3d(pcd, eps=0.1, min_points=10):
    """
    使用 Open3D 进行 DBSCAN 聚类
    Args:
        pcd: Open3D 点云对象
        eps: 邻域半径
        min_points: 最小邻域点数
    Returns:
        labels: 每个点的簇标签（-1 表示噪声）
        clusters: 各簇的点索引列表
    """
    # DBSCAN 聚类
    labels = np.array(pcd.cluster_dbscan(eps=eps, min_points=min_points))

    # 提取各簇索引
    unique_labels = np.unique(labels)
    clusters = []
    for label in unique_labels:
        if label == -1:
            continue
        clusters.append(np.where(labels == label)[0])

    return labels, clusters


# 使用示例
pcd = o3d.io.read_point_cloud("E:\point_distance\data\part_car_head.ply")

# 预处理
pcd = pcd.voxel_down_sample(voxel_size=0.05)
pcd, _ = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)

# DBSCAN 聚类
labels, clusters = dbscan_open3d(pcd, eps=0.1, min_points=50)

print(f"找到 {len(clusters)} 个聚类")
print(f"噪声点: {np.sum(labels == -1)}")

# 为每个聚类着色
colors = np.random.rand(len(clusters), 3)
colored_points = np.zeros((len(pcd.points), 3))
for i, cluster in enumerate(clusters):
    colored_points[cluster] = colors[i]
pcd.colors = o3d.utility.Vector3dVector(colored_points)

o3d.visualization.draw_geometries([pcd])