79 lines
3.0 KiB
Python
79 lines
3.0 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_horizontal_track_edge
|
|
|
|
def process_image(image_path, save_dir=None, show_steps=False):
|
|
"""处理单张图像"""
|
|
print(f"处理图像: {image_path}")
|
|
|
|
# 检测赛道并估算距离
|
|
start_time = time.time()
|
|
edge_point, edge_info = detect_horizontal_track_edge(image_path, observe=show_steps, save_log=True)
|
|
processing_time = time.time() - start_time
|
|
|
|
# 输出结果
|
|
if edge_point is not None and edge_info is not None:
|
|
print(f"处理时间: {processing_time:.3f}秒")
|
|
print(f"边缘点: ({edge_point[0]}, {edge_point[1]})")
|
|
print(f"到中线距离: {edge_info['distance_to_center']}像素")
|
|
print(f"边缘斜率: {edge_info['slope']:.4f}")
|
|
print(f"是否水平: {edge_info['is_horizontal']}")
|
|
print(f"点数量: {edge_info['points_count']}")
|
|
print(f"中线交点: ({edge_info['intersection_point'][0]}, {edge_info['intersection_point'][1]})")
|
|
print(f"交点到底部距离: {edge_info['distance_to_bottom']:.1f}像素")
|
|
print(f"注意: 中线交点是垂直中线与边缘横线的交点")
|
|
print("-" * 30)
|
|
|
|
# 如果指定了保存目录，保存结果
|
|
if save_dir:
|
|
if not os.path.exists(save_dir):
|
|
os.makedirs(save_dir)
|
|
else:
|
|
print("未能检测到黄色赛道")
|
|
|
|
return edge_point, edge_info
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='黄色赛道检测演示程序')
|
|
parser.add_argument('--input', type=str, default='res/path/test-2.jpg', help='输入图像或视频的路径')
|
|
parser.add_argument('--output', type=str, default='res/path/test/result_image_20250514_024313.png', help='输出结果的保存路径')
|
|
parser.add_argument('--type', type=str, choices=['image', 'video'], help='输入类型，不指定会自动检测')
|
|
parser.add_argument('--show', default=True, action='store_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':
|
|
# 获取输出目录
|
|
output_dir = os.path.dirname(args.output)
|
|
process_image(args.input, output_dir, args.show)
|
|
else:
|
|
print("错误：不支持的视频类型")
|
|
|
|
if __name__ == "__main__":
|
|
main() |