from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch

# Load model from Hugging Face (replace with your actual model)
MODEL_NAME = "your-hf-username/solar-panel-fault-detector"

processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
model.eval()

def predict_fault(image: Image.Image) -> str:
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    label = model.config.id2label[predicted_class_idx]
    return label