首页/在线代码库/demo29_3d_segment.py

demo29_3d_segment.py

用合成点云(两个平面 + 噪声)演示区域生长分割: 1. 计算每个点法向 (PCA 邻域) 2. 选种子点, 向法向/距离一致的邻域生长 3. 分割出不同平面区域

Python 99 行 系列:识别系列 🔒 源码仅供在线学习,不提供文件下载
完整源码 · 在线查看demo29_3d_segment.py
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
"""演示2: 三维点云分割 — 区域生长 + 平面拟合 (经典方法).

用合成点云(两个平面 + 噪声)演示区域生长分割:
1. 计算每个点法向 (PCA 邻域)
2. 选种子点, 向法向/距离一致的邻域生长
3. 分割出不同平面区域
"""
import os
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from sklearn.neighbors import NearestNeighbors
from figure_style import figures_dir, configure

configure()
OUT = figures_dir()
os.makedirs(OUT, exist_ok=True)

rng = np.random.default_rng(0)

# 两个相交平面 + 噪声点
def plane(center, normal, n=400):
    u = rng.normal(0, 1, (n, 3))
    v = u - u.dot(normal)[:, None] * normal   # 投影到平面
    return v + center + rng.normal(0, 0.02, (n, 3))

p1 = plane([0, 0, 0], [0, 0, 1])
p2 = plane([0, 0, 0.3], [0, 1, 0.3])
noise = rng.uniform(-1.5, 1.5, (50, 3))
pts = np.vstack([p1, p2, noise])
print(f"点云: {pts.shape[0]} 点 (2 平面 + 噪声)")


def compute_normals(pc, k=12):
    """PCA 法向估计: 邻域协方差最小特征向量."""
    nn = NearestNeighbors(n_neighbors=k).fit(pc)
    _, idx = nn.kneighbors(pc)
    normals = np.zeros_like(pc)
    for i, ni in enumerate(idx):
        cov = np.cov((pc[ni] - pc[i]).T)
        w, v = np.linalg.eigh(cov)
        normals[i] = v[:, np.argmin(w)]
    # 一致朝外
    normals[normals[:, 2] < 0] *= -1
    return normals


def region_grow(pc, normals, th_angle=0.3, th_dist=0.1, min_size=50):
    """区域生长: 法向夹角 + 到平面距离约束."""
    n = len(pc)
    labels = np.full(n, -1)
    nn = NearestNeighbors(n_neighbors=20).fit(pc)
    _, idx_all = nn.kneighbors(pc)
    label = 0
    for seed in range(n):
        if labels[seed] != -1:
            continue
        # BFS 生长
        queue = [seed]; labels[seed] = label
        pts_l = [pc[seed]]; count = 0
        while queue:
            cur = queue.pop()
            count += 1
            for nb in idx_all[cur]:
                if labels[nb] != -1:
                    continue
                if np.abs(np.dot(normals[cur], normals[nb])) > np.cos(th_angle):
                    # 到种子平面距离 (用局部平面近似)
                    d = np.abs((pc[nb] - pc[seed]) @ normals[seed])
                    if d < th_dist:
                        labels[nb] = label
                        queue.append(nb)
        if count < min_size:
            labels[labels == label] = -1   # 丢弃过小区域
        label += 1
    return labels


normals = compute_normals(pts)
labels = region_grow(pts, normals)
n_seg = len(np.unique(labels)) - (1 if -1 in labels else 0)
print(f"分割出 {n_seg} 个主要区域 (噪声点 label=-1)")

fig = plt.figure(figsize=(8, 6.5))
ax = fig.add_subplot(111, projection='3d')
cmap = plt.cm.tab10
for i, l in enumerate(np.unique(labels)):
    sel = pts[labels == l]
    if l == -1:
        ax.scatter(sel[:, 0], sel[:, 1], sel[:, 2], s=3, c='lightgray', alpha=0.4, label='噪声')
    else:
        ax.scatter(sel[:, 0], sel[:, 1], sel[:, 2], s=5, c=cmap(i), label=f'区域{i}')
ax.set_title(f'点云区域生长分割: {n_seg} 个平面区域')
ax.legend(fontsize=8); ax.set_xlabel('X'); ax.set_ylabel('Y'); ax.set_zlabel('Z')
fig.tight_layout()
out = os.path.join(OUT, 'fig3d_segment.png')
fig.savefig(out, dpi=150)
print(f"-> {out}")

输出图

fig3d_segment.png
输出图 · fig3d_segment.png