Update modules/thermal_fault_detection.py

modules/thermal_fault_detection.py

@@ -3,59 +3,62 @@ from torchvision.models.detection import maskrcnn_resnet50_fpn
 from torchvision.transforms import functional as F
 from PIL import Image, ImageDraw, ImageFont
 
-# Load model - assuming you use a pretrained COCO version
+# Load pretrained COCO model
 def load_model():
-    model = maskrcnn_resnet50_fpn(pretrained=True)  # or load your fine-tuned weights
+    model = maskrcnn_resnet50_fpn(pretrained=True)
     model.eval()
     return model
 
 model = load_model()
 
-# Map COCO classes to your custom fault categories
+# Map COCO classes to your fault types
 CLASS_MAPPING = {
     "overheat": ["person", "fire hydrant"],
     "dust": ["bird", "sheep"],
     "breakage": ["bench", "truck"]
 }
 
+# COCO id-to-name mapping
+COCO_LABELS = {
+    1: "person",
+    10: "fire hydrant",
+    16: "bird",
+    20: "sheep",
+    14: "bench",
+    8: "truck",
+    # Add more if needed
+}
+
 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()
+    for fault_type, classes in CLASS_MAPPING.items():
+        if coco_class_name in classes:
+            return fault_type.capitalize()
     return None
 
 def detect_faults(image: Image.Image, threshold: float = 0.7):
-    image_tensor = F.to_tensor(image).unsqueeze(0)  # Convert to tensor and add batch dimension
+    image_tensor = F.to_tensor(image).unsqueeze(0)  # Convert to batch tensor
 
     with torch.no_grad():
-        output = model(image_tensor)[0]
+        outputs = model(image_tensor)[0]  # Get first (and only) result
 
     draw = ImageDraw.Draw(image)
     font = ImageFont.load_default()
     results = []
 
-    for score, label, box in zip(output["scores"], output["labels"], output["boxes"]):
+    for score, label, box in zip(outputs["scores"], outputs["labels"], outputs["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
+            # Draw bounding 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)
+            draw.rectangle(box, outline="red", width=2)
+            draw.text((box[0], box[1] - 10), f"{fault_type} ({score:.2f})", fill="red", font=font)
 
     return results, image