PokemonCardSearch/code/tools/test_card_merge.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
近距 mask 合并逻辑测试（在 249 上跑；本地无模型权重）。

用 data/gallery_images/ 的干净单卡图 PIL 合成多卡场景，只测 CardDetector 层
（检出数 / 合并行为 / 阅读顺序），不加载 DINOv2 不查检索，秒级出结果。

用例（期望数按 CARD_MERGE_DIST_PX=50、IoU 门槛关闭）：
  same@20/40    同一张卡水平偏移粘贴 → 中心距=偏移 → 合并 → 检出 1
  same@50       恰好阈值（<= 判定）   → 合并 → 检出 1
  same@60/100   超阈值              → 并存 → 检出 2
  diff@wide     两张不同卡左右并排（间隔≈卡宽）→ 不合并 → 检出 2
  2x2           四卡田字格 → 检出 4 且阅读顺序=先上后下、同行左→右
  vert          两卡竖排 → 检出 2（上在前）

用法（pytorch env）:
  python tools/test_card_merge.py
  CARD_MERGE_DIST_PX=80 python tools/test_card_merge.py   # 换阈值复跑
"""
import os
import sys
import glob
import random

sys.stdout.reconfigure(encoding="utf-8")
_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, _ROOT)

import numpy as np
from PIL import Image

import config
from modules.yolo_detector import CardDetector

GAP = 60            # 拼接图之间的空白间隔(px)，避免贴边影响检测
MARGIN = 40         # 画布边缘留白


def _fetch_minio_fronts(n, out_dir):
    """gallery/query_images 都空时（如 249），从 MinIO grading/raspi_img_data 拉 N 张正面拍摄图。"""
    import config as _cfg
    from minio import Minio
    cfg = _cfg.CAPP_MINIO
    client = Minio(cfg["endpoint"], access_key=cfg["access_key"],
                   secret_key=cfg["secret_key"], secure=bool(cfg.get("secure")))
    objs = [o for o in client.list_objects(cfg["bucket"], prefix="raspi_img_data/", recursive=True)
            if os.path.basename(o.object_name).startswith("raspi_front_")]
    objs.sort(key=lambda o: o.last_modified or 0, reverse=True)
    os.makedirs(out_dir, exist_ok=True)
    paths = []
    for o in objs:
        if len(paths) >= n:
            break
        dst = os.path.join(out_dir, os.path.basename(o.object_name))
        if not os.path.exists(dst) or os.path.getsize(dst) == 0:
            try:
                client.fget_object(cfg["bucket"], o.object_name, dst)
            except Exception:
                continue
        paths.append(dst)
    return paths


def _load_card_images(n):
    """单卡图源：gallery_images → query_images → MinIO raspi_front（树莓派真实拍摄图）。"""
    files = sorted(glob.glob(os.path.join(config.GALLERY_IMG_DIR, "*.jpg")))
    if not files and os.path.isdir(config.QUERY_IMG_DIR):
        files = [os.path.join(config.QUERY_IMG_DIR, f) for f in os.listdir(config.QUERY_IMG_DIR)]
    random.seed(42)
    random.shuffle(files)
    out = []
    for f in files:
        try:
            img = Image.open(f).convert("RGB")
            if img.width >= 200 and img.height >= 200:
                out.append((os.path.basename(f), img))
        except Exception:
            continue
        if len(out) >= n:
            break
    if len(out) < n:
        cache_dir = os.path.join(config.DATA_DIR, "audit_imgs")
        need = n - len(out)
        fetched = _fetch_minio_fronts(max(need, 6), cache_dir)
        for f in fetched[:need]:
            try:
                out.append((os.path.basename(f), Image.open(f).convert("RGB")))
            except Exception:
                continue
    if len(out) < n:
        raise RuntimeError(f"可用单卡图不足: {len(out)}/{n}（gallery/query_images/MinIO raspi_front 均不够）")
    return out


def _paste_offset(base_img, offset_px):
    """同卡水平偏移粘贴：中心距 = offset_px。"""
    w, h = base_img.size
    canvas = Image.new("RGB", (w + offset_px + 2 * MARGIN, h + 2 * MARGIN), (245, 245, 245))
    canvas.paste(base_img, (MARGIN, MARGIN))
    canvas.paste(base_img, (MARGIN + offset_px, MARGIN))
    return canvas


def _paste_two(img_a, img_b, gap=GAP, vertical=False):
    """两张不同卡并排/竖排粘贴。"""
    if vertical:
        w = max(img_a.width, img_b.width)
        canvas = Image.new("RGB", (w + 2 * MARGIN, img_a.height + gap + img_b.height + 2 * MARGIN),
                           (245, 245, 245))
        canvas.paste(img_a, (MARGIN + (w - img_a.width) // 2, MARGIN))
        canvas.paste(img_b, (MARGIN + (w - img_b.width) // 2, MARGIN + img_a.height + gap))
    else:
        h = max(img_a.height, img_b.height)
        canvas = Image.new("RGB", (img_a.width + gap + img_b.width + 2 * MARGIN, h + 2 * MARGIN),
                           (245, 245, 245))
        canvas.paste(img_a, (MARGIN, MARGIN + (h - img_a.height) // 2))
        canvas.paste(img_b, (MARGIN + img_a.width + gap, MARGIN + (h - img_b.height) // 2))
    return canvas


def _paste_grid(imgs, gap=GAP):
    """2×2 田字格（顺序：左上、右上、左下、右下）。"""
    a, b, c, d = imgs
    cw = max(a.width, b.width, c.width, d.width)
    ch = max(a.height, b.height, c.height, d.height)
    W = 2 * cw + gap + 2 * MARGIN
    H = 2 * ch + gap + 2 * MARGIN
    canvas = Image.new("RGB", (W, H), (245, 245, 245))
    for i, im in enumerate(imgs):
        x = MARGIN + (i % 2) * (cw + gap) + (cw - im.width) // 2
        y = MARGIN + (i // 2) * (ch + gap) + (ch - im.height) // 2
        canvas.paste(im, (x, y))
    return canvas


def _detect(detector, canvas, tag):
    """跑检测，返回 (items, center_dists)。"""
    arr = np.array(canvas)
    items = detector.detect_and_crop_all(arr)
    dists = []
    for i in range(len(items)):
        for j in range(i + 1, len(items)):
            d = float(np.hypot(items[i]["center"][0] - items[j]["center"][0],
                               items[i]["center"][1] - items[j]["center"][1]))
            dists.append(round(d, 1))
    print(f"  [{tag}] 检出 {len(items)} 张, 中心距={dists}, "
          f"conf={[round(it['conf'], 3) for it in items]}", flush=True)
    return items, dists


def _make_meta(i, cx, cy, conf, bw=1000.0, bh=1400.0):
    """构造 _merge_indices 输入用的假实例（box 以中心+固定宽高展开）。"""
    return {"box": (cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2),
            "label": "pokemon", "mask": None, "conf": conf, "center": (float(cx), float(cy))}


def test_merge_indices(detector, thr):
    """_merge_indices 纯算法单测（不依赖 YOLO）：贪心按 conf 保高者、阈值边界、IoU 门槛。"""
    p = f = 0

    def check(tag, got, expect, keep_conf=None):
        nonlocal p, f
        ok = len(got) == expect and (keep_conf is None or
                                     all(abs(m["conf"] - keep_conf) < 1e-9 for m in got))
        print(f"  [{'PASS' if ok else 'FAIL'}] {tag}: 保留 {len(got)} 个 conf="
              f"{[m['conf'] for m in got]}（期望 {expect} 个"
              + (f", conf={keep_conf}" if keep_conf is not None else "") + "）", flush=True)
        p += int(ok)
        f += int(not ok)

    # 近距高低 conf：低者被吸收
    detector.merge_dist_px = thr
    got, dropped = detector._merge_indices([
        _make_meta(0, 1000, 1000, 0.60), _make_meta(1, 1030, 1000, 0.95)])
    check("近距保高conf", got, 1, keep_conf=0.95)
    # 阈值边界：中心距恰 = thr → 合并（<= 判定）
    got, _ = detector._merge_indices([
        _make_meta(0, 1000, 1000, 0.95), _make_meta(1, 1000 + thr, 1000, 0.60)])
    check(f"边界{thr:g}px合并(<=)", got, 1, keep_conf=0.95)
    # 刚超阈值：并存
    got, _ = detector._merge_indices([
        _make_meta(0, 1000, 1000, 0.95), _make_meta(1, 1000 + thr + 1, 1000, 0.60)])
    check(f"{thr + 1:g}px并存", got, 2)
    # 链式：A(0.9)-B(0.8) 近、B-C 近、A-C 远 → 贪心保 A，C 与 A 远所以 C 存活
    got, _ = detector._merge_indices([
        _make_meta(0, 1000, 1000, 0.90), _make_meta(1, 1040, 1000, 0.80),
        _make_meta(2, 1080 + thr, 1000, 0.70)])
    check("链式贪心", got, 2, keep_conf=None)
    # 三个近距聚堆：只留 conf 最高
    got, _ = detector._merge_indices([
        _make_meta(0, 1000, 1000, 0.70), _make_meta(1, 1020, 1020, 0.95),
        _make_meta(2, 1040, 980, 0.80)])
    check("三实例聚堆", got, 1, keep_conf=0.95)
    # IoU 门槛启用：距离近但框几乎不重叠（交叉摆卡）→ 不合并
    detector.merge_min_iou = 0.35
    m_a = _make_meta(0, 1000, 1000, 0.95)
    m_b = _make_meta(1, 1000 + thr * 0.5, 1000 + thr * 0.5, 0.60)
    m_b["box"] = (m_b["box"][0] + 900, m_b["box"][1] + 1300, m_b["box"][2] + 900, m_b["box"][3] + 1300)
    m_b["center"] = (m_a["center"][0] + thr * 0.5, m_a["center"][1] + thr * 0.5)
    got, _ = detector._merge_indices([m_a, m_b])
    check("IoU门槛防误杀", got, 2)
    detector.merge_min_iou = 0.0
    # 阈值 0：关闭合并
    detector.merge_dist_px = 0.0
    got, _ = detector._merge_indices([
        _make_meta(0, 1000, 1000, 0.95), _make_meta(1, 1001, 1001, 0.60)])
    check("阈值0关闭", got, 2)
    detector.merge_dist_px = thr
    return p, f


def _reading_order_ok(items):
    """验证输出顺序 == 阅读顺序期望序列：按行容差聚簇（0.5*中位高），
    行间按行内平均 cy，行内按 cx——与 _sort_reading_order 同一套语义重建期望后逐项比对。"""
    if len(items) <= 1:
        return True
    heights = [max(1.0, it["box"][3] - it["box"][1]) for it in items]
    row_tol = 0.5 * float(np.median(heights))
    ordered = sorted(items, key=lambda it: ((it["box"][1] + it["box"][3]) / 2.0,
                                            (it["box"][0] + it["box"][2]) / 2.0))
    rows, cur_cy = [], None
    for it in ordered:
        cy = (it["box"][1] + it["box"][3]) / 2.0
        if rows and abs(cy - cur_cy) <= row_tol:
            rows[-1].append(it)
            cur_cy = sum((x["box"][1] + x["box"][3]) / 2.0 for x in rows[-1]) / len(rows[-1])
        else:
            rows.append([it])
            cur_cy = cy
    expected = [it for r in rows for it in
                sorted(r, key=lambda x: (x["box"][0] + x["box"][2]) / 2.0)]
    return all(a is b for a, b in zip(items, expected))


def main():
    detector = CardDetector(config.YOLO_MODEL_PATH)
    thr = float(getattr(detector, "merge_dist_px", 50.0))
    print(f"[Test] CARD_MERGE_DIST_PX={thr}  MIN_IOU={getattr(detector, 'merge_min_iou', 0.0)}", flush=True)

    passed = failed = 0

    print("\n[Part1] _merge_indices 纯算法单测（不依赖 YOLO 输出）", flush=True)
    p, f = test_merge_indices(detector, thr)
    passed += p
    failed += f

    print("\n[Part2] 合成多卡图全链路 smoke（YOLO 检出 + 合并不回归）", flush=True)
    print("  注：同卡小偏移粘贴时 NMS 在检出层就只出 1 框（近全重叠），与距离合并无关；", flush=True)
    print("      本组验证的是多卡主流程检出数与阅读顺序不被合并逻辑破坏。", flush=True)
    imgs = _load_card_images(5)
    name_a, img_a = imgs[0]
    name_b, img_b = imgs[1]
    print(f"[Test] 用图 A={name_a}  B={name_b}", flush=True)

    cases = []
    for off in (20, 60):
        cases.append((f"same@{off}", _paste_offset(img_a, off), None))   # None=不卡期望,对照raw
    cases.append(("diff@wide", _paste_two(img_a, img_b), 2))
    cases.append(("vert", _paste_two(img_a, img_b, vertical=True), 2))
    cases.append(("grid2x2", _paste_grid([img_a, img_b, imgs[2][1], imgs[3][1]]), 4))

    for tag, canvas, expect in cases:
        detector.merge_dist_px = 0.0
        raw_items, _ = _detect(detector, canvas, tag + "/raw")
        detector.merge_dist_px = thr
        items, dists = _detect(detector, canvas, tag + "/merge")
        if expect is None:
            ok = len(items) == len(raw_items)     # 近重叠场景：合并前后一致（不误减）
        else:
            ok = len(items) == expect
        print(f"  ==> [{'PASS' if ok else 'FAIL'}] {tag}: 原始检出 {len(raw_items)} → 合并后 {len(items)}"
              f"（期望 {'=raw' if expect is None else expect}）\n", flush=True)
        passed += int(ok)
        failed += int(not ok)

        if tag == "grid2x2":
            order_ok = _reading_order_ok(items)
            print(f"  ==> grid2x2 阅读顺序（行容差语义）centers="
                  f"{[(round(it['center'][0]), round(it['center'][1])) for it in items]} "
                  f"[{'PASS' if order_ok else 'FAIL'}]\n", flush=True)
            passed += int(order_ok)
            failed += int(not order_ok)

    print(f"[Result] PASS {passed} / FAIL {failed}", flush=True)
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    sys.exit(main())