src/yolo_predictions.py

from datetime import datetime

from data_models.bbox_manager import BoundingBoxManager, add_bboxes_to_db
from data_models.image_manager import ImageManager
from ultralytics import YOLO


model = None
bbox_manager = BoundingBoxManager()
image_manager = ImageManager()


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 load_model(model_file):
    """
    Load the YOLO model.

    Args:
        model_file (str): Path to the YOLO model.

    Returns:
        YOLO: The YOLO model.
    """
    global model
    if model:
        return model

    try:
        print(f"Loading YOLO model from file: {model_file}...")
        model = YOLO(f"models/{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(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.

    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...")

    model = load_model(model_file)
    if not model:
        raise Exception("Error loading YOLO model. Aborting processing.")

    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)

        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()