57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
#!/usr/bin/env python3
|
||
|
||
import os
|
||
import sys
|
||
import cv2
|
||
import numpy as np
|
||
import argparse
|
||
|
||
# 添加父目录到路径，以便能够导入utils
|
||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
from utils.decode_arrow import detect_arrow_direction, visualize_arrow_detection
|
||
|
||
def main():
|
||
# 创建参数解析器
|
||
parser = argparse.ArgumentParser(description='测试箭头方向检测')
|
||
parser.add_argument('image_path', help='图像文件路径')
|
||
parser.add_argument('--save', help='保存可视化结果的路径', default=None)
|
||
parser.add_argument('--show', help='显示可视化结果', action='store_true')
|
||
parser.add_argument('--debug', help='输出详细的调试信息', action='store_true')
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 检查文件是否存在
|
||
if not os.path.exists(args.image_path):
|
||
print(f"错误: 文件 '{args.image_path}' 不存在")
|
||
sys.exit(1)
|
||
|
||
print(f"正在处理图像: {args.image_path}")
|
||
|
||
# 加载图像
|
||
img = cv2.imread(args.image_path)
|
||
if img is None:
|
||
print(f"错误: 无法加载图像 '{args.image_path}'")
|
||
sys.exit(1)
|
||
|
||
# 如果需要，显示原始图像
|
||
if args.debug:
|
||
cv2.imshow('Original Image', img)
|
||
cv2.waitKey(0)
|
||
|
||
# 检测箭头方向
|
||
direction = detect_arrow_direction(img)
|
||
print(f"检测到的箭头方向: {direction}")
|
||
|
||
# 如果需要，显示可视化结果
|
||
if args.show or args.save:
|
||
visualize_arrow_detection(img, args.save)
|
||
|
||
# 如果不需要显示可视化结果但需要保存
|
||
if args.save and not args.show:
|
||
print(f"可视化结果已保存到: {args.save}")
|
||
|
||
return direction
|
||
|
||
if __name__ == "__main__":
|
||
main() |