feat: update algorithm integrations and point cloud overlays

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

@@ -0,0 +1,195 @@
+"""可视化工具 - 用于显示点云、平面、ROI等"""
+import numpy as np
+import open3d as o3d
+from pathlib import Path
+from typing import List, Tuple, Optional
+
+
+class Visualizer:
+    """点云和测量结果可视化工具"""
+
+    @staticmethod
+    def visualize_roi_in_context(
+        full_cloud: np.ndarray,
+        roi_cloud: np.ndarray,
+        bbox_min: List[float],
+        bbox_max: List[float],
+        title: str = "ROI in Context",
+        downsample_full: int = 10,
+        save_path: Optional[str] = None
+    ):
+        """
+        在完整点云上下文中可视化ROI区域
+
+        Args:
+            full_cloud: 完整点云 (N, 3)
+            roi_cloud: ROI区域点云 (M, 3)
+            bbox_min: 包围盒最小坐标 [x, y, z]
+            bbox_max: 包围盒最大坐标 [x, y, z]
+            title: 窗口标题
+            downsample_full: 完整点云降采样倍数（每N个点取1个）
+            save_path: 保存截图路径（可选）
+        """
+        # 降采样完整点云以提高显示性能
+        full_cloud_downsampled = full_cloud[::downsample_full]
+
+        # 创建完整点云对象（灰色）
+        pcd_full = o3d.geometry.PointCloud()
+        pcd_full.points = o3d.utility.Vector3dVector(full_cloud_downsampled)
+        pcd_full.paint_uniform_color([0.7, 0.7, 0.7])  # 灰色
+
+        # 创建ROI点云对象（红色）
+        pcd_roi = o3d.geometry.PointCloud()
+        pcd_roi.points = o3d.utility.Vector3dVector(roi_cloud)
+        pcd_roi.paint_uniform_color([1.0, 0.0, 0.0])  # 红色
+
+        # 创建包围盒线框
+        bbox = Visualizer._create_bbox_lineset(bbox_min, bbox_max, color=[0, 1, 0])
+
+        # 可视化
+        geometries = [pcd_full, pcd_roi, bbox]
+        o3d.visualization.draw_geometries(
+            geometries,
+            window_name=title,
+            width=1200,
+            height=800,
+            point_show_normal=False
+        )
+
+        # 保存截图（如果指定）
+        if save_path:
+            vis = o3d.visualization.Visualizer()
+            vis.create_window(window_name=title, width=1200, height=800, visible=False)
+            for geom in geometries:
+                vis.add_geometry(geom)
+            vis.update_renderer()
+            vis.capture_screen_image(save_path)
+            vis.destroy_window()
+            print(f"截图已保存: {save_path}")
+
+    @staticmethod
+    def visualize_plane_measurement(
+        source_cloud: np.ndarray,
+        target_cloud: np.ndarray,
+        source_plane_normal: np.ndarray,
+        source_plane_point: np.ndarray,
+        target_plane_normal: np.ndarray,
+        target_plane_point: np.ndarray,
+        sample_points: np.ndarray,
+        distances: np.ndarray,
+        title: str = "Plane Measurement",
+        save_path: Optional[str] = None
+    ):
+        """
+        可视化平面到平面的距离测量
+
+        Args:
+            source_cloud: 源平面点云
+            target_cloud: 目标平面点云
+            source_plane_normal: 源平面法向量
+            source_plane_point: 源平面上的点
+            target_plane_normal: 目标平面法向量
+            target_plane_point: 目标平面上的点
+            sample_points: 采样点
+            distances: 每个采样点到目标平面的距离
+            title: 窗口标题
+            save_path: 保存截图路径
+        """
+        # 创建源平面点云（蓝色）
+        pcd_source = o3d.geometry.PointCloud()
+        pcd_source.points = o3d.utility.Vector3dVector(source_cloud)
+        pcd_source.paint_uniform_color([0.0, 0.5, 1.0])  # 蓝色
+
+        # 创建目标平面点云（绿色）
+        pcd_target = o3d.geometry.PointCloud()
+        pcd_target.points = o3d.utility.Vector3dVector(target_cloud)
+        pcd_target.paint_uniform_color([0.0, 1.0, 0.0])  # 绿色
+
+        # 创建采样点（红色）
+        pcd_samples = o3d.geometry.PointCloud()
+        pcd_samples.points = o3d.utility.Vector3dVector(sample_points)
+        pcd_samples.paint_uniform_color([1.0, 0.0, 0.0])  # 红色
+
+        # 创建距离线（从采样点到目标平面的投影）
+        lines = []
+        for i, (point, dist) in enumerate(zip(sample_points, distances)):
+            # 计算投影点
+            projection = point - dist * target_plane_normal
+            lines.append([point, projection])
+
+        # 创建线集
+        line_points = []
+        line_indices = []
+        for i, (start, end) in enumerate(lines):
+            line_points.append(start)
+            line_points.append(end)
+            line_indices.append([2*i, 2*i+1])
+
+        lineset = o3d.geometry.LineSet()
+        lineset.points = o3d.utility.Vector3dVector(line_points)
+        lineset.lines = o3d.utility.Vector2iVector(line_indices)
+        lineset.paint_uniform_color([1.0, 0.5, 0.0])  # 橙色
+
+        # 可视化
+        geometries = [pcd_source, pcd_target, pcd_samples, lineset]
+        o3d.visualization.draw_geometries(
+            geometries,
+            window_name=title,
+            width=1200,
+            height=800,
+            point_show_normal=False
+        )
+
+        if save_path:
+            vis = o3d.visualization.Visualizer()
+            vis.create_window(window_name=title, width=1200, height=800, visible=False)
+            for geom in geometries:
+                vis.add_geometry(geom)
+            vis.update_renderer()
+            vis.capture_screen_image(save_path)
+            vis.destroy_window()
+            print(f"截图已保存: {save_path}")
+
+    @staticmethod
+    def _create_bbox_lineset(
+        bbox_min: List[float],
+        bbox_max: List[float],
+        color: List[float] = [0, 1, 0]
+    ) -> o3d.geometry.LineSet:
+        """创建包围盒线框"""
+        # 8个顶点
+        points = [
+            [bbox_min[0], bbox_min[1], bbox_min[2]],  # 0
+            [bbox_max[0], bbox_min[1], bbox_min[2]],  # 1
+            [bbox_max[0], bbox_max[1], bbox_min[2]],  # 2
+            [bbox_min[0], bbox_max[1], bbox_min[2]],  # 3
+            [bbox_min[0], bbox_min[1], bbox_max[2]],  # 4
+            [bbox_max[0], bbox_min[1], bbox_max[2]],  # 5
+            [bbox_max[0], bbox_max[1], bbox_max[2]],  # 6
+            [bbox_min[0], bbox_max[1], bbox_max[2]],  # 7
+        ]
+
+        # 12条边
+        lines = [
+            [0, 1], [1, 2], [2, 3], [3, 0],  # 底面
+            [4, 5], [5, 6], [6, 7], [7, 4],  # 顶面
+            [0, 4], [1, 5], [2, 6], [3, 7],  # 竖边
+        ]
+
+        lineset = o3d.geometry.LineSet()
+        lineset.points = o3d.utility.Vector3dVector(points)
+        lineset.lines = o3d.utility.Vector2iVector(lines)
+        lineset.paint_uniform_color(color)
+
+        return lineset
+
+    @staticmethod
+    def save_roi_pointcloud(
+        roi_cloud: np.ndarray,
+        save_path: str
+    ):
+        """保存ROI点云为PLY文件"""
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(roi_cloud)
+        o3d.io.write_point_cloud(save_path, pcd)
+        print(f"ROI点云已保存: {save_path}")