Update tasks/image.py

tasks/image.py

@@ -2,11 +2,11 @@ from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
 import numpy as np
-from sklearn.metrics import accuracy_score
+from sklearn.metrics import accuracy_score, precision_score, recall_score
 import random
 import os
+from ultralytics import YOLO
 
-from ultralytics import YOLO  # Import YOLO
 from .utils.evaluation import ImageEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
@@ -15,11 +15,10 @@ load_dotenv()
 
 router = APIRouter()
 
-DESCRIPTION = "YOLO Smoke Detection"
+# MODEL_TYPE = "YOLOv11n"
+DESCRIPTION = "YOLOv11"
 ROUTE = "/image"
 
-yolo_model = YOLO("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"""
@@ -70,6 +69,16 @@ def compute_max_iou(true_boxes, pred_box):
         max_iou = max(max_iou, iou)
     return max_iou
 
+def load_model(path_to_model, model_type="YOLO"):
+    if model_type == "YOLO":
+        model = YOLO(path_to_model)
+    else:
+        raise NotImplementedError
+    return model
+
+def get_boxes_list(predictions):
+    return [box.tolist() for box in predictions.boxes.xywhn]
+
 @router.post(ROUTE, tags=["Image Task"],
              description=DESCRIPTION)
 async def evaluate_image(request: ImageEvaluationRequest):
@@ -92,7 +101,7 @@ async def evaluate_image(request: ImageEvaluationRequest):
     
     # Split dataset
     train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
-    test_dataset = dataset["val"]#train_test["test"]
+    test_dataset = train_test["test"]
     
     # Start tracking emissions
     tracker.start()
@@ -102,52 +111,44 @@ async def evaluate_image(request: ImageEvaluationRequest):
     # YOUR MODEL INFERENCE CODE HERE
     # Update the code below to replace the random baseline with your model inference
     #--------------------------------------------------------------------------------------------   
+    
+    
+    PATH_TO_MODEL = f"best.pt"
+    model = load_model(PATH_TO_MODEL)
+    
+    print(f"Model info: {model.info()}")
     predictions = []
     true_labels = []
-    pred_boxes = []  
-    true_boxes_list = []  
-
-    for example in test_dataset:
-        # Extract image and annotations
-        image = example["image"]
-        annotation = example.get("annotations", "").strip()
-        
+    pred_boxes = []
+    true_boxes_list = []  # List of lists, each inner list contains boxes for one image
     
+    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
-        true_labels.append(1 if has_smoke else 0)
-    
-
+        true_labels.append(int(has_smoke))
+        
+        model_preds = model(example['image'])[0]
+        pred_has_smoke = len(model_preds) > 0
+        predictions.append(int(pred_has_smoke))
+        
+        # If there's a true box, parse it and make random box prediction
         if has_smoke:
+            
+            # Parse all true boxes from the annotation
             image_true_boxes = parse_boxes(annotation)
-            if image_true_boxes:
-                true_boxes_list.append(image_true_boxes)
-            else:
-                true_boxes_list.append([])
-        else:
-            true_boxes_list.append([])
-    
-        # results = yolo_model .predict(image, verbose=False)    # INFERENCE - prediction
-        results = yolo_model.predict(image)    # INFERENCE - prediction
+            true_boxes_list.append(image_true_boxes)
+            
+            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)
 
-        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
@@ -156,8 +157,10 @@ async def evaluate_image(request: ImageEvaluationRequest):
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
     
-    # Calculate classification accuracy
+    # Calculate classification metrics
     classification_accuracy = accuracy_score(true_labels, predictions)
+    classification_precision = precision_score(true_labels, predictions)
+    classification_recall = recall_score(true_labels, predictions)
     
     # Calculate mean IoU for object detection (only for images with smoke)
     # For each image, we compute the max IoU between the predicted box and all true boxes
@@ -175,6 +178,8 @@ async def evaluate_image(request: ImageEvaluationRequest):
         "submission_timestamp": datetime.now().isoformat(),
         "model_description": DESCRIPTION,
         "classification_accuracy": float(classification_accuracy),
+        "classification_precision": float(classification_precision),
+        "classification_recall": float(classification_recall),
         "mean_iou": mean_iou,
         "energy_consumed_wh": emissions_data.energy_consumed * 1000,
         "emissions_gco2eq": emissions_data.emissions * 1000,
@@ -186,4 +191,5 @@ async def evaluate_image(request: ImageEvaluationRequest):
             "test_seed": request.test_seed
         }
     }
+    
     return results