CheckCardBoxAndDefectServer/app/services/card_rectify_and_center.py

import cv2
import numpy as np
from ..core.model_loader import get_predictor
from app.utils.defect_inference.img_rectify_and_center import (
    FryCardProcessParams, FryCardProcessor, CenterMode, FillMode)
from app.core.config import settings
from app.core.logger import get_logger
import json

logger = get_logger(__name__)


class CardRectifyAndCenter:
    def __init__(self):
        self.inference_type = "outer_box"
        self.outer_box_model = get_predictor(self.inference_type)

    def _keep_max_prob_shape(self, json_data: dict) -> dict:
        """
        辅助函数：过滤掉低置信度的框，只保留 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

        logger.info(f"检测到 {len(shapes)} 个外框，开始筛选最大置信度目标...")

        # 使用 max 函数找出 probability 最大的 shape
        # x.get('probability', 0) 确保如果字段不存在默认为0
        best_shape = max(shapes, key=lambda x: x.get('probability', 0))

        logger.info(f"筛选结果: 保留置信度为 {best_shape.get('probability')} 的目标，"
                    f"丢弃了其他 {len(shapes) - 1} 个目标。")

        # 重新赋值给 shapes，保持列表结构
        json_data['shapes'] = [best_shape]

        # 同时更新 num 数量，保持数据一致性
        if 'num' in json_data:
            json_data['num'] = 1

        return json_data

    def rectify_and_center(self, img_bgr: np.ndarray) -> np.ndarray:
        # 1. 设置处理参数
        params = FryCardProcessParams(
            debug_level="detail",
            label_name="outer_box",
            center_mode=CenterMode.BOUNDING_RECT,
            fill_mode=FillMode.BLACK
        )

        # 2. 初始化处理器
        processor = FryCardProcessor()

        seg_json = self.outer_box_model.predict_from_image(img_bgr)
        # 过滤噪点，只留最大置信度的外框
        seg_json = self._keep_max_prob_shape(seg_json)

        # 4. 执行处理
        final_image = processor.process_image_with_json(img_bgr, seg_json, params)
        temp_img_path = settings.TEMP_WORK_DIR / "rectify_center_img.jpg"
        cv2.imwrite(temp_img_path, final_image)

        return final_image