Create services/fault_service.py

services/fault_service.py

@@ -0,0 +1,34 @@
+from ultralytics import YOLO
+from ultralytics.nn.tasks import DetectionModel
+import os
+import torch.serialization
+
+# Allowlist Ultralytics DetectionModel to avoid UnpicklingError
+torch.serialization.add_safe_globals([DetectionModel])
+
+def load_fault_model():
+    model_path = "pole_fault_model.pt"
+    if os.path.exists(model_path):
+        print(f"Loading custom model: {model_path}")
+        return YOLO(model_path)
+    else:
+        print(f"Warning: {model_path} not found. Falling back to YOLOv8s.")
+        return YOLO("yolov8s.pt")  # Fallback to pre-trained YOLOv8s
+
+fault_model = load_fault_model()
+
+def detect_pole_faults(image_path):
+    try:
+        results = fault_model(image_path)
+        flagged = []
+        for r in results:
+            # Check if model is custom-trained with fault_type (for custom models)
+            if hasattr(r, 'fault_type') and r.conf > 0.6:
+                flagged.append({"fault_type": r.fault_type, "confidence": r.conf})
+            # Fallback for generic YOLOv8 models (no fault_type)
+            elif r.names and r.conf > 0.6:
+                flagged.append({"fault_type": r.names[int(r.cls)], "confidence": r.conf})
+        return flagged
+    except Exception as e:
+        print(f"Error in fault detection: {e}")
+        return []