import matplotlib.pyplot as plt
import numpy as np
import PIL
import tensorflow as tf
import gradio as gr

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential

import pathlib

from PIL import Image

model = tf.keras.models.load_model('modelA.keras')


img_height = 240
img_width = 240
class_names = ['Healthy', 'Maize Lethal Necrosis', 'Maize Streak Virus']
num_classes = len(class_names)

def classify_image(img):
    # Convert the NumPy array to a PIL Image if necessary
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img)

    # Resize the image
    img = img.resize((img_height, img_width))

    # Convert the PIL Image to a NumPy array for further processing
    img_array = np.array(img)

    # Adding the batch dimension
    img_array = tf.expand_dims(img_array, 0)  # Shape becomes (1, img_height, img_width, 3)

    # Make a prediction
    predictions = model.predict(img_array)
    score = tf.nn.softmax(predictions[0])

    # Return a dictionary of class names and their confidence scores
    return {class_names[i]: float(score[i]) for i in range(num_classes)}


# Set up the Gradio interface
iface = gr.Interface(fn=classify_image, inputs="image", outputs="label")


# Launch the interface
iface.launch()