init
This commit is contained in:
183
test.py
Normal file
183
test.py
Normal file
@@ -0,0 +1,183 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user