Update services/thermal_service.py

services/thermal_service.py

@@ -1,17 +1,29 @@
 # services/thermal_service.py
 
-from ultralytics import YOLO
+import cv2
+from transformers import pipeline
 
-# ✅ Load directly from the Ultralytics model hub (no .pt file)
-thermal_model = YOLO('yolov8n')  # Smallest, fast model
+# ✅ Use DETR model for thermal detection as well
+thermal_detector = pipeline("object-detection", model="facebook/detr-resnet-50")
 
 def detect_thermal_anomalies(image_path):
-    results = thermal_model.predict(image_path)
-    flagged_boxes = []
+    image = cv2.imread(image_path)
+    results = thermal_detector(image)
 
-    for r in results:
-        for box in r.boxes:
-            if box.conf > 0.7:
-                flagged_boxes.append(box)
+    thermal_boxes = []
+    for result in results:
+        if result['score'] > 0.7:
+            box = result['box']
+            thermal_boxes.append(box)
 
-    return flagged_boxes
+    return thermal_boxes
+
+def overlay_thermal_boxes(image_path, boxes):
+    image = cv2.imread(image_path)
+
+    for box in boxes:
+        xmin, ymin, xmax, ymax = box['xmin'], box['ymin'], box['xmax'], box['ymax']
+        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2)
+        cv2.putText(image, 'Hotspot', (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
+
+    return image