Update script.py

script.py

@@ -1,91 +1,94 @@
 import os
-import json
-import glob
+import pandas as pd
 from ultralytics import YOLO
+import glob
 
 # --- CONFIGURATION ---
-# These paths are standard for most competition containers.
-# If testing locally, change 'TEST_IMAGES_DIR' to your local test folder.
-MODEL_PATH = "best.pt"  # The model must be in the same folder as this script
-TEST_IMAGES_DIR = "test/images" # The competition usually puts images here
-OUTPUT_FILE = "submission.json"
+# Use the exact paths from the baseline script to avoid "File Not Found" errors
+# The server puts images in "/tmp/data/test_images"
+TEST_IMAGE_PATH = "/tmp/data/test_images"
+SUBMISSION_SAVE_PATH = "submission.csv"
+MODEL_WEIGHTS = "best.pt"
+CONF_THRESHOLD = 0.30
 
-# Phase 2 Category Mapping (Ensure these match your training!)
-# 0: Large Needle Driver, 1: Prograsp Forceps, 2: Monopolar Curved Scissors
-# We map YOLO ID -> Category ID required by submission (usually 1, 2, 3 or same)
-# If your competition expects IDs 1, 2, 3, use this map:
-ID_MAPPING = {
-    0: 1,  # Large Needle Driver
-    1: 2,  # Prograsp Forceps
-    2: 3   # Monopolar Curved Scissors
-}
+# Mapping: YOLO ID (0,1,2) -> Competition Category ID (likely 1,2,3)
+# If your previous training used 0,1,2, usually competitions want 1-based IDs.
+# We will add +1 to be safe, matching standard Phase 2 rules.
+def get_category_id(cls_id):
+    return int(cls_id) + 1 
 
-def main():
-    print(f"🚀 Loading model from {MODEL_PATH}...")
+def run_inference(model, image_path, conf_threshold, save_path):
+    print(f"🚀 Checking for images in {image_path}...")
     
-    # 1. Load the trained YOLOv8 model
-    try:
-        model = YOLO(MODEL_PATH)
-    except Exception as e:
-        print(f"❌ Error loading model: {e}")
-        return
-
-    # 2. Get all test images
-    # We look for jpg, png, and jpeg
-    image_paths = glob.glob(os.path.join(TEST_IMAGES_DIR, "*.*"))
-    print(f"🔍 Found {len(image_paths)} images in {TEST_IMAGES_DIR}")
+    # 1. Get all images (support multiple extensions)
+    if os.path.exists(image_path):
+        test_images = [f for f in os.listdir(image_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))]
+        test_images.sort()
+    else:
+        print(f"⚠️ Warning: {image_path} not found. Using current folder for testing.")
+        test_images = []
 
-    submission_results = []
+    print(f"🔍 Found {len(test_images)} images.")
 
-    # 3. Run Inference
-    # stream=True prevents crashing on large datasets
-    results = model.predict(
-        source=TEST_IMAGES_DIR, 
-        conf=0.25,      # Confidence threshold (0.25 is standard for mAP)
-        iou=0.45,       # NMS IoU threshold
-        save=False,     # Don't save plotted images (saves time)
-        save_txt=False, 
-        verbose=False,
-        stream=True
-    )
+    # Prepare lists for the dataframe
+    # The baseline wants a specific format: stringified lists of lists
+    df_rows = []
 
-    print("🏃 Processing predictions...")
-    
-    for result in results:
-        # Get filename (e.g., 'frame_001.jpg')
-        file_name = os.path.basename(result.path)
-        
-        # Get filename without extension (often used as image_id)
-        image_id = os.path.splitext(file_name)[0]
-        
-        # Loop through detections
-        for box in result.boxes:
-            # Get data
-            cls_id = int(box.cls[0])
-            score = float(box.conf[0])
-            bbox = box.xywh[0].tolist() # x_center, y_center, width, height (Normalized? No, usually pixels)
+    # 2. Run Inference
+    if len(test_images) > 0:
+        # Load images one by one or in batches
+        # We loop to match the baseline structure exactly
+        for image_name in test_images:
+            full_path = os.path.join(image_path, image_name)
             
-            # YOLOv8 .xywh returns pixels: [x_center, y_center, width, height]
-            # MTEC/COCO usually wants: [x_min, y_min, width, height]
-            x_c, y_c, w, h = bbox
-            x_min = x_c - (w / 2)
-            y_min = y_c - (h / 2)
+            # Run YOLOv8 inference
+            results = model.predict(full_path, conf=conf_threshold, verbose=False)
             
-            # Create annotation entry
-            annotation = {
-                "image_id": file_name,      # Or use image_id variable depending on rules
-                "category_id": ID_MAPPING.get(cls_id, cls_id + 1), # Map to correct ID
-                "bbox": [x_min, y_min, w, h],
-                "score": score
-            }
-            submission_results.append(annotation)
+            bbox_list = []
+            category_list = []
+            
+            for result in results:
+                for box in result.boxes:
+                    # Get box in xywh format (x, y, width, height)
+                    # The baseline calculated width = xmax - xmin, so it wants [x, y, w, h]
+                    x_c, y_c, w, h = box.xywh[0].tolist()
+                    
+                    # Convert Center-XY to Top-Left-XY
+                    x_min = x_c - (w / 2)
+                    y_min = y_c - (h / 2)
+                    
+                    # Append to list
+                    bbox_list.append([x_min, y_min, w, h])
+                    
+                    # Get Class ID
+                    cls_id = int(box.cls[0])
+                    category_list.append(get_category_id(cls_id))
+
+            # 3. Format exactly like the baseline
+            # It expects columns: file_name, bbox (string), category_id (string)
+            df_rows.append({
+                "file_name": image_name,
+                "bbox": str(bbox_list),           # e.g. "[[10, 20, 50, 50]]"
+                "category_id": str(category_list) # e.g. "[1]"
+            })
+
+    # 4. Create DataFrame and Save
+    df_predictions = pd.DataFrame(df_rows, columns=["file_name", "bbox", "category_id"])
+    
+    # Safety check: if empty, create empty CSV with correct headers
+    if df_predictions.empty:
+        df_predictions = pd.DataFrame(columns=["file_name", "bbox", "category_id"])
 
-    # 4. Save to JSON
-    print(f"💾 Saving {len(submission_results)} detections to {OUTPUT_FILE}...")
-    with open(OUTPUT_FILE, 'w') as f:
-        json.dump(submission_results, f, indent=4)
-        
+    print(f"💾 Saving {len(df_predictions)} rows to {save_path}...")
+    df_predictions.to_csv(save_path, index=False)
     print("✅ Done!")
 
 if __name__ == "__main__":
-    main()
+    # Load your YOLOv8 model
+    # Note: We do NOT use torch.hub anymore
+    print(f"🔥 Loading YOLOv8 model: {MODEL_WEIGHTS}...")
+    try:
+        model = YOLO(MODEL_WEIGHTS)
+        run_inference(model, TEST_IMAGE_PATH, CONF_THRESHOLD, SUBMISSION_SAVE_PATH)
+    except Exception as e:
+        print(f"❌ Critical Error: {e}")