score_service.py 3.8 KB

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