import gradio as gr
import numpy as np
from PIL import Image
import tensorflow as tf
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.efficientnet import preprocess_input

# Load the trained model
model = load_model("efficent_net224B0.h5")

# Define the classes
waste_labels = {0: 'Fibres', 1: 'Nanowires', 2: 'Particles', 3: 'Powder'}

# Define the Gradio interface
def classify_image(pil_image):
    # Convert PIL.Image to Numpy array
    img = image.img_to_array(pil_image)
    
    # Resize to the model's expected input size
    img = tf.image.resize(img, (224, 224))
    
    # Expand dimensions to create a batch size of 1
    img = np.expand_dims(img, axis=0)
    
    # Preprocess the input for the EfficientNet model
    img = preprocess_input(img)
    
    # Make prediction
    prediction = model.predict(img)
    
    # Get predicted class and confidence
    predicted_class = np.argmax(prediction)
    predicted_class = waste_labels[predicted_class]
    confidence = prediction[0, np.argmax(prediction)]

    return predicted_class

# Create the Gradio interface
iface = gr.Interface(fn=classify_image, inputs="image", outputs="text")

# Launch the Gradio interface
iface.launch()