Update script.py

script.py

@@ -1,79 +1,79 @@
-import os
-import pandas as pd
-from ultralytics import YOLO
-
-def run_inference(model, image_path, conf_threshold, save_path):
-    test_images = os.listdir(image_path)
-    test_images.sort()
-    
-    bboxes = []
-    category_ids = []
-    test_images_names = []
-    
-    # Iterate through images for inference
-    for image_name in test_images:
-        # Skip any non-image files if they exist in the directory
-        if not image_name.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
-            continue
-
-        full_image_path = os.path.join(image_path, image_name)
-        
-        current_image_bboxes = []
-        current_image_category_ids = []
-        
-        # Perform inference
-        results = model(full_image_path)
-        
-        # Process results
-        # results is a list of Results objects, one for each image
-        # Since we pass one image at a time, results[0] is the relevant object
-        for pred in results[0].boxes:
-            # Bounding box in xyxy format, confidence, class_id
-            xmin, ymin, xmax, ymax = pred.xyxy[0].tolist()
-            conf = pred.conf.item()
-            class_id = int(pred.cls.item())
-
-            if conf >= conf_threshold:
-                width = xmax - xmin
-                height = ymax - ymin
-
-                current_image_bboxes.append([xmin, ymin, width, height])
-                current_image_category_ids.append(class_id)
-            
-        test_images_names.append(image_name)
-        bboxes.append(current_image_bboxes)
-        category_ids.append(current_image_category_ids)
-    
-    # Create DataFrame for predictions
-    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)
-        
-    # Ensure the save directory exists
-    os.makedirs(os.path.dirname(save_path), exist_ok=True)
-    df_predictions.to_csv(save_path, index=False)
-    print(f"Inference results saved to: {save_path}")
-
-
-if __name__ == "__main__":
-    # Define paths
-    # You might need to change TEST_IMAGE_PATH to where your actual test images are stored
-    current_directory = os.path.dirname(os.path.abspath(__file__))
-    TEST_IMAGE_PATH = "/tmp/data/test_images"
-    SUBMISSION_SAVE_PATH = os.path.join(current_directory, "submission.csv")" 
-    
-    # Path to your trained model weights
-    MODEL_WEIGHTS_PATH = os.path.join(current_directory, "best.pt")
-    CONF_THRESHOLD = 0.30 # Confidence threshold for predictions
-    
-    # Load the YOLO model
-    model = YOLO(MODEL_WEIGHTS_PATH) # Using ultralytics.YOLO for loading
-
-    # Run inference
-    run_inference(model, TEST_IMAGE_PATH, CONF_THRESHOLD, SUBMISSION_SAVE_PATH)
+import os
+import pandas as pd
+from ultralytics import YOLO
+
+def run_inference(model, image_path, conf_threshold, save_path):
+    test_images = os.listdir(image_path)
+    test_images.sort()
+    
+    bboxes = []
+    category_ids = []
+    test_images_names = []
+    
+    # Iterate through images for inference
+    for image_name in test_images:
+        # Skip any non-image files if they exist in the directory
+        if not image_name.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
+            continue
+
+        full_image_path = os.path.join(image_path, image_name)
+        
+        current_image_bboxes = []
+        current_image_category_ids = []
+        
+        # Perform inference
+        results = model(full_image_path)
+        
+        # Process results
+        # results is a list of Results objects, one for each image
+        # Since we pass one image at a time, results[0] is the relevant object
+        for pred in results[0].boxes:
+            # Bounding box in xyxy format, confidence, class_id
+            xmin, ymin, xmax, ymax = pred.xyxy[0].tolist()
+            conf = pred.conf.item()
+            class_id = int(pred.cls.item())
+
+            if conf >= conf_threshold:
+                width = xmax - xmin
+                height = ymax - ymin
+
+                current_image_bboxes.append([xmin, ymin, width, height])
+                current_image_category_ids.append(class_id)
+            
+        test_images_names.append(image_name)
+        bboxes.append(current_image_bboxes)
+        category_ids.append(current_image_category_ids)
+    
+    # Create DataFrame for predictions
+    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)
+        
+    # Ensure the save directory exists
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    df_predictions.to_csv(save_path, index=False)
+    print(f"Inference results saved to: {save_path}")
+
+
+if __name__ == "__main__":
+    # Define paths
+    # You might need to change TEST_IMAGE_PATH to where your actual test images are stored
+    current_directory = os.path.dirname(os.path.abspath(__file__))
+    TEST_IMAGE_PATH = "/tmp/data/test_images"
+    SUBMISSION_SAVE_PATH = os.path.join(current_directory, "submission.csv") 
+    
+    # Path to your trained model weights
+    MODEL_WEIGHTS_PATH = os.path.join(current_directory, "best.pt")
+    CONF_THRESHOLD = 0.30 # Confidence threshold for predictions
+    
+    # Load the YOLO model
+    model = YOLO(MODEL_WEIGHTS_PATH) # Using ultralytics.YOLO for loading
+
+    # Run inference
+    run_inference(model, TEST_IMAGE_PATH, CONF_THRESHOLD, SUBMISSION_SAVE_PATH)