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

# Load the saved models
cnn_model = load_model('cnn_image_classifier.h5')
resnet_model = load_model('resnet_image_classifier.h5')

# Load label map (assuming you have it defined or loaded elsewhere)
# Replace with your actual loading method
with open('EuroSAT/label_map.json', 'r') as f:
    label_map = json.load(f)
label_map_inv = {v: k for k, v in label_map.items()}


def predict_image(image):
    # Preprocess the image (resize, normalize, etc.)
    image = cv2.resize(image, (64, 64))
    image = image / 255.0
    image = np.expand_dims(image, axis=0)

    # Make predictions with both models
    cnn_pred = cnn_model.predict(image)[0]
    resnet_pred = resnet_model.predict(image)[0]

    # Get top 5 predictions for CNN
    cnn_top5_indices = np.argsort(cnn_pred)[::-1][:5]
    cnn_top5 = {
        label_map_inv[idx]: float(cnn_pred[idx]) for idx in cnn_top5_indices
    }

    # Get top 5 predictions for ResNet
    resnet_top5_indices = np.argsort(resnet_pred)[::-1][:5]
    resnet_top5 = {
        label_map_inv[idx]: float(resnet_pred[idx]) for idx in resnet_top5_indices
    }

    # Final predictions
    cnn_final_prediction = label_map_inv[np.argmax(cnn_pred)]
    resnet_final_prediction = label_map_inv[np.argmax(resnet_pred)]

    return cnn_top5, cnn_final_prediction, resnet_top5, resnet_final_prediction


iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="numpy"),
    outputs=[
        gr.Label(num_top_classes=5, label="CNN Top 5 Predictions"),
        gr.Textbox(label="CNN Final Prediction"),
        gr.Label(num_top_classes=5, label="ResNet Top 5 Predictions"),
        gr.Textbox(label="ResNet Final Prediction"),
    ],
    title="Image Classification with CNN and ResNet",
    description="Upload an image to classify using two different models.",
)

iface.launch(debug=True)