mi-task/test/visualize_dual_track.py

98 lines
3.2 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import cv2
import os
import glob
import datetime
from utils.detect_dual_track_lines import auto_detect_dual_track_lines
def visualize_detection(img_path, output_dir='results'):
"""可视化轨迹线检测结果并保存"""
# 读取图像
img = cv2.imread(img_path)
if img is None:
print(f"无法读取图像: {img_path}")
return False
# 获取文件名
filename = os.path.basename(img_path)
# 检测轨迹线
result = auto_detect_dual_track_lines(img, observe=False, save_log=False)
if result[0] is None:
print(f"检测失败: {img_path}")
return False
center_info, left_info, right_info = result
# 绘制结果
result_img = img.copy()
height, width = img.shape[:2]
center_x = width // 2
# 获取线段坐标
left_x1, left_y1, left_x2, left_y2 = left_info["line"]
right_x1, right_y1, right_x2, right_y2 = right_info["line"]
center_point = center_info["point"]
deviation = center_info["deviation"]
# 绘制左右轨迹线
cv2.line(result_img, (int(left_x1), int(left_y1)), (int(left_x2), int(left_y2)), (255, 0, 0), 2)
cv2.line(result_img, (int(right_x1), int(right_y1)), (int(right_x2), int(right_y2)), (0, 0, 255), 2)
# 计算中心线
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, (center_line_x1, center_line_y1), (center_line_x2, center_line_y2), (0, 255, 0), 2)
# 绘制图像中心线
cv2.line(result_img, (center_x, 0), (center_x, height), (0, 0, 255), 1)
# 标记中心点
cv2.circle(result_img, center_point, 10, (255, 0, 255), -1)
# 显示偏差信息
cv2.putText(result_img, f"Deviation: {deviation}px", (10, 30),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
# 显示是否为石板路模式
if center_info["stone_path_mode"]:
cv2.putText(result_img, "Stone Path Mode", (10, 60),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
# 创建输出目录
os.makedirs(output_dir, exist_ok=True)
# 保存结果图像
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
output_path = os.path.join(output_dir, f"result_{filename}")
cv2.imwrite(output_path, result_img)
print(f"结果已保存到: {output_path}")
return True
def process_all_images(input_dir, output_dir='results'):
"""处理目录中的所有图像"""
# 获取所有jpg图像
image_paths = glob.glob(os.path.join(input_dir, "*.jpg"))
total = len(image_paths)
success = 0
for img_path in sorted(image_paths):
print(f"\n处理图像: {img_path}")
if visualize_detection(img_path, output_dir):
success += 1
success_rate = (success / total) * 100 if total > 0 else 0
print(f"\n处理完成: 成功率 {success_rate:.2f}% ({success}/{total})")
if __name__ == "__main__":
input_dir = "res/path/dual"
output_dir = "results/dual_track"
process_all_images(input_dir, output_dir)