208 lines
8.2 KiB
Python
208 lines
8.2 KiB
Python
"""train-measurement 算法适配器 (v1.1)。
|
||
|
||
读 input.json → 调 TaskManager（来自 alg/点云测距v1.1）→ 写统一 result.json。
|
||
|
||
v1.1 变更：
|
||
- 路径从 260505列车点云测距v1 → 点云测距v1.1
|
||
- 算法包现在是可安装的 Python 包（有 __init__.py），PYTHONPATH 指向 src/
|
||
- TRAIN_MEAS_ROOT 指向 train_measurement_system/（用于定位 config/measurements.yaml）
|
||
- TaskManager.run_all_tasks() 返回结构不变，data 字段含 mean_distance / std_distance 等
|
||
|
||
约定：
|
||
- 第一个点云路径作为 pointcloud_file 注入到测量配置 YAML 中；
|
||
- 若用户在 UI 选择「自定义 YAML」并提供 customYaml，则使用之；
|
||
否则使用算法仓库自带的 config/measurements.yaml 作为模板。
|
||
- 所有任务结果合成成 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,
|
||
ensure_ascii_ply,
|
||
first_pointcloud,
|
||
load_input,
|
||
write_error_result,
|
||
write_result,
|
||
)
|
||
|
||
|
||
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 v1.1 算法包")
|
||
|
||
# TRAIN_MEAS_ROOT 由 manifest envExtra 注入，指向 train_measurement_system/
|
||
train_root = Path(os.environ.get("TRAIN_MEAS_ROOT") or "")
|
||
if not train_root.exists():
|
||
# 回退：从 ALG_ROOT 推导
|
||
alg_root = Path(os.environ.get("ALG_ROOT") or "/alg")
|
||
train_root = alg_root / "点云测距v1.1" / "train_measurement_system"
|
||
|
||
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)}")
|
||
|
||
# 若是二进制 PLY，先转成 ASCII PLY（算法内置 reader 不依赖 open3d）
|
||
pc_path = ensure_ascii_ply(pc_path, out_dir / "_ply_cache")
|
||
|
||
# ── 准备 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
|
||
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()
|
||
# 先注入点云路径，让 load_config 跳过 YAML 里的 pointcloud_file（开发测试路径）
|
||
tm.pointcloud_path_override = pc_path
|
||
if not tm.load_config(str(yaml_path)):
|
||
raise RuntimeError(f"TaskManager 加载配置失败：{yaml_path}")
|
||
|
||
# 注入点云路径（覆盖 YAML 中的 pointcloud_file）
|
||
tm.config["pointcloud_file"] = pc_path
|
||
|
||
# 注入可视化输出目录
|
||
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.30, "加载点云（已在 load_config 中完成）")
|
||
|
||
emit_progress(0.40, f"执行 {len(tm.tasks)} 个测量任务")
|
||
results = tm.run_all_tasks()
|
||
|
||
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")
|
||
measurements.append(
|
||
{
|
||
"name": name,
|
||
"value": float(value) if isinstance(value, (int, float)) else float("nan"),
|
||
"unit": "mm",
|
||
"status": "info",
|
||
}
|
||
)
|
||
std = clean_data.get("std_distance")
|
||
if isinstance(std, (int, float)):
|
||
measurements.append(
|
||
{
|
||
"name": f"{name} · 标准差",
|
||
"value": float(std),
|
||
"unit": "mm",
|
||
"status": "info",
|
||
}
|
||
)
|
||
# v1.1 新增字段：min/max distance
|
||
for field_name, label in [
|
||
("min_distance", "最小距离"),
|
||
("max_distance", "最大距离"),
|
||
]:
|
||
v = clean_data.get(field_name)
|
||
if isinstance(v, (int, float)):
|
||
measurements.append(
|
||
{
|
||
"name": f"{name} · {label}",
|
||
"value": float(v),
|
||
"unit": "mm",
|
||
"status": "info",
|
||
}
|
||
)
|
||
else:
|
||
measurements.append(
|
||
{
|
||
"name": name,
|
||
"value": float("nan"),
|
||
"unit": "mm",
|
||
"status": "fail",
|
||
}
|
||
)
|
||
|
||
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())
|