defect_service.py 7.2 KB

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