mi-task/test/task-arrow/test_arrow.py

#!/usr/bin/env python3

import os
import sys
import cv2
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.decode_arrow import detect_arrow_direction, visualize_arrow_detection

def main():
    # 创建参数解析器
    parser = argparse.ArgumentParser(description='箭头方向检测测试')
    parser.add_argument('--image', default="res/arrows/left/image-3.png", 
                        help='图像文件路径 (默认: image_20250511_121219.png)')
    parser.add_argument('--save', default="res/arrows/test/arrow_detection_result.jpg",
                        help='保存可视化结果的路径 (默认: res/arrows/test/arrow_detection_result.jpg)')
    parser.add_argument('--show', action='store_true',
                        help='显示可视化结果')
    
    args = parser.parse_args()
    
    # 获取图像路径
    image_path = args.image
    
    # 检查文件是否存在
    if not os.path.exists(image_path):
        print(f"错误: 文件 '{image_path}' 不存在")
        sys.exit(1)
    
    print(f"正在处理图像: {image_path}")
    
    # 加载图像
    img = cv2.imread(image_path)
    if img is None:
        print(f"错误: 无法加载图像 '{image_path}'")
        sys.exit(1)
    
    try:
        # 检测箭头方向
        direction = detect_arrow_direction(img, observe=True)
        print(f"检测到的箭头方向: {direction}")
        
        # 可视化检测过程并保存结果
        visualize_arrow_detection(img, args.save)
        
        print(f"可视化结果已保存到: {args.save}")
        
        # 如果需要显示结果，等待用户按键
        print("按任意键退出...")
        cv2.waitKey(0)
    except Exception as e:
        print(f"发生错误: {e}")
    finally:
        cv2.destroyAllWindows()
    
    return direction

if __name__ == "__main__":
    main()