corner_detector.py

"""棋盘角点检测模块。"""

from __future__ import annotations

import cv2
import numpy as np

from config import BLUR_THRESHOLD, BOARD_COLS, BOARD_ROWS, COVERAGE_THRESHOLD, SQUARE_SIZE

# 亚像素精化参数
_SUBPIX_CRITERIA = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
_PATTERN_SIZE = (BOARD_COLS, BOARD_ROWS)

# 世界坐标模板（只需构建一次）
_OBJP_TEMPLATE = np.zeros((1, BOARD_ROWS * BOARD_COLS, 3), np.float32)
_OBJP_TEMPLATE[0, :, :2] = np.mgrid[0:BOARD_COLS, 0:BOARD_ROWS].T.reshape(-1, 2) * SQUARE_SIZE


def compute_sharpness(gray: np.ndarray) -> float:
    """计算灰度图的清晰度分数（Laplacian 方差，越高越清晰）。"""
    return float(cv2.Laplacian(gray, cv2.CV_64F).var())


def detect_frame(
    frame: np.ndarray,
) -> tuple[bool, np.ndarray | None, np.ndarray | None, float]:
    """检测单帧的棋盘角点，同时进行清晰度和覆盖率过滤。

    返回:
        (有效, objp, 角点坐标, 清晰度分数)
        若无效则 objp 和角点坐标为 None，清晰度分数为 0.0
    """
    frame_h, frame_w = frame.shape[:2]
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # 1. 模糊过滤
    sharpness = compute_sharpness(gray)
    if sharpness < BLUR_THRESHOLD:
        return False, None, None, sharpness

    # 2. 角点检测
    found, corners = cv2.findChessboardCorners(gray, _PATTERN_SIZE, None)
    if not found:
        return False, None, None, sharpness

    # 3. 亚像素精化
    cv2.cornerSubPix(gray, corners, winSize=(11, 11), zeroZone=(-1, -1), criteria=_SUBPIX_CRITERIA)

    # 4. 覆盖率过滤
    _, _, box_w, box_h = cv2.boundingRect(corners)
    if float(box_w * box_h) < float(frame_w * frame_h) * COVERAGE_THRESHOLD:
        return False, None, None, sharpness

    return True, _OBJP_TEMPLATE.copy(), corners, sharpness


def detect_corners(frames: list[np.ndarray]) -> tuple[list, list, tuple[int, int]]:
    """【兼容旧接口】批量检测帧，返回 (objpoints, imgpoints, image_size)。"""
    objpoints: list[np.ndarray] = []
    imgpoints: list[np.ndarray] = []
    image_size: tuple[int, int] = (0, 0)

    for frame_index, frame in enumerate(frames):
        image_size = (frame.shape[1], frame.shape[0])
        valid, objp, corners, sharpness = detect_frame(frame)
        if valid:
            objpoints.append(objp)
            imgpoints.append(corners)
            print(f"第 {frame_index} 帧: 通过 (清晰度={sharpness:.1f})")
        else:
            print(f"第 {frame_index} 帧: 跳过 (清晰度={sharpness:.1f})")

    print(f"最终有效帧数: {len(objpoints)}")
    return objpoints, imgpoints, image_size