score_service.py 4.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677
  1. from app.core.config import settings
  2. from app.core.logger import get_logger
  3. from app.services.defect_service import DefectInferenceService
  4. from app.utils.score_inference.CardScorer import CardScorer
  5. from app.utils.json_data_formate import formate_one_card_result
  6. import json
  7. logger = get_logger(__name__)
  8. class ScoreService:
  9. def __init__(self):
  10. self.scoring_config_path = settings.SCORE_CONFIG_PATH
  11. def score_inference(self, score_type: str, is_reflect_card: bool,
  12. image_bytes: bytes) -> dict:
  13. defect_service = DefectInferenceService()
  14. card_scorer = CardScorer(config_path=self.scoring_config_path)
  15. logger.info("开始进行卡片分数推理")
  16. if score_type == 'front_corner_edge' or score_type == 'front_face':
  17. center_data = defect_service.defect_inference("pokemon_front_card_center", image_bytes)
  18. else:
  19. center_data = defect_service.defect_inference("pokemon_back_card_center", image_bytes)
  20. if is_reflect_card:
  21. if score_type == 'front_corner_edge':
  22. defect_data = defect_service.defect_inference('pokemon_front_corner_reflect_defect', image_bytes)
  23. elif score_type == 'front_face':
  24. defect_data = defect_service.defect_inference('pokemon_front_face_reflect_defect', image_bytes)
  25. elif score_type == 'back_corner_edge':
  26. defect_data = defect_service.defect_inference('pokemon_back_corner_defect', image_bytes)
  27. elif score_type == 'back_face':
  28. defect_data = defect_service.defect_inference('pokemon_back_face_defect', image_bytes)
  29. else:
  30. return {}
  31. else:
  32. if score_type == 'front_corner_edge':
  33. defect_data = defect_service.defect_inference('pokemon_front_corner_no_reflect_defect', image_bytes)
  34. elif score_type == 'front_face':
  35. defect_data = defect_service.defect_inference('pokemon_front_face_no_reflect_defect', image_bytes)
  36. elif score_type == 'back_corner_edge':
  37. defect_data = defect_service.defect_inference('pokemon_back_corner_defect', image_bytes)
  38. elif score_type == 'back_face':
  39. defect_data = defect_service.defect_inference('pokemon_back_face_defect', image_bytes)
  40. else:
  41. return {}
  42. logger.info("模型推理结束, 开始计算分数")
  43. if score_type == 'front_corner_edge' or score_type == 'front_face':
  44. center_score_data = card_scorer.calculate_centering_score('front', center_data, True)
  45. else:
  46. center_score_data = card_scorer.calculate_centering_score('back', center_data, True)
  47. if score_type == 'front_corner_edge':
  48. # 先计算角的分数, 会把分数写入json, 然后传入刚写好的, 继续写边的分数
  49. corner_score_data = card_scorer.calculate_defect_score('corner', 'front', defect_data, True)
  50. defect_score_data = card_scorer.calculate_defect_score('edge', 'front', corner_score_data, True)
  51. elif score_type == 'front_face':
  52. defect_score_data = card_scorer.calculate_defect_score('face', 'front', defect_data, True)
  53. elif score_type == 'back_corner_edge':
  54. corner_score_data = card_scorer.calculate_defect_score('corner', 'back', defect_data, True)
  55. defect_score_data = card_scorer.calculate_defect_score('edge', 'back', corner_score_data, True)
  56. elif score_type == 'back_face':
  57. defect_score_data = card_scorer.calculate_defect_score('face', 'back', defect_data, True)
  58. else:
  59. return {}
  60. result_json = formate_one_card_result(center_score_data, defect_score_data)
  61. temp_score_json_path = settings.TEMP_WORK_DIR / f'{score_type}_score.json'
  62. with open(temp_score_json_path, 'w', encoding='utf-8') as f:
  63. json.dump(result_json, f, ensure_ascii=False, indent=2)
  64. logger.info("分数推理完成 ")
  65. return result_json