vent_agent_2.0/algorithms/characteristic_curve_recognition/worker.py

"""由 service 在可超时终止的独立进程中运行。"""
from __future__ import annotations

from pathlib import Path
import json
import sys


def main() -> int:
    import cv2
    import numpy as np

    from .service import MAX_IMAGE_PIXELS
    from .vendor.chart_axis_detector.current_image_pipeline import (
        CurrentImagePipelineConfig, process_new_image,
    )

    cv2.setNumThreads(1)
    image_path, result_dir = Path(sys.argv[1]), Path(sys.argv[2])
    try:
        image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    except cv2.error:
        return 2
    if image is None or min(image.shape[:2]) < 32 or image.shape[0] * image.shape[1] > MAX_IMAGE_PIXELS:
        return 2
    del image
    def stage(name):
        if len(sys.argv) > 4:
            path = Path(sys.argv[4])
            temporary = path.with_suffix(".tmp")
            temporary.write_text(json.dumps({"stage": name}), encoding="utf-8")
            temporary.replace(path)

    try:
        process_new_image(image_path, result_dir, CurrentImagePipelineConfig(detailed=sys.argv[3] == "1"), on_stage=stage)
    except ValueError as exc:
        print(str(exc), file=sys.stderr)
        return 3
    return 0


if __name__ == "__main__":
    sys.exit(main())