mi-task/test/test_detect_track.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import os
import cv2
import numpy as np
import time
import matplotlib.pyplot as plt
import argparse
from matplotlib.patches import Rectangle

# 添加项目根目录到 Python 路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
sys.path.append(project_root)

# 导入待测试模块
from utils.detect_track import detect_dual_track_lines

def display_image_with_results(image, center_info, left_info, right_info, save_path=None):
    """
    显示检测结果并标记轨道线
    """
    # 创建图像副本
    result_img = image.copy()
    height, width = image.shape[:2]
    center_x = width // 2
    
    # 如果检测到轨道线
    if center_info is not None and left_info is not None and right_info is not None:
        # 获取线段坐标
        left_x1, left_y1, left_x2, left_y2 = left_info["line"]
        right_x1, right_y1, right_x2, right_y2 = right_info["line"]
        
        # 计算中心线坐标
        center_line_x1 = (left_x1 + right_x1) // 2
        center_line_y1 = (left_y1 + right_y1) // 2
        center_line_x2 = (left_x2 + right_x2) // 2
        center_line_y2 = (left_y2 + right_y2) // 2
        
        # 绘制左右轨道线
        cv2.line(result_img, (left_x1, left_y1), (left_x2, left_y2), (255, 0, 0), 3)  # 蓝色
        cv2.line(result_img, (right_x1, right_y1), (right_x2, right_y2), (0, 0, 255), 3)  # 红色
        
        # 绘制中心线
        cv2.line(result_img, (center_line_x1, center_line_y1), (center_line_x2, center_line_y2), (0, 255, 0), 3)  # 绿色
        
        # 绘制中线与图像底部的交点
        bottom_x = center_info["point"][0]
        cv2.circle(result_img, (bottom_x, height), 12, (255, 0, 255), -1)  # 紫色
        
        # 绘制图像中心线
        cv2.line(result_img, (center_x, 0), (center_x, height), (0, 0, 255), 1)  # 红色
        
        # 计算并显示偏差
        deviation = center_info["deviation"]
        track_width = center_info["track_width"]
        
        # 添加文本信息
        cv2.putText(result_img, f"偏差: {deviation}px", (20, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
        cv2.putText(result_img, f"轨道宽度: {track_width:.1f}px", (20, 70), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
        
        # 添加方向指示
        if deviation > 0:
            direction = "向右偏"
        elif deviation < 0:
            direction = "向左偏"
        else:
            direction = "居中"
            
        cv2.putText(result_img, f"方向: {direction}", (20, 110), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
    else:
        # 如果未检测到轨道线，显示错误信息
        cv2.putText(result_img, "未检测到双轨道线", (width // 3, height // 2), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 3)
    
    # 如果指定了保存路径，保存图像
    if save_path:
        cv2.imwrite(save_path, result_img)
        print(f"结果已保存至: {save_path}")
    
    return result_img

def process_image(image_path, save_dir=None, show_steps=False):
    """
    处理单张图像，检测双轨道线
    """
    print(f"\n正在处理图像: {image_path}")
    
    # 读取图像
    try:
        image = cv2.imread(image_path)
        if image is None:
            print(f"无法读取图像: {image_path}")
            return None, None, None
    except Exception as e:
        print(f"读取图像出错: {e}")
        return None, None, None
    
    # 获取图像尺寸
    height, width = image.shape[:2]
    
    # 测试双轨道线检测
    print("\n执行双轨道线检测...")
    start_time = time.time()
    center_info, left_info, right_info = detect_dual_track_lines(image, observe=show_steps, save_log=True)
    processing_time = time.time() - start_time
    
    print(f"处理时间: {processing_time:.4f} 秒")
    
    # 输出结果
    if center_info is not None:
        print(f"检测成功 - 中心偏差: {center_info['deviation']}像素")
        print(f"轨道宽度: {center_info['track_width']:.1f}像素")
        print(f"中心线斜率: {center_info['slope']:.4f}")
        print(f"是否垂直: {'是' if center_info['is_vertical'] else '否'}")
        print("-" * 30)
    else:
        print("未检测到双轨道线")
    
    # 如果指定了保存目录，保存结果
    if save_dir and center_info is not None:
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        
        # 生成结果文件名
        base_name = os.path.basename(image_path)
        result_name = f"result_{base_name}"
        result_path = os.path.join(save_dir, result_name)
        
        # 显示并保存结果
        result_img = display_image_with_results(image, center_info, left_info, right_info, result_path)
    
    return center_info, left_info, right_info

def main():
    parser = argparse.ArgumentParser(description='双轨道线检测演示程序')
    parser.add_argument('--input', type=str, default='res/path/image_20250514_024347.png', help='输入图像路径')
    parser.add_argument('--output', type=str, default='res/path/test', 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
    
    # 处理图像
    process_image(args.input, args.output, args.show)
    
    print("\n处理完成!")

if __name__ == "__main__":
    main()