231 lines
8.7 KiB
Python
231 lines
8.7 KiB
Python
import cv2
|
|
import os
|
|
import sys
|
|
import time
|
|
import argparse
|
|
|
|
# 添加父目录到系统路径
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
project_root = os.path.dirname(os.path.dirname(current_dir))
|
|
sys.path.append(project_root)
|
|
|
|
from utils.detect_track import detect_left_side_track
|
|
|
|
def process_image(image_path, save_dir=None, show_steps=False):
|
|
"""处理单张图像"""
|
|
print(f"处理图像: {image_path}")
|
|
|
|
# 检测左侧轨迹线
|
|
start_time = time.time()
|
|
track_info, tracking_point = detect_left_side_track(image_path, observe=show_steps, save_log=True)
|
|
processing_time = time.time() - start_time
|
|
|
|
# 输出结果
|
|
if track_info is not None and tracking_point is not None:
|
|
print(f"处理时间: {processing_time:.3f}秒")
|
|
print(f"最佳跟踪点: ({tracking_point[0]}, {tracking_point[1]})")
|
|
print(f"距左边界: {track_info['distance_to_left']:.1f}像素")
|
|
print(f"线段斜率: {track_info['slope']:.4f}")
|
|
print(f"是否垂直: {track_info['is_vertical']}")
|
|
print(f"线段中点: ({track_info['mid_x']:.1f}, {track_info['mid_y']:.1f})")
|
|
print(f"地面交点: ({track_info['ground_intersection'][0]}, {track_info['ground_intersection'][1]})")
|
|
|
|
# 提取线段坐标
|
|
x1, y1, x2, y2 = track_info['line']
|
|
print(f"线段端点: ({x1}, {y1}) - ({x2}, {y2})")
|
|
print("-" * 30)
|
|
|
|
# 如果指定了保存目录，加载原始图像并绘制检测结果
|
|
if save_dir:
|
|
if not os.path.exists(save_dir):
|
|
os.makedirs(save_dir)
|
|
|
|
# 构建输出文件路径
|
|
base_name = os.path.basename(image_path)
|
|
out_path = os.path.join(save_dir, f"result_{base_name}")
|
|
|
|
# 加载原始图像
|
|
img = cv2.imread(image_path)
|
|
if img is not None:
|
|
# 绘制检测结果
|
|
height, width = img.shape[:2]
|
|
center_x = width // 2
|
|
|
|
# 绘制左侧区域范围
|
|
cv2.rectangle(img, (0, 0), (center_x, height), (255, 0, 0), 2)
|
|
|
|
# 绘制检测到的线段
|
|
cv2.line(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
|
|
# 绘制跟踪点
|
|
cv2.circle(img, tracking_point, 8, (0, 0, 255), -1)
|
|
|
|
# 绘制地面交点
|
|
cv2.circle(img, track_info['ground_intersection'], 8, (255, 0, 255), -1)
|
|
|
|
# 添加文本信息
|
|
cv2.putText(img, f"斜率: {track_info['slope']:.2f}", (10, 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
|
|
cv2.putText(img, f"距左边界: {track_info['distance_to_left']:.1f}px", (10, 70),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
|
|
|
|
# 保存结果图像
|
|
cv2.imwrite(out_path, img)
|
|
print(f"结果已保存至: {out_path}")
|
|
else:
|
|
print("未能检测到左侧黄色轨迹线")
|
|
|
|
return track_info, tracking_point
|
|
|
|
def process_video(video_path, save_dir=None, show_steps=False):
|
|
"""处理视频"""
|
|
print(f"处理视频: {video_path}")
|
|
|
|
# 打开视频文件
|
|
cap = cv2.VideoCapture(video_path)
|
|
if not cap.isOpened():
|
|
print("错误：无法打开视频文件")
|
|
return
|
|
|
|
# 获取视频信息
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
|
|
print(f"视频信息: {width}x{height}, {fps}fps, 总帧数: {frame_count}")
|
|
|
|
# 如果需要保存结果视频
|
|
video_writer = None
|
|
if save_dir:
|
|
if not os.path.exists(save_dir):
|
|
os.makedirs(save_dir)
|
|
|
|
# 构建输出文件路径
|
|
base_name = os.path.basename(video_path)
|
|
name, ext = os.path.splitext(base_name)
|
|
out_path = os.path.join(save_dir, f"result_{name}{ext}")
|
|
|
|
# 创建视频写入器
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 可以根据需要修改编码器
|
|
video_writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
|
|
|
|
# 帧计数器
|
|
frame_idx = 0
|
|
detect_success_count = 0
|
|
processing_times = []
|
|
|
|
# 处理每一帧
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
frame_idx += 1
|
|
print(f"\r处理帧 {frame_idx}/{frame_count}", end="")
|
|
|
|
# 检测左侧轨迹线
|
|
start_time = time.time()
|
|
track_info, tracking_point = detect_left_side_track(frame, observe=False, save_log=False)
|
|
processing_time = time.time() - start_time
|
|
processing_times.append(processing_time)
|
|
|
|
# 如果检测成功
|
|
if track_info is not None and tracking_point is not None:
|
|
detect_success_count += 1
|
|
|
|
# 提取线段坐标
|
|
x1, y1, x2, y2 = track_info['line']
|
|
|
|
# 在帧上绘制检测结果
|
|
center_x = width // 2
|
|
|
|
# 绘制左侧区域范围
|
|
cv2.rectangle(frame, (0, 0), (center_x, height), (255, 0, 0), 2)
|
|
|
|
# 绘制检测到的线段
|
|
cv2.line(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
|
|
# 绘制跟踪点
|
|
cv2.circle(frame, tracking_point, 8, (0, 0, 255), -1)
|
|
|
|
# 绘制地面交点
|
|
cv2.circle(frame, track_info['ground_intersection'], 8, (255, 0, 255), -1)
|
|
|
|
# 添加文本信息
|
|
cv2.putText(frame, f"斜率: {track_info['slope']:.2f}", (10, 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
|
|
cv2.putText(frame, f"距左边界: {track_info['distance_to_left']:.1f}px", (10, 70),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
|
|
cv2.putText(frame, f"帧: {frame_idx}/{frame_count}", (10, height-30),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
|
|
else:
|
|
# 如果检测失败，显示错误消息
|
|
cv2.putText(frame, "未检测到轨迹线", (10, 50),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
|
|
|
|
# 如果需要显示
|
|
if show_steps:
|
|
cv2.imshow('Left Track Detection', frame)
|
|
key = cv2.waitKey(1)
|
|
if key == 27: # ESC键退出
|
|
break
|
|
|
|
# 如果需要保存
|
|
if video_writer is not None:
|
|
video_writer.write(frame)
|
|
|
|
# 清理资源
|
|
cap.release()
|
|
if video_writer is not None:
|
|
video_writer.release()
|
|
cv2.destroyAllWindows()
|
|
|
|
# 打印统计信息
|
|
print(f"\n视频处理完成")
|
|
print(f"总帧数: {frame_count}")
|
|
print(f"成功检测帧数: {detect_success_count}")
|
|
print(f"检测成功率: {detect_success_count/frame_count*100:.2f}%")
|
|
if processing_times:
|
|
avg_time = sum(processing_times) / len(processing_times)
|
|
print(f"平均处理时间: {avg_time*1000:.2f}ms")
|
|
print(f"处理帧率: {1/avg_time:.2f}fps")
|
|
|
|
if save_dir:
|
|
print(f"结果已保存至: {out_path}")
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='左侧黄色轨迹线检测演示程序')
|
|
parser.add_argument('--input', type=str, default='res/path/image_20250513_162556.png', help='输入图像或视频的路径')
|
|
parser.add_argument('--output', type=str, default='res/path/test/left_track_results/', help='输出结果的保存目录')
|
|
parser.add_argument('--type', type=str, choices=['image', 'video'], help='输入类型，不指定会自动检测')
|
|
parser.add_argument('--show', default=True, help='显示处理步骤')
|
|
|
|
args = parser.parse_args()
|
|
|
|
# 检查输入路径
|
|
if not os.path.exists(args.input):
|
|
print(f"错误：文件 '{args.input}' 不存在")
|
|
return
|
|
|
|
# 如果未指定类型，根据文件扩展名判断
|
|
if args.type is None:
|
|
ext = os.path.splitext(args.input)[1].lower()
|
|
if ext in ['.jpg', '.jpeg', '.png', '.bmp']:
|
|
args.type = 'image'
|
|
elif ext in ['.mp4', '.avi', '.mov']:
|
|
args.type = 'video'
|
|
else:
|
|
print(f"错误：无法确定文件类型 '{ext}'")
|
|
return
|
|
|
|
# 根据类型处理
|
|
if args.type == 'image':
|
|
process_image(args.input, args.output, args.show)
|
|
elif args.type == 'video':
|
|
process_video(args.input, args.output, args.show)
|
|
else:
|
|
print("错误：不支持的类型")
|
|
|
|
if __name__ == "__main__":
|
|
main() |