shudao-main/shudao-chat-py/services/yolo_service.py

"""
YOLO 隐患识别服务
"""
import base64
import json
from typing import Dict, Any

import httpx

from utils.config import settings
from utils.logger import logger


class YoloService:
    def __init__(self):
        self.base_url = settings.yolo.base_url.rstrip("/")

    async def detect_hazards(self, image_bytes: bytes, scene_type: str = "") -> Dict[str, Any]:
        """检测图片中的隐患。

        对齐 Go 版本行为：
        - 调用 `/predict`
        - 请求字段为 `modeltype`、`image`、`conf_threshold`
        - `image` 传 base64 编码后的图片内容
        - 不输出模拟数据；失败直接抛错
        """
        detect_url = f"{self.base_url}/predict"

        if not image_bytes:
            raise ValueError("image_bytes 不能为空")

        image_base64 = base64.b64encode(image_bytes).decode("utf-8")

        logger.info(
            "[YOLO API] 开始请求检测接口: base_url=%s, detect_url=%s, scene_type=%s, image_bytes_len=%s, image_base64_len=%s",
            self.base_url,
            detect_url,
            scene_type or "<empty>",
            len(image_bytes),
            len(image_base64),
        )

        payload = {
            "modeltype": scene_type,
            "image": image_base64,
            "conf_threshold": 0.5,
        }

        payload_json = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))

        headers = {
            "Content-Type": "application/json",
        }

        logger.info("[YOLO API] 场景名称(modeltype): %s", json.dumps(payload["modeltype"], ensure_ascii=False))
        logger.info("[YOLO API] 置信度阈值(conf_threshold): %.2f", payload["conf_threshold"])
        logger.info("[YOLO API] base64前120字符: %s", image_base64[:120])
        logger.info("[YOLO API] YOLO请求JSON长度: %s", len(payload_json))
        logger.info("[YOLO API] YOLO请求JSON前1000字符: %s", payload_json[:1000])

        try:
            async with httpx.AsyncClient(timeout=30.0) as client:
                response = await client.post(
                    detect_url,
                    content=payload_json.encode("utf-8"),
                    headers=headers,
                )

            response_text = response.text
            logger.info(
                "[YOLO API] 响应返回: status_code=%s, content_type=%s, request_url=%s, body_prefix=%s",
                response.status_code,
                response.headers.get("content-type", ""),
                detect_url,
                response_text[:2000],
            )

            if response.status_code != 200:
                raise RuntimeError(
                    f"YOLO服务HTTP错误: {response.status_code}, body={response_text}"
                )

            result = response.json()

            if not isinstance(result, dict):
                raise RuntimeError(
                    f"YOLO服务响应格式错误: {type(result).__name__}"
                )

            return result

        except httpx.TimeoutException as e:
            logger.error(
                "[YOLO API] 请求超时: base_url=%s, detect_url=%s, error=%s",
                self.base_url,
                detect_url,
                repr(e),
            )
            raise RuntimeError("YOLO服务请求超时") from e

        except httpx.RequestError as e:
            logger.error(
                "[YOLO API] 请求异常: base_url=%s, detect_url=%s, error=%s, type=%s",
                self.base_url,
                detect_url,
                str(e),
                type(e).__name__,
            )
            raise RuntimeError(f"YOLO服务请求失败: {type(e).__name__}") from e

        except ValueError as e:
            logger.error(
                "[YOLO API] 响应JSON解析失败: base_url=%s, detect_url=%s, error=%s",
                self.base_url,
                detect_url,
                str(e),
            )
            raise RuntimeError("YOLO服务响应解析失败") from e

        except Exception as e:
            logger.exception(
                "[YOLO API] 未知异常: base_url=%s, detect_url=%s, error=%s",
                self.base_url,
                detect_url,
                str(e),
            )
            raise RuntimeError("YOLO服务异常") from e


# 全局实例
yolo_service = YoloService()