import gradio as gr
import numpy as np
from PIL import Image
import tensorflow as tf
from tensorflow.keras.applications.resnet import preprocess_input
from tensorflow.keras.models import load_model

# Load the model
model = load_model("best_model.h5")

# Class names
class_names = ['Cloudy', 'Rain', 'Shine', 'Sunrise']

# Preprocessing function
def preprocess_image(img):
    img = img.resize((224, 224))
    img_array = np.array(img)
    img_array = preprocess_input(img_array)
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

# Prediction function
def classify_image(image):
    processed_img = preprocess_image(image)
    preds = model.predict(processed_img)[0]
    predicted_class = class_names[np.argmax(preds)]
    confidence = float(np.max(preds))
    return {predicted_class: confidence}

# Gradio Interface
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=4),
    title="Weather Image Classifier",
    description="Upload an image of the weather and get the predicted category (Cloudy, Rain, Shine, Sunrise)"
)

if __name__ == "__main__":
    interface.launch()