Update services/thermal_service.py

services/thermal_service.py

@@ -1,14 +1,12 @@
-import os
 import torch
+import os
 from ultralytics import YOLO
-from torch.serialization import add_safe_globals
-
-# Import all necessary classes explicitly
+from torch.serialization import add_safe_globals, safe_load
+import torch.nn.modules.container as container
 from ultralytics.nn.tasks import DetectionModel
 from ultralytics.nn.modules import Conv
-import torch.nn.modules.container as container
 
-# ✅ Register all trusted classes
+# ✅ Register all necessary classes
 add_safe_globals({
     container.Sequential: "torch.nn.modules.container.Sequential",
     container.ModuleList: "torch.nn.modules.container.ModuleList",
@@ -17,33 +15,45 @@ 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:
     """
-    Safely loads a YOLO model with necessary PyTorch 2.6+ safe globals registration
-    and ensures auto-download if missing.
+    Custom safe load for YOLO models trained before torch 2.6 weight-only enforcement.
     """
     if not os.path.isfile(model_path):
-        print(f"[INFO] Model {model_path} not found locally. Auto-downloading...")
+        print(f"[INFO] Downloading {model_path}...")
+        # Will auto-download internally by Ultralytics YOLO
+
     try:
         model = YOLO(model_path)
-        print(f"[INFO] YOLO model {model_path} loaded successfully.")
         return model
     except Exception as e:
-        print(f"[ERROR] Could not load YOLO model: {e}")
-        raise
+        # 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
 
-# ✅ Load YOLO model globally
+# ✅ Initialize model
 thermal_model = load_yolo_model_safely()
 
-def detect_thermal_anomalies(image_path: str):
+def detect_thermal_anomalies(image_path):
     """
-    Detects thermal anomalies in a given image frame using YOLO.
+    Detect anomalies using YOLO model.
     """
     results = thermal_model(image_path)
     flagged = []
     for r in results:
         for box in r.boxes:
-            # Simulate thermal detection by confidence threshold
             if box.conf > 0.7:
                 flagged.append({
                     "confidence": float(box.conf),