fix(model): fix yolo v8

Dockerfile

@@ -4,6 +4,8 @@
 FROM python:3.9
 
 RUN useradd -m -u 1000 user
+RUN apt-get update && apt-get install ffmpeg libsm6 libxext6  -y
+
 USER user
 ENV PATH="/home/user/.local/bin:$PATH"
 

requirements.txt

@@ -11,4 +11,5 @@ librosa==0.10.2.post1
 numpy==1.26.4
 ultralytics==8.3.68
 ultralytics-thop==2.0.14
-opencv-python==4.11.0.86
+#opencv-python==4.11.0.86
+python-dotenv==1.0.0

tasks/image.py

@@ -5,6 +5,7 @@ import numpy as np
 from sklearn.metrics import accuracy_score, precision_score, recall_score
 import random
 import os
+from ultralytics import YOLO
 
 from .utils.evaluation import ImageEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
@@ -14,7 +15,7 @@ load_dotenv()
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "Simple YOLOv8n model"
 ROUTE = "/image"
 
 def parse_boxes(annotation_string):
@@ -120,7 +121,9 @@ async def evaluate_image(request: ImageEvaluationRequest):
     pred_boxes = []
     true_boxes_list = []  # List of lists, each inner list contains boxes for one image
     
-    for example in test_dataset:
+    n_examples = len(test_dataset)
+    for i, example in enumerate(test_dataset):
+        print(f"Running {i+1} of {n_examples}")
         # Parse true annotation (YOLO format: class_id x_center y_center width height)
         annotation = example.get("annotations", "").strip()
         has_smoke = len(annotation) > 0
@@ -137,9 +140,14 @@ async def evaluate_image(request: ImageEvaluationRequest):
             image_true_boxes = parse_boxes(annotation)
             true_boxes_list.append(image_true_boxes)
             
-            pred_box_list = get_boxes_list(model_preds)[0] # With one bbox to start with (as in the random baseline)
+            try:
+                pred_box_list = get_boxes_list(model_preds)[0] # With one bbox to start with (as in the random baseline)
+            except:
+                print("No boxes found")
+                pred_box_list = [0, 0, 0, 0] # Hacky way to make sure that compute_max_iou doesn't fail
             pred_boxes.append(pred_box_list)
 
+
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE