Update tasks/image.py

tasks/image.py

@@ -5,7 +5,7 @@ import numpy as np
 from sklearn.metrics import accuracy_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
 
@@ -17,6 +17,8 @@ router = APIRouter()
 DESCRIPTION = "Random Baseline"
 ROUTE = "/image"
 
+model = YOLO("test_best.pt")
+
 def parse_boxes(annotation_string):
     """Parse multiple boxes from a single annotation string.
     Each box has 5 values: class_id, x_center, y_center, width, height"""
@@ -102,35 +104,50 @@ async def evaluate_image(request: ImageEvaluationRequest):
     
     predictions = []
     true_labels = []
-    pred_boxes = []
-    true_boxes_list = []  # List of lists, each inner list contains boxes for one image
+    pred_boxes = []  
+    true_boxes_list = []  
     
     for example in test_dataset:
-        # Parse true annotation (YOLO format: class_id x_center y_center width height)
+        # Extract image and annotations
+        image = example["image"]
         annotation = example.get("annotations", "").strip()
-        has_smoke = len(annotation) > 0
-        true_labels.append(int(has_smoke))
-        
-        # Make random classification prediction
-        pred_has_smoke = random.random() > 0.5
-        predictions.append(int(pred_has_smoke))
         
-        # If there's a true box, parse it and make random box prediction
+    
+        has_smoke = len(annotation) > 0
+        true_labels.append(1 if has_smoke else 0)
+    
+    
         if has_smoke:
-            # Parse all true boxes from the annotation
             image_true_boxes = parse_boxes(annotation)
-            true_boxes_list.append(image_true_boxes)
-            
-            # For baseline, make one random box prediction per image
-            # In a real model, you might want to predict multiple boxes
-            random_box = [
-                random.random(),  # x_center
-                random.random(),  # y_center
-                random.random() * 0.5,  # width (max 0.5)
-                random.random() * 0.5   # height (max 0.5)
-            ]
-            pred_boxes.append(random_box)
+            if image_true_boxes:
+                true_boxes_list.append(image_true_boxes)
+            else:
+                true_boxes_list.append([])
+        else:
+            true_boxes_list.append([])
+    
+        results =  model.predict(image, verbose=False)    # INFERENCE - prediction
+        
+        if len(results[0].boxes):
+            pred_box = results[0].boxes.xywhn[0].cpu().numpy().tolist()
+            predictions.append(1)  
+            pred_boxes.append(pred_box)
+        else:
+            predictions.append(0)  
+            pred_boxes.append([])  
+    
+    filtered_true_boxes_list = []
+    filtered_pred_boxes = []
     
+    for true_boxes, pred_boxes_entry in zip(true_boxes_list, pred_boxes):  # Only see when annotation(s) is/are both on true label and prediction
+        if true_boxes and pred_boxes_entry:  
+            filtered_true_boxes_list.append(true_boxes)
+            filtered_pred_boxes.append(pred_boxes_entry)
+    
+    
+    true_boxes_list = filtered_true_boxes_list
+    pred_boxes = filtered_pred_boxes
+
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
     #--------------------------------------------------------------------------------------------