Update modules/thermal_fault_detection.py

modules/thermal_fault_detection.py

@@ -1,33 +1,61 @@
-# modules/thermal_fault_detection.py
-from transformers import DetrImageProcessor, DetrForObjectDetection
-from PIL import Image
 import torch
+from torchvision.models.detection import maskrcnn_resnet50_fpn
+from torchvision.transforms import functional as F
+from PIL import Image, ImageDraw, ImageFont
 
-# Load processor and model
-processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
-model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+# Load model - assuming you use a pretrained COCO version
+def load_model():
+    model = maskrcnn_resnet50_fpn(pretrained=True)  # or load your fine-tuned weights
+    model.eval()
+    return model
 
-# Define class mapping for your use case
-CUSTOM_CLASSES = {
+model = load_model()
+
+# Map COCO classes to your custom fault categories
+CLASS_MAPPING = {
     "overheat": ["person", "fire hydrant"],
     "dust": ["bird", "sheep"],
     "breakage": ["bench", "truck"]
 }
 
+def get_fault_label(coco_class_name: str):
+    for fault, coco_aliases in CLASS_MAPPING.items():
+        if coco_class_name in coco_aliases:
+            return fault.capitalize()
+    return None
+
 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
+    image_tensor = F.to_tensor(image).unsqueeze(0)  # Convert to tensor and add batch dimension
+
+    with torch.no_grad():
+        output = model(image_tensor)[0]
+
+    draw = ImageDraw.Draw(image)
+    font = ImageFont.load_default()
+    results = []
+
+    for score, label, box in zip(output["scores"], output["labels"], output["boxes"]):
+        if score < threshold:
+            continue
+
+        label_id = label.item()
+        coco_label = model.coco_labels[label_id] if hasattr(model, "coco_labels") else label_id
+        class_name = model.coco_labels.get(label_id, f"class_{label_id}") if hasattr(model, "coco_labels") else f"class_{label_id}"
+        
+        # COCO id2label mapping from torchvision
+        COCO_LABELS = {
+            1: "person", 10: "fire hydrant", 16: "bird", 20: "sheep", 14: "bench", 8: "truck",
+            # Add more if needed
+        }
+        class_name = COCO_LABELS.get(label_id, f"class_{label_id}")
+        fault_type = get_fault_label(class_name)
+
+        if fault_type:
+            results.append((fault_type, score.item()))
+
+            # Draw box
+            box = box.tolist()
+            draw.rectangle(box, outline="red", width=3)
+            draw.text((box[0], box[1] - 10), f"{fault_type}: {score:.2f}", fill="red", font=font)
+
+    return results, image