PokemonCardSearch/code/modules/grading_detector.py

"""
评级卡检测模块（评级封装壳上的公司标签识别）

用 YOLO26 训练的 card_v1/best.pt 检测交易图上的评级公司标签：
PSA / BGS / CGC / SGC（含子类 AUTHENTIC / AUTO / OLD），共 11 类。

⚠️ 必须在【原图】上检测，不能在卡牌裁切图上：
   卡牌 YOLO(yolov11n_card_seg) 会把卡面裁出来、丢掉评级壳，
   评级标签（印在壳的边框/正面）会随之被裁掉，所以在裁切图上几乎检不到。
   正确做法：在 match_transactions.py 里，卡牌 YOLO 裁切【之前】对原图跑一次本检测。

定位：本模块的输出是匹配结果的【额外字段】，用于回填/校验交易表的 grade_company，
      不参与 DINOv2 特征提取与向量检索（与语种 detected_lang 同级）。
"""
import os
import numpy as np
from modules.ultralytics_compat import import_yolo

# 11 个类别，顺序必须与 ultralytics/data.yaml 的 names 一致（card_v1 训练时定的）
# 0:BGS 1:BGS-AUTHENTIC 2:BGS-AUTO 3:CGC 4:CGC-AUTO 5:CGC-OLD
# 6:PSA 7:PSA-AUTHENTIC 8:SGC 9:SGC-AUTO 10:SGC-OLD
GRADING_NAMES = [
    "BGS", "BGS-AUTHENTIC", "BGS-AUTO",
    "CGC", "CGC-AUTO", "CGC-OLD",
    "PSA", "PSA-AUTHENTIC",
    "SGC", "SGC-AUTO", "SGC-OLD",
]


def parse_label(name):
    """把检测类名拆成 (公司, 子类)。
    'PSA-AUTHENTIC' -> ('PSA', 'AUTHENTIC')；'PSA' -> ('PSA', None)。
    公司即交易表 grade_company 字段的取值（PSA/BGS/CGC/SGC）。"""
    parts = name.split("-", 1)
    return parts[0], (parts[1] if len(parts) > 1 else None)


class GradingDetector:
    """评级公司标签检测器（YOLO26 detect）"""

    def __init__(self, model_path, conf=0.25, imgsz=640, device=None):
        YOLO = import_yolo()  # 规避框架根下同名目录对 ultralytics 的命名空间包遮蔽
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"评级检测模型不存在: {model_path}")
        self.model = YOLO(model_path, task="detect")
        self.conf = conf
        self.imgsz = imgsz
        self.device = device

    def detect(self, img, return_box=False):
        """
        在【原图】上检测评级公司标签，取置信度最高的一个。
        Args:
            img: 图片路径(str) 或 numpy 数组（路径输入用 BGR，数组按原样，颜色不影响类别）
            return_box: 是否返回标签框
        Returns:
            company:  'PSA'/'BGS'/'CGC'/'SGC'  或 None（未检出）
            subtype:  'AUTHENTIC'/'AUTO'/'OLD' 或 None
            conf:     float 或 None
            box:      (x1,y1,x2,y2) 或 None（仅 return_box=True）
        """
        none = (None, None, None, None) if return_box else (None, None, None)
        is_path = isinstance(img, str)
        results = self.model.predict(
            img if is_path else img,
            conf=self.conf, imgsz=self.imgsz, device=self.device, verbose=False,
        )
        if not results:
            return none
        boxes = results[0].boxes
        if boxes is None or len(boxes) == 0:
            return none

        confs = boxes.conf.cpu().numpy()
        clses = boxes.cls.cpu().numpy().astype(int)
        xyxy = boxes.xyxy.cpu().numpy()
        best = int(np.argmax(confs))
        if clses[best] >= len(GRADING_NAMES):
            return none
        company, subtype = parse_label(GRADING_NAMES[clses[best]])
        c = float(confs[best])
        if return_box:
            return company, subtype, c, tuple(float(v) for v in xyxy[best])
        return company, subtype, c