PokemonCardSearch/code/modules/attribute_detector.py

# -*- coding: utf-8 -*-
"""
attribute_detector.py - 宝可梦卡牌属性检测（火/水/草/电...）

与 cat16 三属性流水线（ebay_sync/infer_cat16_pipeline.py）同款配方：
  裁卡图 → roi_box.json 的 box3 归一化框（右上角元素图标区）裁小图
        → 4x LANCZOS 放大 → yolo26s_det_attribute 检测取 top1 类别名
整卡直接喂模型检不出框（图标太小），必须先裁 ROI 再放大。

轻量常驻，加载 20MB pt，推理 ~10ms。
"""
import os
import json
import numpy as np
from PIL import Image
from modules.ultralytics_compat import import_yolo

_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ATTRIBUTE_MODEL_PATH = os.path.join(
    _ROOT, "ultralytics/runs/detect/yolo26s_det_attribute/weights/best.pt"
)
ROI_JSON_PATH = os.path.join(
    _ROOT, "data_all/卡牌区域/images_pokemon_card_crops/roi_box.json"
)
BOX3_LABEL = "box3"   # 右上角元素图标区（归一化 x1,y1,x2,y2）
UPSCALE = 4           # LANCZOS 放大倍数（同 cat16 / OCR v5 方案）


def _load_box3(path=None):
    """roi_box.json → box3 归一化框 (x1, y1, x2, y2)；缺失返回 None。"""
    try:
        with open(path or ROI_JSON_PATH, encoding="utf-8-sig") as f:
            rois = json.load(f)["rois_normalized"]
        for item in rois:
            if item["label"] == BOX3_LABEL:
                return (item["x1"], item["y1"], item["x2"], item["y2"])
    except Exception:
        pass
    return None


class AttributeDetector:
    def __init__(self, model_path=None, roi_json_path=None):
        self.roi = _load_box3(roi_json_path)
        if self.roi is None:
            print("[AttributeDetector] roi_box.json 缺失或无 box3，attribute 将始终为 None", flush=True)
            self.model = None
            return
        YOLO = import_yolo()
        path = model_path or ATTRIBUTE_MODEL_PATH
        if not os.path.exists(path):
            print("[AttributeDetector] 模型不存在: %s，attribute 将始终为 None" % path, flush=True)
            self.model = None
            return
        self.model = YOLO(path, task="detect")
        self.device = ""  # 空串走auto分支: vendored select_device(device=0)会把CVD覆盖成0从而撞GTX1060
        self.names = self.model.model.names
        print("[AttributeDetector] loaded %d classes: %s" % (len(self.names), list(self.names.values())), flush=True)

    def _crop_roi_4x(self, img_rgb):
        """裁卡 RGB → box3 ROI 裁块 → 4x LANCZOS 放大（cat16 同款）。"""
        h, w = img_rgb.shape[:2]
        x1, y1, x2, y2 = self.roi
        px1, py1 = max(0, int(x1 * w)), max(0, int(y1 * h))
        px2, py2 = min(w, int(x2 * w)), min(h, int(y2 * h))
        if px2 <= px1 or py2 <= py1:
            return None
        patch = Image.fromarray(img_rgb[py1:py2, px1:px2])
        return np.array(patch.resize((patch.size[0] * UPSCALE, patch.size[1] * UPSCALE), Image.LANCZOS))

    def predict(self, img_rgb):
        """输入裁卡 RGB numpy 数组，返回元素标签字符串（如 '草'/'水'）或 None。"""
        if self.model is None or img_rgb is None:
            return None
        roi_img = self._crop_roi_4x(img_rgb)
        if roi_img is None:
            return None
        results = self.model.predict(
            roi_img, conf=0.25, imgsz=640,
            device=self.device, verbose=False,
        )
        if not results:
            return None
        res = results[0]
        boxes = res.boxes
        if boxes is None or len(boxes) == 0:
            return None
        # 取置信度最高的实例类别名
        confs = boxes.conf.cpu().numpy()
        clses = boxes.cls.cpu().numpy()
        best_idx = int(np.argmax(confs))
        cls_id = int(clses[best_idx])
        return self.names.get(cls_id)