feat: update algorithm integrations and point cloud overlays

alg/train_measurement_system_v3.2_0629/src/line_fitter.py

@@ -0,0 +1,169 @@
+"""直线拟合模块 - 使用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