Optimize submission

tasks/image.py

@@ -116,52 +116,42 @@ async def evaluate_image(request: ImageEvaluationRequest):
     model = YOLO(Path(model_path, model_name), task="detect")
     device_name = device("cuda" if is_available() else "cpu")
     
-    # Preprocess annotations before the loop
-    preprocessed_annotations = [parse_boxes(example.get("annotations", "").strip()) for example in test_dataset]
-
-    batch_size = 16  # Define a batch size
-    batch_images = []
-    batch_annotations = []
-
     predictions = []
     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}")
+   # 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
+        
+        logging.info(f"Inference start on device: {device_name}")
         for idx, example in enumerate(test_dataset):
-            batch_images.append(example["image"])
-            batch_annotations.append(preprocessed_annotations[idx])
-
-            # When the batch size is met, or it's the last image, perform inference
-            if (len(batch_images) == batch_size or idx == len(test_dataset) - 1):
-                # Make a prediction for the current batch
-                results = model.predict(batch_images, device=device_name, conf=THRESHOLD, verbose=False, half=True, imgsz=IMGSIZE)[0]
-
-                for batch_idx, result in enumerate(results):
-                    annotation = batch_annotations[batch_idx]
-                    has_smoke = len(annotation) > 0
-                    true_labels.append(int(has_smoke))
-
-                    pred_has_smoke = len(result) > 0
-                    predictions.append(int(pred_has_smoke))
-
-                    if has_smoke:
-                        true_boxes_list.append(annotation)
-
-                    # Handle prediction boxes for each image in the batch
-                    if result.boxes.cls.numel() != 0:
-                        pred_boxes.append(result.boxes[0].xywhn.tolist()[0])
-                    else:
-                        pred_boxes.append([0, 0, 0, 0])
-
-                # Clear the batch after processing
-                batch_images.clear()
-                batch_annotations.clear()
-
+            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
+
+            # Handle prediction boxes: Append first box (or default box if none detected)
+            if results.boxes.cls.numel() != 0:
+                pred_boxes.append(results.boxes[0].xywhn.tolist()[0])
+            else:
+                pred_boxes.append([0, 0, 0, 0])
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE