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

# Load your Keras model
model = load_model('vege_classifier_model.h5')

# Define the updated classes
classes = [
    'Bean', 'Bitter_Gourd', 'Bottle_Gourd', 'Brinjal', 'Broccoli', 
    'Cabbage', 'Capsicum', 'Carrot', 'Cauliflower', 'Cucumber', 
    'Papaya', 'Potato', 'Pumpkin', 'Radish', 'Tomato'
]

# Define a prediction function
def predict_image(img):
    # Preprocess the image to fit the model's input requirements
    img = img.resize((224, 224))  # Match the model's expected input size
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)  # Create a batch
    img_array /= 255.0  # Rescale the image to [0, 1] to match training preprocessing
    
    # Predict with your model
    predictions = model.predict(img_array)
    predicted_class_index = np.argmax(predictions, axis=1)
    
    # Assuming 'classes' is a list of class names in the order they are represented in the model
    return classes[predicted_class_index[0]]

# Create a simplified Gradio interface
iface = gr.Interface(fn=predict_image, 
                     inputs="image", 
                     outputs="label", 
                     title="Vegetable Image Classifier",
                     description="Classify images of various vegetables into 15 categories: Bean, Bitter Gourd, Bottle Gourd, Brinjal, Broccoli, Cabbage, Capsicum, Carrot, Cauliflower, Cucumber, Papaya, Potato, Pumpkin, Radish, Tomato.")

if __name__ == "__main__":
    iface.launch()