CheckCardBoxAndDefectServer/app/services/card_rectify_and_center.py

import cv2
import numpy as np
import copy
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) -> tuple[np.ndarray, dict]:
        """
        修改返回值: (final_image, rectified_json)
        """
        # 1. 设置处理参数
        params = FryCardProcessParams(
            debug_level="detail",
            label_name="outer_box",
            center_mode=CenterMode.BOUNDING_RECT,
            fill_mode=FillMode.BLACK
        )

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

        # 3. 原始外框推理
        seg_json = self.outer_box_model.predict_from_image(img_bgr)
        seg_json = self._keep_max_prob_shape(seg_json)

        # 4. 执行处理，获取图像和变换矩阵
        final_image, matrix = processor.process_image_with_json(img_bgr, seg_json, params)

        # 5. 生成变换后的 JSON 数据
        rectified_json = copy.deepcopy(seg_json)
        if matrix is not None and 'shapes' in rectified_json:
            logger.info("正在将原始外框坐标映射到转正后的坐标系...")
            for shape in rectified_json['shapes']:
                original_points = shape.get('points', [])
                if original_points:
                    # 使用矩阵变换点
                    new_points = processor.transform_points_with_matrix(original_points, matrix)
                    shape['points'] = new_points
                    # 清除可能存在的 rect_box，确保后续逻辑重新计算最小外接矩形
                    if 'rect_box' in shape:
                        del shape['rect_box']

        temp_img_path = settings.TEMP_WORK_DIR / "rectify_center_img.jpg"
        if final_image is not None:
            cv2.imwrite(temp_img_path, final_image)

        return final_image, rectified_json