Update services/services/fault_service.py

services/services/fault_service.py

@@ -1,11 +1,33 @@
 from ultralytics import YOLO
+import os
+import torch.serialization
 
-fault_model = YOLO("pole_fault_model.pt")
+# Allowlist Ultralytics DetectionModel to avoid UnpicklingError
+torch.serialization.add_safe_globals(['ultralytics.nn.tasks.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):
-    results = fault_model(image_path)
-    flagged = []
-    for r in results:
-        if hasattr(r, 'fault_type') and r.conf > 0.6:
-            flagged.append({"fault_type": r.fault_type, "confidence": r.conf})
-    return flagged
+    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 []