card_inference.py 2.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. from fastapi import APIRouter, File, UploadFile, Depends, HTTPException, Path
  2. from fastapi.responses import FileResponse, JSONResponse
  3. from fastapi.concurrency import run_in_threadpool
  4. from enum import Enum
  5. from ..core.config import settings
  6. from app.services.card_service import CardInferenceService, card_service
  7. from app.services.defect_service import DefectInferenceService
  8. from app.core.logger import get_logger
  9. import json
  10. logger = get_logger(__name__)
  11. router = APIRouter()
  12. model_names = list(settings.CARD_MODELS_CONFIG.keys())
  13. defect_names = list(settings.DEFECT_TYPE.keys())
  14. InferenceType = Enum("InferenceType", {name: name for name in model_names})
  15. DefectType = Enum("InferenceType", {name: name for name in defect_names})
  16. @router.post("/model_inference", description="内外框类型输入大图, 其他输入小图")
  17. async def card_model_inference(
  18. inference_type: InferenceType,
  19. service: CardInferenceService = Depends(lambda: card_service),
  20. file: UploadFile = File(...)
  21. ):
  22. """
  23. 接收一张卡片图片,使用指定类型的模型进行推理,并返回JSON结果。
  24. - **inference_type**: 要使用的模型类型(从下拉列表中选择)。
  25. - **file**: 要上传的图片文件。
  26. """
  27. image_bytes = await file.read()
  28. try:
  29. # 3. 传递参数时,使用 .value 获取 Enum 的字符串值
  30. json_result = await run_in_threadpool(
  31. service.predict,
  32. inference_type=inference_type.value, # 使用 .value
  33. image_bytes=image_bytes
  34. )
  35. return json_result
  36. except ValueError as e:
  37. raise HTTPException(status_code=400, detail=str(e))
  38. except Exception as e:
  39. raise HTTPException(status_code=500, detail=f"服务器内部错误: {e}")
  40. @router.post("/defect_inference",
  41. description="环形光居中计算, 环形光正反边角缺陷, 同轴光正反表面缺陷")
  42. async def card_model_inference(
  43. defect_type: DefectType,
  44. # service: DefectInferenceService = Depends(lambda: defect_service),
  45. file: UploadFile = File(...),
  46. is_draw_image: bool = False,
  47. ):
  48. service = DefectInferenceService()
  49. image_bytes = await file.read()
  50. try:
  51. # 3. 传递参数时,使用 .value 获取 Enum 的字符串值
  52. json_result = await run_in_threadpool(
  53. service.defect_inference,
  54. inference_type=defect_type.value,
  55. image_bytes=image_bytes,
  56. is_draw_image=is_draw_image
  57. )
  58. return json_result
  59. except ValueError as e:
  60. logger.error(e)
  61. raise HTTPException(status_code=400, detail=str(e))
  62. except Exception as e:
  63. logger.error(e)
  64. raise HTTPException(status_code=500, detail=f"服务器内部错误: {e}")