Update gully_drs_core/model_utils.py

gully_drs_core/model_utils.py

@@ -1,17 +1,6 @@
-import torch
-import logging
+# gully_drs_core/model_utils.py
+from ultralytics import YOLO
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-def load_yolo_model():
-    try:
-        # Load YOLOv5s model (use pre-trained or fine-tuned weights)
-        model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
-        model.eval()
-        logger.info("YOLOv5 model loaded successfully")
-        return model
-    except Exception as e:
-        logger.error(f"Error loading YOLOv5 model: {str(e)}")
-        raise
+def load_model():
+    model = YOLO('yolov5s.pt')  # Replace with custom model if available
+    return model