from ultralytics import YOLO
from PIL import Image
import numpy as np

# Load the YOLOv8 model (make sure 'solar_fault.pt' exists in your root directory)
model = YOLO("solar_fault.pt")

def detect_faults(image, threshold=0.3):
    # Convert image to numpy array
    results = model.predict(np.array(image), conf=threshold)
    detections = []

    for r in results:
        for box in r.boxes:
            label_id = int(box.cls[0])
            confidence = float(box.conf[0])
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            label = model.names[label_id]
            detections.append((label, confidence, (x1, y1, x2, y2)))

    return detections

CUSTOM_CLASSES = ["crack", "burn", "hotspot", "dust"]