Update script.py

script.py

@@ -1,70 +1,91 @@
-import torch
-from PIL import Image
-import matplotlib.pyplot as plt
-import cv2
-import numpy as np
 import os
-import pandas as pd
-import ultralytics
-from torchvision import transforms
+import json
+import glob
+from ultralytics import YOLO
 
+# --- 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"
 
+# 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
+}
 
-def run_inference(model, image_path, conf_threshold, save_path):
-
-    test_images = os.listdir(image_path)
-    test_images.sort()
+def main():
+    print(f"🚀 Loading model from {MODEL_PATH}...")
     
-    bboxes = []
-    category_ids = []
-    test_images_names = []
+    # 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}")
+
+    submission_results = []
+
+    # 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
+    )
+
+    print("🏃 Processing predictions...")
     
-    for image in test_images:
-        
-        test_images_names.append(image)
-        bbox = []
-        category_id = []
+    for result in results:
+        # Get filename (e.g., 'frame_001.jpg')
+        file_name = os.path.basename(result.path)
         
-        results = model(os.path.join(image_path, image))
+        # Get filename without extension (often used as image_id)
+        image_id = os.path.splitext(file_name)[0]
         
-        for pred in results.pred[0]:
-            xmin, ymin, xmax, ymax, conf, class_id = pred.tolist()
-            if conf >= conf_threshold:
+        # 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)
             
-                width = xmax - xmin
-                height = ymax - ymin
-
-                bbox.append([xmin, ymin, width, height])
-                category_id.append(int(class_id))
+            # 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)
             
-        bboxes.append(bbox)
-        category_ids.append(category_id)  # Convert class_id to int
-    
-    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])
-        df_predictions = pd.concat([df_predictions, new_row], ignore_index=True)
-        
-    df_predictions.to_csv(save_path, index=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)
 
+    # 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("✅ Done!")
 
 if __name__ == "__main__":
-
-
-    current_directory = os.path.dirname(os.path.abspath(__file__))
-    # print(current_directory)
-    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")
-    CONF_THRESHOLD = 0.30
-    
-    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)
+    main()