import gradio as gr
import tensorflow as tf
import tensorflow_hub as hub
from PIL import Image
import numpy as np
from tensorflow.keras.preprocessing.image import img_to_array, load_img

# Function to load the model with custom objects
def load_model_with_hub(model_path):
    # Load the model architecture without weights
    model = tf.keras.models.load_model(model_path, compile=False)

    # Define the KerasLayer from TensorFlow Hub
    keras_layer = hub.KerasLayer("https://tfhub.dev/google/imagenet/resnet_v2_101/feature_vector/5", trainable=False)

    # Add the KerasLayer to the model
    model.add(keras_layer)

    return model

# Loading saved model with custom objects
model = load_model_with_hub('model_cat_dog.h5')

def predict(input_image):
    try:
        # Convert PIL Image to Numpy array
        input_image = img_to_array(input_image)
        # Resize the Numpy array
        input_image = np.resize(input_image, (224, 224, 3))
        input_image = np.array(input_image).astype(np.float32) / 255.0
        input_image = np.expand_dims(input_image, axis=0) 


        # Making prediction
        prediction = model.predict(input_image)

        # Postprocess prediction
        labels = ['Cat', 'Dog']
        threshold = 0.5  # threshold for classifying as 'Dog'
        predicted_class = 'Dog' if prediction[0] > threshold else 'Cat'
        prediction_probability = prediction[0] if predicted_class == 'Dog' else 1 - prediction[0]

        cat_emoji = "\U0001F408"  # Cat emoji
        dog_emoji = "\U0001F415"  # Dog emoji

        selected_emoji = dog_emoji if predicted_class == 'Dog' else cat_emoji

        # Combine the predicted class and the probability into a single string
        output = f"{selected_emoji} {predicted_class}"

        return output
    except Exception as e:
        return str(e)

examples = ["dog1.jpeg",
            "cat1.jpg"]

# Creating Gradio interface
iface = gr.Interface(
    fn=predict, 
    inputs=gr.inputs.Image(shape=(224, 224)), 
    outputs="text",
    title = 'Image Recognition - Cats vs Dogs',
    examples = examples
)

iface.launch()