YOLO bulk image prediction

bulk_image_prediction.py

from ultralytics import YOLO
import os
import cv2

# Load the model
model = YOLO("model.pt")

# Specify the folder containing the images
image_folder = "./images"

# Create a folder to save the results
output_folder = "./output"
os.makedirs(output_folder, exist_ok=True)

# Iterate over the images in the folder
for filename in os.listdir(image_folder):
    # Adjust file extensions if needed
    if filename.endswith(".jpeg") or filename.endswith(".png"):
        image_path = os.path.join(image_folder, filename)

        # Perform prediction on the image
        results = model(image_path)

        # Get the annotated image from the results
        annotated_image = results[0].plot()

        # Save the annotated image
        results_filename = os.path.splitext(filename)[0] + "_results.jpg"
        results_path = os.path.join(output_folder, results_filename)
        cv2.imwrite(results_path, annotated_image)

print("Bulk prediction and saving completed.")