import gradio as gr
import tensorflow as tf
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess_input, decode_predictions
from PIL import Image
import numpy as np

# Load pre-trained ResNet50 model + higher level layers
model = ResNet50(weights='imagenet')


#chameleon = load_img("example_1.jpeg", output_type="pil")


def classify_image(img):
    # Resize the image to 224x224 pixels (required input size for ResNet50)
    img = img.resize((224, 224))
    # Convert the image to array
    img_array = np.array(img)
    # Expand dimensions to match the shape required by the model
    img_array = np.expand_dims(img_array, axis=0)
    # Preprocess the image
    img_array = preprocess_input(img_array)
    # Predict the classification
    predictions = model.predict(img_array)
    # Decode predictions into readable labels
    decoded_predictions = decode_predictions(predictions, top=3)[0]
    
    # Format the output
    return {label: float(probability) for (_, label, probability) in decoded_predictions}

with gr.Blocks(theme="Hev832/Applio") as iface:
    gr.Interface(
        fn=classify_image,  # Function to call for predictions
        inputs=gr.Image(type="pil"),  # Input is an image        
        outputs=gr.Label(num_top_classes=3),  # Output is a label with top 3 predictions
        title="Contextual Image Classification",
        description="Upload an image, and the model will classify it based on the context.",
        #examples=[chameleon],
    )

# Launch the interface
iface.launch()