CheckCardBoxAndDefectServer/app/services/score_service.py

import cv2

from app.core.config import settings
from app.core.logger import get_logger
from app.services.defect_service import DefectInferenceService
from app.services.card_rectify_and_center import CardRectifyAndCenter
from app.utils.score_inference.CardScorer import CardScorer
from app.utils.json_data_formate import formate_one_card_result
import numpy as np
import json

logger = get_logger(__name__)


class ScoreService:
    def __init__(self):
        self.scoring_config_path = settings.SCORE_CONFIG_PATH

    def score_inference(self, score_type: str, is_reflect_card: bool,
                        img_bgr: np.ndarray) -> dict:
        defect_service = DefectInferenceService()
        rectify_center_service = CardRectifyAndCenter()
        card_scorer = CardScorer(config_path=self.scoring_config_path)

        logger.info("开始进行卡片居中和转正")
        img_bgr = rectify_center_service.rectify_and_center(img_bgr)


        logger.info("开始进行卡片分数推理")
        if score_type == 'front_corner_edge' or score_type == 'front_face':
            center_data = defect_service.defect_inference("pokemon_front_card_center", img_bgr.copy())
        else:
            center_data = defect_service.defect_inference("pokemon_back_card_center", img_bgr.copy())

        if is_reflect_card:
            if score_type == 'front_corner_edge':
                defect_data = defect_service.defect_inference('pokemon_front_corner_reflect_defect', img_bgr.copy())
            elif score_type == 'front_face':
                defect_data = defect_service.defect_inference('pokemon_front_face_reflect_defect', img_bgr.copy())
            elif score_type == 'back_corner_edge':
                defect_data = defect_service.defect_inference('pokemon_back_corner_defect', img_bgr.copy())
            elif score_type == 'back_face':
                defect_data = defect_service.defect_inference('pokemon_back_face_defect', img_bgr.copy())
            else:
                return {}
        else:
            if score_type == 'front_corner_edge':
                defect_data = defect_service.defect_inference('pokemon_front_corner_no_reflect_defect', img_bgr.copy())
            elif score_type == 'front_face':
                defect_data = defect_service.defect_inference('pokemon_front_face_no_reflect_defect', img_bgr.copy())
            elif score_type == 'back_corner_edge':
                defect_data = defect_service.defect_inference('pokemon_back_corner_defect', img_bgr.copy())
            elif score_type == 'back_face':
                defect_data = defect_service.defect_inference('pokemon_back_face_defect', img_bgr.copy())

            else:
                return {}

        logger.info("模型推理结束, 开始计算分数")

        if score_type == 'front_corner_edge' or score_type == 'front_face':
            center_score_data = card_scorer.calculate_centering_score('front', center_data, True)
        else:
            center_score_data = card_scorer.calculate_centering_score('back', center_data, True)

        if score_type == 'front_corner_edge':
            # 先计算角的分数, 会把分数写入json, 然后传入刚写好的, 继续写边的分数
            corner_score_data = card_scorer.calculate_defect_score('corner', 'front', defect_data, True)
            defect_score_data = card_scorer.calculate_defect_score('edge', 'front', corner_score_data, True)
        elif score_type == 'front_face':
            defect_score_data = card_scorer.calculate_defect_score('face', 'front', defect_data, True)
        elif score_type == 'back_corner_edge':
            corner_score_data = card_scorer.calculate_defect_score('corner', 'back', defect_data, True)
            defect_score_data = card_scorer.calculate_defect_score('edge', 'back', corner_score_data, True)
        elif score_type == 'back_face':
            defect_score_data = card_scorer.calculate_defect_score('face', 'back', defect_data, True)
        else:
            return {}

        result_json = formate_one_card_result(center_score_data, defect_score_data)

        temp_score_json_path = settings.TEMP_WORK_DIR / f'{score_type}_score.json'
        with open(temp_score_json_path, 'w', encoding='utf-8') as f:
            json.dump(result_json, f, ensure_ascii=False, indent=2)
        logger.info("分数推理完成 ")
        return result_json