Added model evaluation code

tasks/image.py

@@ -1,20 +1,24 @@
-from fastapi import APIRouter
+import os
+import random
 from datetime import datetime
-from datasets import load_dataset
+
 import numpy as np
+import torch
+from datasets import load_dataset
+from dotenv import load_dotenv
+from fastapi import APIRouter
 from sklearn.metrics import accuracy_score, precision_score, recall_score
-import random
-import os
 
+from ultralytics import YOLOv10
+
+from .utils.emissions import clean_emissions_data, get_space_info, tracker
 from .utils.evaluation import ImageEvaluationRequest
-from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
-from dotenv import load_dotenv
 load_dotenv()
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "YOLO-EFD"
 ROUTE = "/image"
 
 def parse_boxes(annotation_string):
@@ -90,6 +94,10 @@ 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 = train_test["test"]
+
+    model = YOLOv10("best.pt")
+    device = "cuda"
+    model = model.to(device)
     
     # Start tracking emissions
     tracker.start()
@@ -111,25 +119,20 @@ async def evaluate_image(request: ImageEvaluationRequest):
         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))
+        # Make prediction with model
+        image = example['image']
+        with torch.inference_mode():
+            pred = model(image)[0]
+        smoke_detected = pred.boxes.xywhn.shape[0] > 0
+        predictions.append(int(smoke_detected))
         
-        # If there's a true box, parse it and make random box prediction
-        if has_smoke:
+        # If there's a true box, and at least one box is predicted, parse them
+        # If one of the two boxes is empty, mIoU computation fails
+        if has_smoke and smoke_detected:
             # 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)
+            pred_boxes.append(pred.boxes.xywhn.tolist()[0])
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
@@ -137,7 +140,7 @@ async def evaluate_image(request: ImageEvaluationRequest):
     
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
-    
+
     # Calculate classification metrics
     classification_accuracy = accuracy_score(true_labels, predictions)
     classification_precision = precision_score(true_labels, predictions)