import cv2
import numpy as np
import os
import sys
from pathlib import Path
from datetime import datetime
class EnhancedLaneDetector:
def __init__(self):
# Canny边缘检测参数
self.canny_low_threshold = 50
self.canny_high_threshold = 150
# 霍夫变换参数
self.rho = 2
self.theta = np.pi / 180
self.hough_threshold = 100
self.min_line_length = 40
self.max_line_gap = 5
# 车道线颜色定义
self.left_lane_color = (0, 0, 255) # 红色 - 左车道线
self.right_lane_color = (255, 0, 0) # 蓝色 - 右车道线
self.lane_area_color = (0, 255, 0) # 绿色 - 车道区域
self.center_line_color = (0, 255, 255) # 黄色 - 中心线
# 跟踪状态
self.previous_left_lane = None
self.previous_right_lane = None
def canny_edge_detection(self, image):
"""Canny边缘检测"""
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
blur = cv2.GaussianBlur(gray, (5, 5), 0)
edges = cv2.Canny(blur, self.canny_low_threshold, self.canny_high_threshold)
return edges
def region_of_interest(self, image):
"""定义感兴趣区域(ROI)- 自适应大小"""
height, width = image.shape
# 自适应ROI,根据图像大小调整
bottom_width = 0.85
top_width = 0.45
height_ratio = 0.6
bottom_left = (width * (1 - bottom_width) // 2, height)
bottom_right = (width * (1 + bottom_width) // 2, height)
top_left = (width * (1 - top_width) // 2, height * height_ratio)
top_right = (width * (1 + top_width) // 2, height * height_ratio)
polygons = np.array([[bottom_left, top_left, top_right, bottom_right]])
mask = np.zeros_like(image)
cv2.fillPoly(mask, polygons.astype(np.int32), 255)
masked_image = cv2.bitwise_and(image, mask)
return masked_image
def hough_lines_detection(self, image):
"""霍夫变换检测直线"""
lines = cv2.HoughLinesP(image, self.rho, self.theta, self.hough_threshold,
np.array([]), minLineLength=self.min_line_length,
maxLineGap=self.max_line_gap)
return lines
def separate_lines(self, lines, image):
"""将检测到的线分为左右车道线"""
left_lines = []
right_lines = []
if lines is not None:
for line in lines:
x1, y1, x2, y2 = line[0]
# 计算斜率
if x2 - x1 == 0:
continue
slope = (y2 - y1) / (x2 - x1)
# 过滤掉水平线和异常斜率的线
if abs(slope) < 0.3 or abs(slope) > 2.0:
continue
# 计算线的长度
length = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
if slope < 0: # 左车道线
left_lines.append({'points': line[0], 'slope': slope, 'length': length})
else: # 右车道线
right_lines.append({'points': line[0], 'slope': slope, 'length': length})
return left_lines, right_lines
def average_lines(self, lines, image, side):
"""对检测到的多条线进行加权平均处理"""
if not lines:
# 如果没有检测到线,使用上一帧的结果(如果存在)
if side == 'left' and self.previous_left_lane:
return self.previous_left_lane
elif side == 'right' and self.previous_right_lane:
return self.previous_right_lane
return None
x_coords = []
y_coords = []
weights = []
for line_info in lines:
line = line_info['points']
x1, y1, x2, y2 = line
x_coords.extend([x1, x2])
y_coords.extend([y1, y2])
# 使用线的长度作为权重
weights.extend([line_info['length'], line_info['length']])
if len(x_coords) > 1:
# 加权线性回归
coefficients = np.polyfit(x_coords, y_coords, 1, w=weights)
slope, intercept = coefficients
# 计算线在图像底部和指定高度的x坐标
y1 = image.shape[0] # 图像底部
y2 = int(y1 * 0.65) # 图像高度的65%处
x1 = int((y1 - intercept) / slope)
x2 = int((y2 - intercept) / slope)
result = [x1, y1, x2, y2, slope]
# 保存当前结果供下一帧使用
if side == 'left':
self.previous_left_lane = result
else:
self.previous_right_lane = result
return result
return None
def calculate_center_line(self, left_line, right_line, image):
"""计算中心线"""
if left_line is None or right_line is None:
return None
# 提取左右车道线的端点
left_x1, left_y1, left_x2, left_y2, left_slope = left_line
right_x1, right_y1, right_x2, right_y2, right_slope = right_line
# 计算中心线的端点
center_x1 = (left_x1 + right_x1) // 2
center_y1 = left_y1 # 左右y1相同
center_x2 = (left_x2 + right_x2) // 2
center_y2 = left_y2 # 左右y2相同
return [center_x1, center_y1, center_x2, center_y2]
def draw_lanes(self, image, left_line, right_line, center_line):
"""在图像上绘制车道线和车道区域"""
result = image.copy()
height, width = image.shape[:2]
# 绘制车道区域
if left_line is not None and right_line is not None:
left_x1, left_y1, left_x2, left_y2, _ = left_line
right_x1, right_y1, right_x2, right_y2, _ = right_line
# 创建车道区域的顶点
pts = np.array([
[left_x1, left_y1], # 左线底部
[left_x2, left_y2], # 左线顶部
[right_x2, right_y2], # 右线顶部
[right_x1, right_y1] # 右线底部
], np.int32)
# 创建半透明覆盖层
overlay = result.copy()
cv2.fillPoly(overlay, [pts], self.lane_area_color)
# 合并原图和覆盖层
cv2.addWeighted(overlay, 0.3, result, 0.7, 0, result)
# 绘制左车道线(红色)
if left_line is not None:
x1, y1, x2, y2, _ = left_line
cv2.line(result, (x1, y1), (x2, y2), self.left_lane_color, 8)
cv2.putText(result, "Left Lane", (x1 - 100, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, self.left_lane_color, 2)
# 绘制右车道线(蓝色)
if right_line is not None:
x1, y1, x2, y2, _ = right_line
cv2.line(result, (x1, y1), (x2, y2), self.right_lane_color, 8)
cv2.putText(result, "Right Lane", (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, self.right_lane_color, 2)
# 绘制中心线(黄色)
if center_line is not None:
x1, y1, x2, y2 = center_line
cv2.line(result, (x1, y1), (x2, y2), self.center_line_color, 4)
cv2.putText(result, "Center", (x2 - 50, y2 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, self.center_line_color, 2)
# 在中心线底部添加方向指示器
cv2.circle(result, (x1, y1), 10, self.center_line_color, -1)
return result
def add_info_panel(self, image, left_detected, right_detected, frame_count):
"""添加信息面板"""
height, width = image.shape[:2]
# 创建信息面板背景
info_panel = np.zeros((150, width, 3), dtype=np.uint8)
# 添加系统标题
cv2.putText(info_panel, "LANE DETECTION SYSTEM", (20, 30),
cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
# 添加车道状态
left_status = "DETECTED" if left_detected else "NOT DETECTED"
right_status = "DETECTED" if right_detected else "NOT DETECTED"
left_color = (0, 255, 0) if left_detected else (0, 0, 255)
right_color = (0, 255, 0) if right_detected else (0, 0, 255)
cv2.putText(info_panel, f"Left Lane: {left_status}", (20, 70),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, left_color, 2)
cv2.putText(info_panel, f"Right Lane: {right_status}", (20, 100),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, right_color, 2)
# 添加帧计数
cv2.putText(info_panel, f"Frame: {frame_count}", (width - 200, 70),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
# 添加时间戳
timestamp = datetime.now().strftime("%H:%M:%S")
cv2.putText(info_panel, f"Time: {timestamp}", (width - 200, 100),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
# 合并信息面板和主图像
result = np.vstack([info_panel, image])
return result
def add_legend(self, image):
"""添加图例"""
legend_height = 80
legend = np.zeros((legend_height, image.shape[1], 3), dtype=np.uint8)
# 图例项目
items = [
("Left Lane", self.left_lane_color),
("Right Lane", self.right_lane_color),
("Lane Area", self.lane_area_color),
("Center Line", self.center_line_color)
]
for i, (text, color) in enumerate(items):
x_pos = 20 + i * 200
# 颜色方块
cv2.rectangle(legend, (x_pos, 20), (x_pos + 30, 50), color, -1)
# 文本
cv2.putText(legend, text, (x_pos + 40, 40),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)
# 添加操作说明
cv2.putText(legend, "Press 'Q' to Quit | 'S' to Save | 'R' to Reset",
(image.shape[1] - 400, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
# 合并图例和图像
result = np.vstack([image, legend])
return result
def process_frame(self, frame, frame_count):
"""处理单帧图像"""
# 边缘检测
edges = self.canny_edge_detection(frame)
# 区域掩码
roi_edges = self.region_of_interest(edges)
# 霍夫变换检测直线
lines = self.hough_lines_detection(roi_edges)
# 分离左右车道线
left_lines, right_lines = self.separate_lines(lines, frame)
# 平均化处理
left_lane = self.average_lines(left_lines, frame, 'left')
right_lane = self.average_lines(right_lines, frame, 'right')
# 计算中心线
center_line = self.calculate_center_line(left_lane, right_lane, frame)
# 绘制车道
result = self.draw_lanes(frame, left_lane, right_lane, center_line)
# 添加信息面板
result = self.add_info_panel(result, left_lane is not None, right_lane is not None, frame_count)
# 添加图例
result = self.add_legend(result)
return result, left_lane is not None, right_lane is not None
def main():
print("=" * 50)
print(" Enhanced Lane Detection System")
print("=" * 50)
print("Color Coding:")
print(" 🔴 Red: Left Lane")
print(" 🔵 Blue: Right Lane")
print(" 🟢 Green: Lane Area")
print(" 🟡 Yellow: Center Line")
print("\nControls:")
print(" Press 'Q' to quit")
print(" Press 'S' to save current frame")
print(" Press 'R' to reset lane tracking")
print("=" * 50)
detector = EnhancedLaneDetector()
# 尝试打开视频文件或摄像头
cap = None
video_files = ['zyh.mp4', 'road.mp4', 'highway.mp4', 'drive.mp4']
for video_file in video_files:
if os.path.exists(video_file):
cap = cv2.VideoCapture(video_file)
print(f"📹 Using video file: {video_file}")
break
if cap is None or not cap.isOpened():
cap = cv2.VideoCapture(0)
if cap.isOpened():
print("📹 Using default camera")
else:
print("❌ Error: Cannot open camera or find video file!")
print("Please ensure:")
print("1. Camera is connected")
print("2. Or place a test video file in the same directory")
input("Press Enter to exit...")
return
# 创建输出目录
output_dir = Path("output")
output_dir.mkdir(exist_ok=True)
frame_count = 0
saved_count = 0
try:
while True:
ret, frame = cap.read()
if not ret:
print("🎬 Video ended or cannot read frame")
break
frame_count += 1
# 处理帧
processed_frame, left_detected, right_detected = detector.process_frame(frame, frame_count)
# 显示处理结果
cv2.imshow('Enhanced Lane Detection System', processed_frame)
# 键盘输入处理
key = cv2.waitKey(1) & 0xFF
if key == ord('q') or key == ord('Q'):
break
elif key == ord('s') or key == ord('S'):
# 保存当前帧
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = output_dir / f"lane_detection_{timestamp}_{saved_count:04d}.jpg"
cv2.imwrite(str(filename), processed_frame)
print(f"💾 Frame saved: {filename}")
saved_count += 1
elif key == ord('r') or key == ord('R'):
# 重置车道跟踪
detector.previous_left_lane = None
detector.previous_right_lane = None
print("🔄 Lane tracking reset")
except KeyboardInterrupt:
print("\n⏹️ Program interrupted by user")
finally:
# 释放资源
cap.release()
cv2.destroyAllWindows()
print(f"✅ Program exited. Processed {frame_count} frames.")
print(f"💾 Saved {saved_count} frames to 'output' directory.")
if __name__ == "__main__":
main() 帮我改进一下
最新发布