PokemonCardSearch/code/yolo_crop.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""YOLO 裁剪卡牌区域"""
import os, sys, csv, hashlib, json
import numpy as np
sys.stdout.reconfigure(encoding="utf-8")
sys.path.insert(0, "/home/martin/顾工交接/wzj")
import config
import cv2

DATA = config.DATA_DIR
def md5n(u): return hashlib.md5(u.encode()).hexdigest() + ".jpg"

print("[YOLO] 初始化裁剪器...", flush=True)
from modules.ultralytics_compat import import_yolo   # 规避框架根下同名目录遮蔽 ultralytics
YOLO = import_yolo()
yolo = YOLO(config.YOLO_MODEL_PATH, task="segment")
print("[YOLO] 加载完成", flush=True)

CROP_DIR = os.path.join(DATA, "query_crops")
os.makedirs(CROP_DIR, exist_ok=True)

txs = list(csv.DictReader(open(os.path.join(DATA, "transactions_ebay.csv"), encoding="utf-8-sig")))
print("[YOLO] 处理 %d 张图片" % len(txs), flush=True)

cropped = {}
no_crop = 0
for i, tx in enumerate(txs):
    tid = tx["tx_id"]
    img_path = os.path.join(config.QUERY_IMG_DIR, md5n(tx["img_url"]))
    if not os.path.exists(img_path):
        continue
    
    try:
        results = yolo.predict(img_path, conf=config.YOLO_CONF_THRESHOLD,
                               imgsz=config.YOLO_IMG_SIZE, verbose=False)
        if not results:
            no_crop += 1
            continue
        res = results[0]
        orig = res.orig_img
        boxes = res.boxes
        if boxes is None or len(boxes) == 0:
            # 没检测到卡牌，用原图
            crop_path = os.path.join(CROP_DIR, md5n(tx["img_url"]))
            cv2.imwrite(crop_path, orig)
            cropped[tid] = crop_path
            continue
        
        xyxy = boxes.xyxy.cpu().numpy()
        areas = (xyxy[:, 2] - xyxy[:, 0]) * (xyxy[:, 3] - xyxy[:, 1])
        max_idx = int(np.argmax(areas))
        x1, y1, x2, y2 = xyxy[max_idx]
        h, w = orig.shape[:2]
        x1, y1 = max(0, int(x1)), max(0, int(y1))
        x2, y2 = min(w, int(x2)), min(h, int(y2))
        if x2 <= x1 or y2 <= y1:
            cv2.imwrite(os.path.join(CROP_DIR, md5n(tx["img_url"])), orig)
            cropped[tid] = os.path.join(CROP_DIR, md5n(tx["img_url"]))
            continue
        
        crop = orig[y1:y2, x1:x2]
        crop_path = os.path.join(CROP_DIR, md5n(tx["img_url"]))
        cv2.imwrite(crop_path, crop)
        cropped[tid] = crop_path
        
        if (i+1) % 10 == 0:
            print("[YOLO] %d/%d 裁剪完成" % (i+1, len(txs)), flush=True)
    except Exception as e:
        no_crop += 1

print("[YOLO] 完成! 成功裁剪 %d 张, 失败 %d 张" % (len(cropped), no_crop), flush=True)
print("[YOLO] 裁剪图保存于: %s" % CROP_DIR, flush=True)