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

# Load the saved model
model = load_model('fruit_veg_classifier.h5')

# Create a class mapping (based on your dataset)
class_names = list(train_generator.class_indices.keys())

# Define the prediction function
def classify_image(image):
    # Preprocess the image
    image = image.resize((150, 150))
    image = img_to_array(image) / 255.0
    image = np.expand_dims(image, axis=0)
    
    # Make prediction
    predictions = model.predict(image)
    predicted_class = np.argmax(predictions)
    return class_names[predicted_class]

# Create the Gradio interface
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.inputs.Image(shape=(150, 150)),
    outputs=gr.outputs.Label(),
    title="Fruit & Vegetable Classifier",
    description="Upload an image of a fruit or vegetable, and the model will predict what it is!"
)

# Launch the app
interface.launch()