feat: make image loop in predict.py

src/predict.py

@@ -1,10 +1,58 @@
 import os
 from tqdm import tqdm
 from yolo_predictions import process_YOLO_predictions
+from openai_predictions import process_agent_predictions
+from data_models.park_manager import ParkManager
+
+
+def load_images_from_folder(folder):
+    """
+    Load images from a folder.
+
+    Args:
+        folder (str): Path to the folder.
+
+    Returns:
+        list: List of image paths.
+    """
+    for file in os.listdir(folder):
+        if file.endswith(".jpg") or file.endswith(".png"):
+            yield os.path.join(folder, file)
+
+
+def process_predictions(list_folder, model_file):
+    """
+    Process predictions for a list of image folders, adding images and bounding boxes to the database.
+
+    Args:
+        list_folder (list): List of image folders.
+        model_file (str): Path to the YOLO model.
+
+    Raises:
+        Exception: If an error occurs during processing.
+    """
+    park_manager = ParkManager()
+
+    for folder in list_folder:
+        print(f"Loading images from folder: {folder}...")
+
+        park_name = folder.split("/")[-1]
+
+        print(f"Retrieving park ID for '{park_name}'...")
+        park_id = park_manager.get_park_id(park_name)
+        if not park_id:
+            print(f"Park '{park_name}' not found in the database. Skipping folder.")
+            continue
+
+        for image in tqdm(load_images_from_folder(folder)):
+            process_YOLO_predictions(park_id, image, model_file)
+            # process_agent_predictions(park_id, image)
+        print(f"Folders {folder} processed successfully!")
+    park_manager.close_connection()
 
 
 if __name__ == "__main__":
     list_folder = ["./images/Al Khaldiyah Park", "./images/Family Park"]
     model_file = "yolo11s.pt"
 
-    process_YOLO_predictions(list_folder, model_file)
+    process_predictions(list_folder, model_file)

src/yolo_predictions.py

@@ -49,6 +49,10 @@ def load_model(model_file):
     Returns:
         YOLO: The YOLO model.
     """
+    global model
+    if model:
+        return model
+
     try:
         print(f"Loading YOLO model from file: {model_file}...")
         model = YOLO(model_file)
@@ -59,7 +63,7 @@ def load_model(model_file):
         return None
 
 
-def process_YOLO_predictions(list_folder: list[str], model_file: str):
+def process_YOLO_predictions(park_id: int, image: str, model_file: str):
     """
     Process predictions for a list of image folders, adding images and bounding boxes to the database using the YOLO model.
 
@@ -71,35 +75,21 @@ def process_YOLO_predictions(list_folder: list[str], model_file: str):
         Exception: If an error occurs during processing.
     """
     print(f"Initializing processing of predictions...")
-    bbox_manager = BoundingBoxManager()
-    image_manager = ImageManager()
-    park_manager = ParkManager()
 
     model = load_model(model_file)
     if not model:
         raise Exception("Error loading YOLO model. Aborting processing.")
 
-    for folder in list_folder:
-        print(f"Processing folder: {folder}...")
-        park_name = folder.split("/")[-1]
-
-        print(f"Retrieving park ID for '{park_name}'...")
-        park_id = park_manager.get_park_id(park_name)
-        if not park_id:
-            print(f"Park '{park_name}' not found in the database. Skipping folder.")
-            continue
-
-        predictions = model(folder, stream=True, show=False)
+    predictions = model(image, stream=True, show=False)
 
-        for result in predictions:
-            print(f"Processing result for image: {result.path.split('/')[-1]}...")
-            image_id = add_image_to_db(result, image_manager, park_id)
+    for result in predictions:
+        print(f"Processing result for image: {result.path.split('/')[-1]}...")
+        image_id = add_image_to_db(result, image_manager, park_id)
 
-            if image_id:
-                print(f"Adding bounding boxes to the database...")
-                add_bboxes_to_db(result, bbox_manager, image_id)
+        if image_id:
+            print(f"Adding bounding boxes to the database...")
+            add_bboxes_to_db(result, bbox_manager, image_id)
 
     print(f"Processing complete. Closing database connections.")
     bbox_manager.close_connection()
     image_manager.close_connection()
-    park_manager.close_connection()