Upload script.py

script.py

@@ -8,8 +8,6 @@ import pandas as pd
 import ultralytics
 from torchvision import transforms
 
-
-
 def run_inference(model, image_path, conf_threshold, save_path):
 
     test_images = os.listdir(image_path)
@@ -19,52 +17,94 @@ def run_inference(model, image_path, conf_threshold, save_path):
     category_ids = []
     test_images_names = []
     
+    print(f"Starting inference on {len(test_images)} images...")
+    
     for image in test_images:
         
+        # Skip non-image files if any
+        if not (image.endswith('.jpg') or image.endswith('.png')):
+            continue
+            
         test_images_names.append(image)
         bbox = []
         category_id = []
         
+        # Run inference
         results = model(os.path.join(image_path, image))
         
+        # Parse detections
         for pred in results.pred[0]:
             xmin, ymin, xmax, ymax, conf, class_id = pred.tolist()
-            if conf >= conf_threshold:
             
+            if conf >= conf_threshold:
+                
+                # --- CALCULATE WIDTH & HEIGHT ---
                 width = xmax - xmin
                 height = ymax - ymin
 
-                bbox.append([xmin, ymin, width, height])
-                category_id.append(int(class_id))
+                # --- FORMATTING FIXES FOR LEADERBOARD ---
+                # Round coordinates to avoid messy floats (e.g. 100.0 instead of 100.43242)
+                final_box = [round(xmin, 1), round(ymin, 1), round(width, 1), round(height, 1)]
+                
+                # --- CRITICAL FIX: SHIFT CLASS ID ---
+                # YOLO predicts 0, 1, 2.
+                # Your JSON categories are 1 (Needle), 2 (Grasper), 3 (Scissors).
+                # We simply add +1 to align them.
+                final_id = int(class_id) + 1 
+
+                bbox.append(final_box)
+                category_id.append(final_id)
             
         bboxes.append(bbox)
-        category_ids.append(category_id)  # Convert class_id to int
+        category_ids.append(category_id)
     
+    # Create Submission DataFrame
     df_predictions = pd.DataFrame(columns=["file_name", "bbox", "category_id"])
     
     for i in range(len(test_images_names)):
         file_name = test_images_names[i]
-        new_row = pd.DataFrame({"file_name": file_name,
-                                "bbox": str(bboxes[i]),
-                                "category_id": str(category_ids[i]),
-                                }, index=[0])
+        
+        # Format the list exactly how standard submission files look
+        # e.g. "[[10, 10, 50, 50]]" instead of just list objects
+        str_bbox = str(bboxes[i])
+        str_cat_id = str(category_ids[i])
+        
+        new_row = pd.DataFrame({
+            "file_name": file_name,
+            "bbox": str_bbox,
+            "category_id": str_cat_id
+        }, index=[0])
+        
         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__":
 
-
     current_directory = os.path.dirname(os.path.abspath(__file__))
-    # print(current_directory)
+    
+    # Default paths (These usually don't change on the server)
     TEST_IMAGE_PATH = "/tmp/data/test_images"
     SUBMISSION_SAVE_PATH = os.path.join(current_directory, "submission.csv")
     
-    RUN_NAME = "instrument_detection"
-    MODEL_WEIGHTS_PATH = os.path.join(current_directory, "yolov5", "runs", "train", RUN_NAME, "weights", "best.pt")
+    # --- MODEL LOADING CONFIG ---
+    # We look for 'best.pt' in the SAME directory as this script
+    MODEL_WEIGHTS_PATH = os.path.join(current_directory, "best.pt")
+    
+    # Increase threshold slightly to avoid garbage predictions
     CONF_THRESHOLD = 0.30
     
-    model = torch.hub.load(os.path.join(current_directory, 'yolov5'), 'custom', path=MODEL_WEIGHTS_PATH, source="local")
+    print(f"Loading model from: {MODEL_WEIGHTS_PATH}")
+    
+    # Load YOLOv5 from the local folder to ensure it works offline
+    # 'source=local' means it looks for the 'yolov5' folder in this directory
+    model = torch.hub.load(
+        os.path.join(current_directory, 'yolov5'), 
+        'custom', 
+        path=MODEL_WEIGHTS_PATH, 
+        source='local'
+    )
     
     run_inference(model, TEST_IMAGE_PATH, CONF_THRESHOLD, SUBMISSION_SAVE_PATH)