Update services/thermal_service.py

services/thermal_service.py

@@ -1,31 +1,29 @@
-from transformers import pipeline
+# services/thermal_service.py
+
+from ultralytics import YOLO
 import cv2
 
-# Load DETR model from Hugging Face
-object_detector = pipeline("object-detection", model="facebook/detr-resnet-50")
+# ✅ Load a pre-trained YOLO model (small model)
+thermal_model = YOLO('yolov8n.pt')  # assuming model auto-download
 
 def detect_thermal_anomalies(image_path):
     """
-    Simulate thermal anomalies using DETR object detection.
-    If an object has a high confidence and occupies a big bounding box, simulate a "hotspot."
+    Detect thermal anomalies using YOLOv8 model.
+    - Accepts a file path to an image.
+    - Processes it through YOLO.
+    - Flags boxes above a confidence threshold.
     """
+    # ✅ Read image properly using OpenCV
     image = cv2.imread(image_path)
-    results = object_detector(image)
 
-    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
+    # ✅ Perform detection
+    results = thermal_model.predict(source=image, save=False, conf=0.5, imgsz=640)
 
-        # 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
-            })
+    flagged = []
+    for r in results:
+        if hasattr(r, 'boxes'):
+            for box in r.boxes:
+                if box.conf > 0.7:
+                    flagged.append(box)
 
-    return anomalies
+    return flagged