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

print(tf.__version__)

print(f"Current working directory: {os.getcwd()}")

model_path = 'transferlearning_pokemon.h5'

# Check if the model exists
if os.path.exists(model_path):
    print(f"Model found at {model_path}")
    try:
        # Load the trained model
        model = tf.keras.models.load_model(model_path)
        print("Model loaded successfully.")
    except Exception as e:
        print(f"Error loading model: {e}")
else:
    print(f"Model not found at {model_path}. Please check the path.")


# Define class names (make sure this matches the classes used during training)
class_names = ['Machamp', 'Raichu', 'Vulpix']

# Define the prediction function
def predict(image):
    image = Image.fromarray(image.astype('uint8'))  # Convert numpy array to PIL image
    image = image.resize((150, 150)) #resize the image to 28x28 and converts it to gray scale
    image = np.array(image)
    image = np.expand_dims(image, axis=0) # same as image[None, ...]
    predictions = model.predict(image)
    predicted_class = np.argmax(predictions, axis=1)[0]
    confidence = np.max(predictions)
    return {class_names[predicted_class]: float(confidence)}

# Create a Gradio interfac
input_image = gr.Image()
output_text = gr.Textbox(label="Predicted Value")
interface = gr.Interface(fn=predict, 
                         inputs=input_image, 
                         outputs=gr.Label(),
                         examples=["00000000.jpg", "00000001.jpg", "00000010.png", "00000017.jpg", "00000021.jpg", "00000067.jpg"],   
                         description="A simple mlp classification model for image classification using the mnist dataset.")

# Launch the Gradio interface
interface.launch()