score_service.py 4.2 KB

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