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

fix: replace absolute path with relative path and update JSON key fro…

src/predict.py

@@ -1,6 +1,7 @@
 """Predict the class of a list of images, and return a json file with the prediction"""
 
 import json
+
 from ultralytics import YOLO
 
 # [{
@@ -16,15 +17,24 @@ from ultralytics import YOLO
 
 
 def add_prediction_to_json(predictions) -> None:
+    """Add the prediction to the json file.
+
+    Args:
+        predictions (generator): generator of predictions from the YOLO model
+    Returns:
+        list: list of predictions to add to the json result
+    """
     formatted_predictions = []
     for result in predictions:
         for box in result.boxes:
-            formatted_predictions.append({
-                "class": int(box.cls[0].numpy().item()),
-                "confidence": box.conf[0].numpy().item(),
-                "box": box.xyxy[0].numpy().tolist()
-            })
-
+            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(),
+                }
+            )
     return formatted_predictions
 
 
@@ -43,13 +53,15 @@ def get_prediction(list_folder, model_file, output_file) -> None:
         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 with the prediction
+            # show=False to not show the image when getting the prediction
             predictions = model(folder, stream=True, show=False)
             try:
-                dict_predictions.append({
-                    "image": folder,
-                    "prediction": add_prediction_to_json(predictions)
-                })
+                dict_predictions.append(
+                    {
+                        "park": folder.split("/")[-1],
+                        "prediction": add_prediction_to_json(predictions),
+                    }
+                )
             except AttributeError as att_error:
                 dict_predictions.append(
                     {
@@ -63,7 +75,7 @@ def get_prediction(list_folder, model_file, output_file) -> None:
 
 if __name__ == "__main__":
     get_prediction(
-        ["images/ChIJ_yEAweFlXj4Reo-x-AghMRM"],
+        ["../images/Al Khaldiyah Park", "../images/Family Park"],
         "models/yolov5s.pt",
         "./predictions.json",
     )