183 lines
6.8 KiB
Python
183 lines
6.8 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
test.py
|
||
摄像头畸变校正效果测试 - 仅显示校正前后对比画面
|
||
"""
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import argparse
|
||
import os
|
||
|
||
def is_rtsp_url(input_str):
|
||
"""检查输入是否为RTSP URL"""
|
||
return input_str.startswith('rtsp://')
|
||
|
||
def load_calibration_params(yaml_file):
|
||
"""从YAML文件加载标定参数"""
|
||
if not os.path.exists(yaml_file):
|
||
raise FileNotFoundError(f"标定文件 {yaml_file} 不存在")
|
||
|
||
fs = cv2.FileStorage(yaml_file, cv2.FILE_STORAGE_READ)
|
||
if not fs.isOpened():
|
||
raise IOError(f"无法打开标定文件 {yaml_file}")
|
||
|
||
# 读取参数
|
||
camera_matrix = fs.getNode("camera_matrix").mat()
|
||
dist_coeffs = fs.getNode("dist_coeffs").mat()
|
||
|
||
fs.release()
|
||
|
||
print(f"✓ 成功加载标定参数")
|
||
print(f" 相机矩阵:\n{camera_matrix}")
|
||
print(f" 畸变系数: {dist_coeffs.flatten()}")
|
||
|
||
return camera_matrix, dist_coeffs
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description='摄像头畸变校正效果测试')
|
||
parser.add_argument("-i", "--input", type=str, required=True,
|
||
help="输入源（摄像头设备索引、视频文件路径或RTSP URL）")
|
||
parser.add_argument("-c", "--calibration", type=str, required=True,
|
||
help="标定参数YAML文件路径")
|
||
parser.add_argument("-f", "--fisheye", action="store_true",
|
||
help="是否为鱼眼相机")
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 加载标定参数
|
||
try:
|
||
camera_matrix, dist_coeffs = load_calibration_params(args.calibration)
|
||
except Exception as e:
|
||
print(f"❌ 加载标定参数失败: {e}")
|
||
return
|
||
|
||
# 初始化视频源
|
||
if is_rtsp_url(args.input):
|
||
# RTSP流
|
||
cap = cv2.VideoCapture(args.input, cv2.CAP_FFMPEG)
|
||
try:
|
||
cap.set(cv2.CAP_PROP_FFMPEG_OPTION, "rtsp_transport", "tcp")
|
||
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
|
||
except:
|
||
pass
|
||
elif args.input.isdigit():
|
||
# 摄像头设备索引
|
||
cap = cv2.VideoCapture(int(args.input))
|
||
else:
|
||
# 视频文件
|
||
cap = cv2.VideoCapture(args.input)
|
||
|
||
if not cap.isOpened():
|
||
print(f"❌ 无法打开视频源: {args.input}")
|
||
return
|
||
|
||
# 获取实际分辨率
|
||
actual_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
actual_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
print(f"视频源分辨率: {actual_w}x{actual_h}")
|
||
|
||
print("✓ 准备开始校正")
|
||
print("按 'q' 键退出")
|
||
print("按 's' 键切换校正方法")
|
||
|
||
window_name = "畸变校正对比 - 原始(左) vs 校正后(右)"
|
||
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
|
||
|
||
# 校正方法选择
|
||
use_remap = True # 默认使用remap方法
|
||
|
||
frame_count = 0
|
||
while True:
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
print("读取帧失败")
|
||
# 对于RTSP流，尝试重新连接
|
||
if is_rtsp_url(args.input):
|
||
cap.release()
|
||
cap = cv2.VideoCapture(args.input, cv2.CAP_FFMPEG)
|
||
continue
|
||
break
|
||
|
||
frame_count += 1
|
||
|
||
# 应用畸变校正
|
||
try:
|
||
if args.fisheye:
|
||
if use_remap:
|
||
# 使用映射方法
|
||
new_camera_matrix = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(
|
||
camera_matrix, dist_coeffs, (actual_w, actual_h), np.eye(3), balance=0.8)
|
||
map1, map2 = cv2.fisheye.initUndistortRectifyMap(
|
||
camera_matrix, dist_coeffs, np.eye(3), new_camera_matrix,
|
||
(actual_w, actual_h), cv2.CV_16SC2)
|
||
undistorted = cv2.remap(frame, map1, map2,
|
||
interpolation=cv2.INTER_LINEAR,
|
||
borderMode=cv2.BORDER_CONSTANT)
|
||
else:
|
||
# 使用直接校正方法
|
||
undistorted = cv2.fisheye.undistortImage(frame, camera_matrix, dist_coeffs)
|
||
else:
|
||
# 普通相机校正
|
||
undistorted = cv2.undistort(frame, camera_matrix, dist_coeffs)
|
||
|
||
except Exception as e:
|
||
print(f"校正失败: {e}")
|
||
undistorted = frame.copy() # 失败时使用原图
|
||
|
||
# 检查校正后的图像是否有效
|
||
if undistorted is None or undistorted.size == 0:
|
||
print(f"警告: 第 {frame_count} 帧校正失败")
|
||
undistorted = frame.copy()
|
||
|
||
# 调整大小以便显示
|
||
max_height = 600
|
||
scale_orig = max_height / frame.shape[0]
|
||
orig_resized = cv2.resize(frame,
|
||
(int(frame.shape[1] * scale_orig),
|
||
int(frame.shape[0] * scale_orig)))
|
||
|
||
scale_undist = max_height / undistorted.shape[0]
|
||
undistorted_resized = cv2.resize(undistorted,
|
||
(int(undistorted.shape[1] * scale_undist),
|
||
int(undistorted.shape[0] * scale_undist)))
|
||
|
||
# 确保两个图像高度相同
|
||
min_height = min(orig_resized.shape[0], undistorted_resized.shape[0])
|
||
orig_resized = orig_resized[:min_height, :]
|
||
undistorted_resized = undistorted_resized[:min_height, :]
|
||
|
||
# 水平拼接显示对比
|
||
comparison = np.hstack((orig_resized, undistorted_resized))
|
||
|
||
# 添加文本说明
|
||
cv2.putText(comparison, "原始图像 (有畸变)",
|
||
(10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
||
cv2.putText(comparison, "校正后图像",
|
||
(orig_resized.shape[1] + 10, 30),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
||
|
||
# 显示相机类型和校正方法
|
||
camera_type = "鱼眼相机" if args.fisheye else "普通相机"
|
||
method_type = "映射方法" if use_remap else "直接方法"
|
||
cv2.putText(comparison, f"类型: {camera_type} | 方法: {method_type}",
|
||
(10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
|
||
cv2.putText(comparison, "按 's' 切换方法 | 按 'q' 退出",
|
||
(10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
|
||
|
||
cv2.imshow(window_name, comparison)
|
||
|
||
key = cv2.waitKey(1) & 0xFF
|
||
if key == ord('q'):
|
||
break
|
||
elif key == ord('s'):
|
||
use_remap = not use_remap
|
||
method_name = "映射方法" if use_remap else "直接方法"
|
||
print(f"切换到 {method_name}")
|
||
|
||
cap.release()
|
||
cv2.destroyAllWindows()
|
||
print("测试完成")
|
||
|
||
if __name__ == "__main__":
|
||
main() |