94 lines
3.6 KiB
Python
94 lines
3.6 KiB
Python
|
import cv2
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
# 添加父目录到系统路径
|
||
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||
|
|
project_root = os.path.dirname(current_dir)
|
||
|
|
sys.path.append(project_root)
|
||
|
|
|
||
|
|
from utils.detect_bar import detect_bar_top_distance
|
||
|
|
|
||
|
|
def process_image(image_path, save_dir=None, show_steps=False, observe=False, delay=800):
|
||
|
|
"""处理单张图像"""
|
||
|
|
print(f"处理图像: {image_path}")
|
||
|
|
|
||
|
|
# 读取图片
|
||
|
|
image = cv2.imread(image_path)
|
||
|
|
if image is None:
|
||
|
|
print(f"错误：无法读取图像 '{image_path}'")
|
||
|
|
return None
|
||
|
|
|
||
|
|
# 检测栏杆到顶部的距离
|
||
|
|
start_time = time.time()
|
||
|
|
top_distance = detect_bar_top_distance(image, observe=observe, delay=delay)
|
||
|
|
processing_time = time.time() - start_time
|
||
|
|
|
||
|
|
# 输出结果
|
||
|
|
if top_distance is not None:
|
||
|
|
print(f"处理时间: {processing_time:.3f}秒")
|
||
|
|
print(f"栏杆顶部到图像顶部的距离: {top_distance}像素")
|
||
|
|
|
||
|
|
# 如果要显示结果
|
||
|
|
if show_steps:
|
||
|
|
# 在图像上标记栏杆顶部位置
|
||
|
|
result_image = image.copy()
|
||
|
|
cv2.line(result_image, (0, top_distance), (result_image.shape[1], top_distance), (0, 255, 0), 2)
|
||
|
|
cv2.putText(result_image, f"Distance: {top_distance}px", (10, 30),
|
||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
||
|
|
|
||
|
|
# 显示结果
|
||
|
|
cv2.imshow("Bar Detection Result", result_image)
|
||
|
|
cv2.waitKey(delay)
|
||
|
|
|
||
|
|
# 如果指定了保存目录，保存结果
|
||
|
|
if save_dir:
|
||
|
|
if not os.path.exists(save_dir):
|
||
|
|
os.makedirs(save_dir)
|
||
|
|
|
||
|
|
# 在图像上标记栏杆顶部位置（如果之前没有标记）
|
||
|
|
if not show_steps:
|
||
|
|
result_image = image.copy()
|
||
|
|
cv2.line(result_image, (0, top_distance), (result_image.shape[1], top_distance), (0, 255, 0), 2)
|
||
|
|
cv2.putText(result_image, f"Distance: {top_distance}px", (10, 30),
|
||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
||
|
|
|
||
|
|
# 构建保存路径
|
||
|
|
base_name = os.path.basename(image_path)
|
||
|
|
save_path = os.path.join(save_dir, f"result_{base_name}")
|
||
|
|
cv2.imwrite(save_path, result_image)
|
||
|
|
print(f"结果已保存至: {save_path}")
|
||
|
|
else:
|
||
|
|
print("未能检测到栏杆")
|
||
|
|
|
||
|
|
# 关闭所有窗口
|
||
|
|
if show_steps:
|
||
|
|
cv2.destroyAllWindows()
|
||
|
|
|
||
|
|
return top_distance
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(description='灰色栏杆检测演示程序')
|
||
|
|
parser.add_argument('--input', default='./res/bar/image_20250525_085240.png', type=str, help='输入图像的路径')
|
||
|
|
parser.add_argument('--output', default='./res/bar/result', type=str, help='输出结果的保存目录')
|
||
|
|
parser.add_argument('--show', default=True, action='store_true', help='显示处理结果')
|
||
|
|
parser.add_argument('--observe', default=True, action='store_true', help='显示中间处理步骤')
|
||
|
|
parser.add_argument('--delay', type=int, default=800, help='显示中间结果的延迟时间(ms)')
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
# 检查输入路径
|
||
|
|
if not os.path.exists(args.input):
|
||
|
|
print(f"错误：文件 '{args.input}' 不存在")
|
||
|
|
return
|
||
|
|
|
||
|
|
# 获取输出目录
|
||
|
|
output_dir = args.output if args.output else os.path.dirname(args.input)
|
||
|
|
|
||
|
|
# 处理图像
|
||
|
|
process_image(args.input, output_dir, args.show, args.observe, args.delay)
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|