Update services/thermal_service.py

services/thermal_service.py

@@ -1,12 +1,12 @@
-import torch
 import os
+import torch
 from ultralytics import YOLO
-from torch.serialization import add_safe_globals, safe_load
+from torch.serialization import add_safe_globals
 import torch.nn.modules.container as container
 from ultralytics.nn.tasks import DetectionModel
 from ultralytics.nn.modules import Conv
 
-# ✅ Register all necessary classes
+# ✅ Register all necessary safe classes
 add_safe_globals({
     container.Sequential: "torch.nn.modules.container.Sequential",
     container.ModuleList: "torch.nn.modules.container.ModuleList",
@@ -15,40 +15,23 @@ add_safe_globals({
     Conv: "ultralytics.nn.modules.Conv"
 })
 
-def custom_safe_load(filepath):
-    """
-    Force torch to load YOLO weights without weights_only=True limitation.
-    """
-    with open(filepath, 'rb') as f:
-        return torch.load(f, map_location='cpu', weights_only=False)
-
-def load_yolo_model_safely(model_path: str = 'yolov8n.pt') -> YOLO:
+def load_yolo_model_safely():
     """
-    Custom safe load for YOLO models trained before torch 2.6 weight-only enforcement.
+    Use direct pretrained YOLOv8n model from Ultralytics Hub (no local weights download needed).
     """
-    if not os.path.isfile(model_path):
-        print(f"[INFO] Downloading {model_path}...")
-        # Will auto-download internally by Ultralytics YOLO
-
     try:
-        model = YOLO(model_path)
+        model = YOLO('yolov8n.pt')  # pretrained small model directly from Ultralytics hub
+        print("[INFO] YOLOv8 model loaded successfully.")
         return model
     except Exception as e:
-        # If default load fails, force fallback load
-        print(f"[WARNING] Normal YOLO load failed: {e}")
-        print(f"[INFO] Trying manual safe load...")
-
-        # Manual fallback load
-        weights = custom_safe_load(model_path)
-        model = YOLO(model=weights)  # Load model from raw weights
-        return model
+        print(f"[ERROR] Failed to load YOLO model: {e}")
+        raise
 
-# ✅ Initialize model
 thermal_model = load_yolo_model_safely()
 
 def detect_thermal_anomalies(image_path):
     """
-    Detect anomalies using YOLO model.
+    Detect anomalies in an image using the loaded YOLO model.
     """
     results = thermal_model(image_path)
     flagged = []