implementation of inference

tasks/image.py

@@ -6,6 +6,7 @@ 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("models/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"""
@@ -99,37 +102,34 @@ async def evaluate_image(request: ImageEvaluationRequest):
     # YOUR MODEL INFERENCE CODE HERE
     # Update the code below to replace the random baseline with your model inference
     #--------------------------------------------------------------------------------------------   
-    
     predictions = []
     true_labels = []
     pred_boxes = []
-    true_boxes_list = []  # List of lists, each inner list contains boxes for one image
-    
+    true_boxes_list = []
+
     for example in test_dataset:
-        # Parse true annotation (YOLO format: class_id x_center y_center width height)
+        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
         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)
+
+            # Perform YOLO inference
+            results = yolo_model.predict(image)
+            if len(results[0].boxes):  # If predictions exist
+                pred_box = results[0].boxes.xywh[0].cpu().numpy().tolist()  # First box in YOLO format
+                pred_boxes.append(pred_box)
+            else:
+                pred_boxes.append([])  # No prediction for this image
+        else:
+            true_boxes_list.append([])
+            pred_boxes.append([])
+
+        # Classification: If predictions exist, assume smoke is present
+        predictions.append(1 if len(results[0].boxes) > 0 else 0)
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE