59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import os
|
|
import sys
|
|
import cv2
|
|
import argparse
|
|
|
|
# 添加工作目录到路径
|
|
sys.path.append(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', default="image_20250511_121219.png",
|
|
help='图像文件路径 (默认: image_20250511_121219.png)')
|
|
parser.add_argument('--save', default="arrow_detection_result.jpg",
|
|
help='保存可视化结果的路径 (默认: 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)
|
|
|
|
# 检测箭头方向
|
|
direction = detect_arrow_direction(img)
|
|
print(f"检测到的箭头方向: {direction}")
|
|
|
|
# 可视化检测过程并保存结果
|
|
visualize_arrow_detection(img, args.save)
|
|
|
|
print(f"可视化结果已保存到: {args.save}")
|
|
|
|
# 如果需要显示结果，等待用户按键
|
|
if args.show:
|
|
print("按任意键退出...")
|
|
cv2.waitKey(0)
|
|
cv2.destroyAllWindows()
|
|
|
|
return direction
|
|
|
|
if __name__ == "__main__":
|
|
main() |