39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
|
import os
|
||
|
|
from ultralytics import YOLO
|
||
|
|
import cv2
|
||
|
|
|
||
|
|
def predict_and_save_images(model, image_folder, output_folder):
|
||
|
|
"""使用YOLOv8模型预测图像并保存结果。"""
|
||
|
|
if not os.path.exists(output_folder):
|
||
|
|
os.makedirs(output_folder)
|
||
|
|
print(f"Created output directory: {output_folder}")
|
||
|
|
|
||
|
|
image_files = os.listdir(image_folder)
|
||
|
|
print(f"Found {len(image_files)} images in {image_folder}")
|
||
|
|
|
||
|
|
for image_file in image_files:
|
||
|
|
image_path = os.path.join(image_folder, image_file)
|
||
|
|
image = cv2.imread(image_path)
|
||
|
|
if image is None:
|
||
|
|
print(f"Could not read image: {image_path}")
|
||
|
|
continue
|
||
|
|
|
||
|
|
results = model(image_path)
|
||
|
|
annotated_image = results[0].plot()
|
||
|
|
|
||
|
|
output_path = os.path.join(output_folder, image_file)
|
||
|
|
success = cv2.imwrite(output_path, annotated_image)
|
||
|
|
if not success:
|
||
|
|
print(f"Could not write image: {output_path}")
|
||
|
|
else:
|
||
|
|
print(f"Image saved: {output_path}")
|
||
|
|
|
||
|
|
# 示例用法
|
||
|
|
image_folder = '/root/catkin_ws/src/ultralytics/ours_15000/renders' # 输入图像文件夹路径
|
||
|
|
output_folder = '/root/catkin_ws/src/ultralytics/ours_15000/renders_box' # 输出结果文件夹路径
|
||
|
|
|
||
|
|
# 加载YOLOv8模型 (例如使用预训练的 YOLOv8n 模型)
|
||
|
|
model = YOLO('yolov8n.pt')
|
||
|
|
|
||
|
|
predict_and_save_images(model, image_folder, output_folder)
|