vent_agent_2.0/algorithms/characteristic_curve_recognition/service.py

"""与 HTTP 无关的单图调用封装；每次调用独立目录与计算进程。"""
from __future__ import annotations

import importlib.util
import json
import logging
import os
from pathlib import Path
import subprocess
import sys
import tempfile
import threading
from typing import Literal
import zipfile
import io

import config

MAX_UPLOAD_BYTES = 20 * 1024 * 1024
MAX_IMAGE_PIXELS = 16_000_000
TIMEOUT_SECONDS = 180
_slots = threading.BoundedSemaphore(config.ALGORITHM_MAX_PARALLEL)
_ROOT = Path(__file__).resolve().parents[2]
logger = logging.getLogger(__name__)


class AlgorithmError(Exception):
    def __init__(self, code: str, message: str, status_code: int):
        super().__init__(message)
        self.code = code
        self.status_code = status_code


def validate_image(content: bytes) -> str:
    if len(content) > MAX_UPLOAD_BYTES:
        raise AlgorithmError("image_too_large", "图片不能超过 20 MiB", 413)
    if content.startswith(b"\x89PNG\r\n\x1a\n"):
        suffix = ".png"
    elif content.startswith(b"\xff\xd8\xff"):
        suffix = ".jpg"
    else:
        raise AlgorithmError("invalid_image", "请上传有效的 PNG 或 JPEG 图片", 422)
    return suffix


def check_dependencies() -> None:
    if any(importlib.util.find_spec(name) is None for name in ("numpy", "cv2")):
        raise AlgorithmError("algorithm_unavailable", "算法依赖未安装，请按项目锁文件安装依赖", 503)


def execute_image(image_path: Path, result_dir: Path, *, detailed: bool = False,
                  wait_for_slot: bool = False, stage_path: Path | None = None) -> dict:
    """在共享并行闸下计算单张图；批量任务可等待空位。"""
    check_dependencies()
    if not _slots.acquire(blocking=wait_for_slot):
        raise AlgorithmError("algorithm_busy", "算法正在处理其他图片，请稍后重试", 429)
    try:
        env = dict(os.environ, OPENCV_IO_MAX_IMAGE_PIXELS=str(MAX_IMAGE_PIXELS),
                   OPENBLAS_NUM_THREADS="1", OMP_NUM_THREADS="1")
        command = [sys.executable, "-B", "-m", "algorithms.characteristic_curve_recognition.worker",
                   str(image_path), str(result_dir), "1" if detailed else "0"]
        if stage_path is not None:
            command.append(str(stage_path))
        try:
            process = subprocess.run(command, cwd=_ROOT, env=env, capture_output=True, timeout=TIMEOUT_SECONDS)
        except subprocess.TimeoutExpired as exc:
            raise AlgorithmError("algorithm_timeout", "图片处理超过 180 秒，请缩小图片后重试", 504) from exc
        if process.returncode:
            logger.error("特性曲线识别失败：%s", process.stderr.decode("utf-8", errors="replace")[-4000:])
            if process.returncode == 2:
                raise AlgorithmError("invalid_image", "图片损坏、尺寸超限或小于 32×32 像素", 422)
            if process.returncode == 3:
                raise AlgorithmError("detection_failed", "无法识别图表，请确认图片含清晰的坐标轴和曲线", 422)
            raise AlgorithmError("algorithm_failed", "算法处理失败，请通过服务日志排查", 500)
        result = {
            "algorithm": "characteristic_curve_recognition",
            "model": json.loads((result_dir / "final_model.json").read_text()),
            "geometry": json.loads((result_dir / "diagnostics/geometry.json").read_text()),
            "timing": json.loads((result_dir / "diagnostics/timing.json").read_text()),
        }
        (result_dir / "result.json").write_text(json.dumps(result, ensure_ascii=False, allow_nan=False), encoding="utf-8")
        return result
    finally:
        _slots.release()


def analyze_image(content: bytes, *, detailed: bool = False,
                  output: Literal["json", "zip"] = "json") -> dict | bytes:
    """上传 PNG/JPEG，返回模型与几何结果，或可离线查看的 ZIP。"""
    if output not in {"json", "zip"}:
        raise ValueError("output 必须为 json 或 zip")
    suffix = validate_image(content)
    with tempfile.TemporaryDirectory(prefix="vent-characteristic-curve-") as temporary:
        root = Path(temporary)
        image_path = root / ("input" + suffix)
        image_path.write_bytes(content)
        result_dir = root / "result"
        result = execute_image(image_path, result_dir, detailed=detailed)
        if output == "json":
            return result
        archive = io.BytesIO()
        with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as zip_file:
            for path in sorted(result_dir.rglob("*")):
                if path.is_file():
                    zip_file.write(path, path.relative_to(result_dir).as_posix())
        return archive.getvalue()