Update script.py

script.py

@@ -4,44 +4,52 @@ from ultralytics import YOLO
 import glob
 
 # --- CONFIGURATION ---
-# Use the exact paths from the baseline script to avoid "File Not Found" errors
-# The server puts images in "/tmp/data/test_images"
+# The evaluation server usually puts test images here
 TEST_IMAGE_PATH = "/tmp/data/test_images"
 SUBMISSION_SAVE_PATH = "submission.csv"
 MODEL_WEIGHTS = "best.pt"
-CONF_THRESHOLD = 0.30
 
-# 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.
+# You can tune this. 0.25 - 0.30 is usually a safe balance for mAP.
+CONF_THRESHOLD = 0.25 
+
 def get_category_id(cls_id):
-    return int(cls_id) + 1 
+    # ❌ OLD INCORRECT LOGIC: return int(cls_id) + 1
+    # ✅ CORRECT LOGIC: Your JSON annotations use IDs 0, 1, 2.
+    # Since your model was trained on 0, 1, 2, we return the class ID directly.
+    return int(cls_id)
 
 def run_inference(model, image_path, conf_threshold, save_path):
     print(f"🚀 Checking for images in {image_path}...")
     
     # 1. Get all images (support multiple extensions)
+    test_images = []
     if os.path.exists(image_path):
-        test_images = [f for f in os.listdir(image_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))]
+        # Recursive search just in case, though usually they are flat in the folder
+        extensions = ['*.png', '*.jpg', '*.jpeg', '*.bmp']
+        for ext in extensions:
+            test_images.extend(glob.glob(os.path.join(image_path, ext)))
+        
+        # Sort to ensure consistent order if needed
         test_images.sort()
+        # Just keep filenames for the loop to match your specific structure requirements
+        test_images = [os.path.basename(x) for x in test_images]
     else:
-        print(f"⚠️ Warning: {image_path} not found. Using current folder for testing.")
-        test_images = []
+        print(f"⚠️ Warning: {image_path} not found. Testing with empty list (or verify path).")
+        # If running locally for debug, you might want to point to a local folder here
+        # test_images = os.listdir(".") 
 
     print(f"🔍 Found {len(test_images)} images.")
 
     # Prepare lists for the dataframe
-    # The baseline wants a specific format: stringified lists of lists
     df_rows = []
 
     # 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)
             
             # Run YOLOv8 inference
+            # verbose=False keeps the logs clean
             results = model.predict(full_path, conf=conf_threshold, verbose=False)
             
             bbox_list = []
@@ -49,34 +57,34 @@ def run_inference(model, image_path, conf_threshold, save_path):
             
             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]
+                    # YOLO format is Center-X, Center-Y, Width, Height
                     x_c, y_c, w, h = box.xywh[0].tolist()
                     
-                    # Convert Center-XY to Top-Left-XY
+                    # Convert to COCO format: Top-Left-X, Top-Left-Y, Width, Height
                     x_min = x_c - (w / 2)
                     y_min = y_c - (h / 2)
                     
-                    # Append to list
+                    # Store as list [x, y, w, h]
                     bbox_list.append([x_min, y_min, w, h])
                     
-                    # Get Class ID
+                    # Get predicted class ID (0, 1, or 2)
                     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)
+            # 3. Format Prediction String
+            # The baseline requires the bbox and category_id columns to be Strings of Lists
             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]"
+                "bbox": str(bbox_list),           # e.g. "[[10.5, 20.1, 50.0, 50.0]]"
+                "category_id": str(category_list) # e.g. "[0]" or "[0, 2]"
             })
 
     # 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
+    # Safety check: if no images found, create an empty CSV with headers to prevent crash
     if df_predictions.empty:
+        print("⚠️ No predictions made (no images found?). Creating empty CSV.")
         df_predictions = pd.DataFrame(columns=["file_name", "bbox", "category_id"])
 
     print(f"💾 Saving {len(df_predictions)} rows to {save_path}...")
@@ -84,11 +92,15 @@ def run_inference(model, image_path, conf_threshold, save_path):
     print("✅ Done!")
 
 if __name__ == "__main__":
-    # Load your YOLOv8 model
-    # Note: We do NOT use torch.hub anymore
+    # Ensure ultralytics is installed via requirements.txt before this runs
     print(f"🔥 Loading YOLOv8 model: {MODEL_WEIGHTS}...")
+    
     try:
+        # Load the trained model
         model = YOLO(MODEL_WEIGHTS)
+        
+        # Execute inference
         run_inference(model, TEST_IMAGE_PATH, CONF_THRESHOLD, SUBMISSION_SAVE_PATH)
+        
     except Exception as e:
-        print(f"❌ Critical Error: {e}")
+        print(f"❌ Critical Error in main execution: {e}")