feat: update algorithm integrations and point cloud overlays

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

@@ -0,0 +1,510 @@
+"""任务管理器 - 负责加载配置、调度任务、汇总结果"""
+import yaml
+import numpy as np
+import time
+from pathlib import Path
+from typing import List, Dict, Any
+from datetime import datetime
+
+from pointcloud_loader import PointCloudLoader
+from measurements.base import MeasurementResult
+from measurements.plane_distance import PlaneToPlaneDistance
+from measurements.point_to_plane import PointToPlaneDistance
+from measurements.point_to_point import PointToPointDistance
+from measurements.plane_with_holes import PlaneToPlaneWithHoles
+from measurements.point_to_plane_with_holes import PointToPlaneWithHoles
+from measurements.point_to_point_with_holes import PointToPointWithHoles
+from measurements.line_to_line import LineToLineDistance
+from measurements.line_to_plane import LineToPlaneDistance
+from excel_exporter import ExcelExporter
+from performance_logger import PerformanceLogger
+from measurement_logger import MeasurementLogger
+from spatial_index import LinearScanIndex, GridIndex, GridIndexOptimized, OctreeIndex, KDTreeIndex
+
+
+class TaskManager:
+    """测量任务管理器"""
+
+    def __init__(self):
+        """初始化任务管理器"""
+        self.config = None
+        self.point_cloud = None
+        self.spatial_index = None  # 空间索引
+        self.tasks = []
+        self.results = []
+        self.pointcloud_path_override = None
+        self.output_dir_override = None
+
+    def load_config(self, config_path: str) -> bool:
+        """
+        加载配置文件
+
+        Args:
+            config_path: 配置文件路径
+
+        Returns:
+            bool: 是否加载成功
+        """
+        try:
+            with open(config_path, 'r', encoding='utf-8') as f:
+                self.config = yaml.safe_load(f)
+
+            # 加载点云（支持相对路径）
+            pointcloud_path = self.pointcloud_path_override or self.config.get('pointcloud_file')
+            if not pointcloud_path:
+                print("错误: 配置文件中未指定点云文件路径")
+                return False
+
+            # 处理相对路径
+            pointcloud_path = self._resolve_path(pointcloud_path, config_path)
+
+            # 加载点云
+            print(f"\n{'='*70}")
+            print("阶段1: 加载点云")
+            print("="*70)
+            print(f"  文件: {Path(pointcloud_path).name}")
+
+            load_start = time.time()
+            self.point_cloud = PointCloudLoader.load_ply(pointcloud_path)
+            load_time = time.time() - load_start
+
+            print(f"  点云数量: {len(self.point_cloud):,} 个点")
+            print(f"  加载耗时: {load_time:.2f} 秒")
+
+            # 保存点云加载信息供性能日志使用
+            self.pointcloud_load_time = load_time
+            self.pointcloud_filename = Path(pointcloud_path).name
+
+            # 构建空间索引
+            print(f"\n{'='*70}")
+            print("阶段2: 构建空间索引")
+            print("="*70)
+            index_start = time.time()
+            self._build_spatial_index()
+            self.index_build_time = time.time() - index_start
+
+            # 解析任务
+            print(f"\n{'='*70}")
+            print("阶段3: 解析测量任务")
+            print("="*70)
+            tasks_config = self.config.get('tasks', [])
+            if not tasks_config:
+                print("警告: 配置文件中没有定义任务")
+                return True
+
+            print(f"\n找到 {len(tasks_config)} 个测量任务")
+
+            for task_config in tasks_config:
+                task = self._create_task(task_config)
+                if task:
+                    self.tasks.append(task)
+                    print(f"  - {task.name} ({task.task_type})")
+
+            return True
+
+        except FileNotFoundError:
+            print(f"错误: 配置文件不存在: {config_path}")
+            return False
+        except yaml.YAMLError as e:
+            print(f"错误: 配置文件格式错误: {e}")
+            return False
+        except Exception as e:
+            print(f"错误: 加载配置失败: {e}")
+            return False
+
+    def _resolve_path(self, path: str, config_path: str) -> str:
+        """
+        解析路径（支持相对路径和绝对路径）
+
+        Args:
+            path: 待解析的路径
+            config_path: 配置文件路径（用于相对路径的基准）
+
+        Returns:
+            str: 解析后的绝对路径
+        """
+        path_obj = Path(path)
+        if path_obj.is_absolute():
+            return str(path_obj)
+        else:
+            # 相对于配置文件所在目录
+            config_dir = Path(config_path).parent
+            return str(config_dir / path_obj)
+
+    def override_pointcloud_path(self, path: str):
+        """
+        覆盖点云文件路径（用于命令行参数）
+
+        Args:
+            path: 点云文件路径
+        """
+        self.pointcloud_path_override = path
+
+    def override_output_dir(self, output_dir: str):
+        """
+        覆盖输出目录（用于命令行参数）
+
+        Args:
+            output_dir: 输出目录路径
+        """
+        self.output_dir_override = output_dir
+
+    def _create_task(self, task_config: Dict[str, Any]):
+        """
+        根据配置创建测量任务
+
+        Args:
+            task_config: 任务配置字典
+
+        Returns:
+            BaseMeasurement: 测量任务实例
+        """
+        task_type = task_config.get('type')
+        task_name = task_config.get('name', 'Unnamed Task')
+
+        # 处理带有孔中心检测的任务类型
+        if '4_hole_center' in task_type:
+            if 'plane_to_plane' in task_type:
+                return PlaneToPlaneWithHoles(
+                    name=task_name,
+                    point_cloud=self.point_cloud,
+                    config=task_config,
+                    spatial_index=self.spatial_index
+                )
+            elif 'point_to_plane' in task_type:
+                return PointToPlaneWithHoles(
+                    name=task_name,
+                    point_cloud=self.point_cloud,
+                    config=task_config,
+                    spatial_index=self.spatial_index
+                )
+        elif 'hole_center' in task_type and 'point_to_point' in task_type:
+            return PointToPointWithHoles(
+                name=task_name,
+                point_cloud=self.point_cloud,
+                config=task_config,
+                spatial_index=self.spatial_index
+            )
+        elif task_type == 'plane_to_plane':
+            return PlaneToPlaneDistance(
+                name=task_name,
+                point_cloud=self.point_cloud,
+                config=task_config,
+                spatial_index=self.spatial_index
+            )
+        elif task_type == 'point_to_plane':
+            return PointToPlaneDistance(
+                name=task_name,
+                point_cloud=self.point_cloud,
+                config=task_config,
+                spatial_index=self.spatial_index
+            )
+        elif task_type == 'point_to_point':
+            return PointToPointDistance(
+                name=task_name,
+                point_cloud=self.point_cloud,
+                config=task_config,
+                spatial_index=self.spatial_index
+            )
+        elif task_type == 'line_to_line':
+            return LineToLineDistance(
+                name=task_name,
+                point_cloud=self.point_cloud,
+                config=task_config,
+                spatial_index=self.spatial_index
+            )
+        elif task_type == 'line_to_plane':
+            return LineToPlaneDistance(
+                name=task_name,
+                point_cloud=self.point_cloud,
+                config=task_config,
+                spatial_index=self.spatial_index
+            )
+        else:
+            print(f"警告: 未知的任务类型: {task_type}")
+            return None
+
+    def run_all_tasks(self) -> List[Dict[str, Any]]:
+        """
+        执行所有测量任务
+
+        Returns:
+            List[Dict]: 任务结果列表
+        """
+        import time
+
+        self.results = []
+
+        # 创建性能日志记录器
+        if self.output_dir_override:
+            output_dir = Path(self.output_dir_override)
+        else:
+            output_dir = Path(__file__).parent.parent / "results"
+
+        perf_logger = PerformanceLogger(str(output_dir))
+
+        # 记录初始化阶段性能
+        if hasattr(self, 'pointcloud_load_time'):
+            perf_logger.log_pointcloud_loading(
+                self.pointcloud_load_time,
+                len(self.point_cloud),
+                self.pointcloud_filename
+            )
+        if hasattr(self, 'index_build_time'):
+            index_info = self.spatial_index.get_info() if self.spatial_index else {}
+            index_method = self.config.get('roi_extraction', {}).get('method', 'unknown')
+            perf_logger.log_spatial_index(
+                self.index_build_time,
+                index_method,
+                index_info
+            )
+
+        # 创建详细测量日志记录器
+        detail_logger = MeasurementLogger(str(output_dir))
+
+        for i, task in enumerate(self.tasks, 1):
+            print(f"\n[{i}/{len(self.tasks)}] 执行任务: {task.name}")
+            print("-" * 60)
+
+            # 记录任务开始到详细日志
+            detail_logger.log_task_start(i, len(self.tasks), task.name)
+
+            # 记录开始时间
+            start_time = time.time()
+
+            # 将详细日志记录器传递给任务
+            task.detail_logger = detail_logger
+
+            # 执行任务
+            result = task.execute()
+
+            # 记录结束时间
+            end_time = time.time()
+            execution_time = end_time - start_time
+
+            # 将执行时间添加到结果数据中
+            if result.data is None:
+                result.data = {}
+            result.data['execution_time'] = execution_time
+
+            # 获取详细的时间分布
+            timing_breakdown = result.data.get('timing_breakdown', {})
+
+            # 记录到性能日志
+            perf_logger.log_task(task.name, execution_time, timing_breakdown)
+
+            # 记录结果
+            result_dict = {
+                'name': result.task_name,
+                'type': result.task_type,
+                'status': 'success' if result.success else 'failed',
+                'data': result.data if result.success else None,
+                'error': result.error_message if not result.success else None,
+                'execution_time': execution_time
+            }
+            self.results.append(result_dict)
+
+            # 检查是否有宽度差值检查结果，如果有则生成额外的任务19结果
+            if result.success and result.data and 'width_variation' in result.data:
+                width_var = result.data['width_variation']
+
+                # 创建任务19的结果
+                task19_name = task.name.replace("任务18", "任务19").replace("门口宽度", "门口宽度差值")
+                task19_result_dict = {
+                    'name': task19_name,
+                    'type': result.task_type,
+                    'status': 'success',
+                    'data': {
+                        'upper_width': width_var['upper_width'],
+                        'middle_width': width_var['middle_width'],
+                        'lower_width': width_var['lower_width'],
+                        'diff_upper_middle': width_var['diff_upper_middle'],
+                        'diff_middle_lower': width_var['diff_middle_lower'],
+                        'diff_upper_lower': width_var['diff_upper_lower'],
+                        'max_diff': max(abs(width_var['diff_upper_middle']),
+                                       abs(width_var['diff_middle_lower']),
+                                       abs(width_var['diff_upper_lower'])),
+                        'tolerance': width_var['tolerance'],
+                        'pass': width_var['pass'],
+                        'execution_time': execution_time
+                    },
+                    'error': None,
+                    'execution_time': execution_time
+                }
+                self.results.append(task19_result_dict)
+
+                # 打印任务19结果
+                print(f"\n  [自动生成] {task19_name}")
+                print(f"    上部宽度: {width_var['upper_width']:.4f} mm")
+                print(f"    中部宽度: {width_var['middle_width']:.4f} mm")
+                print(f"    下部宽度: {width_var['lower_width']:.4f} mm")
+                print(f"    最大差值: {task19_result_dict['data']['max_diff']:.4f} mm")
+                print(f"    公差要求: ±{width_var['tolerance']:.1f} mm")
+                status = "✅ 通过" if width_var['pass'] else "❌ 超差"
+                print(f"    检查结果: {status}")
+
+            # 打印结果
+            if result.success:
+                print(f"[成功] 任务完成 (用时: {execution_time:.2f} 秒)")
+                if result.data:
+                    print(f"  平均距离: {result.data.get('mean_distance', 0):.2f} mm")
+                    print(f"  标准差: {result.data.get('std_distance', 0):.2f} mm")
+                    print(f"  范围: [{result.data.get('min_distance', 0):.2f}, {result.data.get('max_distance', 0):.2f}] mm")
+                    print(f"  测量点数: {result.data.get('point_count', result.data.get('sample_count', 0))}")
+                    if 'source_inlier_ratio' in result.data:
+                        print(f"  源平面内点率: {result.data.get('source_inlier_ratio', 0):.1%}")
+                    if 'target_inlier_ratio' in result.data:
+                        print(f"  目标平面内点率: {result.data.get('target_inlier_ratio', 0):.1%}")
+                    if 'plane_inlier_ratio' in result.data:
+                        print(f"  平面内点率: {result.data.get('plane_inlier_ratio', 0):.1%}")
+
+                # 记录成功到详细日志
+                detail_logger.log_task_success(execution_time)
+                if result.data:
+                    # 记录测量结果到详细日志
+                    if 'mean_distance' in result.data:
+                        detail_logger.log_distance_result(
+                            result.data.get('mean_distance', 0),
+                            result.data.get('std_distance', 0),
+                            result.data.get('min_distance', 0),
+                            result.data.get('max_distance', 0)
+                        )
+            else:
+                print(f"[失败] {result.error_message} (用时: {execution_time:.2f} 秒)")
+                # 记录失败到详细日志
+                detail_logger.log_task_failure(result.error_message, execution_time)
+
+            # 可视化（如果配置中启用）
+            if result.success and task.config.get('visualization', False):
+                self._visualize_result(task, result)
+
+        # 保存性能日志
+        perf_logger.save()
+
+        # 关闭详细测量日志
+        detail_logger.close()
+
+        # 导出Excel结果
+        self._export_results()
+
+        return self.results
+
+    def _visualize_result(self, task, result: MeasurementResult):
+        """
+        可视化测量结果
+
+        Args:
+            task: 测量任务
+            result: 测量结果
+        """
+        try:
+            import open3d as o3d
+            import matplotlib.pyplot as plt
+
+            print("  生成可视化...")
+
+            data = result.data
+            if not data:
+                return
+
+            # 创建可视化窗口
+            vis = o3d.visualization.Visualizer()
+            vis.create_window(window_name=f"测量结果: {task.name}", width=1200, height=800)
+
+            # 显示源平面点云（蓝色）
+            source_pcd = o3d.geometry.PointCloud()
+            source_pcd.points = o3d.utility.Vector3dVector(data['_source_points'])
+            source_pcd.paint_uniform_color([0.2, 0.4, 0.8])
+            vis.add_geometry(source_pcd)
+
+            # 显示目标平面点云（绿色）
+            target_pcd = o3d.geometry.PointCloud()
+            target_pcd.points = o3d.utility.Vector3dVector(data['_target_points'])
+            target_pcd.paint_uniform_color([0.2, 0.8, 0.4])
+            vis.add_geometry(target_pcd)
+
+            # 显示采样点（红色）
+            sampled_pcd = o3d.geometry.PointCloud()
+            sampled_pcd.points = o3d.utility.Vector3dVector(data['_sampled_points'])
+            sampled_pcd.paint_uniform_color([0.9, 0.2, 0.2])
+            vis.add_geometry(sampled_pcd)
+
+            # 运行可视化
+            vis.run()
+            vis.destroy_window()
+
+        except Exception as e:
+            print(f"  可视化失败: {e}")
+
+    def _export_results(self):
+        """导出结果到Excel"""
+        try:
+            # 转换为MeasurementResult对象列表
+            measurement_results = []
+            for result_dict in self.results:
+                measurement_results.append(MeasurementResult(
+                    task_name=result_dict['name'],
+                    task_type=result_dict['type'],
+                    success=result_dict['status'] == 'success',
+                    data=result_dict.get('data'),
+                    error_message=result_dict.get('error')
+                ))
+
+            # 导出Excel（支持命令行覆盖输出目录）
+            if self.output_dir_override:
+                output_dir = Path(self.output_dir_override)
+            else:
+                output_dir = Path(__file__).parent.parent / "results"
+
+            exporter = ExcelExporter(str(output_dir))
+
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            filename = f"measurement_results_{timestamp}"
+
+            # 传递任务配置信息到导出器
+            excel_path = exporter.export(measurement_results, filename, self.config.get('tasks', []))
+            print(f"\n结果已导出到: {excel_path}")
+
+        except Exception as e:
+            print(f"\n导出Excel失败: {e}")
+
+    def _build_spatial_index(self):
+        """根据配置构建空间索引"""
+        index_config = self.config.get('roi_extraction', {})
+        index_method = index_config.get('method', 'linear').lower()
+
+        if index_method == 'linear':
+            print("  方法: 线性扫描（无索引）")
+            self.spatial_index = LinearScanIndex(self.point_cloud)
+
+        elif index_method == 'grid':
+            grid_size = index_config.get('grid_size', 100.0)
+            use_optimized = index_config.get('use_optimized', True)  # 默认使用优化版
+
+            if use_optimized:
+                print(f"  方法: Grid网格索引（优化版 - numpy向量化）")
+                self.spatial_index = GridIndexOptimized(self.point_cloud, grid_size=grid_size)
+            else:
+                print(f"  方法: Grid网格索引（原始版 - Python循环）")
+                self.spatial_index = GridIndex(self.point_cloud, grid_size=grid_size)
+
+        elif index_method == 'octree':
+            max_depth = index_config.get('max_depth', 8)
+            max_points = index_config.get('max_points_per_node', 1000)
+            print(f"  方法: Octree八叉树索引")
+            self.spatial_index = OctreeIndex(self.point_cloud, max_depth=max_depth, max_points_per_node=max_points)
+
+        elif index_method == 'kdtree':
+            print(f"  方法: KDTree索引")
+            self.spatial_index = KDTreeIndex(self.point_cloud)
+
+        else:
+            print(f"  警告: 未知的索引方法 '{index_method}'，使用线性扫描")
+            self.spatial_index = LinearScanIndex(self.point_cloud)
+
+        # 显示索引信息
+        info = self.spatial_index.get_info()
+        if 'grid_count' in info:
+            print(f"  详细信息: {info['grid_count']:,} 个网格, 平均 {info['avg_points_per_grid']:.1f} 点/格")
+        elif 'node_count' in info:
+            print(f"  详细信息: {info['node_count']:,} 个节点")