Create modules/thermal_fault_detection.py

modules/thermal_fault_detection.py

@@ -0,0 +1,33 @@
+# modules/thermal_fault_detection.py
+from transformers import DetrImageProcessor, DetrForObjectDetection
+from PIL import Image
+import torch
+
+# Load processor and model
+processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+
+# Define class mapping for your use case
+CUSTOM_CLASSES = {
+    "overheat": ["person", "fire hydrant"],
+    "dust": ["bird", "sheep"],
+    "breakage": ["bench", "truck"]
+}
+
+def detect_faults(image: Image.Image, threshold: float = 0.7):
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=threshold)[0]
+
+    faults = []
+
+    for score, label in zip(results["scores"], results["labels"]):
+        class_name = model.config.id2label[label.item()].lower()
+        for fault_type, tags in CUSTOM_CLASSES.items():
+            if class_name in tags:
+                faults.append((fault_type.capitalize(), score.item()))
+                break
+    
+    return faults