import gradio as gr
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array
import matplotlib.pyplot as plt

# Load the model
model_path = './model-0.9980.hdf5'
try:
    model = load_model(model_path)
except OSError:
    print("Error: Could not load the model. Please check the path.")
    exit()

# Define class labels
class_labels = {
    0: 'Stage = I, Situation = Tumor, Tumour Type = Benign, Tumour Class = Adenoma',
    1: 'Stage = I, Situation = Tumor, Tumour Type = Benign, Tumour Class = Angiomyolipoma',
    2: 'Stage = I, Situation = Tumor, Tumour Type = Benign, Tumour Class = Angiomyolipoma ',
    3: 'Stage = I, Situation = Tumor, Tumour Type = Benign, Tumour Class = Angiomyolipoma and Adenomas',
    4: 'Stage = I, Situation = Tumor, Tumour Type = Benign, Tumour Class = Lipomas',
    5: 'Stage = I, Situation = Tumor, Tumour Type = Malignant, Tumour Class = RCC',
    6: 'Stage = I, Situation = Tumor, Tumour Type = Malignant, Tumour Class = RCC ',
    7: 'Stage = I, Situation = Tumor, Tumour Type = Malignant, Tumour Class = Secondary',
    8: 'Stage = II, Situation = Tumor, Tumour Type = Malignant, Tumour Class = RCC',
    9: 'Stage = II, Situation = Tumor, Tumour Type = Malignant, Tumour Class = Secondary',
    10: 'Stage = III, Situation = Tumor, Tumour Type = Malignant, Tumour Class = RCC',
    11: 'Stage = IV, Situation = Tumor, Tumour Type = Malignant, Tumour Class = Secondary',
    12: 'Stage = Null, Situation = Normal case with cyst, Tumour Type = Null, Tumour Class = Null',
    13: 'Stage = Null, Situation = Normal case, Tumour Type = Null, Tumour Class = Null'
}

def predict_class(image):
    # Convert image to numpy array and preprocess
    img_array = img_to_array(image)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = img_array / 255.0
    
    predictions = model.predict(img_array)
    predicted_class_indices = np.argsort(predictions[0])[::-1][:5]  # Get top 5 indices

    # Create label_confidences dictionary
    label_confidences = {
        class_labels[idx]: float(predictions[0][idx])
        for idx in predicted_class_indices
    }

    # Get top predicted class and probability
    top_predicted_class = max(label_confidences, key=label_confidences.get)
    top_probability = label_confidences[top_predicted_class]
    top_predicted_class = top_predicted_class.replace(","," :: ")

    # Return top class and probability
    return top_predicted_class, top_probability


# Create Gradio interface
input = gr.Image(label="Upload Kidney CT Image")
output = [gr.Label(label="Top Predicted Class", type="text"), gr.Label(label="Probabilities", type="text")]
interface = gr.Interface(
    fn=predict_class,
    inputs=gr.Image(shape=(256,256),label="Upload Kidney CT Image"),
    outputs=[gr.Label(label="Top Predicted Class", type="text"), gr.Label(label="Probabilities", type="text")],
    title="ResNext-Powered Kidney Tumour Detection",
    description="Upload a CT image to classify the tumor and get the top predicted class with its probabilities.",
)

# Launch the interface
interface.launch()