Merge pull request #3 from ia-labo/feat/prediction

src/predict.py

@@ -1,81 +1,115 @@
-"""Predict the class of a list of images, and return a json file with the prediction"""
-
-import json
+from datetime import datetime
 
 from ultralytics import YOLO
 
-# [{
-#     "image": "images/ChIJ_yEAweFlXj4Reo-x-AghMRM",
-#     "prediction": [
-#         {
-#             "class": "person",
-#             "confidence": 0.99,
-#             "box": [0.0, 0.0, 0.0, 0.0]
-#         }
-#     ]
-# }]
+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_prediction_to_json(predictions) -> None:
-    """Add the prediction to the json file.
+def add_bboxes_to_db(result, bbox_manager, image_id):
+    """
+    Adds bounding boxes from a YOLO result to the database.
 
     Args:
-        predictions (generator): generator of predictions from the YOLO model
+        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:
-        list: list of predictions to add to the json result
+        None
     """
-    formatted_predictions = []
-    for result in predictions:
-        for box in result.boxes:
-            formatted_predictions.append(
-                {
-                    "image": result.path.split("/")[-1],
-                    "class": int(box.cls[0].numpy().item()),
-                    "confidence": box.conf[0].numpy().item(),
-                    "box": box.xyxy[0].numpy().tolist(),
-                }
+    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(),
             )
-    return formatted_predictions
+            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 get_prediction(list_folder, model_file, output_file) -> None:
-    """Get the prediction for a list of images
+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_path (list): list of images
-        model_file (str): path to the model
-        output_file (str): path to the output file
+        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)
 
-    with open(output_file, "w+", encoding="utf8") as file:
-        dict_predictions = []
-        for folder in list_folder[:2]:
-            # stream=True to get the prediction for each image
-            # instead of trying to get all the predictions at once
-            # show=False to not show the image when getting the prediction
-            predictions = model(folder, stream=True, show=False)
-            try:
-                dict_predictions.append(
-                    {
-                        "park": folder.split("/")[-1],
-                        "prediction": add_prediction_to_json(predictions),
-                    }
-                )
-            except AttributeError as att_error:
-                dict_predictions.append(
-                    {
-                        "status": "error",
-                        "error_name": str(att_error.__class__.__name__),
-                        "message": str(att_error),
-                    }
-                )
-        json.dump(dict_predictions, file, indent=4)
+    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()
 
 
 if __name__ == "__main__":
-    get_prediction(
-        ["../images/Al Khaldiyah Park", "../images/Family Park"],
-        "models/yolov5s.pt",
-        "./predictions.json",
-    )
+    list_folder = ["../images/Al Khaldiyah Park", "../images/Family Park"]
+    model_file = "yolov5s.pt"
+    process_predictions(list_folder, model_file)