algonacci · May 22, 2024 03:20 · May 22, 2024
diff --git a/bulk_image_prediction.py b/bulk_image_prediction.py
@@ -0,0 +1,32 @@
+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.")