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

# Define the class names
class_names = ['AnnualCrop', 'Forest', 'HerbaceousVegetation', 'Highway', 'Industrial', 'Pasture', 'PermanentCrop', 'Residential', 'River', 'SeaLake']

# Function to load the selected model
def load_selected_model(model_name):
    if model_name == "ResNet50":
        return load_model('resnet.h5')
    elif model_name == "CNN":
        return load_model('cnn.h5')
    else:
        raise ValueError("Invalid model name")

# Function to predict the image
def predict_image(image, model_name):
    model = load_selected_model(model_name)
    img = cv2.resize(image, (64, 64))
    img = img / 255.0
    img = np.expand_dims(img, axis=0)

    predictions = model.predict(img)[0]

    predicted_class_index = np.argmax(predictions)
    predicted_class_name = class_names[predicted_class_index]

    confidence = predictions[predicted_class_index]

    # Get the top 3 predictions
    top_indices = predictions.argsort()[-3:][::-1]
    top_classes = [class_names[i] for i in top_indices]
    top_confidences = [predictions[i] for i in top_indices]

    # Format the output
    result = "Top Predictions:\n"
    for i in range(3):
        result += f"{top_classes[i]} (Confidence: {top_confidences[i]:.2f})\n"

    return result

with gr.Blocks() as demo:
    gr.Markdown("# Satellite Image Classifier (CaseOh edition)")

    with gr.Row():
        with gr.Column():
            model_selector = gr.Dropdown(choices=["ResNet50", "CNN"], label="Select Model")
            image_input = gr.Image(type="numpy", label="Upload Image")
            predict_button = gr.Button("Predict")

        image_output = gr.Image(type="numpy", label="Uploaded Image")

    prediction_output = gr.Textbox(label="Prediction")

    predict_button.click(fn=predict_image, inputs=[image_input, model_selector], outputs=prediction_output)
    image_input.change(fn=lambda x: x, inputs=image_input, outputs=image_output)

demo.launch(share=True)