import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
model_path = "recycling-model_transferlearning.keras"
model = tf.keras.models.load_model(model_path)
# Define the core prediction function
def predict_recycling(image):
    # Preprocess image
    print(type(image))
    image = Image.fromarray(image.astype('uint8'))  # Convert numpy array to PIL image
    image = image.resize((150, 150))  # Resize the image to 150x150
    image = np.array(image)
    image = np.expand_dims(image, axis=0)  # Expand dimensions to create batch size of 1

    # Predict
    prediction = model.predict(image)
    
    # Assuming the model's output layer uses softmax activation and there are three outputs
    prediction = prediction.flatten()
    predictions = np.round(prediction, 2)  # Flatten the predictions and round them

    # Separate the probabilities for each class
    p_battery = predictions[0]  # Probability for battery
    p_biological = predictions[1]   # Probability for biological
    p_brownglass = predictions[2]  # Probability for brown-glass
    p_cardboard = predictions[3]  # Probability for cardboard
    p_clothes = predictions[4]  # Probability for clothes
    p_greenglass = predictions[5]  # Probability for green-glass
    p_metal = predictions[6]  # Probability for metal
    p_paper = predictions[7]  # Probability for paper
    p_plastic = predictions[8]  # Probability for plastic
    p_shoes = predictions[9]  # Probability for shoes
    p_whiteglass = predictions[10]  # Probability for white-glass

    return {
        'battery': p_battery,
        'biological': p_biological,
        'brown-glass': p_brownglass,
        'cardboard': p_cardboard,
        'clothes': p_clothes,
        'green-glass': p_greenglass,
        'metal': p_metal,
        'paper': p_paper,
        'plastic': p_plastic,
        'shoes': p_shoes,
        'white-glass': p_whiteglass
    }
# Create the Gradio interface
input_image = gr.Image()
interface = gr.Interface(
    fn=predict_recycling,
    inputs=input_image,
    outputs=gr.Label(),
    examples=["test/battery1.jpg", "test/battery2.jpg", "test/battery3.jpg",
              "test/biological1.jpg", "test/biological2.jpg", "test/biological3.jpg",
              "test/brown-glass1.jpg", "test/brown-glass2.jpg", "test/brown-glass3.jpg",
              "test/cardboard1.jpg", "test/cardboard2.jpg", "test/cardboard3.jpg",
              "test/clothes1.jpg", "test/clothes2.jpg", "test/clothes3.jpg",
              "test/green-glass1.jpg", "test/green-glass2.jpg", "test/green-glass3.jpg",
              "test/metal1.jpg", "test/metal2.jpg", "test/metal3.jpg",
              "test/paper1.jpg", "test/paper2.jpg", "test/paper3.jpg",
              "test/plastic1.jpg", "test/plastic2.jpg", "test/plastic3.jpg",
              "test/shoes1.jpg", "test/shoes2.jpg", "test/shoes3.jpg",
              "test/white-glass1.jpg", "test/white-glass2.jpg", "test/white-glass3.jpg"],
    title="Bildklassifikation für Recycling-Materialien",
    description="Dieses Tool klassifiziert Bilder in verschiedene Recycling-Kategorien. Bitte lade ein Bild hoch, benutze die Kamera oder verwende ein Beispiel von unten, um die Klassifikation zu sehen.",
    theme = gr.themes.Soft(
    primary_hue="emerald",
    secondary_hue="emerald",
    ).set(
        background_fill_primary='*neutral_100'
    ))
interface.launch()