from transformers import AutoModelForObjectDetection, AutoProcessor
import torch

# Load the processor for the model
processor = AutoProcessor.from_pretrained("ultralytics/yolov5s")  # Ensure this matches the model you're using

def detect_faults_solar(model, image):
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)

    results = outputs.logits.detach().cpu().numpy()
    faults = []
    for result in results:
        class_id = int(result[5])  # Get the class ID from the model's output
        if class_id == 0:  # Cracks (example)
            faults.append({"type": "Crack", "location": (result[0], result[1])})
        elif class_id == 1:  # Damage
            faults.append({"type": "Damage", "location": (result[0], result[1])})

    return faults

def detect_faults_windmill(model, image):
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)

    results = outputs.logits.detach().cpu().numpy()
    faults = []
    for result in results:
        class_id = int(result[5])  # Get the class ID from the model's output
        if class_id == 2:  # Blade damage (example)
            faults.append({"type": "Blade Damage", "location": (result[0], result[1])})
        elif class_id == 3:  # Motor fault
            faults.append({"type": "Motor Fault", "location": (result[0], result[1])})

    return faults