Upload script.py

script.py

@@ -0,0 +1,78 @@
+import os
+import torch
+import pandas as pd
+from rfdetr import RFDETRBase
+
+def run_inference(model, image_path, conf_threshold, save_path):
+
+    test_images = sorted([
+        f for f in os.listdir(image_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}")
+        test_images_names.append(image_name)
+
+        image_file = os.path.join(image_path, image_name)
+
+        # Run prediction
+        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 score >= conf_threshold:
+                    xmin, ymin, xmax, ymax = 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}.")
+
+        bboxes.append(image_bboxes)
+        category_ids.append(image_categories)
+
+    # 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]),
+        ]
+
+    df_predictions.to_csv(save_path, index=False)
+    print(f"\nInference complete. Predictions saved to {save_path}")
+
+
+if __name__ == "__main__":
+
+    TEST_IMAGE_PATH = r"rf-detr\dataset\test"
+    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)