import gradio as gr
import tensorflow as tf
from tensorflow.keras import models, layers


# Define model architecture
image_size=256
channels=3
input_shape = (None, image_size, image_size, channels)
n_classes = 3
#preprocessing during model creation RESCALING & RESIZING
resize_and_rescale = tf.keras.Sequential([
    layers.Resizing(image_size, image_size),
    layers.Rescaling(1.0/255)
])

#DATA AUGMENTATION
data_augmentation = tf.keras.Sequential([
    layers.RandomFlip("horizontal_and_vertical"),
    layers.RandomRotation(0.2)
])
model = models.Sequential([
    resize_and_rescale,
    data_augmentation,
    layers.Conv2D(64, kernel_size=3, activation='relu', input_shape=input_shape),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(64, kernel_size=(3, 3), activation='relu'),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(64, kernel_size=(3, 3), activation='relu'),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(64, kernel_size=(3, 3), activation='relu'),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(64, kernel_size=(3, 3), activation='relu'),
    layers.MaxPooling2D((2, 2)),
    layers.Flatten(),
    layers.Dense(64, activation='relu'),
    layers.Dense(n_classes, activation='softmax')
])

# Load pre-trained weights
model.load_weights('model911.h5')

# Function to make predictions
def classify_image(image):
    # Preprocess image if necessary
    # Make prediction
    prediction = model.predict(image)
    classes = ['Potato___Early_blight', 'Potato___Late_blight', 'Potato___healthy']
    return {classes[i]: float(prediction[0][i]) for i in range(len(classes))}


# Input component
inputs = gr.Image(shape=(image_size, image_size,channels))

# Output component
outputs = gr.outputs.Label(num_top_classes=3)

# Create Gradio interface
gr.Interface(fn=classify_image, inputs=inputs, outputs=outputs, title='Potato Plant Diseases Classifier').launch()