import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

# 1. Load the trained model
try:
    model = tf.keras.models.load_model("best_dp_cnn_model (1).keras")
except Exception as e:
    print(f"Error loading model: {e}")
    model = None

# 2. Define the real class names
CLASS_NAMES = [
    "Food Waste",
    "Animal Dead Body",
    "Cardboard",
    "Newspaper",
    "Paper Cups",
    "Papers",
    "Brown Glass",
    "Porcelain",
    "Green Glass",
    "White Glass",
    "Beverage Cans",
    "Construction Scrap",
    "Metal Containers",
    "Plastic Bag",
    "Plastic Bottle",
    "Plastic Containers",
    "Plastic Cups",
    "Tetra Pak",
    "Clothes",
    "Shoes",
    "Gloves",
    "Masks",
    "Bandaids",
    "Medicine",
    "Syringe",
    "Diaper",
    "Electrical Cable",
    "Electronic Chip",
    "Laptops",
    "Small Appliances",
    "Smartphones",
    "Battery",
    "Thermometer",
    "Cigarette Butt",
    "Pesticide Bottle",
    "Spray Cans"
]

def predict_image(img):
    if model is None:
        return {"Error: Model not loaded properly": 1.0}
        
    # Resize image to the size the EfficientNetB3 model expects (IMG_SIZE = 300)
    img = img.resize((384, 384))
    
    # Convert to numpy array
    img_array = np.array(img)
    
    # Ensure the image has 3 channels (RGB) in case a grayscale or RGBA image is uploaded
    if len(img_array.shape) == 2:
        img_array = np.stack((img_array,)*3, axis=-1)
    elif img_array.shape[-1] == 4:
        img_array = img_array[..., :3]
        
    # Expand dimensions to create a batch of size 1: (1, 300, 300, 3)
    img_array = np.expand_dims(img_array, axis=0)
    
    # Preprocess the image exactly as done during training
    img_array = tf.keras.applications.efficientnet.preprocess_input(img_array)
    
    # Make prediction
    predictions = model.predict(img_array)[0]
    
    # Create a dictionary of class names and their corresponding probabilities for Gradio
    confidences = {CLASS_NAMES[i]: float(predictions[i]) for i in range(len(CLASS_NAMES))}
    
    return confidences

# 3. Create the Gradio interface
interface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=5), # Shows the top 5 predictions
    title="Waste Classification Model",
    description="Upload an image of waste, and the Resnet50 model will classify it into one of the 36 categories."
)

# 4. Launch the app
if __name__ == "__main__":
    interface.launch()