Update services/thermal_service.py

services/thermal_service.py

@@ -1,45 +1,31 @@
-import os
-import torch
-from ultralytics import YOLO
-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
+from transformers import pipeline
+import cv2
 
-# ✅ Register all necessary safe 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",
-    Conv: "ultralytics.nn.modules.Conv"
-})
+# Load DETR model from Hugging Face
+object_detector = pipeline("object-detection", model="facebook/detr-resnet-50")
 
-def load_yolo_model_safely():
+def detect_thermal_anomalies(image_path):
     """
-    Use direct pretrained YOLOv8n model from Ultralytics Hub (no local weights download needed).
+    Simulate thermal anomalies using DETR object detection.
+    If an object has a high confidence and occupies a big bounding box, simulate a "hotspot."
     """
-    try:
-        model = YOLO('yolov8n.pt')  # pretrained small model directly from Ultralytics hub
-        print("[INFO] YOLOv8 model loaded successfully.")
-        return model
-    except Exception as e:
-        print(f"[ERROR] Failed to load YOLO model: {e}")
-        raise
+    image = cv2.imread(image_path)
+    results = object_detector(image)
 
-thermal_model = load_yolo_model_safely()
+    anomalies = []
+    for result in results:
+        score = result.get('score', 0)
+        box = result.get('box', {})
+        width = box.get('width', 0)
+        height = box.get('height', 0)
+        area = width * height
 
-def detect_thermal_anomalies(image_path):
-    """
-    Detect anomalies in an image using the loaded YOLO model.
-    """
-    results = thermal_model(image_path)
-    flagged = []
-    for r in results:
-        for box in r.boxes:
-            if box.conf > 0.7:
-                flagged.append({
-                    "confidence": float(box.conf),
-                    "bbox": box.xyxy.tolist() if hasattr(box, 'xyxy') else []
-                })
-    return flagged
+        # Simulate: large bounding boxes + high score = thermal anomaly
+        if score > 0.7 and area > 5000:  # adjust thresholds if needed
+            anomalies.append({
+                "label": result.get('label', ''),
+                "score": score,
+                "bbox": box
+            })
+
+    return anomalies