253 lines
8.4 KiB
Python
253 lines
8.4 KiB
Python
"""算法插件适配器共用 I/O 工具。
|
||
|
||
约定（与 server/plugins/types.ts 对齐）：
|
||
- 子进程入口的唯一参数是 ``input.json`` 的绝对路径；
|
||
- 子进程通过 stdout 单行 JSON 上报进度：``{"type":"progress","percent":42,"message":"..."}``
|
||
- 子进程在 ``outDir/result.json`` 写最终结果；
|
||
- stderr 仅做日志（runner 会原样转发）。
|
||
|
||
所有适配器都很薄：解析 input.json → 调真实算法 API → 写 result.json。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
import time
|
||
import traceback
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
|
||
def load_input(argv: list[str]) -> dict[str, Any]:
|
||
"""读取 input.json。"""
|
||
if len(argv) < 2:
|
||
emit_error("usage: <adapter>.py <input.json>", exit_code=2)
|
||
path = Path(argv[1])
|
||
if not path.exists():
|
||
emit_error(f"input.json 不存在: {path}", exit_code=2)
|
||
return json.loads(path.read_text(encoding="utf-8"))
|
||
|
||
|
||
def emit_progress(percent: float, message: str = "") -> None:
|
||
"""以单行 JSON 上报进度。runner 解析后写入数据库。
|
||
|
||
percent 接受 0~1 或 0~100，runner 端会归一化。
|
||
"""
|
||
line = {"type": "progress", "percent": float(percent), "message": message}
|
||
sys.stdout.write(json.dumps(line, ensure_ascii=False) + "\n")
|
||
sys.stdout.flush()
|
||
|
||
|
||
def write_result(out_dir: str | os.PathLike, payload: dict[str, Any]) -> Path:
|
||
"""写 result.json 到 outDir，并强制 schemaVersion / status 字段存在。
|
||
NaN/Infinity 会被替换为 null，确保输出合法 JSON。
|
||
"""
|
||
import math
|
||
|
||
def _sanitize(obj: Any) -> Any:
|
||
if isinstance(obj, float) and (math.isnan(obj) or math.isinf(obj)):
|
||
return None
|
||
if isinstance(obj, dict):
|
||
return {k: _sanitize(v) for k, v in obj.items()}
|
||
if isinstance(obj, list):
|
||
return [_sanitize(v) for v in obj]
|
||
return obj
|
||
|
||
out_path = Path(out_dir) / "result.json"
|
||
payload.setdefault("schemaVersion", 1)
|
||
payload.setdefault("status", "completed")
|
||
out_path.write_text(json.dumps(_sanitize(payload), ensure_ascii=False, indent=2), encoding="utf-8")
|
||
return out_path
|
||
|
||
|
||
def write_error_result(out_dir: str | os.PathLike, message: str, extra: dict | None = None) -> None:
|
||
"""统一错误写出：result.json 仍然存在，runner 据此把数据库改成 error。"""
|
||
payload: dict[str, Any] = {
|
||
"schemaVersion": 1,
|
||
"status": "error",
|
||
"message": message,
|
||
}
|
||
if extra:
|
||
payload["extra"] = extra
|
||
Path(out_dir).mkdir(parents=True, exist_ok=True)
|
||
(Path(out_dir) / "result.json").write_text(
|
||
json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
|
||
)
|
||
|
||
|
||
def emit_error(message: str, exit_code: int = 1) -> None:
|
||
sys.stderr.write(f"[adapter-error] {message}\n")
|
||
sys.exit(exit_code)
|
||
|
||
|
||
def first_pointcloud(input_payload: dict[str, Any]) -> str:
|
||
"""取 inputs.pointCloudPaths 的第一个；无则报错。"""
|
||
paths = (input_payload.get("inputs") or {}).get("pointCloudPaths") or []
|
||
if not paths:
|
||
raise RuntimeError("inputs.pointCloudPaths 为空，请确认任务已上传点云。")
|
||
return str(paths[0])
|
||
|
||
|
||
def now() -> float:
|
||
return time.time()
|
||
|
||
|
||
# ─── Binary PLY → ASCII PLY 转换 ─────────────────────────────────────────────
|
||
|
||
def ensure_ascii_ply(src: str | os.PathLike, cache_dir: str | os.PathLike) -> str:
|
||
"""如果 src 是二进制 PLY，把它转成 ASCII PLY 写到 cache_dir 并返回新路径。
|
||
如果已经是 ASCII PLY 或其他格式，直接返回原路径。
|
||
|
||
不依赖 open3d / plyfile，只用 numpy + struct（标准库）。
|
||
"""
|
||
src = Path(src)
|
||
suffix = src.suffix.lower()
|
||
if suffix != ".ply":
|
||
return str(src)
|
||
|
||
# 读文件头（最多 16 KB）
|
||
HEADER_MAX = 16 * 1024
|
||
with open(src, "rb") as fh:
|
||
raw = fh.read(HEADER_MAX)
|
||
|
||
try:
|
||
header_text = raw.decode("ascii", errors="replace")
|
||
except Exception:
|
||
return str(src) # 无法解析，原样返回
|
||
|
||
if not header_text.startswith("ply"):
|
||
return str(src)
|
||
|
||
end_marker = "end_header"
|
||
end_idx = header_text.find(end_marker)
|
||
if end_idx == -1:
|
||
return str(src)
|
||
|
||
# 找 end_header 后的换行
|
||
data_start = end_idx + len(end_marker)
|
||
if header_text[data_start : data_start + 2] == "\r\n":
|
||
data_start += 2
|
||
elif header_text[data_start] in ("\r", "\n"):
|
||
data_start += 1
|
||
|
||
header_lines = header_text[:end_idx].splitlines()
|
||
|
||
# 解析 format
|
||
fmt = "ascii"
|
||
for line in header_lines:
|
||
parts = line.strip().split()
|
||
if parts and parts[0] == "format":
|
||
fmt = parts[1] if len(parts) > 1 else "ascii"
|
||
break
|
||
|
||
if fmt == "ascii":
|
||
return str(src) # 已经是 ASCII，无需转换
|
||
|
||
# 解析 vertex 属性
|
||
import struct as _struct
|
||
import numpy as _np
|
||
|
||
_PLY_DTYPE: dict[str, tuple[str, int]] = {
|
||
"char": ("b", 1), "int8": ("b", 1),
|
||
"uchar": ("B", 1), "uint8": ("B", 1),
|
||
"short": ("h", 2), "int16": ("h", 2),
|
||
"ushort": ("H", 2), "uint16": ("H", 2),
|
||
"int": ("i", 4), "int32": ("i", 4),
|
||
"uint": ("I", 4), "uint32": ("I", 4),
|
||
"float": ("f", 4), "float32": ("f", 4),
|
||
"double": ("d", 8), "float64": ("d", 8),
|
||
}
|
||
|
||
vertex_count = 0
|
||
props: list[tuple[str, str, int]] = [] # (name, fmt_char, byte_size)
|
||
in_vertex = False
|
||
for line in header_lines:
|
||
parts = line.strip().split()
|
||
if not parts:
|
||
continue
|
||
if parts[0] == "element":
|
||
in_vertex = parts[1] == "vertex" if len(parts) > 1 else False
|
||
if in_vertex and len(parts) > 2:
|
||
vertex_count = int(parts[2])
|
||
elif parts[0] == "property" and in_vertex:
|
||
if parts[1] == "list":
|
||
in_vertex = False
|
||
continue
|
||
ptype = parts[1] if len(parts) > 1 else "float"
|
||
pname = parts[2] if len(parts) > 2 else ""
|
||
fc, bs = _PLY_DTYPE.get(ptype, ("f", 4))
|
||
props.append((pname, fc, bs))
|
||
|
||
if vertex_count == 0 or not props:
|
||
return str(src)
|
||
|
||
# Find x, y, z offsets
|
||
x_off = y_off = z_off = -1
|
||
x_fc = y_fc = z_fc = "f"
|
||
byte_off = 0
|
||
for pname, fc, bs in props:
|
||
n = pname.lower()
|
||
if n == "x":
|
||
x_off, x_fc = byte_off, fc
|
||
elif n == "y":
|
||
y_off, y_fc = byte_off, fc
|
||
elif n == "z":
|
||
z_off, z_fc = byte_off, fc
|
||
byte_off += bs
|
||
|
||
if x_off < 0 or y_off < 0 or z_off < 0:
|
||
return str(src) # 找不到 xyz，原样返回
|
||
|
||
bytes_per_vertex = sum(bs for _, _, bs in props)
|
||
little_endian = fmt != "binary_big_endian"
|
||
endian = "<" if little_endian else ">"
|
||
|
||
# 读整个数据区
|
||
file_size = src.stat().st_size
|
||
data_bytes = file_size - data_start
|
||
expected = vertex_count * bytes_per_vertex
|
||
|
||
with open(src, "rb") as fh:
|
||
fh.seek(data_start)
|
||
data = fh.read(expected)
|
||
|
||
if len(data) < expected:
|
||
return str(src) # 文件截断，原样返回
|
||
|
||
# 用 numpy structured array 解析
|
||
dtype_fields = []
|
||
for pname, fc, bs in props:
|
||
np_type = endian + fc
|
||
dtype_fields.append((pname or f"_f{len(dtype_fields)}", np_type))
|
||
|
||
try:
|
||
arr = _np.frombuffer(data, dtype=_np.dtype(dtype_fields))
|
||
xs = arr[props[0][0] if props[0][0] else "_f0"] # fallback
|
||
# 直接按名字取
|
||
x_name = props[[i for i, (n, _, _) in enumerate(props) if n.lower() == "x"][0]][0]
|
||
y_name = props[[i for i, (n, _, _) in enumerate(props) if n.lower() == "y"][0]][0]
|
||
z_name = props[[i for i, (n, _, _) in enumerate(props) if n.lower() == "z"][0]][0]
|
||
xs = arr[x_name].astype(_np.float64)
|
||
ys = arr[y_name].astype(_np.float64)
|
||
zs = arr[z_name].astype(_np.float64)
|
||
except Exception:
|
||
return str(src)
|
||
|
||
# 写 ASCII PLY
|
||
cache_dir = Path(cache_dir)
|
||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||
out_path = cache_dir / (src.stem + "_ascii.ply")
|
||
|
||
with open(out_path, "w", encoding="ascii") as fh:
|
||
fh.write("ply\nformat ascii 1.0\n")
|
||
fh.write(f"element vertex {vertex_count}\n")
|
||
fh.write("property double x\nproperty double y\nproperty double z\n")
|
||
fh.write("end_header\n")
|
||
for x, y, z in zip(xs, ys, zs):
|
||
fh.write(f"{x:.10g} {y:.10g} {z:.10g}\n")
|
||
|
||
return str(out_path)
|
||
|