feat: split prediction into YOLO predictions

src/{models → data_models}/bbox_manager.py

@@ -2,7 +2,7 @@
 
 from sqlalchemy import text
 
-from models.sql_connection import get_db_connection
+from data_models.sql_connection import get_db_connection
 
 
 class BoundingBoxManager:
@@ -62,3 +62,32 @@ class BoundingBoxManager:
     def close_connection(self):
         """Close the connection."""
         self.session.close()
+
+
+def add_bboxes_to_db(result, bbox_manager, image_id):
+    """
+    Adds bounding boxes from a YOLO result to the database.
+
+    Args:
+        result: YOLO result object containing the bounding box details.
+        bbox_manager (BoundingBoxManager): Instance of BoundingBoxManager to interact with the database.
+        image_id (int): The ID of the associated image.
+
+    Returns:
+        None
+    """
+    for box in result.boxes:
+        try:
+            print(f"Adding bounding box for image ID: {image_id}...")
+            bbox_manager.add_bbox(
+                confidence=box.conf[0].numpy().item(),
+                class_id=int(box.cls[0].numpy().item()),
+                img_id=image_id,
+                x_min=box.xyxy[0][0].numpy().item(),
+                y_min=box.xyxy[0][1].numpy().item(),
+                x_max=box.xyxy[0][2].numpy().item(),
+                y_max=box.xyxy[0][3].numpy().item(),
+            )
+            print(f"Bounding box added to DB for image ID: {image_id}")
+        except Exception as e:
+            print(f"Error inserting bounding box into DB: {e}")

src/{models → data_models}/image_manager.py

@@ -1,8 +1,9 @@
-"""Mnages images table in the database."""
+"""Manages images table in the database."""
 
+from datetime import datetime
 from sqlalchemy import text
 
-from models.sql_connection import get_db_connection
+from data_models.sql_connection import get_db_connection
 
 
 class ImageManager:
@@ -97,3 +98,32 @@ class ImageManager:
     def close_connection(self):
         """Close the connection."""
         self.session.close()
+
+
+def add_image_to_db(result, image_manager, park_id):
+    """
+    Adds an image to the database.
+
+    Args:
+        result: YOLO result object containing the image details.
+        image_manager (ImageManager): Instance of ImageManager to interact with the database.
+        park_id (int): The ID of the associated park.
+
+    Returns:
+        int: The ID of the added image.
+    """
+    image_name = result.path.split("/")[-1]
+    current_datetime = datetime.now()
+    created_at = current_datetime.strftime("%Y-%m-%d %H:%M:%S")
+
+    try:
+        print(f"Adding image '{image_name}' to the database...")
+        image = image_manager.add_image(
+            name=image_name, created_at=created_at, park_id=park_id
+        )
+        image_id = image_manager.get_image_id(image_name)
+        print(f"Image '{image_name}' added with ID: {image_id}")
+        return image_id
+    except Exception as e:
+        print(f"Error adding image '{image_name}' to the database: {e}")
+        return None

src/{models → data_models}/park_manager.py

@@ -2,7 +2,7 @@
 
 from sqlalchemy import text
 
-from models.sql_connection import get_db_connection
+from data_models.sql_connection import get_db_connection
 
 
 class ParkManager:

src/predict.py

@@ -1,115 +1,10 @@
-from datetime import datetime
-
-from ultralytics import YOLO
-
-from models.bbox_manager import BoundingBoxManager
-from models.image_manager import ImageManager
-from models.park_manager import ParkManager
-
-
-def add_image_to_db(result, image_manager, park_id):
-    """
-    Adds an image to the database.
-
-    Args:
-        result: YOLO result object containing the image details.
-        image_manager (ImageManager): Instance of ImageManager to interact with the database.
-        park_id (int): The ID of the associated park.
-
-    Returns:
-        int: The ID of the added image.
-    """
-    image_name = result.path.split("/")[-1]
-    current_datetime = datetime.now()
-    created_at = current_datetime.strftime("%Y-%m-%d %H:%M:%S")
-
-    try:
-        print(f"Adding image '{image_name}' to the database...")
-        image = image_manager.add_image(
-            name=image_name, created_at=created_at, park_id=park_id
-        )
-        image_id = image_manager.get_image_id(image_name)
-        print(f"Image '{image_name}' added with ID: {image_id}")
-        return image_id
-    except Exception as e:
-        print(f"Error adding image '{image_name}' to the database: {e}")
-        return None
-
-
-def add_bboxes_to_db(result, bbox_manager, image_id):
-    """
-    Adds bounding boxes from a YOLO result to the database.
-
-    Args:
-        result: YOLO result object containing the bounding box details.
-        bbox_manager (BoundingBoxManager): Instance of BoundingBoxManager to interact with the database.
-        image_id (int): The ID of the associated image.
-
-    Returns:
-        None
-    """
-    for box in result.boxes:
-        try:
-            print(f"Adding bounding box for image ID: {image_id}...")
-            bbox_manager.add_bbox(
-                confidence=box.conf[0].numpy().item(),
-                class_id=int(box.cls[0].numpy().item()),
-                img_id=image_id,
-                x_min=box.xyxy[0][0].numpy().item(),
-                y_min=box.xyxy[0][1].numpy().item(),
-                x_max=box.xyxy[0][2].numpy().item(),
-                y_max=box.xyxy[0][3].numpy().item(),
-            )
-            print(f"Bounding box added to DB for image ID: {image_id}")
-        except Exception as e:
-            print(f"Error inserting bounding box into DB: {e}")
-
-
-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.
-    """
-    print(f"Initializing processing of predictions...")
-    bbox_manager = BoundingBoxManager()
-    image_manager = ImageManager()
-    park_manager = ParkManager()
-
-    model = YOLO(model_file)
-
-    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)
-
-        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)
-
-    print(f"Processing complete. Closing database connections.")
-    bbox_manager.close_connection()
-    image_manager.close_connection()
-    park_manager.close_connection()
+import os
+from tqdm import tqdm
+from yolo_predictions import process_YOLO_predictions
 
 
 if __name__ == "__main__":
-    list_folder = ["../images/Al Khaldiyah Park", "../images/Family Park"]
-    model_file = "yolov5s.pt"
-    process_predictions(list_folder, model_file)
+    list_folder = ["./images/Al Khaldiyah Park", "./images/Family Park"]
+    model_file = "yolo11s.pt"
+
+    process_YOLO_predictions(list_folder, model_file)

src/yolo_predictions.py

@@ -0,0 +1,70 @@
+from data_models.bbox_manager import BoundingBoxManager, add_bboxes_to_db
+from data_models.image_manager import ImageManager, add_image_to_db
+from data_models.park_manager import ParkManager
+from ultralytics import YOLO
+
+
+def load_model(model_file):
+    """
+    Load the YOLO model.
+
+    Args:
+        model_file (str): Path to the YOLO model.
+
+    Returns:
+        YOLO: The YOLO model.
+    """
+    try:
+        print(f"Loading YOLO model from file: {model_file}...")
+        model = YOLO(model_file)
+        print(f"Model loaded successfully!")
+        return model
+    except Exception as e:
+        print(f"Error loading YOLO model: {e}")
+        return None
+
+
+def process_YOLO_predictions(list_folder: list[str], model_file: str):
+    """
+    Process predictions for a list of image folders, adding images and bounding boxes to the database using the YOLO model.
+
+    Args:
+        list_folder (list): List of image folders.
+        model_file (str): Path to the YOLO model.
+
+    Raises:
+        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)
+
+        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)
+
+    print(f"Processing complete. Closing database connections.")
+    bbox_manager.close_connection()
+    image_manager.close_connection()
+    park_manager.close_connection()