Update services/thermal_service.py

services/thermal_service.py

@@ -1,48 +1,52 @@
-import torch
 import os
+import torch
 from ultralytics import YOLO
 from torch.serialization import add_safe_globals
+
+# Import all necessary classes explicitly
 from ultralytics.nn.tasks import DetectionModel
+from ultralytics.nn.modules import Conv
 import torch.nn.modules.container as container
 
-# ✅ Register all required globals
+# ✅ Register all trusted classes
 add_safe_globals({
     container.Sequential: "torch.nn.modules.container.Sequential",
     container.ModuleList: "torch.nn.modules.container.ModuleList",
     container.ModuleDict: "torch.nn.modules.container.ModuleDict",
-    DetectionModel: "ultralytics.nn.tasks.DetectionModel"
+    DetectionModel: "ultralytics.nn.tasks.DetectionModel",
+    Conv: "ultralytics.nn.modules.Conv"
 })
 
 def load_yolo_model_safely(model_path: str = 'yolov8n.pt') -> YOLO:
     """
-    Safely loads a YOLO model with necessary PyTorch 2.6+ fixes
-    and ensures the model file exists.
+    Safely loads a YOLO model with necessary PyTorch 2.6+ safe globals registration
+    and ensures auto-download if missing.
     """
-    # Auto-download yolov8n.pt if not present
     if not os.path.isfile(model_path):
-        print(f"[INFO] Downloading {model_path}...")
-        # Ultralytics automatically handles download inside YOLO if path missing
-        return YOLO(model_path)
-
+        print(f"[INFO] Model {model_path} not found locally. Auto-downloading...")
     try:
         model = YOLO(model_path)
+        print(f"[INFO] YOLO model {model_path} loaded successfully.")
         return model
     except Exception as e:
-        print(f"[ERROR] Failed loading YOLO model: {e}")
+        print(f"[ERROR] Could not load YOLO model: {e}")
         raise
 
-# ✅ Load model once at module load
+# ✅ Load YOLO model globally
 thermal_model = load_yolo_model_safely()
 
-def detect_thermal_anomalies(image_path):
+def detect_thermal_anomalies(image_path: str):
     """
-    Detects thermal-like anomalies (placeholder simulation).
+    Detects thermal anomalies in a given image frame using YOLO.
     """
     results = thermal_model(image_path)
     flagged = []
     for r in results:
         for box in r.boxes:
-            # Simulate flagging if confidence > 70%
+            # Simulate thermal detection by confidence threshold
             if box.conf > 0.7:
-                flagged.append(box)
+                flagged.append({
+                    "confidence": float(box.conf),
+                    "bbox": box.xyxy.tolist() if hasattr(box, 'xyxy') else []
+                })
     return flagged