170 lines
5.4 KiB
Python
170 lines
5.4 KiB
Python
"""直线拟合模块 - 使用RANSAC算法拟合3D直线"""
|
||
import numpy as np
|
||
from dataclasses import dataclass
|
||
from typing import Tuple, Optional
|
||
|
||
|
||
@dataclass
|
||
class LineModel:
|
||
"""直线模型"""
|
||
point: np.ndarray # 直线上的点 (3,)
|
||
direction: np.ndarray # 方向向量（单位向量） (3,)
|
||
inliers: np.ndarray # 内点索引
|
||
inlier_ratio: float # 内点比例
|
||
|
||
def distance_to_point(self, point: np.ndarray) -> float:
|
||
"""计算点到直线的距离"""
|
||
vector = point - self.point
|
||
projection = np.dot(vector, self.direction) * self.direction
|
||
perpendicular = vector - projection
|
||
return np.linalg.norm(perpendicular)
|
||
|
||
def distance_to_points(self, points: np.ndarray) -> np.ndarray:
|
||
"""计算多个点到直线的距离"""
|
||
vectors = points - self.point
|
||
projections = np.dot(vectors, self.direction)[:, np.newaxis] * self.direction
|
||
perpendicular_vectors = vectors - projections
|
||
distances = np.linalg.norm(perpendicular_vectors, axis=1)
|
||
return distances
|
||
|
||
def sample_points_on_line(self, n_samples: int, point_cloud: np.ndarray) -> np.ndarray:
|
||
"""
|
||
从拟合的直线对应的真实点云中采样点
|
||
|
||
Args:
|
||
n_samples: 采样点数
|
||
point_cloud: 用于拟合的原始点云（内点）
|
||
|
||
Returns:
|
||
采样的点 (n_samples, 3)
|
||
"""
|
||
if len(point_cloud) == 0:
|
||
return np.array([]).reshape(0, 3)
|
||
|
||
n_samples = min(n_samples, len(point_cloud))
|
||
indices = np.random.choice(len(point_cloud), size=n_samples, replace=False)
|
||
return point_cloud[indices]
|
||
|
||
|
||
class LineFitter:
|
||
"""RANSAC直线拟合器"""
|
||
|
||
def __init__(
|
||
self,
|
||
threshold: float = 2.0,
|
||
max_iterations: int = 1000,
|
||
min_inlier_ratio: float = 0.6
|
||
):
|
||
"""
|
||
Args:
|
||
threshold: 内点距离阈值(mm)
|
||
max_iterations: 最大迭代次数
|
||
min_inlier_ratio: 最小内点比例
|
||
"""
|
||
self.threshold = threshold
|
||
self.max_iterations = max_iterations
|
||
self.min_inlier_ratio = min_inlier_ratio
|
||
|
||
def fit(self, points: np.ndarray) -> Optional[LineModel]:
|
||
"""
|
||
使用RANSAC拟合3D直线
|
||
|
||
Args:
|
||
points: 点云数组 (N, 3)
|
||
|
||
Returns:
|
||
LineModel或None（拟合失败）
|
||
"""
|
||
if len(points) < 2:
|
||
return None
|
||
|
||
best_inliers = None
|
||
best_point = None
|
||
best_direction = None
|
||
best_inlier_count = 0
|
||
|
||
for _ in range(self.max_iterations):
|
||
# 随机选择2个点
|
||
sample_indices = np.random.choice(len(points), 2, replace=False)
|
||
p1, p2 = points[sample_indices]
|
||
|
||
# 计算直线方向
|
||
direction = p2 - p1
|
||
if np.linalg.norm(direction) < 1e-6:
|
||
continue
|
||
direction = direction / np.linalg.norm(direction)
|
||
|
||
# 计算所有点到直线的距离
|
||
vectors = points - p1
|
||
projections = np.dot(vectors, direction)[:, np.newaxis] * direction
|
||
distances = np.linalg.norm(vectors - projections, axis=1)
|
||
|
||
# 找到内点
|
||
inliers = np.where(distances < self.threshold)[0]
|
||
|
||
if len(inliers) > best_inlier_count:
|
||
best_inlier_count = len(inliers)
|
||
best_inliers = inliers
|
||
best_point = p1
|
||
best_direction = direction
|
||
|
||
if best_inliers is None or len(best_inliers) < 2:
|
||
return None
|
||
|
||
# 使用所有内点重新拟合直线（最小二乘）
|
||
inlier_points = points[best_inliers]
|
||
final_point, final_direction = self._fit_line_least_squares(inlier_points)
|
||
|
||
if final_direction is None:
|
||
return None
|
||
|
||
# 重新计算内点
|
||
vectors = points - final_point
|
||
projections = np.dot(vectors, final_direction)[:, np.newaxis] * final_direction
|
||
distances = np.linalg.norm(vectors - projections, axis=1)
|
||
final_inliers = np.where(distances < self.threshold)[0]
|
||
|
||
inlier_ratio = len(final_inliers) / len(points)
|
||
|
||
# 检查内点比例
|
||
if inlier_ratio < self.min_inlier_ratio:
|
||
return None
|
||
|
||
return LineModel(
|
||
point=final_point,
|
||
direction=final_direction,
|
||
inliers=final_inliers,
|
||
inlier_ratio=inlier_ratio
|
||
)
|
||
|
||
def _fit_line_least_squares(self, points: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
||
"""
|
||
使用最小二乘法拟合直线
|
||
|
||
Returns:
|
||
(point, direction): 直线上的点和方向向量
|
||
"""
|
||
if len(points) < 2:
|
||
return None, None
|
||
|
||
# 计算质心
|
||
centroid = np.mean(points, axis=0)
|
||
|
||
# 中心化点云
|
||
centered = points - centroid
|
||
|
||
# PCA：找主方向
|
||
try:
|
||
# 协方差矩阵
|
||
cov = np.dot(centered.T, centered) / len(points)
|
||
# 特征值分解
|
||
eigenvalues, eigenvectors = np.linalg.eig(cov)
|
||
# 最大特征值对应的特征向量就是直线方向
|
||
max_idx = np.argmax(eigenvalues)
|
||
direction = eigenvectors[:, max_idx].real
|
||
direction = direction / np.linalg.norm(direction)
|
||
|
||
return centroid, direction
|
||
except np.linalg.LinAlgError:
|
||
return None, None
|