import gradio as gr
import numpy as np
from tensorflow.keras.models import load_model
from PIL import Image

# Load the trained model
model = load_model("m32.h5")

# Class labels
class_labels = [
    'Cell', 'Cell-Multi', 'Cracking', 'Diode', 'Diode-Multi',
    'Hot-Spot', 'Hot-Spot-Multi', 'No-Anomaly', 'Offline-Module',
    'Shadowing', 'Soiling', 'Vegetation'
]

# Preprocessing and prediction function
def predict_image(img: Image.Image):
    img = img.convert("RGB")
    img = img.resize((24, 40))  # Resize to (width=24, height=40)
    img_array = np.array(img)
    img_input = img_array.reshape(1, 40, 24, 3)

    prediction = model.predict(img_input)
    max_prob = np.max(prediction[0])
    class_ind = np.argmax(prediction[0])
    predicted_class = class_labels[class_ind]

    return f"Predicted class: {predicted_class} with probability {max_prob:.4f}"

# Gradio interface
interface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Solar Module Anomaly Classifier",
    description="Upload an infrared image to detect anomalies in solar modules."
)

interface.launch()