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

# Ensure TensorFlow does not allocate any GPU
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

# Define the data augmentation pipeline
data_augmentation = tf.keras.Sequential([
    tf.keras.layers.RandomFlip("horizontal"),
    tf.keras.layers.RandomRotation(0.2),
    tf.keras.layers.RandomZoom(0.2),
    tf.keras.layers.RandomHeight(0.2),
    tf.keras.layers.RandomWidth(0.2),
], name="data_augmentation")

# Load your trained model
model_path = 'garbage-classification.h5'
model = tf.keras.models.load_model(model_path, custom_objects={'data_augmentation': data_augmentation})
class_names = ['battery', 'biological', 'cardboard', 'clothes', 'glass', 'metal', 'paper', 'plastic', 'shoes', 'trash']
IMG_SIZE = (400, 400)  # replace with your image size, same as used during training

def classify_image(image):
    img = Image.fromarray(image.astype('uint8'), 'RGB')
    img = img.resize(IMG_SIZE)
    
    # Convert image to tensor
    img_tensor = tf.convert_to_tensor(img)
    img_tensor = tf.cast(img_tensor, tf.float32)  # Ensure float32 cast if not already
    
    # Expand dimensions to match the model's expected input
    img_tensor = tf.expand_dims(img_tensor, axis=0)

    # Make prediction
    predictions = model.predict(img_tensor)
    predicted_class = class_names[np.argmax(predictions)]
    probability = float(np.max(predictions))
    
    return predicted_class, probability

# Create a Gradio interface
iface = gr.Interface(fn=classify_image,
                     inputs=gr.Image(label="Upload an Image"),
                     outputs=[gr.Label(num_top_classes=1, label="Prediction"),
                              gr.Textbox(label="Probability")],
                     title="Garbage Classification",
                     description="Upload an image of garbage, and the model will classify it.")

iface.launch(share=True)