import gradio as gr
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

# Load model (ignore compile warning — you're only predicting)
model = load_model("waste_classification(Mobilenetv2).h5", compile=False)
class_names = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']

# Prediction function: outputs label + chart
def predict_with_chart(image):
    if image is None:
        return "No image received", None

    image = image.resize((224, 224))
    img_array = img_to_array(image) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    prediction = model.predict(img_array)[0]
    pred_index = np.argmax(prediction)
    pred_label = class_names[pred_index]
    confidence = float(np.max(prediction))

    # Create bar chart
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.bar(class_names, prediction, color='skyblue')
    ax.set_ylabel('Probability')
    ax.set_ylim(0, 1)
    ax.set_title('Class Probabilities')
    plt.xticks(rotation=45)
    plt.tight_layout()

    return f"Prediction: {pred_label} ({confidence*100:.1f}%)", fig

# Gradio Interface (Gradio 4.x compatible)
with gr.Blocks() as demo:
    gr.Markdown("## 🗑️ Waste Classifier — Upload or Capture an Image")
    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload or Webcam (click camera icon)")
    with gr.Row():
        label_output = gr.Textbox(label="Predicted Class")
        plot_output = gr.Plot(label="Class Probability Chart")

    image_input.change(fn=predict_with_chart, inputs=image_input, outputs=[label_output, plot_output])

demo.launch()