defect_service.py 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148
  1. import cv2
  2. import numpy as np
  3. from ..core.model_loader import get_predictor
  4. from app.utils.CardDefectAggregator import CardDefectAggregator
  5. from app.utils.arean_anylize_draw import DefectProcessor, DrawingParams
  6. from app.utils.AnalyzeCenter import analyze_centering_rotated
  7. from app.utils.json_data_formate import formate_center_data
  8. from app.core.config import settings
  9. from app.core.logger import logger
  10. import json
  11. class DefectInferenceService:
  12. def defect_inference(self, inference_type: str, image_bytes: bytes,
  13. is_draw_image=False) -> dict:
  14. """
  15. 执行卡片识别推理。
  16. Args:
  17. inference_type: 模型类型 (e.g., 'outer_box').
  18. image_bytes: 从API请求中获得的原始图像字节。
  19. Returns:
  20. 一个包含推理结果的字典。
  21. """
  22. # 2. 将字节流解码为OpenCV图像
  23. # 将字节数据转换为numpy数组
  24. np_arr = np.frombuffer(image_bytes, np.uint8)
  25. # 从numpy数组中解码图像
  26. img_bgr = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
  27. if img_bgr is None:
  28. logger.error("无法解码图像,请确保上传的是有效的图片格式 (JPG, PNG, etc.)")
  29. return {}
  30. # 面
  31. if (inference_type == "pokemon_front_face_no_reflect_defect"
  32. or inference_type == "pokemon_front_face_reflect_defect"
  33. or inference_type == "pokemon_back_face_defect"):
  34. # 1. 获取对应的预测器实例
  35. predictor = get_predictor(inference_type)
  36. # 3. 调用我们新加的 predict_from_image 方法进行推理
  37. # result = predictor.predict_from_image(img_bgr)
  38. # 3. 实例化我们聚合器,传入预测器
  39. aggregator = CardDefectAggregator(
  40. predictor=predictor,
  41. tile_size=512,
  42. overlap_ratio=0.1, # 10% 重叠
  43. )
  44. json_data = aggregator.process_image(
  45. image=img_bgr,
  46. mode='face'
  47. )
  48. # merge_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-merge.json'
  49. # with open(merge_json_path, 'w', encoding='utf-8') as f:
  50. # json.dump(json_data, f, ensure_ascii=False, indent=4)
  51. # logger.info(f"合并结束")
  52. processor = DefectProcessor(pixel_resolution=settings.pixel_resolution)
  53. area_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-area_result.json'
  54. if is_draw_image:
  55. drawing_params_with_rect = DrawingParams(draw_min_rect=True)
  56. drawn_image, area_json = processor.analyze_and_draw(img_bgr, json_data,
  57. drawing_params_with_rect)
  58. temp_img_path = settings.TEMP_WORK_DIR / f'{inference_type}-area_result.jpg'
  59. cv2.imwrite(temp_img_path, drawn_image)
  60. else:
  61. area_json = processor.analyze_from_json(json_data)
  62. with open(area_json_path, 'w', encoding='utf-8') as f:
  63. json.dump(area_json, f, ensure_ascii=False, indent=4)
  64. logger.info("面积计算结束")
  65. return area_json
  66. # 边角
  67. elif (inference_type == "pokemon_front_corner_no_reflect_defect"
  68. or inference_type == "pokemon_front_corner_reflect_defect"
  69. or inference_type == "pokemon_back_corner_defect"):
  70. predictor = get_predictor(inference_type)
  71. aggregator = CardDefectAggregator(
  72. predictor=predictor,
  73. tile_size=512,
  74. overlap_ratio=0.1, # 10% 重叠
  75. )
  76. json_data = aggregator.process_image(
  77. image=img_bgr,
  78. mode='edge'
  79. )
  80. # merge_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-merge.json'
  81. # with open(merge_json_path, 'w', encoding='utf-8') as f:
  82. # json.dump(json_data, f, ensure_ascii=False, indent=4)
  83. # logger.info(f"合并结束")
  84. processor = DefectProcessor(pixel_resolution=settings.pixel_resolution)
  85. area_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-area_result.json'
  86. if is_draw_image:
  87. drawing_params_with_rect = DrawingParams(draw_min_rect=True)
  88. drawn_image, area_json = processor.analyze_and_draw(img_bgr, json_data,
  89. drawing_params_with_rect)
  90. temp_img_path = settings.TEMP_WORK_DIR / f'{inference_type}-area_result.jpg'
  91. cv2.imwrite(temp_img_path, drawn_image)
  92. else:
  93. area_json = processor.analyze_from_json(json_data)
  94. with open(area_json_path, 'w', encoding='utf-8') as f:
  95. json.dump(area_json, f, ensure_ascii=False, indent=4)
  96. logger.info("面积计算结束")
  97. return area_json
  98. elif inference_type == "pokemon_card_center":
  99. predictor_inner = get_predictor(settings.DEFECT_TYPE[inference_type]['inner_box'])
  100. predictor_outer = get_predictor(settings.DEFECT_TYPE[inference_type]['outer_box'])
  101. inner_result = predictor_inner.predict_from_image(img_bgr)
  102. outer_result = predictor_outer.predict_from_image(img_bgr)
  103. # temp_inner_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-inner_result.json'
  104. # temp_outer_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-outer_result.json'
  105. # with open(temp_inner_json_path, 'w', encoding='utf-8') as f:
  106. # json.dump(inner_result, f, ensure_ascii=False, indent=4)
  107. # with open(temp_outer_json_path, 'w', encoding='utf-8') as f:
  108. # json.dump(outer_result, f, ensure_ascii=False, indent=4)
  109. inner_points = inner_result['shapes'][0]['points']
  110. outer_points = outer_result['shapes'][0]['points']
  111. center_result = analyze_centering_rotated(inner_points, outer_points)
  112. logger.info("格式化居中数据")
  113. center_result = formate_center_data(center_result, inner_result, outer_result)
  114. temp_center_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-center_result.json'
  115. with open(temp_center_json_path, 'w', encoding='utf-8') as f:
  116. json.dump(center_result, f, ensure_ascii=False, indent=2)
  117. return center_result
  118. else:
  119. return {}
  120. # 创建一个单例服务
  121. defect_service = DefectInferenceService()