Update modules/thermal_fault_detection.py

modules/thermal_fault_detection.py

@@ -1,35 +1,23 @@
-from transformers import DetrImageProcessor, DetrForObjectDetection
-import torch
+from ultralytics import YOLO
 from PIL import Image
+import numpy as np
 
-# Load processor and model
-processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
-model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+# Load the YOLOv8 model (make sure 'solar_fault.pt' exists in your root directory)
+model = YOLO("solar_fault.pt")
 
-CUSTOM_CLASSES = ["Overheating", "Dust", "Breakage"]
+def detect_faults(image, threshold=0.3):
+    # Convert image to numpy array
+    results = model.predict(np.array(image), conf=threshold)
+    detections = []
 
-def detect_faults(image, threshold=0.1):
-    # Preprocess the image
-    inputs = processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
+    for r in results:
+        for box in r.boxes:
+            label_id = int(box.cls[0])
+            confidence = float(box.conf[0])
+            x1, y1, x2, y2 = map(int, box.xyxy[0])
+            label = model.names[label_id]
+            detections.append((label, confidence, (x1, y1, x2, y2)))
 
-    # ✅ Define probs from logits
-    probs = outputs.logits.softmax(-1)[0]         # [num_predictions, num_classes]
-    boxes = outputs.pred_boxes[0]                 # [num_predictions, 4]
+    return detections
 
-    results = []
-
-    for i, prob in enumerate(probs):
-        label_id = prob.argmax().item()
-        confidence = prob[label_id].item()
-
-        # Get label safely
-        label = model.config.id2label.get(label_id, None)
-        if label is None or confidence < threshold:
-            continue
-
-        if label in CUSTOM_CLASSES:
-            box = boxes[i].tolist()
-            results.append((label, confidence, box))
-
-    return results
+CUSTOM_CLASSES = ["crack", "burn", "hotspot", "dust"]