# -*- coding: utf-8 -*-
import gradio as gr
import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
import cv2

# Define a dictionary to map the custom layer to its implementation
custom_objects = {'KerasLayer': hub.KerasLayer}

# Load your model (ensure the path to the model file is correct)
model = tf.keras.models.load_model('bird_model4.h5', custom_objects=custom_objects)

# Load the class labels from the text file
train_info = []
with open('labelwithspace.txt', 'r') as file:
    train_info = [line.strip() for line in file.readlines()]

# Function to preprocess the image and make predictions
def predict_image(image):
    # Resize the image to the model's input size
    img = cv2.resize(image, (224, 224))
    img = img / 255.0  # Normalize the image
    # Make predictions using the model
    predictions = model.predict(img[np.newaxis, ...])[0]
    top_classes = np.argsort(predictions)[-3:][::-1]  # Get indices of top 3 predictions
    top_class = top_classes[0]  # Index of the top prediction
    label = train_info[top_class]  # Retrieve the label using the index
    return label

# Define the Gradio interface
input_image = gr.inputs.Image(shape=(224, 224))
output_label = gr.outputs.Label()

# Launch the Gradio app
gr.Interface(fn=predict_image, inputs=input_image, outputs=output_label, capture_session=True).launch()