from ultralytics import YOLO
import os

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 []