Update script.py

script.py

@@ -3,6 +3,7 @@ import torch
 import pandas as pd
 from rfdetr import RFDETRBase
 
+
 def run_inference(model, image_path, conf_threshold, save_path):
 
     test_images = sorted([
@@ -10,56 +11,51 @@ def run_inference(model, image_path, conf_threshold, save_path):
         if f.lower().endswith((".jpg", ".jpeg", ".png"))
     ])
 
-    print(f"Found {len(test_images)} images for inference.")
-
     bboxes = []
     category_ids = []
     test_images_names = []
 
-    for idx, image_name in enumerate(test_images):
-        print(f"\nProcessing image {idx+1}/{len(test_images)}: {image_name}")
+    for image_name in test_images:
         test_images_names.append(image_name)
 
         image_file = os.path.join(image_path, image_name)
 
-        # Run prediction
+        bbox = []
+        category_id = []
+
         preds = model.predict(image_file)
 
-        # Debug print
-        print(f"Raw prediction output: {preds}")
-
-        # Handle empty predictions
-        if preds is None or preds.xyxy is None or len(preds.xyxy) == 0:
-            print("No predictions returned for this image.")
-            image_bboxes = []
-            image_categories = []
-        else:
-            image_bboxes = []
-            image_categories = []
-            for box, score, label in zip(preds.xyxy, preds.confidence, preds.class_id):
+        if preds is not None and preds.xyxy is not None and len(preds.xyxy) > 0:
+            for box, score, label in zip(
+                preds.xyxy,
+                preds.confidence,
+                preds.class_id
+            ):
+                score = float(score)
                 if score >= conf_threshold:
-                    xmin, ymin, xmax, ymax = box
+                    xmin, ymin, xmax, ymax = map(float, box)
+
                     width = xmax - xmin
                     height = ymax - ymin
-                    image_bboxes.append([float(xmin), float(ymin), float(width), float(height)])
-                    image_categories.append(int(label))
 
-        print(f"Detected {len(image_bboxes)} objects above threshold {conf_threshold}.")
+                    bbox.append([xmin, ymin, width, height])
+                    category_id.append(int(label))
 
-        bboxes.append(image_bboxes)
-        category_ids.append(image_categories)
+        bboxes.append(bbox)
+        category_ids.append(category_id)
 
-    # Prepare DataFrame
     df_predictions = pd.DataFrame(columns=["file_name", "bbox", "category_id"])
+
     for i in range(len(test_images_names)):
-        df_predictions.loc[i] = [
-            test_images_names[i],
-            str(bboxes[i]),
-            str(category_ids[i]),
-        ]
+        new_row = pd.DataFrame({
+            "file_name": test_images_names[i],
+            "bbox": str(bboxes[i]),
+            "category_id": str(category_ids[i])
+        }, index=[0])
+
+        df_predictions = pd.concat([df_predictions, new_row], ignore_index=True)
 
     df_predictions.to_csv(save_path, index=False)
-    print(f"\nInference complete. Predictions saved to {save_path}")
 
 
 if __name__ == "__main__":
@@ -68,11 +64,9 @@ if __name__ == "__main__":
     SUBMISSION_SAVE_PATH = "submission.csv"
     CONF_THRESHOLD = 0.30
 
-    print("Loading RF-DETR model...")
     model = RFDETRBase(
         checkpoint_path="checkpoint_best_ema.pth",
         device="cuda" if torch.cuda.is_available() else "cpu"
     )
 
-    print("Starting inference...")
     run_inference(model, TEST_IMAGE_PATH, CONF_THRESHOLD, SUBMISSION_SAVE_PATH)