Update tasks/image.py

tasks/image.py

@@ -99,37 +99,50 @@ 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
+    yolov8_model = YOLO("./../yolov8n.pt")
+
     predictions = []
     true_labels = []
     pred_boxes = []
     true_boxes_list = []  # List of lists, each inner list contains boxes for one image
     
+    
+    # Inference loop
     for example in test_dataset:
         # 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(int(has_smoke))
-        
-        # Make random classification prediction
-        pred_has_smoke = random.random() > 0.5
+    
+        # 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))
-        
-        # If there's a true box, parse it and make random box prediction
+    
+        # 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)
-            
-            # 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)
+    
+            # 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)
+        else:
+            true_boxes_list.append([])  # No true boxes for this image
+            pred_boxes.append([])  # No predicted boxes for this image
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE