Update tasks/image.py

tasks/image.py

@@ -2,10 +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, precision_score, recall_score
+from sklearn.metrics import accuracy_score
 import random
 import os
 
+from ultralytics import YOLO  # Import YOLO
 from .utils.evaluation import ImageEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
@@ -14,9 +15,11 @@ load_dotenv()
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "YOLO Smoke Detection"
 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"""
@@ -89,7 +92,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 = train_test["test"]
+    test_dataset = dataset["val"]#train_test["test"]
     
     # Start tracking emissions
     tracker.start()
@@ -99,50 +102,51 @@ async def evaluate_image(request: ImageEvaluationRequest):
     # YOUR MODEL INFERENCE CODE HERE
     # Update the code below to replace the random baseline with your model inference
     #--------------------------------------------------------------------------------------------   
-    from ultralytics import YOLO
-
-    # Load the trained YOLOv8 model
-    model = YOLO("./../best.pt")
-
     predictions = []
     true_labels = []
-    pred_boxes = []
-    true_boxes_list = []  # List of lists, each inner list contains boxes for one image
-    
-    
-    # Inference loop
+    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))
-    
-        # Load the image
-        image_path = example.get("image_path")  # Assuming access to image file path
+        
     
-        # Perform YOLOv8 inference
-        results = yolov8_model.predict(image_path)
-    
-        # Classification: Check if the model predicted any boxes (smoke presence)
-        pred_has_smoke = len(results[0].boxes) > 0
-        predictions.append(int(pred_has_smoke))
+        has_smoke = len(annotation) > 0
+        true_labels.append(1 if has_smoke else 0)
     
-        # Parse YOLO format true boxes
+
         if has_smoke:
-            # Parse all true boxes from the annotation
             image_true_boxes = parse_boxes(annotation)
-            true_boxes_list.append(image_true_boxes)
-    
-            # Get predicted bounding boxes (YOLO format)
-            image_pred_boxes = []
-            for box in results[0].boxes:
-                x_center, y_center, width, height = box.xywh[0].tolist()
-                image_pred_boxes.append([x_center, y_center, width, height])
-    
-            pred_boxes.append(image_pred_boxes)
+            if image_true_boxes:
+                true_boxes_list.append(image_true_boxes)
+            else:
+                true_boxes_list.append([])
         else:
-            true_boxes_list.append([])  # No true boxes for this image
-            pred_boxes.append([])  # No predicted boxes for this image
+            true_boxes_list.append([])
+    
+        results = yolo_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
@@ -151,10 +155,8 @@ async def evaluate_image(request: ImageEvaluationRequest):
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
     
-    # Calculate classification metrics
+    # Calculate classification accuracy
     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
@@ -172,8 +174,6 @@ 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,
@@ -185,5 +185,4 @@ async def evaluate_image(request: ImageEvaluationRequest):
             "test_seed": request.test_seed
         }
     }
-    
     return results