CheckCardBoxAndDefectServer/app/api/card_inference.py

from fastapi import APIRouter, File, UploadFile, Depends, HTTPException, Path, Response
from fastapi.responses import FileResponse, JSONResponse
from fastapi.concurrency import run_in_threadpool
from enum import Enum
from ..core.config import settings
from app.services.card_rectify_and_center import CardRectifyAndCenter
from app.services.card_service import CardInferenceService, card_service
from app.services.defect_service import DefectInferenceService
from app.core.logger import get_logger
import cv2
import numpy as np
import json

logger = get_logger(__name__)

router = APIRouter()

model_names = list(settings.CARD_MODELS_CONFIG.keys())
defect_names = list(settings.DEFECT_TYPE.keys())
InferenceType = Enum("InferenceType", {name: name for name in model_names})
DefectType = Enum("InferenceType", {name: name for name in defect_names})


@router.post("/model_inference", description="内外框类型输入大图, 其他输入小图")
async def card_model_inference(
        inference_type: InferenceType,
        service: CardInferenceService = Depends(lambda: card_service),
        file: UploadFile = File(...)
):
    """
    接收一张卡片图片，使用指定类型的模型进行推理，并返回JSON结果。

    - **inference_type**: 要使用的模型类型（从下拉列表中选择）。
    - **file**: 要上传的图片文件。
    """
    image_bytes = await file.read()

    try:
        # 3. 传递参数时，使用 .value 获取 Enum 的字符串值
        json_result = await run_in_threadpool(
            service.predict,
            inference_type=inference_type.value,  # 使用 .value
            image_bytes=image_bytes
        )
        return json_result
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"服务器内部错误: {e}")


@router.post("/defect_inference",
             description="环形光居中计算, 环形光正反边角缺陷, 同轴光正反表面缺陷")
async def card_model_inference(
        defect_type: DefectType,
        file: UploadFile = File(...),
        is_draw_image: bool = False,
):
    service = DefectInferenceService()
    image_bytes = await file.read()

    # 将字节数据转换为numpy数组
    np_arr = np.frombuffer(image_bytes, np.uint8)
    # 从numpy数组中解码图像
    img_bgr = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

    if img_bgr is None:
        raise ValueError("无法解码图像，请确保上传的是有效的图片格式 (JPG, PNG, etc.)")

    try:
        # 3. 传递参数时，使用 .value 获取 Enum 的字符串值
        json_result = await run_in_threadpool(
            service.defect_inference,
            inference_type=defect_type.value,
            img_bgr=img_bgr,
            is_draw_image=is_draw_image
        )
        return json_result
    except ValueError as e:
        logger.error(e)
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(e)
        raise HTTPException(status_code=500, detail=f"服务器内部错误: {e}")


@router.post("/card_rectify_and_center",
             description="对卡片图像进行转正和居中处理")
async def card_rectify_and_center(
        file: UploadFile = File(...)
):
    service = CardRectifyAndCenter()

    image_bytes = await file.read()
    # 将字节数据转换为numpy数组
    np_arr = np.frombuffer(image_bytes, np.uint8)
    # 从numpy数组中解码图像
    img_bgr = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

    if img_bgr is None:
        raise ValueError("无法解码图像，请确保上传的是有效的图片格式 (JPG, PNG, etc.)")

    try:
        # 3. 传递参数时，使用 .value 获取 Enum 的字符串值
        img_result = await run_in_threadpool(
            service.rectify_and_center,
            img_bgr=img_bgr
        )

        is_success, buffer = cv2.imencode(".jpg", img_result)
        jpeg_bytes = buffer.tobytes()

        return Response(content=jpeg_bytes, media_type="image/jpeg")
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"服务器内部错误: {e}")