PokemonCardSearch/code/modules/rectifier.py

"""
卡牌透视矫正模块 (rectifier)
作用：把 YOLO 检出的卡牌区域，用分割 mask 的 4 角点做透视拉正（梯形→矩形），
      并按 PaddleOCR 给出的方向 angle 把卡牌转到正立竖向。
设计要点：
  - 透视拉正只用 YOLO mask（凸包→approxPolyDP→4角点），与 OCR 无关；
  - 正立(防180°倒置)用 PaddleOCR doc_preprocessor_res.angle（0/90/180/270），
    该 angle 由 paddleocr 环境的语种脚本一并产出，主流程读 JSON 获得；
  - mask 可用时直接在原图 warpPerspective（裁切+去透视一步到位）；
    mask 不可用时回退到 YOLO bbox 矩形裁切；再不行用原图。
  - 只影响下游 DINOv2 匹配，不影响 OCR / 语种判断。
纯 cv2 + numpy，运行在 pytorch 环境，无 OCR 依赖。
"""
import cv2
import numpy as np

# 卡牌标准宽高比（宝可梦≈0.716，球星卡≈0.714），用于约束 warp 输出比例，避免 mask 抖动导致畸变
CARD_ASPECT = 0.714

# PaddleOCR angle = 输入图相对正立顺时针旋转的角度；要转正就反向旋转
# 90  → 输入顺时针倒 → 逆时针转90° (ROTATE_90_COUNTERCLOCKWISE)
# 180 → 转180°
# 270 → 输入逆时针倒 → 顺时针转90° (ROTATE_90_CLOCKWISE)
_ANGLE_TO_CV = {
    0: None,
    90: cv2.ROTATE_90_COUNTERCLOCKWISE,
    180: cv2.ROTATE_180,
    270: cv2.ROTATE_90_CLOCKWISE,
}


def mask_to_corners(mask, min_area=30.0):
    """
    二值 mask → 凸包 → approxPolyDP 取 4 角点。
    Args:
        mask: uint8 HxW，非零区域为卡牌
    Returns:
        corners (4,2) float32；失败返回 None
    """
    if mask is None:
        return None
    m = (mask > 0).astype(np.uint8)
    cnts, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not cnts:
        return None
    c = max(cnts, key=cv2.contourArea)
    if cv2.contourArea(c) < min_area:
        return None
    hull = cv2.convexHull(c)
    peri = cv2.arcLength(hull, True)
    for eps in (0.01, 0.02, 0.03, 0.04, 0.06, 0.08, 0.10, 0.15):
        approx = cv2.approxPolyDP(hull, eps * peri, True)
        if len(approx) == 4:
            return approx.reshape(4, 2).astype(np.float32)
    # 兜底：最小外接矩形 4 角点
    box = cv2.boxPoints(cv2.minAreaRect(hull))
    return box.astype(np.float32)


def order_corners(pts):
    """
    4 点排序为 TL, TR, BR, BL（按图像坐标，与卡牌真实上下无关；上下由 angle 后续修正）。
    """
    pts = np.asarray(pts, dtype=np.float32).reshape(4, 2)
    rect = np.zeros((4, 2), dtype=np.float32)
    s = pts.sum(axis=1)
    d = np.diff(pts, axis=1)  # y - x
    rect[0] = pts[np.argmin(s)]   # TL: 最小 x+y
    rect[2] = pts[np.argmax(s)]   # BR: 最大 x+y
    rect[1] = pts[np.argmin(d)]   # TR: 最小 y-x
    rect[3] = pts[np.argmax(d)]   # BL: 最大 y-x
    return rect


def _angle_at(a, b, c):
    """角 abc 的度数"""
    v1, v2 = a - b, c - b
    n = np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-9
    return float(np.degrees(np.arccos(np.clip(np.dot(v1, v2) / n, -1, 1))))


def valid_quad(ordered, angle_tol=40.0, edge_ratio_max=3.0):
    """4 角点四边形是否合理：每个内角近 90°、边长比不极端。
    用于挡住 mask 抖动/角点崩坏（如 demo 第6张那种被标偏的情况），异常则回退 bbox。
    阈值宽松：内角 50~130°、边长 max/min<3。第6张(87-93°,1.32)轻松通过。"""
    tl, tr, br, bl = ordered
    quad = [tl, tr, br, bl]
    for i in range(4):
        if abs(_angle_at(quad[i - 1], quad[i], quad[(i + 1) % 4]) - 90) > angle_tol:
            return False
    edges = [np.linalg.norm(tr - tl), np.linalg.norm(br - tr),
             np.linalg.norm(bl - br), np.linalg.norm(tl - bl)]
    emin = min(edges)
    if emin < 1 or max(edges) / emin > edge_ratio_max:
        return False
    return True


def has_perspective(ordered, trap_thresh=0.10, maxdev_thresh=3.5):
    """四边形是否有"值得矫正"的真实透视形变。
    两个指标：
      trap = 上下边长差的相对量（梯形度，0=对边完全平行）；
      maxdev = 四个内角偏离 90° 的最大值（倾斜/剪切）。
    任一超阈值才算有透视，否则视为"本就端正"——这种卡 warp 只会引入重采样噪声
    （实测 tx13/25/86 透视极小，硬 warp 反而掉 0.06~0.36 相似度）。
    阈值由 tx13/25/86(无,≤0.083/≤2.4) 与 tx22/55/2(有,≥0.18 或 ≥4.3) 的间隙标定。"""
    tl, tr, br, bl = ordered
    wt, wb = np.linalg.norm(tr - tl), np.linalg.norm(br - bl)
    trap = abs(wt - wb) / (min(wt, wb) + 1e-9)
    angs = [_angle_at(ordered[(i - 1) % 4], ordered[i], ordered[(i + 1) % 4]) for i in range(4)]
    maxdev = max(abs(a - 90) for a in angs)
    return trap > trap_thresh or maxdev > maxdev_thresh


def angle_geometry_disagree(mask, angle):
    """OCR 文字方向(angle) 与 卡牌几何方向(mask) 是否不一致。
    不一致 = 模型对"卡牌到底正不正"没把握 → 应走人工判断(manual_judgment)。
    判定：
      几何：mask 4角点的 W/H，W>H 视为"侧放"(landscape)；
      OCR：angle∈{90,270} 视为"侧放"。
      两者不一致(一个说侧放、一个说正立) → 返回 True。
    实例：tx86(mask横向+angle0)、tx38/65(mask竖向+angle90) 均不一致→人工；
          tx2(mask横向+angle90) 一致→自动转正。"""
    if mask is None:
        return False
    corners = mask_to_corners(mask)
    if corners is None:
        return False
    o = order_corners(corners)
    tl, tr, br, bl = o
    W = (np.linalg.norm(tr - tl) + np.linalg.norm(br - bl)) / 2.0
    H = (np.linalg.norm(bl - tl) + np.linalg.norm(br - tr)) / 2.0
    if W < 1 or H < 1:
        return False
    geo_landscape = W > H
    angle_sideways = int(angle or 0) % 360 in (90, 270)
    return geo_landscape != angle_sideways


def _dst_size(ordered):
    """由角点边长定 warp 输出尺寸（忠实保留实测比例，不强制 CARD_ASPECT）。
    不强制比例的原因：侧放卡的 mask 是横向(w>h)，必须忠实保留，后续 rotate_to_portrait
    才能靠几何判断出"侧放"。若强制 0.714 会把侧放卡压成方形/竖向，破坏方向判断。"""
    tl, tr, br, bl = ordered
    w_top = np.linalg.norm(tr - tl)
    w_bot = np.linalg.norm(br - bl)
    h_left = np.linalg.norm(bl - tl)
    h_right = np.linalg.norm(br - tr)
    W = max(int(round((w_top + w_bot) / 2.0)), 1)
    H = max(int(round((h_left + h_right) / 2.0)), 1)
    return W, H


def warp_perspective(img, corners):
    """
    用 4 角点把卡牌拉正为竖向矩形（裁切+去透视）。
    Returns: 拉正后的图（与输入同通道顺序）；角点无效返回 None。
    """
    ordered = order_corners(corners)
    size = _dst_size(ordered)
    if size is None:
        return None
    W_out, H_out = size
    dst = np.array([[0, 0], [W_out - 1, 0], [W_out - 1, H_out - 1], [0, H_out - 1]],
                   dtype=np.float32)
    M = cv2.getPerspectiveTransform(ordered, dst)
    return cv2.warpPerspective(img, M, (W_out, H_out))


def rotate_upright(img, angle):
    """按 PaddleOCR angle 把图转到正立（angle=输入相对正立顺时针的旋转度数）。"""
    code = _ANGLE_TO_CV.get(int(angle) % 360)
    if code is None:
        return img
    return cv2.rotate(img, code)


def ensure_portrait(img):
    """保证竖向：宽>高则逆时针转90°（仅当无 angle 信息或兜底时生效）。"""
    h, w = img.shape[:2]
    if w > h:
        return cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
    return img


def rotate_to_portrait(img, angle=0):
    """正立旋转：只在 OCR 明确判 90/270（侧放且给了方向）且几何确为横向时才转。
    - 已竖向(w<=h)：不转。覆盖 angle=180 不可靠情况，及 tx38 这类 angle=90 但实际竖向的卡。
    - 横向(w>h)：仅当 angle∈{90,270} 才按方向转；angle=0/180 一律不转——
      实测 tx86(angle=0 但横向) 盲转会倒置(0.45)，而不转保持横向反而(0.86)。
      即：方向不明时不赌，保持原样交给 DINOv2。"""
    h, w = img.shape[:2]
    if w <= h:
        return img
    a = int(angle or 0) % 360
    if a == 90:
        return cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
    if a == 270:
        return cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
    return img


def crop_bbox(img, box):
    """按 YOLO xyxy 框裁切（兜底路径）。box=(x1,y1,x2,y2)。"""
    if box is None:
        return img
    h, w = img.shape[:2]
    x1, y1, x2, y2 = box
    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:
        return img
    return img[y1:y2, x1:x2]


class CardRectifier:
    """
    卡牌矫正器。rectify() 输入原图 + (mask 或 box) + angle，输出正立竖向卡牌图。
    """

    def __init__(self, use_perspective=True):
        self.use_perspective = use_perspective

    def rectify(self, img, mask=None, box=None, angle=0):
        """
        Args:
            img: ndarray 原图（BGR 或 RGB，本模块不转换颜色，原样返回同顺序）
            mask: 卡牌二值 mask，与 img 同坐标（来自 YOLO masks，已缩放到原图尺寸）
            box:  (x1,y1,x2,y2) YOLO bbox，mask 不可用时兜底
            angle: PaddleOCR doc orientation angle (0/90/180/270)
        Returns:
            矫正后的 ndarray（正立竖向）；任何失败回退到 ensure_portrait(原图)
        """
        out = None
        if self.use_perspective:
            corners = mask_to_corners(mask) if mask is not None else None
            if corners is not None:
                ordered = order_corners(corners)
                # 同时满足：四边形合理 + 确有透视形变，才 warp；否则用 bbox（安全）
                if valid_quad(ordered) and has_perspective(ordered):
                    warped = warp_perspective(img, corners)
                    if warped is not None and warped.size > 0:
                        out = warped
        if out is None and box is not None:
            out = crop_bbox(img, box)
        if out is None:
            out = img
        # 正立：几何(横向)主判 + angle(90/270)选方向；竖向卡(含误报180)不碰
        out = rotate_to_portrait(out, angle)
        return out