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

# Load the pre-trained model
model = tf.keras.models.load_model('pokemon_transferlearning.keras')

def classify_image(img):
    
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img.astype('uint8'), 'RGB')

    # Preprocess the image to fit the model's input requirements
    img = img.resize((150, 150))  # Resize the image using PIL, which is intended here
    img_array = np.array(img)  # Convert the resized PIL image to a numpy array
    img_array = img_array / 255.0  # Normalize pixel values to [0, 1]
    img_array = np.expand_dims(img_array, axis=0)  # Expand dimensions to fit model input shape

    # Make prediction
    prediction = model.predict(img_array)

    # prediction = np.round(float(tf.sigmoid(prediction)), 2)
    # p_cat = (1 - prediction)
    # p_dog = prediction
    # return {'cat': p_cat, 'dog': p_dog}

    print(prediction)

    probabilities = tf.sigmoid(prediction).numpy()  # Convert tensor to numpy array if using 

    # Formatting the probabilities
    class_names = ['Hitchoman', 'Pikachu', 'Charmeleon']
    results = {class_names[i]: float(prediction[0][i]) for i in range(3)}  # Convert each probability to float
    
    return results

# Create Gradio interface
iface = gr.Interface(fn=classify_image,
                     inputs=gr.Image(),
                     outputs=gr.Label(num_top_classes=3),
                     title="Pokemon Classifier",
                     description="Upload an image of a pokemon classify.")

# Launch the application
iface.launch(share=True)