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)