# Load the model
model = tf.keras.models.load_model('pokemon_classifier_model.keras')

def predict(image):
    img = tf.keras.preprocessing.image.img_to_array(image)
    img = tf.keras.preprocessing.image.smart_resize(img, (224, 224))
    img = tf.expand_dims(img, 0)  # Make batch of one
    
    pred = model.predict(img)
    pred_label = tf.argmax(pred, axis=1).numpy()[0]  # get the index of the max logit
    pred_class = class_names[pred_label]  # use the index to get the corresponding class name
    confidence = tf.nn.softmax(pred)[0][pred_label]  # softmax to get the confidence
    
    print(f"Predicted: {pred_class}, Confidence: {confidence:.4f}")
    return pred_class


# Setup Gradio interface
iface = gr.Interface(fn=predict, inputs=gr.Image(), outputs="text", title="Pokémon Classifier")

# Run the interface
iface.launch()