from __future__ import annotations import importlib.util import posixpath import time from pathlib import Path from typing import Any from zipfile import ZIP_DEFLATED, ZipFile from xml.sax.saxutils import escape from xml.etree import ElementTree as ET import numpy as np from cad_converter import CADConverter from config_loader import ConfigLoader from evaluation import Evaluator from io_handler import IOHandler from preprocessing.curvature import CurvatureEstimator from preprocessing.filters import PreprocessingRunner from preprocessing.noise import GaussianNoiseAdder from preprocessing.plane import PlaneEstimator from registration import PointCloudRegistration from visualization import Visualizer DEFAULT_RESULT_TEMPLATE = Path("excel/测量结果.xlsx") DEFAULT_COMBINED_CONTEXT_POINTCLOUD = Path("datasets/train/train_2_10M.ply") class MeasurementPipeline: """Main workflow entry class for geometry measurement tasks.""" def __init__(self, config_path: str | None = None) -> None: self.config = ConfigLoader().load(config_path) self.io = IOHandler() self.cad_converter = CADConverter() self.noise_adder = GaussianNoiseAdder() self.filters = PreprocessingRunner() self.registration = PointCloudRegistration() self.evaluator = Evaluator() self.curvature = CurvatureEstimator() self.plane = PlaneEstimator() self.visualizer = Visualizer() def run_cad2mesh(self, cad_path: str | None = None, mesh_path: str | None = None, precision: float | None = None) -> Path: paths = self.config["paths"] cad_config = self.config["cad"] input_step = self.io.require_cad(cad_path or paths["cad"]) output_stl = self.io.ensure_parent_dir(mesh_path or paths["output_mesh"] or paths["mesh"]) return self.cad_converter.convert_step_to_stl( input_step=input_step, output_stl=output_stl, precision=float(precision if precision is not None else cad_config["precision"]), ) def run_denoise( self, pointcloud_path: str | None = None, noise_std: float | None = None, no_visualize: bool = False, ) -> dict[str, Any]: path = pointcloud_path or self.config["paths"]["pointcloud"] clean = self.io.load_point_cloud(path) noise_config = self.config["noise"] adder = GaussianNoiseAdder( noise_std=float(noise_std if noise_std is not None else noise_config["std"]), seed=noise_config.get("seed"), ) noisy = adder.apply(clean) noise_level = adder.noise_level(clean, noisy) filter_results = {} if self.config["filters"].get("enabled", True): filter_results = self.filters.run(noisy, self.config["filters"]) print(f"Clean points: {len(clean.points)}") print(f"Noisy points: {len(noisy.points)}") print(f"Mean noise offset: {noise_level:.6f}") for name, result in filter_results.items(): cd, _, _ = self.evaluator.chamfer_distance(clean, result.pointcloud) result.metrics["chamfer_distance"] = cd print(f"{name}: keep={result.metrics['keep_ratio']:.2f}% CD={cd:.6f}") cd_noisy, _, _ = self.evaluator.chamfer_distance(clean, noisy) print(f"Noisy without filtering: CD={cd_noisy:.6f}") if self._should_visualize(no_visualize): first_result = next(iter(filter_results.values()), None) self.visualizer.show_denoise( clean, noisy, first_result.pointcloud if first_result else None, first_result.name if first_result else "Filtered", ) return { "clean": clean, "noisy": noisy, "noise_level": noise_level, "filter_results": filter_results, "noisy_chamfer_distance": cd_noisy, } def run_registration( self, pointcloud_path: str | None = None, mesh_path: str | None = None, voxel_size: float | None = None, roll: float | None = None, pitch: float | None = None, yaw: float | None = None, trans_range: float | None = None, no_visualize: bool = False, ) -> dict[str, Any]: paths = self.config["paths"] reg_config = self.config["registration"] source = self.io.load_point_cloud(pointcloud_path or paths["noisy_pointcloud"] or paths["pointcloud"]) mesh_file = str(self.io.require_file(mesh_path or paths["mesh"], "Mesh")) return self._run_registration_with_pointcloud( source=source, mesh_file=mesh_file, voxel_size=voxel_size, roll=roll, pitch=pitch, yaw=yaw, trans_range=trans_range, no_visualize=no_visualize, reg_config=reg_config, ) def _run_registration_with_pointcloud( self, source: Any, mesh_file: str, voxel_size: float | None, roll: float | None, pitch: float | None, yaw: float | None, trans_range: float | None, no_visualize: bool, reg_config: dict[str, Any], ) -> dict[str, Any]: result = self.registration.run_simulated_registration( source=source, voxel_size=float(voxel_size if voxel_size is not None else reg_config["voxel_size"]), roll=float(roll if roll is not None else reg_config["roll"]), pitch=float(pitch if pitch is not None else reg_config["pitch"]), yaw=float(yaw if yaw is not None else reg_config["yaw"]), trans_range=float(trans_range if trans_range is not None else reg_config["trans_range"]), random_seed=reg_config.get("random_seed"), ) points = self.io.point_cloud_to_numpy(result.recovered) mesh = self.io.load_mesh(mesh_file) _, distance_stats = self.evaluator.point_to_mesh_distance_torch( points, mesh, chunk_size=int(self.config["evaluation"]["chunk_size"]), triangle_chunk_size=int(self.config["evaluation"]["triangle_chunk_size"]), ) print("Registration result") print(f"Fitness: {result.fitness:.4f}") print(f"RMSE: {result.rmse:.6f}") print(f"Transformation:\n{result.transformation}") self._print_distance_stats(distance_stats) if self._should_visualize(no_visualize): self.visualizer.show_registration(source, result.target, result.recovered) return { "source": source, "registration": result, "distance_stats": distance_stats, } def run_contour( self, pointcloud_path: str | None = None, cad_path: str | None = None, mesh_path: str | None = None, mesh_deflection: float | None = None, point_chunk_size: int | None = None, triangle_chunk_size: int | None = None, save_path: str | None = None, excel_output_path: str | None = None, result_template_path: str | Path | None = DEFAULT_RESULT_TEMPLATE, combined_context_pointcloud_path: str | Path | None = DEFAULT_COMBINED_CONTEXT_POINTCLOUD, no_visualize: bool = False, ) -> dict[str, Any]: paths = self.config["paths"] eval_config = self.config["evaluation"] if mesh_path: pointcloud = self.io.load_point_cloud(pointcloud_path or paths["pointcloud"]) mesh = self.io.load_mesh(mesh_path) points = self.io.point_cloud_to_numpy(pointcloud) distances, stats = self.evaluator.point_to_mesh_distance_torch( points, mesh, chunk_size=int(point_chunk_size if point_chunk_size is not None else eval_config["chunk_size"]), triangle_chunk_size=int( triangle_chunk_size if triangle_chunk_size is not None else eval_config["triangle_chunk_size"] ), ) self._print_distance_stats(stats) heatmap = self.evaluator.colorize_distances(points, distances) if save_path: self.io.save_point_cloud(heatmap, save_path) print(f"Saved contour heatmap: {save_path}") excel_output = None if excel_output_path: excel_output = self._save_contour_excel( excel_output_path, rows=[{ "task_id": self._contour_task_id_from_sample(Path(mesh_path).stem), "mean": stats.mean, "status": "成功", }], result_template_path=result_template_path, ) print(f"Saved contour Excel: {excel_output}") if self._should_visualize(no_visualize): self.visualizer.show_curvature(heatmap, window_name="Point to Mesh Distance Heatmap") return { "pointcloud": heatmap, "mesh": mesh, "distances": distances, "distance_stats": stats, "excel_output": excel_output, } return self.run_contour_distance( cad_path=cad_path or paths["cad"], pointcloud_path=pointcloud_path or paths["pointcloud"], mesh_deflection=mesh_deflection, point_chunk_size=point_chunk_size, triangle_chunk_size=triangle_chunk_size, save_path=save_path, excel_output_path=excel_output_path, result_template_path=result_template_path, no_visualize=no_visualize, ) def run_contour_distance( self, cad_path: str | None = None, pointcloud_path: str | None = None, mesh_deflection: float | None = None, point_chunk_size: int | None = None, triangle_chunk_size: int | None = None, save_path: str | None = None, excel_output_path: str | None = None, result_template_path: str | Path | None = DEFAULT_RESULT_TEMPLATE, combined_context_pointcloud_path: str | Path | None = DEFAULT_COMBINED_CONTEXT_POINTCLOUD, no_visualize: bool = False, ) -> dict[str, Any]: paths = self.config["paths"] eval_config = self.config["evaluation"] input_step = self.io.require_cad(cad_path or paths["cad"]) pointcloud = self.io.load_point_cloud(pointcloud_path or paths["noisy_pointcloud"] or paths["pointcloud"]) points = self.io.point_cloud_to_numpy(pointcloud) vertices, faces = self.cad_converter.step_to_numpy_mesh( input_step, mesh_deflection=float( mesh_deflection if mesh_deflection is not None else eval_config["cad_mesh_deflection"] ), ) print(f"CAD mesh vertices: {vertices.shape[0]}") print(f"CAD mesh triangles: {faces.shape[0]}") distances, stats = self.evaluator.point_to_triangles_distance_torch( points, vertices, faces, point_chunk_size=int(point_chunk_size if point_chunk_size is not None else eval_config["chunk_size"]), triangle_chunk_size=int( triangle_chunk_size if triangle_chunk_size is not None else eval_config["triangle_chunk_size"] ), ) self._print_distance_stats(stats) heatmap = self.evaluator.colorize_distances(points, distances) output = save_path if output: self.io.save_point_cloud(heatmap, output) print(f"Saved contour heatmap: {output}") excel_output = None if excel_output_path: excel_output = self._save_contour_excel( excel_output_path, cad_path=input_step, stats=stats, result_template_path=result_template_path, ) print(f"Saved contour Excel: {excel_output}") if self._should_visualize(no_visualize): self.visualizer.show_curvature(heatmap, window_name="Point to CAD Distance Heatmap") return { "pointcloud": heatmap, "distances": distances, "distance_stats": stats, "cad_vertices": vertices, "cad_faces": faces, "excel_output": excel_output, } def run_contour_batch( self, task_ids: str | list[int | str] | None = None, all_tasks: bool = False, data_root: str | Path = "datasets", pointcloud_kind: str = "noise", mesh_deflection: float | None = None, point_chunk_size: int | None = None, triangle_chunk_size: int | None = None, save_path: str | None = None, excel_output_path: str | None = None, result_template_path: str | Path | None = DEFAULT_RESULT_TEMPLATE, combined_context_pointcloud_path: str | Path | None = DEFAULT_COMBINED_CONTEXT_POINTCLOUD, save_combined_visualization: bool = True, no_visualize: bool = False, ) -> dict[str, Any]: batch_start_time = time.perf_counter() data_root = Path(data_root) task_inputs = self._resolve_contour_task_inputs(task_ids, all_tasks, data_root, pointcloud_kind) task_summaries: dict[str, dict[str, Any]] = {} results: dict[str, Any] = {} combined_visualizations: list[tuple[str, Any]] = [] for task_input in task_inputs: task_start_time = time.perf_counter() task_id = task_input["task_id"] sample_id = task_input["sample_id"] cad_path = task_input.get("cad_path") mesh_path = task_input.get("mesh_path") pointcloud_path = task_input["pointcloud_path"] task_save_path = self._contour_task_save_path(save_path, sample_id) summary = task_summaries.setdefault( task_id, { "task_id": task_id, "sample_ids": [], "point_count": 0, "distance_sum": 0.0, "max": None, "elapsed_sec": 0.0, "failures": [], "sample_results": {}, }, ) summary["sample_ids"].append(sample_id) print(f"===== contour task {task_id} / sample {sample_id} =====") model_path = mesh_path or cad_path missing = [str(path) for path in (model_path, pointcloud_path) if path and not path.exists()] if missing: print("Missing input: " + "; ".join(missing)) elapsed_sec = time.perf_counter() - task_start_time print(f"任务 {task_id} / 样本 {sample_id} 运行时间: {elapsed_sec:.2f} 秒") summary["elapsed_sec"] += elapsed_sec summary["failures"].append(f"{sample_id}: Missing input: " + "; ".join(missing)) continue try: if mesh_path: result = self.run_contour( mesh_path=str(mesh_path), pointcloud_path=str(pointcloud_path), point_chunk_size=point_chunk_size, triangle_chunk_size=triangle_chunk_size, save_path=str(task_save_path) if task_save_path else None, excel_output_path=None, no_visualize=True, ) else: result = self.run_contour_distance( cad_path=cad_path, pointcloud_path=pointcloud_path, mesh_deflection=mesh_deflection, point_chunk_size=point_chunk_size, triangle_chunk_size=triangle_chunk_size, save_path=str(task_save_path) if task_save_path else None, excel_output_path=None, no_visualize=True, ) except Exception as exc: print(f"Task {task_id} sample {sample_id} failed: {exc}") elapsed_sec = time.perf_counter() - task_start_time print(f"任务 {task_id} / 样本 {sample_id} 运行时间: {elapsed_sec:.2f} 秒") summary["elapsed_sec"] += elapsed_sec summary["failures"].append(f"{sample_id}: {exc}") continue stats = result["distance_stats"] distances = result["distances"] elapsed_sec = time.perf_counter() - task_start_time print(f"任务 {task_id} / 样本 {sample_id} 运行时间: {elapsed_sec:.2f} 秒") summary["point_count"] += len(distances) summary["distance_sum"] += float(np.sum(distances)) summary["max"] = stats.max if summary["max"] is None else max(summary["max"], stats.max) summary["elapsed_sec"] += elapsed_sec summary["sample_results"][sample_id] = result combined_visualizations.append((sample_id, result["pointcloud"])) rows: list[dict[str, Any]] = [] for task_id in sorted(task_summaries, key=self._contour_sort_key): summary = task_summaries[task_id] sample_count = len(summary["sample_results"]) failure_count = len(summary["failures"]) if sample_count: status = "成功" if failure_count == 0 else "部分成功" mean = summary["distance_sum"] / summary["point_count"] max_value = summary["max"] else: status = "失败" mean = None max_value = None rows.append({ "task_id": task_id, "sample_count": sample_count, "samples": ",".join(summary["sample_ids"]), "max": max_value, "mean": mean, "status": status, "elapsed_sec": summary["elapsed_sec"], }) results[task_id] = summary["sample_results"] excel_output = None if excel_output_path: excel_output = self._save_contour_excel( excel_output_path, rows=rows, result_template_path=result_template_path, ) print(f"Saved contour Excel: {excel_output}") combined_visualization_output = None combined_pointclouds = [] if save_combined_visualization and combined_visualizations: combined_visualization_output = self._contour_combined_save_path(save_path) self._save_contour_combined_pointcloud_stream( combined_visualizations, combined_context_pointcloud_path, combined_visualization_output, ) print(f"Saved combined contour heatmap: {combined_visualization_output}") if self._should_visualize(no_visualize) and combined_visualizations: combined_pointclouds = self._contour_combined_pointclouds( combined_visualizations, combined_context_pointcloud_path, ) self.visualizer.show_combined_pointclouds( combined_pointclouds, labels=[str(index) for index, _ in enumerate(combined_pointclouds, start=1)], window_name="Contour Batch Distance Heatmaps", ) total_elapsed_sec = time.perf_counter() - batch_start_time print(f"Contour batch total execution time: {total_elapsed_sec:.2f} seconds") return { "rows": rows, "results": results, "excel_output": excel_output, "combined_visualization_output": combined_visualization_output, "total_elapsed_sec": total_elapsed_sec, } def run_curvature( self, pointcloud_path: str | None = None, radius: float | None = None, max_nn: int | None = None, color_scale: float | None = None, save_path: str | None = None, no_visualize: bool = False, ) -> dict[str, Any]: config = self.config["curvature"] pointcloud = self.io.load_point_cloud(pointcloud_path or self.config["paths"]["pointcloud"]) curvatures = self.curvature.compute( pointcloud, radius=float(radius if radius is not None else config["radius"]), max_nn=int(max_nn if max_nn is not None else config["max_nn"]), ) colored = self.curvature.colorize( pointcloud, curvatures, color_scale=float(color_scale if color_scale is not None else config["color_scale"]), ) output = save_path or self.config["paths"].get("curvature_output") if output: self.io.save_point_cloud(colored, output) print(f"Saved curvature point cloud: {output}") print(f"Curvature mean={np.mean(curvatures):.8f}, max={np.max(curvatures):.8f}") if self._should_visualize(no_visualize): self.visualizer.show_curvature(colored) return {"pointcloud": colored, "curvatures": curvatures} def run_plane( self, pointcloud_path: str | None = None, threshold_mm: float | None = None, max_iterations: int | None = None, no_visualize: bool = False, ) -> dict[str, Any]: config = self.config["plane"] if pointcloud_path: pointcloud = self.io.load_point_cloud(pointcloud_path) else: pointcloud = self.plane.create_square_plane_with_noise( size_mm=float(config["size_mm"]), noise_std_mm=float(config["noise_std_mm"]), num_points=int(config["num_points"]), ) result = self.plane.estimate( pointcloud, threshold_mm=float(threshold_mm if threshold_mm is not None else config["threshold_mm"]), max_iterations=int(max_iterations if max_iterations is not None else config["max_iterations"]), ) print(f"Plane model: {result.model}") print(f"Inlier ratio: {result.inlier_ratio:.2f}%") print(f"Normal angle error: {result.normal_angle_deg:.6f} deg") print(f"Plane RMSE: {result.rmse:.8f} m") if self._should_visualize(no_visualize): self.visualizer.show_plane(result) return {"plane_result": result} def run_full(self, no_visualize: bool = False) -> dict[str, Any]: outputs: dict[str, Any] = {} if self.config["cad"].get("enabled", False): outputs["cad2mesh"] = self.run_cad2mesh() if self.config["noise"].get("enabled", True): outputs["denoise"] = self.run_denoise(no_visualize=True) mesh_file = str(self.io.require_file(self.config["paths"]["mesh"], "Mesh")) outputs["registration"] = self._run_registration_with_pointcloud( source=outputs["denoise"]["noisy"], mesh_file=mesh_file, voxel_size=None, roll=None, pitch=None, yaw=None, trans_range=None, no_visualize=no_visualize, reg_config=self.config["registration"], ) else: registration_input = self.config["paths"].get("noisy_pointcloud") or self.config["paths"]["pointcloud"] outputs["registration"] = self.run_registration(pointcloud_path=registration_input, no_visualize=no_visualize) outputs["contour"] = self.run_contour() return outputs def _should_visualize(self, no_visualize: bool) -> bool: return bool(self.config["visualization"].get("enabled", True)) and not no_visualize def _print_distance_stats(self, stats: Any) -> None: print("Distance stats") print(f"Mean: {stats.mean:.8f}") print(f"Max: {stats.max:.8f}") print(f"RMSE: {stats.rmse:.8f}") print(f"Min: {stats.min:.8f}") print(f"Std: {stats.std:.8f}") def _save_contour_excel( self, output_path: str | Path, *, rows: list[dict[str, Any]] | None = None, cad_path: str | Path | None = None, stats: Any | None = None, result_template_path: str | Path | None = DEFAULT_RESULT_TEMPLATE, ) -> Path: output = self._normalize_excel_output_path(output_path) if rows is None: if cad_path is None or stats is None: raise ValueError("cad_path and stats are required when rows is not provided.") rows = [{ "task_id": self._contour_task_id_from_sample(Path(cad_path).stem), "max": stats.max, "mean": stats.mean, "status": "成功", }] if result_template_path and Path(result_template_path).exists(): sheet_name, table = self._build_measurement_table_from_template(rows, Path(result_template_path)) else: sheet_name = "contour" table = [ [ "任务号", "样本数", "样本", "最大值", "平均值", "状态", "elapsed_sec", ] ] for row in rows: table.append([ row.get("task_id"), row.get("sample_count"), row.get("samples"), row.get("max"), row.get("mean"), row.get("status"), row.get("elapsed_sec"), ]) self._write_xlsx(table, output, sheet_name=sheet_name) return output def _build_measurement_table_from_template( self, rows: list[dict[str, Any]], template_path: Path, ) -> tuple[str, list[list[Any]]]: template_sheet_name, headers, template_rows = self._read_measurement_template(template_path) task_col = self._find_header_index(headers, ["任务号", "任务编号", "task_id"]) result_col = self._find_header_index(headers, ["测量结果"]) error_col = self._find_header_index(headers, ["误差结果"]) conclusion_col = self._find_header_index(headers, ["测量结论", "结论", "状态"]) result_by_task = { self._normalize_template_task_id(row.get("task_id")): row for row in rows if row.get("task_id") is not None } used_task_ids = set() table = [headers] for template_row in template_rows: output_row = list(template_row) if len(output_row) < len(headers): output_row.extend([""] * (len(headers) - len(output_row))) task_id = self._normalize_template_task_id(output_row[task_col]) result_row = result_by_task.get(task_id) if result_row is None: continue output_row[result_col] = self._row_measurement_mean(result_row) output_row[error_col] = result_row.get("max") output_row[conclusion_col] = self._row_measurement_conclusion(result_row) table.append(output_row) used_task_ids.add(task_id) for task_id, result_row in result_by_task.items(): if task_id in used_task_ids: continue output_row = [""] * len(headers) output_row[task_col] = task_id output_row[result_col] = self._row_measurement_mean(result_row) output_row[error_col] = result_row.get("max") output_row[conclusion_col] = self._row_measurement_conclusion(result_row) table.append(output_row) return template_sheet_name, table def _read_measurement_template(self, template_path: Path) -> tuple[str, list[Any], list[list[Any]]]: ns = { "a": "http://schemas.openxmlformats.org/spreadsheetml/2006/main", "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships", } with ZipFile(template_path) as xlsx: names = xlsx.namelist() shared_strings = [] if "xl/sharedStrings.xml" in names: shared_root = ET.fromstring(xlsx.read("xl/sharedStrings.xml")) for item in shared_root.findall("a:si", ns): shared_strings.append( "".join((text.text or "") for text in item.findall(".//a:t", ns)) ) workbook_root = ET.fromstring(xlsx.read("xl/workbook.xml")) first_sheet = workbook_root.find("a:sheets/a:sheet", ns) if first_sheet is None: raise RuntimeError(f"模板中没有工作表: {template_path}") sheet_name = first_sheet.attrib.get("name", "Sheet1") relation_id = first_sheet.attrib[ "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}id" ] rels_root = ET.fromstring(xlsx.read("xl/_rels/workbook.xml.rels")) relation_targets = { rel.attrib["Id"]: rel.attrib["Target"] for rel in rels_root } sheet_path = self._xlsx_rel_target_to_path(relation_targets[relation_id]) sheet_root = ET.fromstring(xlsx.read(sheet_path)) parsed_rows = [] max_col_count = 0 for row in sheet_root.findall("a:sheetData/a:row", ns): values_by_col = {} for cell in row.findall("a:c", ns): col_index = self._xlsx_col_to_index(cell.attrib["r"]) values_by_col[col_index] = self._xlsx_cell_value(cell, shared_strings, ns) max_col_count = max(max_col_count, col_index) if values_by_col: parsed_rows.append(values_by_col) table = [] for values_by_col in parsed_rows: table.append([ values_by_col.get(col_index, "") for col_index in range(1, max_col_count + 1) ]) if not table: raise RuntimeError(f"模板为空: {template_path}") return sheet_name, table[0], table[1:] def _row_measurement_mean(self, row: dict[str, Any]) -> Any: for key in ("mean", "mean_dist", "average_distance", "测量结果"): if key in row: return row.get(key) return None def _row_measurement_conclusion(self, row: dict[str, Any]) -> str: status = str(row.get("status", "")).strip() if status in {"成功", "部分成功", "失败"}: return status if status: return "失败" value = self._row_measurement_mean(row) if value is None: return "失败" try: if not np.isfinite(float(value)): return "失败" except (TypeError, ValueError): return "失败" return "成功" def _normalize_template_task_id(self, value: Any) -> str: text = str(value).strip() if text.endswith(".0") and text[:-2].isdigit(): return text[:-2] return text def _find_header_index(self, headers: list[Any], candidates: list[str]) -> int: normalized = [str(header).strip() for header in headers] for candidate in candidates: if candidate in normalized: return normalized.index(candidate) raise ValueError(f"模板缺少列: {'/'.join(candidates)}") def _resolve_contour_task_inputs( self, task_ids: str | list[int | str] | None, all_tasks: bool, data_root: Path, pointcloud_kind: str, ) -> list[dict[str, Any]]: if all_tasks: task_values = self._discover_contour_task_values(data_root) elif task_ids is None: raise ValueError("Batch contour requires --tasks or --all.") elif isinstance(task_ids, str): task_values = [item.strip() for item in task_ids.split(",") if item.strip()] else: task_values = [str(item).strip() for item in task_ids if str(item).strip()] task_inputs: list[dict[str, Any]] = [] for task_value in task_values: task_inputs.extend( self._contour_inputs_for_task_value(data_root, task_value, pointcloud_kind) ) return sorted( task_inputs, key=lambda item: (self._contour_sort_key(item["task_id"]), item["sample_id"]), ) def _discover_contour_task_values(self, data_root: Path) -> list[str]: mesh_dir = data_root / "Mesh" if mesh_dir.exists(): sample_ids = [] for pattern in ("*.stl", "*.obj", "*.ply"): sample_ids.extend(path.stem for path in mesh_dir.glob(pattern)) if sample_ids: return sorted( {self._contour_task_id_from_sample(sample_id) for sample_id in sample_ids}, key=self._contour_sort_key, ) cad_dir = data_root / "CAD" if not cad_dir.exists(): raise FileNotFoundError(f"CAD or Mesh directory not found under: {data_root}") return sorted((path.stem for path in cad_dir.glob("*.step")), key=self._contour_sort_key) def _contour_inputs_for_task_value( self, data_root: Path, task_value: str, pointcloud_kind: str, ) -> list[dict[str, Any]]: mesh_paths = self._contour_mesh_paths_for_task_value(data_root, task_value) if mesh_paths: return [ { "task_id": self._contour_task_id_from_sample(mesh_path.stem), "sample_id": mesh_path.stem, "mesh_path": mesh_path, "pointcloud_path": self._contour_pointcloud_path( data_root, mesh_path.stem, pointcloud_kind, ), } for mesh_path in mesh_paths ] sample_id = task_value return [ { "task_id": self._contour_task_id_from_sample(sample_id), "sample_id": sample_id, "cad_path": data_root / "CAD" / f"{sample_id}.step", "pointcloud_path": self._contour_pointcloud_path( data_root, sample_id, pointcloud_kind, ), } ] def _contour_mesh_paths_for_task_value(self, data_root: Path, task_value: str) -> list[Path]: mesh_dir = data_root / "Mesh" if not mesh_dir.exists(): return [] mesh_paths = [] for pattern in ("*.stl", "*.obj", "*.ply"): mesh_paths.extend(mesh_dir.glob(pattern)) if "_" in task_value: matched = [path for path in mesh_paths if path.stem == task_value] else: matched = [ path for path in mesh_paths if self._contour_task_id_from_sample(path.stem) == task_value ] return sorted(matched, key=lambda path: (self._contour_sort_key(path.stem), path.stem)) def _contour_task_id_from_sample(self, sample_id: str) -> str: return str(sample_id).split("_", 1)[0] def _contour_sort_key(self, value: str) -> tuple[int, int | str]: text = str(value) base = self._contour_task_id_from_sample(text) if base.isdigit(): return (0, int(base)) return (1, base) def _contour_pointcloud_path(self, data_root: Path, task_id: str, pointcloud_kind: str) -> Path: if pointcloud_kind == "clean": return data_root / "PointCloud" / f"{task_id}.ply" if pointcloud_kind == "noise": return data_root / "Noise_PointCloud" / f"{task_id}_noisy.ply" raise ValueError("pointcloud_kind must be 'noise' or 'clean'.") def _contour_task_save_path(self, save_path: str | None, task_id: str) -> Path | None: if not save_path: return None output = Path(save_path) if output.suffix.lower() != ".ply": output.mkdir(parents=True, exist_ok=True) return output / f"{task_id}_contour.ply" output.parent.mkdir(parents=True, exist_ok=True) if "{task_id}" in output.name: return output.with_name(output.name.replace("{task_id}", task_id)) return output.with_name(f"{output.stem}_{task_id}{output.suffix}") def _contour_combined_save_path(self, save_path: str | None) -> Path: if not save_path: return Path("vis") / "combined_contour.ply" output = Path(save_path) if output.suffix.lower() != ".ply": return output / "combined_contour.ply" if "{task_id}" in output.name: return output.with_name(output.name.replace("{task_id}", "combined")) return output.with_name(f"{output.stem}_combined{output.suffix}") def _save_contour_combined_pointcloud_stream( self, task_pointclouds: list[tuple[str, Any]], context_pointcloud_path: str | Path | None, output_path: str | Path, ) -> Path: output = self.io.ensure_parent_dir(output_path) context_path = Path(context_pointcloud_path) if context_pointcloud_path else None context_point_count = 0 context_header_size = None if context_path: if context_path.exists(): context_point_count, context_header_size = self._read_binary_xyz_ply_header(context_path) print(f"Streaming context point cloud into combined contour: {context_path}") else: print(f"Context point cloud not found, skipped: {context_path}") context_path = None task_point_count = sum(len(pointcloud.points) for _, pointcloud in task_pointclouds) total_point_count = context_point_count + task_point_count if total_point_count == 0: raise ValueError("No point clouds to save in combined contour.") header = ( "ply\n" "format binary_little_endian 1.0\n" f"element vertex {total_point_count}\n" "property float x\n" "property float y\n" "property float z\n" "property uchar red\n" "property uchar green\n" "property uchar blue\n" "end_header\n" ) with open(output, "wb") as output_file: output_file.write(header.encode("ascii")) if context_path: self._write_context_xyz_ply_stream( context_path, output_file, context_header_size, context_point_count, ) self._write_contour_pointclouds_stream(task_pointclouds, output_file) return output def _read_binary_xyz_ply_header(self, ply_path: Path) -> tuple[int, int]: with open(ply_path, "rb") as file: header_lines = [] while True: line = file.readline() if not line: raise ValueError(f"Invalid PLY header: {ply_path}") header_lines.append(line) if line.strip() == b"end_header": break header_size = file.tell() text_lines = [ line.decode("ascii", errors="replace").strip() for line in header_lines ] if len(text_lines) < 3 or text_lines[0] != "ply": raise ValueError(f"Not a PLY file: {ply_path}") if text_lines[1] != "format binary_little_endian 1.0": raise ValueError( "Only binary_little_endian PLY context point clouds can be streamed." ) vertex_count = None vertex_properties = [] in_vertex = False for line in text_lines[2:]: parts = line.split() if not parts: continue if parts[0] == "element": in_vertex = parts[1] == "vertex" if in_vertex: vertex_count = int(parts[2]) continue if in_vertex and parts[0] == "property": if len(parts) >= 2 and parts[1] == "list": raise ValueError("List properties are not supported in streamed PLY.") if len(parts) >= 3: vertex_properties.append((parts[1], parts[2])) expected = [("float", "x"), ("float", "y"), ("float", "z")] if vertex_count is None or vertex_properties[:3] != expected: raise ValueError( "Context PLY must start with vertex properties: float x, float y, float z." ) if len(vertex_properties) != 3: raise ValueError( "Context PLY must contain only float x/y/z vertices for streamed merge." ) return vertex_count, header_size def _write_context_xyz_ply_stream( self, input_path: Path, output_file: Any, header_size: int, point_count: int, ) -> None: record_size = 12 chunk_points = 1_000_000 gray = np.array([158, 158, 158], dtype=np.uint8) with open(input_path, "rb") as input_file: input_file.seek(header_size) remaining = point_count while remaining > 0: current = min(chunk_points, remaining) data = input_file.read(current * record_size) if len(data) != current * record_size: raise IOError(f"Unexpected end of context PLY: {input_path}") points = np.frombuffer(data, dtype=" None: chunk_points = 1_000_000 gray = np.array([158, 158, 158], dtype=np.uint8) for _, pointcloud in task_pointclouds: points = np.asarray(pointcloud.points) colors = np.asarray(pointcloud.colors) if pointcloud.has_colors() else gray for start in range(0, len(points), chunk_points): end = min(start + chunk_points, len(points)) self._write_xyz_rgb_chunk(output_file, points[start:end], colors if colors.ndim == 1 else colors[start:end]) def _write_xyz_rgb_chunk(self, output_file: Any, points: np.ndarray, colors: np.ndarray) -> None: points = np.asarray(points) colors = np.asarray(colors) count = len(points) if count == 0: return vertex_dtype = np.dtype([ ("x", " list[Any]: pointclouds = [] if context_pointcloud_path: context_path = Path(context_pointcloud_path) if context_path.exists(): context = self.io.load_point_cloud(context_path) if not context.has_colors(): context.paint_uniform_color([0.62, 0.62, 0.62]) pointclouds.append(context) print(f"Added context point cloud to combined contour: {context_path}") else: print(f"Context point cloud not found, skipped: {context_path}") pointclouds.extend(item[1] for item in task_pointclouds) return pointclouds def run_translate_align_pointcloud_to_cad( self, cad_path: str | None = None, pointcloud_path: str | None = None, output_dir: str = "results", axis: str = "x", crop_ratio: float = 0.01, mesh_deflection: float = 0.1, cad_sample_points: int = 30000, local_sample_points: int = 1000, scale_pointcloud: float = 1.0, no_visualize: bool = False, ) -> dict[str, Any]: """Run the archived front/tail centroid translation alignment workflow.""" paths = self.config["paths"] input_step = self.io.require_cad(cad_path or paths["cad"]) input_pointcloud = self._resolve_translate_align_pointcloud( pointcloud_path or paths["noisy_pointcloud"] or paths["pointcloud"] ) module = self._load_archive_module("translate_align_pointcloud_to_cad.py") return module.translate_align_by_front_tail( cad_path=str(input_step), pointcloud_path=str(input_pointcloud), output_dir=output_dir, axis=axis, crop_ratio=float(crop_ratio), mesh_deflection=float(mesh_deflection), cad_sample_points=int(cad_sample_points), local_sample_points=int(local_sample_points), scale_pointcloud=float(scale_pointcloud), visualize=not no_visualize, ) def _resolve_translate_align_pointcloud(self, pointcloud_path: str | Path) -> Path: path = Path(pointcloud_path) if path.is_dir(): preferred = path / "shifted_whole_pointcloud.ply" if preferred.exists(): return preferred candidates = sorted(path.glob("*.ply")) if len(candidates) == 1: return candidates[0] if not candidates: raise FileNotFoundError(f"No PLY point cloud found in directory: {path}") names = ", ".join(candidate.name for candidate in candidates) raise ValueError(f"Multiple PLY files found in {path}; specify one file explicitly. Found: {names}") return self.io.require_file(path, "PointCloud") def _load_archive_module(self, filename: str) -> Any: script_path = Path(__file__).resolve().parent / "archive" / filename if not script_path.exists(): raise FileNotFoundError(f"Archive script not found: {script_path}") module_name = f"_archive_{script_path.stem}" spec = importlib.util.spec_from_file_location(module_name, script_path) if spec is None or spec.loader is None: raise ImportError(f"Cannot load archive script: {script_path}") module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module def _normalize_excel_output_path(self, output_path: str | Path) -> Path: output = Path(output_path) if output.suffix.lower() != ".xlsx": output = output / "contour.xlsx" elif output.parent == Path("."): output = Path("excel") / output output.parent.mkdir(parents=True, exist_ok=True) return output def _write_xlsx(self, rows: list[list[Any]], output_path: Path, sheet_name: str) -> None: sheet_rows = [] for row_index, row_values in enumerate(rows, start=1): cells = [] for col_index, value in enumerate(row_values, start=1): cell_ref = self._xlsx_cell_ref(row_index, col_index) if isinstance(value, (int, float, np.integer, np.floating)) and not isinstance(value, bool): cells.append(f'{value}') else: text = escape("" if value is None else str(value)) cells.append(f'{text}') sheet_rows.append(f'{"".join(cells)}') workbook_xml = ( '' '' f'' '' ) worksheet_xml = ( '' '' '' f'{"".join(sheet_rows)}' '' '' ) content_types_xml = ( '' '' '' '' '' '' '' ) root_rels_xml = ( '' '' '' '' ) workbook_rels_xml = ( '' '' '' '' ) with ZipFile(output_path, "w", ZIP_DEFLATED) as xlsx: xlsx.writestr("[Content_Types].xml", content_types_xml) xlsx.writestr("_rels/.rels", root_rels_xml) xlsx.writestr("xl/workbook.xml", workbook_xml) xlsx.writestr("xl/_rels/workbook.xml.rels", workbook_rels_xml) xlsx.writestr("xl/worksheets/sheet1.xml", worksheet_xml) def _xlsx_cell_ref(self, row_index: int, col_index: int) -> str: letters = "" col = col_index while col: col, remainder = divmod(col - 1, 26) letters = chr(65 + remainder) + letters return f"{letters}{row_index}" def _xlsx_col_to_index(self, cell_ref: str) -> int: letters = "".join(ch for ch in cell_ref if ch.isalpha()) index = 0 for letter in letters: index = index * 26 + ord(letter.upper()) - ord("A") + 1 return index def _xlsx_rel_target_to_path(self, target: str) -> str: target = target.replace("\\", "/").lstrip("/") if target.startswith("xl/"): return target return posixpath.normpath(posixpath.join("xl", target)) def _xlsx_cell_value(self, cell: ET.Element, shared_strings: list[str], ns: dict[str, str]) -> str: value_node = cell.find("a:v", ns) inline_node = cell.find("a:is/a:t", ns) cell_type = cell.attrib.get("t") if cell_type == "s" and value_node is not None: return shared_strings[int(value_node.text)] if cell_type == "inlineStr" and inline_node is not None: return inline_node.text or "" if value_node is not None: return value_node.text or "" return ""