import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Load the trained model
model = tf.keras.models.load_model("MobileNet_model.h5")  # Ensure the model file is in the same directory

# Define class names from your dataset
class_names = ["Fake", "Low", "Medium", "High"]  # Update based on test_generator.class_indices.keys()

# Image Preprocessing
img_size = (128, 128)  # Same as used in test_generator

def preprocess_image(image):
    image = image.resize(img_size)  # Resize to (128,128)
    image = np.array(image) / 255.0  # Normalize as done in ImageDataGenerator (rescale=1./255)
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    return image

# Prediction Function
def predict(image):
    image = preprocess_image(image)
    predictions = model.predict(image)
    predicted_class = np.argmax(predictions, axis=1)[0]  # Get the predicted class index
    confidence_scores = {class_names[i]: float(predictions[0][i]) for i in range(len(class_names))}  # Get probability scores
    
    return {"Predicted Class": class_names[predicted_class], "Confidence Scores": confidence_scores}

# Gradio Interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(),  # Returns class and confidence scores
    title="Waste Classification Model",
    description="Upload an image to classify it into one of four categories: Fake, Low, Medium, or High."
)

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