Update submission/script.py

submission/script.py

@@ -1,6 +1,7 @@
 import os
 import torch
 import pandas as pd
+from PIL import Image
 from rfdetr import RFDETRBase
 
 
@@ -17,12 +18,15 @@ def run_inference(model, image_path, conf_threshold, save_path):
 
     for image_name in test_images:
         test_images_names.append(image_name)
-
         image_file = os.path.join(image_path, image_name)
 
         bbox = []
         category_id = []
 
+        # Load image to get dimensions (IMPORTANT)
+        with Image.open(image_file) as img:
+            img_w, img_h = img.size
+
         preds = model.predict(image_file)
 
         if preds is not None and preds.xyxy is not None and len(preds.xyxy) > 0:
@@ -32,14 +36,26 @@ def run_inference(model, image_path, conf_threshold, save_path):
                 preds.class_id
             ):
                 score = float(score)
-                if score >= conf_threshold:
-                    xmin, ymin, xmax, ymax = map(float, box)
+                if score < conf_threshold:
+                    continue
+
+                xmin, ymin, xmax, ymax = map(float, box)
+
+                # ---- CLAMP TO IMAGE BOUNDARIES ----
+                xmin = max(0.0, xmin)
+                ymin = max(0.0, ymin)
+                xmax = min(float(img_w), xmax)
+                ymax = min(float(img_h), ymax)
+
+                width = xmax - xmin
+                height = ymax - ymin
 
-                    width = xmax - xmin
-                    height = ymax - ymin
+                # ---- FILTER INVALID BOXES ----
+                if width <= 0 or height <= 0:
+                    continue
 
-                    bbox.append([xmin, ymin, width, height])
-                    category_id.append(int(label))
+                bbox.append([xmin, ymin, width, height])
+                category_id.append(int(label))
 
         bboxes.append(bbox)
         category_ids.append(category_id)
@@ -56,13 +72,14 @@ def run_inference(model, image_path, conf_threshold, save_path):
         df_predictions = pd.concat([df_predictions, new_row], ignore_index=True)
 
     df_predictions.to_csv(save_path, index=False)
+    print(f"Submission saved to {save_path}")
 
 
 if __name__ == "__main__":
 
     TEST_IMAGE_PATH = "/tmp/data/test_images"
     SUBMISSION_SAVE_PATH = "submission.csv"
-    CONF_THRESHOLD = 0.30
+    CONF_THRESHOLD = 0.30  # you may lower to 0.15 if recall is poor
 
     model = RFDETRBase(
         checkpoint_path="checkpoint_best_total.pth",