306 lines
12 KiB
Python
306 lines
12 KiB
Python
"""train-measurement 算法适配器 (v2.2)。
|
||
|
||
读 input.json → 调 TaskManager（来自 alg/train_measurement_system_v2.2）→ 写统一 result.json。
|
||
|
||
v2.2 变更：
|
||
- 算法路径从 alg/点云测距v1.1 → alg/train_measurement_system_v2.2
|
||
- 新增孔中心拟合（hole_center_fitter）、测量日志、性能日志
|
||
- 内置配置文件改为 config/measurements.yaml（v3.2）
|
||
- TRAIN_MEAS_ROOT 指向 train_measurement_system_v3.2_0629/（用于定位 config/）
|
||
- 算法内置 _read_binary_ply，直接支持二进制 PLY，无需 ASCII 转换
|
||
|
||
约定：
|
||
- 第一个点云路径作为 pointcloud_file 注入到测量配置 YAML 中；
|
||
- 若用户在 UI 选择「自定义 YAML」并提供 customYaml，则使用之；
|
||
否则使用算法仓库自带的 config/measurements.yaml（v3）作为模板。
|
||
- 所有任务结果合成成 measurements[] 列表，附带可视化 PLY 路径到 extra。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import sys
|
||
import traceback
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, str(Path(__file__).parent))
|
||
from _plugin_io import (
|
||
emit_progress,
|
||
first_pointcloud,
|
||
load_input,
|
||
write_error_result,
|
||
write_result,
|
||
)
|
||
|
||
|
||
def _to_float(value):
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, (int, float)):
|
||
return float(value)
|
||
try:
|
||
return float(str(value).strip())
|
||
except (TypeError, ValueError):
|
||
return None
|
||
|
||
|
||
def _load_excel_measurement_rows(out_dir: Path) -> dict[str, dict]:
|
||
"""读取算法导出的 Excel 结论，作为前端 measurements 的权威来源。"""
|
||
try:
|
||
from openpyxl import load_workbook
|
||
|
||
excel_files = sorted(
|
||
out_dir.glob("measurement_results_*.xlsx"),
|
||
key=lambda p: p.stat().st_mtime,
|
||
)
|
||
if not excel_files:
|
||
return {}
|
||
|
||
wb = load_workbook(excel_files[-1], data_only=True, read_only=True)
|
||
ws = wb.active
|
||
rows: dict[str, dict] = {}
|
||
for row in ws.iter_rows(min_row=2, values_only=True):
|
||
name = str(row[0] or "").strip()
|
||
if not name:
|
||
continue
|
||
|
||
conclusion = str(row[6] or "").strip()
|
||
status = None
|
||
if conclusion == "通过":
|
||
status = "pass"
|
||
elif conclusion in ("不通过", "失败"):
|
||
status = "fail"
|
||
elif conclusion == "未配置":
|
||
status = "info"
|
||
|
||
rows[name] = {
|
||
"value": _to_float(row[3]),
|
||
"toleranceLabel": str(row[2]) if row[2] is not None else None,
|
||
"status": status,
|
||
}
|
||
return rows
|
||
except Exception as exc:
|
||
print(f"读取 Excel 结论失败，继续使用适配器判定：{exc}", file=sys.stderr)
|
||
return {}
|
||
|
||
|
||
def main() -> int:
|
||
payload = load_input(sys.argv)
|
||
inputs = payload.get("inputs") or {}
|
||
config = inputs.get("config") or {}
|
||
out_dir = Path(payload["outDir"])
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
try:
|
||
emit_progress(0.05, "加载 train_measurement v3.2 算法包")
|
||
|
||
# TRAIN_MEAS_ROOT 由 manifest envExtra 注入，指向 train_measurement_system_v3.2_0629/
|
||
train_root = Path(os.environ.get("TRAIN_MEAS_ROOT") or "")
|
||
if not train_root.exists():
|
||
alg_root = Path(os.environ.get("ALG_ROOT") or "/alg")
|
||
train_root = alg_root / "train_measurement_system_v3.2_0629"
|
||
|
||
if not train_root.exists():
|
||
raise RuntimeError(
|
||
f"找不到 train_measurement_system 目录：{train_root}。"
|
||
"请确认 TRAIN_MEAS_ROOT 或 ALG_ROOT 环境变量已正确设置。"
|
||
)
|
||
|
||
from task_manager import TaskManager # type: ignore
|
||
|
||
pc_path = first_pointcloud(payload)
|
||
emit_progress(0.15, f"读取点云：{os.path.basename(pc_path)}")
|
||
|
||
# ── 准备 YAML 配置 ────────────────────────────────────────────────────
|
||
config_preset = str(config.get("configPreset") or "builtin")
|
||
custom_yaml_str = str(config.get("customYaml") or "").strip()
|
||
|
||
if config_preset == "custom" and custom_yaml_str:
|
||
# 支持 JSON 格式（前端传来）和 YAML 格式
|
||
yaml_path = out_dir / "custom_measurements.yaml"
|
||
try:
|
||
import json as _json
|
||
parsed = _json.loads(custom_yaml_str)
|
||
# JSON 格式：转成 YAML 写盘
|
||
import yaml as _yaml
|
||
yaml_path.write_text(_yaml.dump(parsed, allow_unicode=True, default_flow_style=False), encoding="utf-8")
|
||
except Exception:
|
||
# 已经是 YAML 格式，直接写盘
|
||
yaml_path.write_text(custom_yaml_str, encoding="utf-8")
|
||
else:
|
||
# 使用算法仓库内置的 config/measurements.yaml（v3）
|
||
yaml_path = train_root / "config" / "measurements.yaml"
|
||
if not yaml_path.exists():
|
||
raise RuntimeError(f"内置配置文件不存在：{yaml_path}")
|
||
|
||
emit_progress(0.20, f"加载配置：{yaml_path.name}")
|
||
|
||
# ── 初始化 TaskManager ────────────────────────────────────────────────
|
||
tm = TaskManager()
|
||
tm.pointcloud_path_override = pc_path
|
||
tm.output_dir_override = str(out_dir)
|
||
|
||
emit_progress(0.25, f"读取点云（{os.path.basename(pc_path)}，可能需要 1~2 分钟）")
|
||
if not tm.load_config(str(yaml_path)):
|
||
raise RuntimeError(f"TaskManager 加载配置失败：{yaml_path}")
|
||
|
||
# 注入点云路径（覆盖 YAML 中的 pointcloud_file）
|
||
tm.config["pointcloud_file"] = pc_path
|
||
|
||
# 过滤 enabled=False 的任务（v3 TaskManager 不处理 enabled 字段）
|
||
import yaml as _yaml_check
|
||
_raw = _yaml_check.safe_load(yaml_path.read_text(encoding="utf-8"))
|
||
_enabled_names = {
|
||
t["name"] for t in (_raw.get("tasks") or _raw.get("measurements") or [])
|
||
if t.get("enabled", True)
|
||
}
|
||
tm.tasks = [t for t in tm.tasks if t.name in _enabled_names]
|
||
|
||
# 注入可视化输出目录
|
||
viz_dir = out_dir / "visualizations"
|
||
viz_dir.mkdir(parents=True, exist_ok=True)
|
||
for task in tm.tasks:
|
||
if hasattr(task, "visualization_dir"):
|
||
task.visualization_dir = str(viz_dir)
|
||
|
||
emit_progress(0.40, f"点云已加载，开始执行 {len(tm.tasks)} 个测量任务")
|
||
|
||
# ── 从 YAML 构建公差映射（name → {theoretical, upper, lower}） ─────────
|
||
import yaml as _yaml
|
||
try:
|
||
raw_yaml = _yaml.safe_load(yaml_path.read_text(encoding="utf-8"))
|
||
# v3 用 'tasks' key，v2 用 'measurements' key，兼容两者
|
||
task_cfgs = raw_yaml.get("tasks") or raw_yaml.get("measurements") or []
|
||
tol_map: dict[str, dict] = {}
|
||
for tc in task_cfgs:
|
||
tname = tc.get("name")
|
||
if not tname:
|
||
continue
|
||
tol = tc.get("tolerance")
|
||
theoretical = tc.get("theoretical_dimension")
|
||
if isinstance(tol, dict):
|
||
tol_map[tname] = {
|
||
"theoretical": theoretical,
|
||
"upper": float(tol.get("upper", 0)),
|
||
"lower": float(tol.get("lower", 0)),
|
||
}
|
||
elif tol is not None:
|
||
tol_map[tname] = {
|
||
"theoretical": theoretical,
|
||
"upper": float(tol),
|
||
"lower": float(tol),
|
||
}
|
||
except Exception:
|
||
tol_map = {}
|
||
results = tm.run_all_tasks()
|
||
report_rows = _load_excel_measurement_rows(out_dir)
|
||
|
||
emit_progress(0.92, "整理结果")
|
||
measurements: list[dict] = []
|
||
clean_raw: list[dict] = []
|
||
|
||
for r in results:
|
||
name = str(r.get("name") or "未命名任务")
|
||
status_raw = r.get("status")
|
||
data = r.get("data") if isinstance(r.get("data"), dict) else {}
|
||
# 剥掉以 '_' 开头的私有 numpy 数组（不可 JSON 序列化）
|
||
clean_data = {k: v for k, v in data.items() if not k.startswith("_")} if data else {}
|
||
clean_raw.append(
|
||
{
|
||
"name": name,
|
||
"type": r.get("type"),
|
||
"status": status_raw,
|
||
"data": clean_data,
|
||
"error": r.get("error"),
|
||
}
|
||
)
|
||
|
||
if status_raw == "success":
|
||
value = clean_data.get("mean_distance")
|
||
if value is None:
|
||
value = clean_data.get("distance")
|
||
fval = float(value) if isinstance(value, (int, float)) else None
|
||
|
||
# v3.2 data 里直接有 is_qualified 判定结果（line_to_plane/line_to_line 也支持）
|
||
data_pass = clean_data.get("is_qualified")
|
||
if data_pass is None:
|
||
data_pass = clean_data.get("pass") # 兼容旧字段名
|
||
|
||
tol_cfg = tol_map.get(name)
|
||
tolerance = None
|
||
meas_status = "info"
|
||
if data_pass is not None:
|
||
# 算法已判定，直接用
|
||
meas_status = "pass" if data_pass else "fail"
|
||
if tol_cfg:
|
||
tolerance = tol_cfg.get("upper")
|
||
elif tol_cfg and fval is not None:
|
||
theoretical = tol_cfg.get("theoretical")
|
||
upper = tol_cfg["upper"]
|
||
lower = tol_cfg["lower"]
|
||
tolerance = upper
|
||
if theoretical is not None:
|
||
error = fval - float(theoretical)
|
||
meas_status = "pass" if (-lower <= error <= upper) else "fail"
|
||
|
||
measurements.append(
|
||
{
|
||
"name": name,
|
||
"value": fval if fval is not None else float("nan"),
|
||
"unit": "mm",
|
||
"tolerance": tolerance,
|
||
"status": meas_status,
|
||
"passLabel": "通过",
|
||
"failLabel": "不通过",
|
||
}
|
||
)
|
||
else:
|
||
measurements.append(
|
||
{
|
||
"name": name,
|
||
"value": float("nan"),
|
||
"unit": "mm",
|
||
"status": "fail",
|
||
}
|
||
)
|
||
|
||
for measurement in measurements:
|
||
report_row = report_rows.get(measurement["name"])
|
||
if not report_row:
|
||
continue
|
||
if report_row.get("value") is not None:
|
||
measurement["value"] = report_row["value"]
|
||
if report_row.get("toleranceLabel"):
|
||
measurement["toleranceLabel"] = report_row["toleranceLabel"]
|
||
if report_row.get("status"):
|
||
measurement["status"] = report_row["status"]
|
||
|
||
success_count = sum(1 for r in results if r.get("status") == "success")
|
||
payload_out = {
|
||
"schemaVersion": 1,
|
||
"status": "completed",
|
||
"message": f"共 {len(results)} 个任务，成功 {success_count} 个。",
|
||
"measurements": measurements,
|
||
"overlay3d": {"kind": "none"},
|
||
"extra": {
|
||
"rawTaskResults": clean_raw,
|
||
"visualizationDir": "visualizations",
|
||
},
|
||
}
|
||
write_result(out_dir, payload_out)
|
||
emit_progress(1.0, "完成")
|
||
return 0
|
||
|
||
except Exception as exc:
|
||
traceback.print_exc()
|
||
write_error_result(
|
||
out_dir,
|
||
f"train-measurement 计算失败：{exc}",
|
||
extra={"traceback": traceback.format_exc()},
|
||
)
|
||
return 1
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|