import cv2
import numpy as np
import os
import time
class StereoVision:
def __init__(self, use_wls=True):
self.use_wls = use_wls
# SGBM参数
self.min_disparity = 0
self.num_disparities = 64
self.block_size = 7
self.uniqueness_ratio = 15
self.speckle_window_size = 100
self.speckle_range = 32
self.disp12_max_diff = 1
self.p1 = 8 * 3 * self.block_size ** 2
self.p2 = 32 * 3 * self.block_size ** 2
# WLS滤波参数
self.wls_lambda = 8000.0
self.wls_sigma = 1.5
# 初始化匹配器
self.left_matcher = cv2.StereoSGBM_create(
minDisparity=self.min_disparity,
numDisparities=self.num_disparities,
blockSize=self.block_size,
uniquenessRatio=self.uniqueness_ratio,
speckleWindowSize=self.speckle_window_size,
speckleRange=self.speckle_range,
disp12MaxDiff=self.disp12_max_diff,
P1=self.p1,
P2=self.p2
)
if self.use_wls:
self.right_matcher = cv2.ximgproc.createRightMatcher(self.left_matcher)
self.wls_filter = cv2.ximgproc.createDisparityWLSFilter(self.left_matcher)
self.wls_filter.setLambda(self.wls_lambda)
self.wls_filter.setSigmaColor(self.wls_sigma)
# 相机参数
self.focal_length = 700.0
self.baseline = 0.12
self.cx = 320.0
self.cy = 240.0
def compute_disparity(self, left_img, right_img):
"""计算视差图(带WLS滤波)"""
# 转换为灰度图
gray_left = cv2.cvtColor(left_img, cv2.COLOR_BGR2GRAY)
gray_right = cv2.cvtColor(right_img, cv2.COLOR_BGR2GRAY)
if self.use_wls:
# 计算左右视差图
left_disp = self.left_matcher.compute(gray_left, gray_right).astype(np.float32) / 16.0
right_disp = self.right_matcher.compute(gray_right, gray_left).astype(np.float32) / 16.0
# WLS滤波
filtered_disp = self.wls_filter.filter(left_disp, gray_left, None, right_disp)
# 后处理
filtered_disp = self.post_process_disparity(filtered_disp)
else:
# 普通SGBM
filtered_disp = self.left_matcher.compute(gray_left, gray_right).astype(np.float32) / 16.0
filtered_disp = self.post_process_disparity(filtered_disp)
return filtered_disp
def post_process_disparity(self, disparity):
"""视差图后处理"""
# 过滤无效值
disparity[disparity <= 0] = 0.1
# 中值滤波去噪
disparity = cv2.medianBlur(disparity, 5)
# 形态学操作填充空洞
kernel = np.ones((3, 3), np.uint8)
disparity = cv2.morphologyEx(disparity, cv2.MORPH_CLOSE, kernel)
return disparity
def compute_depth_map(self, disparity):
"""从视差图计算深度图"""
# 避免除以零
disparity[disparity == 0] = 0.1
# 计算深度
depth = (self.focal_length * self.baseline) / disparity
# 深度图后处理
depth = self.post_process_depth(depth)
return depth
def post_process_depth(self, depth_map):
"""深度图后处理"""
# 限制最大深度
max_depth = 10.0 # 10米
depth_map[depth_map > max_depth] = max_depth
# 高斯平滑
depth_map = cv2.GaussianBlur(depth_map, (5, 5), 0)
return depth_map
def get_colormap(self, data):
"""生成彩色图"""
normalized = cv2.normalize(data, None, 0, 255, cv2.NORM_MINMAX)
normalized = np.uint8(normalized)
colormap = cv2.applyColorMap(normalized, cv2.COLORMAP_JET)
return colormap, normalized
def get_3d_point(self, depth_map, x, y):
"""获取指定像素点的3D坐标"""
if 0 <= y < depth_map.shape[0] and 0 <= x < depth_map.shape[1]:
depth = depth_map[y, x]
if depth > 0:
Z = depth
X = (x - self.cx) * Z / self.focal_length
Y = (y - self.cy) * Z / self.focal_length
return (X, Y, Z)
return None
class StereoDemo:
def __init__(self):
self.stereo_vision = StereoVision(use_wls=True)
self.cap = None
self.camera_width = 1280
self.camera_height = 480
self.use_stitched_mode = True
self.current_depth_map = None
def initialize_camera(self):
"""初始化相机"""
print("正在初始化相机...")
# 尝试不同的后端
backends = [
cv2.CAP_DSHOW, # DirectShow (Windows)
cv2.CAP_ANY, # 自动选择
]
for backend in backends:
self.cap = cv2.VideoCapture(0, backend)
if self.cap.isOpened():
# 设置相机参数
self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.camera_width)
self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.camera_height)
self.cap.set(cv2.CAP_PROP_FPS, 30)
# 验证设置是否生效
actual_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
actual_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
print(f"相机初始化成功! 分辨率: {actual_width}x{actual_height}")
return True
print("无法打开相机,将使用虚拟摄像头模式")
return False
def capture_frame(self):
"""捕获一帧并分割为左右图像"""
if self.cap and self.cap.isOpened():
ret, frame = self.cap.read()
if ret and frame is not None:
# 分割左右图像
height, width = frame.shape[:2]
if self.use_stitched_mode and width > height:
left_img = frame[:, :width // 2]
right_img = frame[:, width // 2:]
return left_img, right_img
else:
return frame, frame.copy()
# 如果相机不可用,创建虚拟图像
return self.create_virtual_camera_frames(int(time.time() * 10))
def create_virtual_camera_frames(self, frame_count):
"""创建虚拟摄像头帧"""
height, width = 480, 640
# 左图
imgL = np.zeros((height, width, 3), dtype=np.uint8)
cv2.rectangle(imgL, (100, 100), (200, 200), (255, 0, 0), -1)
cv2.circle(imgL, (400, 150), 50, (0, 255, 0), -1)
cv2.rectangle(imgL, (300, 300), (500, 400), (0, 0, 255), -1)
# 右图(模拟视差)
imgR = np.zeros((height, width, 3), dtype=np.uint8)
cv2.rectangle(imgR, (90, 100), (190, 200), (255, 0, 0), -1)
cv2.circle(imgR, (385, 150), 50, (0, 255, 0), -1)
cv2.rectangle(imgR, (285, 300), (485, 400), (0, 0, 255), -1)
# 添加纹理
for i in range(0, height, 20):
cv2.line(imgL, (0, i), (width, i), (50, 50, 50), 1)
cv2.line(imgR, (0, i), (width, i), (50, 50, 50), 1)
# 添加动态效果
offset_x = int(10 * np.sin(frame_count * 0.05))
offset_y = int(5 * np.sin(frame_count * 0.03))
M = np.float32([[1, 0, offset_x], [0, 1, offset_y]])
imgL_shifted = cv2.warpAffine(imgL, M, (imgL.shape[1], imgL.shape[0]))
M_right = np.float32([[1, 0, offset_x - 15], [0, 1, offset_y]])
imgR_shifted = cv2.warpAffine(imgR, M_right, (imgR.shape[1], imgR.shape[0]))
# 添加噪声
noise = np.random.normal(0, 5, imgL_shifted.shape).astype(np.uint8)
imgL_shifted = cv2.add(imgL_shifted, noise)
imgR_shifted = cv2.add(imgR_shifted, noise)
return imgL_shifted, imgR_shifted
def mouse_callback(self, event, x, y, flags, param):
"""鼠标回调函数"""
if event == cv2.EVENT_LBUTTONDOWN and self.current_depth_map is not None:
depth_value = self.current_depth_map[y, x] if y < self.current_depth_map.shape[0] and x < \
self.current_depth_map.shape[1] else 0
point_3d = self.stereo_vision.get_3d_point(self.current_depth_map, x, y)
if point_3d is not None:
print(
f"点击位置 ({x}, {y}) - 深度: {depth_value:.2f}米, 3D坐标: X={point_3d[0]:.2f}, Y={point_3d[1]:.2f}, Z={point_3d[2]:.2f}")
else:
print(f"点击位置 ({x}, {y}) - 无法计算深度信息")
def camera_mode(self):
"""实时摄像头测距模式"""
print("=== 双目测距演示 - 实时摄像头模式 ===")
# 初始化相机
camera_available = self.initialize_camera()
# 创建窗口
window_names = ['Left Camera', 'Right Camera', 'Disparity Map', 'Depth Map', 'Controls']
for name in window_names:
cv2.namedWindow(name, cv2.WINDOW_NORMAL)
cv2.setMouseCallback('Depth Map', self.mouse_callback)
# 创建控制窗口
control_img = np.zeros((350, 400, 3), dtype=np.uint8)
cv2.putText(control_img, "实时双目测距系统", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
cv2.putText(control_img, "Controls:", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)
cv2.putText(control_img, "q - 退出", (10, 85), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
cv2.putText(control_img, "s - 保存当前帧", (10, 105), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
cv2.putText(control_img, "p - 暂停/继续", (10, 125), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
cv2.putText(control_img, "c - 调整参数", (10, 145), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
cv2.putText(control_img, "w - 切换WLS滤波", (10, 165), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
cv2.putText(control_img, "点击深度图获取3D坐标", (10, 195), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
if camera_available:
cv2.putText(control_img, "模式: 真实摄像头", (10, 225), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
cv2.putText(control_img, f"分辨率: {self.camera_width}x{self.camera_height}", (10, 245),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
else:
cv2.putText(control_img, "模式: 虚拟摄像头", (10, 225), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
cv2.putText(control_img, f"WLS滤波: {'开启' if self.stereo_vision.use_wls else '关闭'}",
(10, 265), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
(0, 255, 0) if self.stereo_vision.use_wls else (0, 0, 255), 1)
frame_count = 0
paused = False
print("实时测距已启动!")
print("按 'q' 退出,'s' 保存,'p' 暂停,'c' 调整参数,'w' 切换WLS滤波")
try:
while True:
if not paused:
# 捕获帧
frame_left, frame_right = self.capture_frame()
# 计算视差图和深度图
start_time = time.time()
disparity = self.stereo_vision.compute_disparity(frame_left, frame_right)
depth_map = self.stereo_vision.compute_depth_map(disparity)
depth_colormap, _ = self.stereo_vision.get_colormap(depth_map)
processing_time = (time.time() - start_time) * 1000
# 保存当前深度图用于鼠标回调
self.current_depth_map = depth_map
# 显示结果
cv2.imshow('Left Camera', frame_left)
cv2.imshow('Right Camera', frame_right)
# 显示视差图
disparity_display, _ = self.stereo_vision.get_colormap(disparity)
cv2.imshow('Disparity Map', disparity_display)
# 显示深度图
depth_display = depth_colormap.copy()
h, w = depth_display.shape[:2]
# 添加信息
cv2.putText(depth_display, "Red: Close", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255),
2)
cv2.putText(depth_display, "Blue: Far", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
cv2.putText(depth_display, f"FPS: {1000 / processing_time:.1f}", (10, 90), cv2.FONT_HERSHEY_SIMPLEX,
0.6, (255, 255, 255), 2)
cv2.putText(depth_display, f"WLS: {'ON' if self.stereo_vision.use_wls else 'OFF'}",
(10, 120), cv2.FONT_HERSHEY_SIMPLEX, 0.6,
(0, 255, 0) if self.stereo_vision.use_wls else (0, 0, 255), 2)
cv2.putText(depth_display, "Click to get 3D coordinates", (10, h - 20),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
cv2.imshow('Depth Map', depth_display)
# 显示控制窗口
cv2.imshow('Controls', control_img)
# 处理键盘输入
key = cv2.waitKey(30) & 0xFF
if key == ord('q'):
break
elif key == ord('s'):
timestamp = int(time.time())
cv2.imwrite(f'capture_{timestamp}_left.jpg', frame_left)
cv2.imwrite(f'capture_{timestamp}_right.jpg', frame_right)
cv2.imwrite(f'capture_{timestamp}_disparity.jpg', disparity_display)
cv2.imwrite(f'capture_{timestamp}_depth.jpg', depth_colormap)
print(f"已保存第 {frame_count} 帧")
elif key == ord('p'):
paused = not paused
print("暂停" if paused else "继续")
elif key == ord('c'):
self.adjust_parameters()
elif key == ord('w'):
self.stereo_vision.use_wls = not self.stereo_vision.use_wls
status = "开启" if self.stereo_vision.use_wls else "关闭"
print(f"WLS滤波: {status}")
frame_count += 1
except KeyboardInterrupt:
print("\n用户中断")
except Exception as e:
print(f"运行时错误: {e}")
finally:
if self.cap:
self.cap.release()
for name in window_names:
cv2.destroyWindow(name)
print("系统已关闭")
def adjust_parameters(self):
"""调整参数"""
print("\n当前参数:")
print(f"num_disparities: {self.stereo_vision.num_disparities}")
print(f"block_size: {self.stereo_vision.block_size}")
print(f"WLS lambda: {self.stereo_vision.wls_lambda}")
print(f"WLS sigma: {self.stereo_vision.wls_sigma}")
try:
new_num_disp = int(input("输入新的num_disparities值(16的倍数): ") or self.stereo_vision.num_disparities)
new_block_size = int(input("输入新的block_size值(奇数): ") or self.stereo_vision.block_size)
if new_num_disp % 16 == 0 and new_block_size % 2 == 1:
self.stereo_vision.num_disparities = new_num_disp
self.stereo_vision.block_size = new_block_size
self.stereo_vision.p1 = 8 * 3 * self.stereo_vision.block_size ** 2
self.stereo_vision.p2 = 32 * 3 * self.stereo_vision.block_size ** 2
# 重新初始化匹配器
self.stereo_vision.left_matcher = cv2.StereoSGBM_create(
minDisparity=self.stereo_vision.min_disparity,
numDisparities=self.stereo_vision.num_disparities,
blockSize=self.stereo_vision.block_size,
uniquenessRatio=self.stereo_vision.uniqueness_ratio,
speckleWindowSize=self.stereo_vision.speckle_window_size,
speckleRange=self.stereo_vision.speckle_range,
disp12MaxDiff=self.stereo_vision.disp12_max_diff,
P1=self.stereo_vision.p1,
P2=self.stereo_vision.p2
)
if self.stereo_vision.use_wls:
self.stereo_vision.right_matcher = cv2.ximgproc.createRightMatcher(self.stereo_vision.left_matcher)
self.stereo_vision.wls_filter = cv2.ximgproc.createDisparityWLSFilter(
self.stereo_vision.left_matcher)
self.stereo_vision.wls_filter.setLambda(self.stereo_vision.wls_lambda)
self.stereo_vision.wls_filter.setSigmaColor(self.stereo_vision.wls_sigma)
print("参数已更新")
else:
print("参数格式错误,保持原值")
except ValueError:
print("输入无效,保持原参数")
def run(self):
"""运行演示程序"""
print("双目立体视觉测距系统")
print("=" * 50)
print("本系统实现了双目测距功能:")
print("1. 立体图像采集")
print("2. 视差计算(SGBM + WLS滤波)")
print("3. 深度图生成")
print("4. 3D坐标计算")
print("5. 后处理优化")
print("=" * 50)
while True:
print("\n请选择运行模式:")
print("1. 相机模式 - 实时测距")
print("2. 退出")
choice = input("请输入选择 (1-2): ").strip()
if choice == '1':
self.camera_mode()
elif choice == '2':
print("感谢使用双目测距系统!")
break
else:
print("无效选择,请重新输入")
if __name__ == "__main__":
demo = StereoDemo()
demo.run()修改这个代码
最新发布