Optimize submission

tasks/image.py

@@ -18,7 +18,6 @@ from pathlib import Path
 from ultralytics import YOLO
 from torch import device
 from torch.cuda import is_available
-from torch import no_grad
 
 router = APIRouter()
 
@@ -120,38 +119,31 @@ async def evaluate_image(request: ImageEvaluationRequest):
     true_labels = []
     pred_boxes = []
     true_boxes_list = []  # List of lists, each inner list contains boxes for one image
+    
+    logging.info(f"Inference start on device: {device_name}")
+    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))
 
-   # Preprocess annotations before the loop
-    preprocessed_annotations = [parse_boxes(example.get("annotations", "").strip()) for example in test_dataset]
-
-    # Use torch.no_grad() to disable gradient tracking during inference
-    with no_grad():
-        predictions = []
-        true_labels = []
-        pred_boxes = []
-        true_boxes_list = []  # List of lists, each inner list contains boxes for one image
+        # Make prediction
+        results = model.predict(example["image"], device=device_name, conf=THRESHOLD, verbose=False, imgsz=IMGSIZE)[0]
+        pred_has_smoke = len(results) > 0
+        predictions.append(int(pred_has_smoke))
         
-        logging.info(f"Inference start on device: {device_name}")
-        for idx, example in enumerate(test_dataset):
-            annotation = preprocessed_annotations[idx]
-            has_smoke = len(annotation) > 0
-            true_labels.append(int(has_smoke))
-
-            # Make prediction for the current image
-            results = model.predict(example["image"], device=device_name, conf=THRESHOLD, verbose=False, imgsz=IMGSIZE)[0]
-            
-            pred_has_smoke = len(results) > 0
-            predictions.append(int(pred_has_smoke))
-            
-            # If there's a true box, add it to the list
-            if has_smoke:
-                true_boxes_list.append(annotation)  # True boxes are already preprocessed
+        # If there's a true box, parse it and add 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)
 
-            # Handle prediction boxes: Append first box (or default box if none detected)
-            if results.boxes.cls.numel() != 0:
+            # Append only one bounding box if at least one fire is detected
+            # Note that multiple boxes could be appended
+            if results.boxes.cls.numel()!=0:
                 pred_boxes.append(results.boxes[0].xywhn.tolist()[0])
             else:
-                pred_boxes.append([0, 0, 0, 0])
+                pred_boxes.append([0,0,0,0])
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE