CheckCardBoxAndDefectServer/app/api/card_inference.py

from fastapi import APIRouter, File, UploadFile, Depends, HTTPException, Path
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_service import CardInferenceService, card_service
from app.services.defect_service import DefectInferenceService
from app.core.logger import logger
import json

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,
        # service: DefectInferenceService = Depends(lambda: defect_service),
        file: UploadFile = File(...),
        is_draw_image: bool = False,
):
    service = DefectInferenceService()
    image_bytes = await file.read()

    try:
        # 3. 传递参数时，使用 .value 获取 Enum 的字符串值
        json_result = await run_in_threadpool(
            service.defect_inference,
            inference_type=defect_type.value,
            image_bytes=image_bytes,
            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("/mock_query")
async def mock_query(img_id: int):
    # json_data = {"img_id": img_id}
    with open("_temp_work/mock_result.json", "r") as f:
        json_data = json.load(f)

    return json_data