first try

tasks/image.py

@@ -10,6 +10,8 @@ from .utils.evaluation import ImageEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 from dotenv import load_dotenv
+from ultralytics import YOLO
+import torch
 load_dotenv()
 
 router = APIRouter()
@@ -96,43 +98,51 @@ async def evaluate_image(request: ImageEvaluationRequest):
     tracker.start_task("inference")
     
     #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE CODE HERE
-    # Update the code below to replace the random baseline with your model inference
+    # YOLO Model Inference
     #--------------------------------------------------------------------------------------------   
     
+    
+    # Load your trained YOLO model
+    model = YOLO('yolotrained.pt')  # Replace with your model path
+    
     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)
+        # Parse true annotation
         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))
+        # Run YOLO inference
+        image_path = example['image_path']  # Adjust according to your dataset structure
+        results = model(image_path)
+        
+        # Process YOLO predictions
+        if len(results[0].boxes) > 0:
+            predictions.append(1)  # Smoke detected
+            # Get the box with highest confidence
+            best_box = results[0].boxes[0]
+            # Convert box to YOLO format (x_center, y_center, width, height)
+            box_xywh = best_box.xywh[0].cpu().numpy()
+            pred_boxes.append([
+                float(box_xywh[0]),  # x_center
+                float(box_xywh[1]),  # y_center
+                float(box_xywh[2]),  # width
+                float(box_xywh[3])   # height
+            ])
+        else:
+            predictions.append(0)  # No smoke detected
         
-        # If there's a true box, parse it and make random box prediction
+        # Process 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)
     
     #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE STOPS HERE
+    # Model Inference Ends
     #--------------------------------------------------------------------------------------------   
     
     # Stop tracking emissions

yolotrained.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47e01d84756c88c9282d1751ec137ea555e54b82e1450cbd7f40eda2e52c225a
+size 15201538