visualizer.py

"""标定调试结果可视化模块。"""

from __future__ import annotations

import os

import cv2
import matplotlib
import numpy as np

matplotlib.use("Agg")
import matplotlib.pyplot as plt

from config import BOARD_COLS, BOARD_ROWS


def save_corner_images(
    frames_with_corners: list[tuple[np.ndarray, np.ndarray]],
    debug_dir: str = "debug",
):
    """保存绘制角点后的调试图像。"""
    os.makedirs(debug_dir, exist_ok=True)
    pattern_size = (BOARD_COLS, BOARD_ROWS)

    for idx, (frame, corners) in enumerate(frames_with_corners):
        canvas = frame.copy()
        cv2.drawChessboardCorners(canvas, pattern_size, corners, True)
        output_path = os.path.join(debug_dir, f"frame_{idx:04d}_corners.png")
        cv2.imwrite(output_path, canvas)


def save_reprojection_plot(per_frame_rms: list[float], debug_dir: str = "debug"):
    """保存每帧重投影误差柱状图。"""
    os.makedirs(debug_dir, exist_ok=True)
    fig, ax = plt.subplots(figsize=(10, 4))
    indices = np.arange(len(per_frame_rms))
    ax.bar(indices, per_frame_rms)
    ax.set_xlabel("帧序号")
    ax.set_ylabel("RMS")
    ax.set_title("Per-frame Reprojection RMS")
    fig.tight_layout()
    fig.savefig(os.path.join(debug_dir, "reproj_error.png"), dpi=150)
    plt.close(fig)


def save_undistort_compare(frame: np.ndarray, K, D, debug_dir: str = "debug"):
    """保存原图与去畸变图的对比图。"""
    os.makedirs(debug_dir, exist_ok=True)
    undistorted = cv2.fisheye.undistortImage(frame, K, D=D, Knew=K)
    compare = np.hstack([frame, undistorted])
    cv2.imwrite(os.path.join(debug_dir, "undistorted_compare.png"), compare)