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}")
|