#!/usr/bin/env python3 import os import sys import cv2 import argparse import time from tqdm import tqdm import pandas as pd import matplotlib.pyplot as plt import matplotlib as mpl # 设置中文字体支持 plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei', 'Microsoft YaHei', 'WenQuanYi Micro Hei', 'sans-serif'] plt.rcParams['axes.unicode_minus'] = False # 添加项目根目录到路径 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 batch_test_arrows(data_dir="res/arrows", save_dir="res/arrows/test", show_results=False): """ 批量测试箭头方向检测算法 参数: data_dir: 包含箭头图像的目录 save_dir: 保存结果的目录 show_results: 是否显示结果 返回: results_df: 包含测试结果的DataFrame """ # 确保保存目录存在 os.makedirs(save_dir, exist_ok=True) # 保存结果的列表 results = [] # 处理左右箭头子目录 for direction in ["left", "right"]: dir_path = os.path.join(data_dir, direction) if not os.path.exists(dir_path): print(f"警告: 目录 '{dir_path}' 不存在") continue # 获取该方向的所有图像文件 image_files = [f for f in os.listdir(dir_path) if f.endswith(('.png', '.jpg', '.jpeg'))] print(f"处理 {direction} 方向的 {len(image_files)} 个图像...") # 处理每个图像 for img_file in tqdm(image_files): img_path = os.path.join(dir_path, img_file) # 读取图像 img = cv2.imread(img_path) if img is None: print(f"错误: 无法加载图像 '{img_path}'") continue # 开始计时 start_time = time.time() # 检测箭头方向 detected_direction = detect_arrow_direction(img) # 结束计时 end_time = time.time() processing_time = end_time - start_time # 确定检测是否正确 is_correct = detected_direction == direction # 保存可视化结果 result_filename = f"{direction}_{img_file.split('.')[0]}_result.jpg" result_path = os.path.join(save_dir, result_filename) # visualize_arrow_detection(img, result_path) # 保存结果 results.append({ "图像文件": img_file, "真实方向": direction, "检测方向": detected_direction, "是否正确": is_correct, "处理时间(秒)": processing_time, "结果文件": result_filename }) # 创建结果DataFrame results_df = pd.DataFrame(results) # 保存结果到CSV csv_path = os.path.join(save_dir, "arrow_detection_results.csv") results_df.to_csv(csv_path, index=False, encoding='utf-8-sig') # 生成统计报告 generate_report(results_df, save_dir) return results_df def generate_report(results_df, save_dir): """生成统计报告和可视化""" # 计算总体准确率 accuracy = results_df["是否正确"].mean() * 100 # 按箭头方向分组计算准确率 direction_accuracy = results_df.groupby("真实方向")["是否正确"].mean() * 100 # 计算平均处理时间 avg_time = results_df["处理时间(秒)"].mean() * 1000 # 转换为毫秒 # 创建报告文件 report_path = os.path.join(save_dir, "arrow_detection_report.txt") with open(report_path, "w", encoding="utf-8") as f: f.write("箭头方向检测 - 测试报告\n") f.write("=======================\n\n") f.write(f"测试日期: {time.strftime('%Y-%m-%d %H:%M:%S')}\n") f.write(f"测试图像总数: {len(results_df)}\n\n") f.write(f"总体准确率: {accuracy:.2f}%\n") f.write("各方向准确率:\n") for direction, acc in direction_accuracy.items(): f.write(f" - {direction}: {acc:.2f}%\n") f.write(f"\n平均处理时间: {avg_time:.2f} 毫秒\n\n") # 错误案例分析 if not results_df["是否正确"].all(): f.write("错误检测案例:\n") error_cases = results_df[~results_df["是否正确"]] for _, row in error_cases.iterrows(): f.write(f" - 文件: {row['图像文件']}, 真实方向: {row['真实方向']}, 错误检测为: {row['检测方向']}\n") # 创建可视化图表 plt.figure(figsize=(12, 6)) # 准确率条形图 plt.subplot(1, 2, 1) # 将中文索引转为英文避免字体问题 direction_accuracy_en = direction_accuracy.copy() direction_accuracy_en.index = direction_accuracy.index.map(lambda x: "Left" if x == "left" else "Right") direction_accuracy_en.plot(kind='bar', color=['blue', 'green']) plt.title('各方向检测准确率') plt.ylabel('准确率 (%)') plt.ylim(0, 100) plt.grid(True, linestyle='--', alpha=0.7) # 处理时间箱线图 plt.subplot(1, 2, 2) # 将中文列名改为英文再制图，避免字体问题 temp_df = results_df.copy() temp_df.rename(columns={"处理时间(秒)": "processing_time", "真实方向": "direction"}, inplace=True) temp_df.boxplot(column=['processing_time'], by='direction') plt.title('处理时间分布') plt.ylabel('时间 (秒)') plt.suptitle('') # 保存图表 plt.tight_layout() plt.savefig(os.path.join(save_dir, "arrow_detection_stats.png")) print(f"测试报告已保存到: {report_path}") print(f"统计图表已保存到: {os.path.join(save_dir, 'arrow_detection_stats.png')}") def main(): # 创建参数解析器 parser = argparse.ArgumentParser(description='箭头方向检测批量测试') parser.add_argument('--data-dir', default="res/arrows", help='箭头图像数据目录 (默认: res/arrows)') parser.add_argument('--save-dir', default="res/arrows/test", help='保存结果的目录 (默认: res/arrows/test)') parser.add_argument('--show', action='store_true', help='显示结果图像') args = parser.parse_args() # 运行批量测试 results = batch_test_arrows(args.data_dir, args.save_dir, args.show) # 输出总体结果 correct = results["是否正确"].sum() total = len(results) print(f"\n测试完成! 总共测试了 {total} 张图像，正确检测了 {correct} 张") print(f"总体准确率: {(correct/total*100):.2f}%") # 按真实方向打印准确率 for direction in ["left", "right"]: dir_results = results[results["真实方向"] == direction] if len(dir_results) > 0: dir_correct = dir_results["是否正确"].sum() dir_total = len(dir_results) print(f"{direction} 方向准确率: {(dir_correct/dir_total*100):.2f}% ({dir_correct}/{dir_total})") if __name__ == "__main__": main()