CheckCardBoxAndDefectServer/app/services/defect_service.py

import cv2
import numpy as np
from ..core.model_loader import get_predictor
from app.utils.defect_inference.CardDefectAggregator import CardDefectAggregator
from app.utils.defect_inference.arean_anylize_draw import DefectProcessor, DrawingParams
from app.utils.defect_inference.AnalyzeCenter import (
    analyze_centering_rotated, analyze_centering_rect, formate_center_data)
from app.utils.defect_inference.DrawCenterInfo import draw_boxes_and_center_info
from app.utils.defect_inference.ClassifyEdgeCorner import ClassifyEdgeCorner
from app.utils.json_data_formate import formate_face_data, formate_add_edit_type
from app.utils.defect_inference.FilterOutsideDefects import FilterOutsideDefects
from app.utils.simplify_points import SimplifyPoints
from app.core.config import settings
from app.core.logger import get_logger
import json

logger = get_logger(__name__)


class DefectInferenceService:
    def _filter_max_prob_shape(self, json_data: dict, tag: str = "unknown") -> dict:
        """
        通用过滤函数：只保留 shapes 中 probability 最大的那一个。
        """
        if not json_data or 'shapes' not in json_data or not json_data['shapes']:
            return json_data

        shapes = json_data['shapes']
        if len(shapes) <= 1:
            return json_data

        # 按 probability 降序排列，取第一个
        best_shape = max(shapes, key=lambda x: x.get('probability', 0))

        logger.info(f"[{tag}] 过滤多余框: 原有 {len(shapes)} 个, 保留最大置信度 {best_shape.get('probability'):.4f}")

        json_data['shapes'] = [best_shape]
        return json_data

    def defect_inference(self, inference_type: str, img_bgr: np.ndarray,
                         is_draw_image=True) -> dict:
        """
        执行卡片识别推理。

        Args:
            inference_type: 模型类型 (e.g., 'outer_box').
            img_bgr: 图像。

        Returns:
            一个包含推理结果的字典。
        """
        simplifyPoints = SimplifyPoints()
        outside_filter = FilterOutsideDefects(expansion_pixel=30)

        # 面
        if (inference_type == "pokemon_front_face_no_reflect_defect"
                or inference_type == "pokemon_front_face_reflect_defect"
                or inference_type == "pokemon_back_face_defect"
                or inference_type == "pokemon_back_face_ring_light_defect"
                or inference_type == "pokemon_front_face_reflect_ring_light_defect"):
            # 1. 获取对应的预测器实例
            predictor = get_predictor(inference_type)

            # 3. 调用我们新加的 predict_from_image 方法进行推理
            # result = predictor.predict_from_image(img_bgr)

            # 3. 实例化我们聚合器，传入预测器
            aggregator = CardDefectAggregator(
                predictor=predictor,
                tile_size=512,
                overlap_ratio=0.1,  # 10% 重叠
            )

            json_data = aggregator.process_image(
                image=img_bgr,
                mode='face'
            )

            # 简化点数
            for shapes in json_data["shapes"]:
                points = shapes["points"]
                num1 = len(points)
                simplify_points = simplifyPoints.simplify_points(points)
                shapes["points"] = simplify_points
                new_num1 = len(simplify_points)
                logger.info(f"num: {num1}, new_num1: {new_num1}")

            logger.info("开始执行外框过滤...")
            try:
                # 1. 因为面原本没有外框推理，这里需要专门加载并推理一次
                predictor_outer = get_predictor("outer_box")
                outer_result = predictor_outer.predict_from_image(img_bgr)

                # 2. 执行过滤
                json_data = outside_filter.execute(json_data, outer_result)
            except Exception as e:
                logger.error(f"面流程外框过滤失败: {e}")

            merge_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-merge.json'
            with open(merge_json_path, 'w', encoding='utf-8') as f:
                json.dump(json_data, f, ensure_ascii=False, indent=4)
            logger.info(f"合并结束")

            processor = DefectProcessor(pixel_resolution=settings.PIXEL_RESOLUTION)
            area_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-face_result.json'
            if is_draw_image:
                drawing_params_with_rect = DrawingParams(draw_min_rect=True)
                drawn_image, area_json = processor.analyze_and_draw(img_bgr, json_data,
                                                                    drawing_params_with_rect)
                temp_img_path = settings.TEMP_WORK_DIR / f'{inference_type}-face_result.jpg'
                cv2.imwrite(temp_img_path, drawn_image)
            else:
                area_json = processor.analyze_from_json(json_data)

            face_json_result = formate_face_data(area_json)
            face_json_result = formate_add_edit_type(face_json_result)

            with open(area_json_path, 'w', encoding='utf-8') as f:
                json.dump(face_json_result, f, ensure_ascii=False, indent=2)
            logger.info("面的面积计算结束")

            return face_json_result

        # 边角
        elif (inference_type == "pokemon_front_corner_no_reflect_defect"
              or inference_type == "pokemon_front_corner_reflect_defect"
              or inference_type == "pokemon_back_corner_defect"):
            predictor = get_predictor(inference_type)

            aggregator = CardDefectAggregator(
                predictor=predictor,
                tile_size=512,
                overlap_ratio=0.1,  # 10% 重叠
            )

            json_data = aggregator.process_image(
                image=img_bgr,
                mode='edge'
            )

            # 简化点数
            for shapes in json_data["shapes"]:
                points = shapes["points"]
                num1 = len(points)
                simplify_points = simplifyPoints.simplify_points(points)
                shapes["points"] = simplify_points
                new_num1 = len(simplify_points)
                logger.info(f"num: {num1}, new_num1: {new_num1}")

            # merge_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-merge.json'
            # with open(merge_json_path, 'w', encoding='utf-8') as f:
            #     json.dump(json_data, f, ensure_ascii=False, indent=4)
            # logger.info(f"合并结束")

            logger.info("开始执行外框过滤...")
            # 外框推理
            predictor_outer = get_predictor("outer_box")
            outer_result = predictor_outer.predict_from_image(img_bgr)

            # 过滤外框，只留一个
            outer_result = self._filter_max_prob_shape(outer_result, tag="Corner流程-外框")

            # 2. 执行过滤, 去掉外框外的缺陷
            json_data = outside_filter.execute(json_data, outer_result)

            processor = DefectProcessor(pixel_resolution=settings.PIXEL_RESOLUTION)
            area_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-corner_result.json'
            if is_draw_image:
                drawing_params_with_rect = DrawingParams(draw_min_rect=True)
                drawn_image, area_json = processor.analyze_and_draw(img_bgr, json_data,
                                                                    drawing_params_with_rect)
                temp_img_path = settings.TEMP_WORK_DIR / f'{inference_type}-corner_result.jpg'
                cv2.imwrite(temp_img_path, drawn_image)
            else:
                area_json: dict = processor.analyze_from_json(json_data)
            logger.info("边角缺陷面积计算结束")

            classifier = ClassifyEdgeCorner(settings.PIXEL_RESOLUTION, settings.CORNER_SIZE_MM)
            edge_corner_data = classifier.classify_defects_location(area_json, outer_result)

            edge_corner_data = formate_add_edit_type(edge_corner_data)
            with open(area_json_path, 'w', encoding='utf-8') as f:
                json.dump(edge_corner_data, f, ensure_ascii=False, indent=2)
            logger.info("边角面积计算结束")

            return edge_corner_data

        elif inference_type == "pokemon_front_card_center" \
                or inference_type == "pokemon_back_card_center":

            predictor_inner = get_predictor(settings.DEFECT_TYPE[inference_type]['inner_box'])
            predictor_outer = get_predictor(settings.DEFECT_TYPE[inference_type]['outer_box'])

            inner_result = predictor_inner.predict_from_image(img_bgr)
            outer_result = predictor_outer.predict_from_image(img_bgr)

            # 过滤内框和外框，只留最大置信度的框
            inner_result = self._filter_max_prob_shape(inner_result, tag="Center流程-内框")
            outer_result = self._filter_max_prob_shape(outer_result, tag="Center流程-外框")

            # temp_inner_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-inner_result.json'
            # temp_outer_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-outer_result.json'
            # with open(temp_inner_json_path, 'w', encoding='utf-8') as f:
            #     json.dump(inner_result, f, ensure_ascii=False, indent=4)
            # with open(temp_outer_json_path, 'w', encoding='utf-8') as f:
            #     json.dump(outer_result, f, ensure_ascii=False, indent=4)

            inner_points = inner_result['shapes'][0]['points']
            outer_points = outer_result['shapes'][0]['points']
            center_result, inner_rect_box, outer_rect_box = analyze_centering_rotated(inner_points, outer_points)

            # logger.info(f"inner_rect_box: {type(inner_rect_box)}, {inner_rect_box}")
            # logger.info(f"outer_rect_box: , {outer_rect_box}")
            logger.info("格式化居中数据")
            center_result = formate_center_data(center_result,
                                                inner_result, outer_result,
                                                inner_rect_box, outer_rect_box)

            draw_img = draw_boxes_and_center_info(img_bgr, center_result)
            temp_center_img_path = settings.TEMP_WORK_DIR / f'{inference_type}-center_result.jpg'
            cv2.imwrite(temp_center_img_path, draw_img)

            temp_center_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-center_result.json'
            with open(temp_center_json_path, 'w', encoding='utf-8') as f:
                json.dump(center_result, f, ensure_ascii=False, indent=2)
            return center_result

        else:
            return {}

    # inference_type: center, face, corner_edge
    def re_inference_from_json(self, inference_type: str, center_json: dict, defect_json: dict) -> dict:
        inference_type_list = ["center", "face", "corner_edge"]
        if inference_type not in inference_type_list:
            logger.error(f"inference_type 只能为{inference_type_list}, 输入为{inference_type}")
            raise ValueError(f"inference_type 只能为{inference_type_list}, 输入为{inference_type}")

        processor = DefectProcessor(pixel_resolution=settings.PIXEL_RESOLUTION)

        # 对于面的图, 不计算居中相关, 这里得到的center_json 应该为 {}
        if inference_type == "face":
            area_json = processor.re_analyze_from_json(defect_json)
            face_json_result = formate_face_data(area_json)
            logger.info("面缺陷面积计算结束")

            return face_json_result

        inner_result = center_json['box_result']['inner_box']
        outer_result = center_json['box_result']['outer_box']

        if inference_type == "center":
            logger.info("居中重计算")
            inner_rect = inner_result['shapes'][0]['rect_box']
            outer_rect = outer_result['shapes'][0]['rect_box']
            center_result, inner_rect_box, outer_rect_box = analyze_centering_rect(inner_rect, outer_rect)

            center_result = formate_center_data(center_result,
                                                inner_result, outer_result,
                                                inner_rect_box, outer_rect_box)
            return center_result
        elif inference_type == "corner_edge":

            area_json: dict = processor.re_analyze_from_json(defect_json)
            logger.info("边角缺陷面积计算结束")

            # 根据外框区分边和角
            classifier = ClassifyEdgeCorner(settings.PIXEL_RESOLUTION, settings.CORNER_SIZE_MM)
            edge_corner_data = classifier.classify_defects_location(area_json, outer_result)

            logger.info("边角面积计算结束")

            return edge_corner_data
        else:
            return {}

# 创建一个单例服务
# defect_service = DefectInferenceService()