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,
                         pre_calculated_outer_result: dict = None) -> dict:
        """
        执行卡片识别推理。

        Args:
            inference_type: 模型类型 (e.g., 'outer_box').
            img_bgr: 图像。
            pre_calculated_outer_result: 如果提供了转正后的外框数据，则不再重新推理外框

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

        # 辅助函数：获取外框结果（优先使用传入的，否则现场推理）
        def get_outer_result(tag=""):
            if pre_calculated_outer_result:
                logger.info(f"[{tag}] 使用预计算的转正外框数据")
                return pre_calculated_outer_result
            else:
                logger.info(f"[{tag}] 未提供外框数据，开始推理外框...")
                p_outer = get_predictor("outer_box")
                res = p_outer.predict_from_image(img_bgr)
                # 如果是现场推理的，需要过滤一下
                if inference_type in ["pokemon_front_card_center", "pokemon_back_card_center"]:
                    res = self._filter_max_prob_shape(res, tag=f"{tag}-过滤")
                return res

        # 面
        if (inference_type == "pokemon_back_face_coaxial_light_defect"
                or inference_type == "pokemon_front_face_reflect_coaxial_light_defect"
                or inference_type == "pokemon_front_face_no_reflect_coaxial_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'
            )

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

            logger.info("开始执行外框过滤...")
            try:
                # 获取外框
                outer_result = get_outer_result(tag="面流程")

                # 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.copy(), 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_back_face_ring_light_defect"
              or inference_type == "pokemon_front_face_reflect_ring_light_defect"
              or inference_type == "pokemon_front_face_no_reflect_ring_light_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='face'
            )

            # 简化点数
            logger.info("开始进行点数简化")
            for shapes in json_data["shapes"]:
                shapes["points"] = simplifyPoints.simplify_points(shapes["points"])

            # 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("开始执行外框过滤...")
            # 外框推理
            outer_result = get_outer_result(tag="Corner流程")

            # 过滤外框，只留一个
            if not pre_calculated_outer_result:
                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.copy(), 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,
                                            settings.EDGE_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 = get_outer_result(tag="Center流程")

            # 过滤内框和外框，只留最大置信度的框
            inner_result = self._filter_max_prob_shape(inner_result, tag="Center流程-内框")
            # 如果是传入的预计算结果，它已经过滤过了；如果是现算的，需要过滤
            if not pre_calculated_outer_result:
                outer_result = self._filter_max_prob_shape(outer_result, tag="Center流程-外框")

            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.copy(), 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, card_light_type: str, center_json: dict, defect_json: dict) -> dict:
        light_type_list = ["center", "coaxial", "ring"]
        if card_light_type not in light_type_list:
            logger.error(f"inference_type 只能为{light_type_list}, 输入为{card_light_type}")
            raise ValueError(f"inference_type 只能为{light_type_list}, 输入为{card_light_type}")

        processor = DefectProcessor(pixel_resolution=settings.PIXEL_RESOLUTION)

        # 对于面的图, 不计算居中相关, 这里得到的center_json 应该为 {}
        if card_light_type == "coaxial":
            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 card_light_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 card_light_type == "ring":

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

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

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

            return edge_corner_data
        else:
            return {}

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