645 lines
27 KiB
Python
645 lines
27 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from zipfile import ZIP_DEFLATED, ZipFile
|
|
from xml.sax.saxutils import escape
|
|
|
|
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
|
|
|
|
|
|
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,
|
|
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)
|
|
return {"pointcloud": pointcloud, "mesh": mesh, "distances": distances, "distance_stats": stats}
|
|
|
|
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,
|
|
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,
|
|
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,
|
|
)
|
|
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,
|
|
no_visualize: bool = False,
|
|
) -> dict[str, Any]:
|
|
data_root = Path(data_root)
|
|
task_id_list = self._resolve_contour_task_ids(task_ids, all_tasks, data_root)
|
|
rows: list[dict[str, Any]] = []
|
|
results: dict[str, Any] = {}
|
|
|
|
for task_id in task_id_list:
|
|
cad_path = data_root / "CAD" / f"{task_id}.step"
|
|
pointcloud_path = self._contour_pointcloud_path(data_root, task_id, pointcloud_kind)
|
|
task_save_path = self._contour_task_save_path(save_path, task_id)
|
|
|
|
print(f"===== contour task {task_id} =====")
|
|
missing = [str(path) for path in (cad_path, pointcloud_path) if not path.exists()]
|
|
if missing:
|
|
print("Missing input: " + "; ".join(missing))
|
|
rows.append({"task_id": task_id, "max": None, "mean": None, "status": "失败"})
|
|
continue
|
|
|
|
try:
|
|
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 if len(task_id_list) > 1 else no_visualize,
|
|
)
|
|
except Exception as exc:
|
|
print(f"Task {task_id} failed: {exc}")
|
|
rows.append({"task_id": task_id, "max": None, "mean": None, "status": "失败"})
|
|
continue
|
|
|
|
stats = result["distance_stats"]
|
|
rows.append({"task_id": task_id, "max": stats.max, "mean": stats.mean, "status": "成功"})
|
|
results[task_id] = result
|
|
|
|
excel_output = None
|
|
if excel_output_path:
|
|
excel_output = self._save_contour_excel(excel_output_path, rows=rows)
|
|
print(f"Saved contour Excel: {excel_output}")
|
|
|
|
return {"rows": rows, "results": results, "excel_output": excel_output}
|
|
|
|
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,
|
|
) -> 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": Path(cad_path).stem, "max": stats.max, "mean": stats.mean, "status": "成功"}]
|
|
|
|
table = [
|
|
[
|
|
"任务号",
|
|
"最大值",
|
|
"平均值",
|
|
"状态",
|
|
]
|
|
]
|
|
for row in rows:
|
|
table.append([
|
|
row.get("task_id"),
|
|
row.get("max"),
|
|
row.get("mean"),
|
|
row.get("status"),
|
|
])
|
|
self._write_xlsx(table, output, sheet_name="contour")
|
|
return output
|
|
|
|
def _resolve_contour_task_ids(
|
|
self,
|
|
task_ids: str | list[int | str] | None,
|
|
all_tasks: bool,
|
|
data_root: Path,
|
|
) -> list[str]:
|
|
if all_tasks:
|
|
cad_dir = data_root / "CAD"
|
|
if not cad_dir.exists():
|
|
raise FileNotFoundError(f"CAD directory not found: {cad_dir}")
|
|
return sorted(path.stem for path in cad_dir.glob("*.step"))
|
|
|
|
if task_ids is None:
|
|
raise ValueError("Batch contour requires --tasks or --all.")
|
|
|
|
if isinstance(task_ids, str):
|
|
values = [item.strip() for item in task_ids.split(",")]
|
|
else:
|
|
values = [str(item).strip() for item in task_ids]
|
|
return [item for item in values if item]
|
|
|
|
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 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'<c r="{cell_ref}"><v>{value}</v></c>')
|
|
else:
|
|
text = escape("" if value is None else str(value))
|
|
cells.append(f'<c r="{cell_ref}" t="inlineStr"><is><t>{text}</t></is></c>')
|
|
sheet_rows.append(f'<row r="{row_index}">{"".join(cells)}</row>')
|
|
|
|
workbook_xml = (
|
|
'<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
|
|
'<workbook xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" '
|
|
'xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships">'
|
|
f'<sheets><sheet name="{escape(sheet_name[:31])}" sheetId="1" r:id="rId1"/></sheets>'
|
|
'</workbook>'
|
|
)
|
|
worksheet_xml = (
|
|
'<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
|
|
'<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">'
|
|
'<sheetData>'
|
|
f'{"".join(sheet_rows)}'
|
|
'</sheetData>'
|
|
'</worksheet>'
|
|
)
|
|
content_types_xml = (
|
|
'<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
|
|
'<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">'
|
|
'<Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>'
|
|
'<Default Extension="xml" ContentType="application/xml"/>'
|
|
'<Override PartName="/xl/workbook.xml" '
|
|
'ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"/>'
|
|
'<Override PartName="/xl/worksheets/sheet1.xml" '
|
|
'ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml"/>'
|
|
'</Types>'
|
|
)
|
|
root_rels_xml = (
|
|
'<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
|
|
'<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
|
|
'<Relationship Id="rId1" '
|
|
'Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" '
|
|
'Target="xl/workbook.xml"/>'
|
|
'</Relationships>'
|
|
)
|
|
workbook_rels_xml = (
|
|
'<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
|
|
'<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
|
|
'<Relationship Id="rId1" '
|
|
'Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet" '
|
|
'Target="worksheets/sheet1.xml"/>'
|
|
'</Relationships>'
|
|
)
|
|
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}"
|