Update detect_bubbles.py

detect_bubbles.py

@@ -1,28 +1,19 @@
 from ultralytics import YOLO
-import torch
 
-torch.serialization.add_safe_globals([YOLO, 'DetectionModel'])
-
-def safe_load_model(model_path):
-    """
-    Manually loads the YOLO model weights with 'weights_only' set to False
-    to avoid the security check issue with PyTorch.
-    """
-    model = YOLO(model_path, weights_only=False)
-    return model
 
 def detect_bubbles(model_path, image_path):
     """
     Detects bubbles in an image using a YOLOv8 model.
+
     Args:
         model_path (str): The file path to the YOLO model.
         image_path (str): The file path to the input image.
+
     Returns:
-        list: A list containing the coordinates, score, and class_id of 
+        list: A list containing the coordinates, score and class_id of 
               the detected bubbles.
     """
-    model = safe_load_model(model_path) 
-    results = model(image_path)
-    bubbles = results[0].boxes.data.tolist()
+    model = YOLO(model_path)
+    bubbles = model(image_path)[0]
 
-    return bubbles
+    return bubbles.boxes.data.tolist()