import tensorflow as tf
from tensorflow.keras import models, layers
import gradio as gr
import numpy as np

# Define and name the model
bat_classifier_model = models.Sequential(name="BatClassifier")
bat_classifier_model.add(layers.Conv2D(20, (5,5), activation='relu', input_shape=(232, 154, 3)))
bat_classifier_model.add(layers.Dropout(0.2))
bat_classifier_model.add(layers.Conv2D(20, (5,5), activation='relu'))
bat_classifier_model.add(layers.Dropout(0.2))
bat_classifier_model.add(layers.MaxPooling2D(3,3))
bat_classifier_model.add(layers.Conv2D(20, (5,5), activation='relu'))
bat_classifier_model.add(layers.Dropout(0.2))
bat_classifier_model.add(layers.Conv2D(10, (5,5), activation='relu'))
bat_classifier_model.add(layers.Dropout(0.2))
bat_classifier_model.add(layers.MaxPooling2D(3,3))
bat_classifier_model.add(layers.Flatten())
bat_classifier_model.add(layers.Dense(4, activation='softmax'))

optimizer = tf.keras.optimizers.Adam(learning_rate=0.02)
bat_classifier_model.compile(optimizer=optimizer, loss='mse', metrics=['accuracy'])

# Load the saved model
bat_classifier_model = tf.keras.models.load_model("bat_classifier_model")

# Preprocess the image
def preprocess_image(image):
    processed_image = np.expand_dims(image, axis=0)
    return processed_image

# Gradio interface function
def classify_bat(image):
    processed_image = preprocess_image(image)
    prediction = bat_classifier_model.predict(processed_image)[0]
    class_names = ["Pipistrellus pygmaus", "Noctula nyctalus "]
    return {class_names[i]: float(prediction[i]) for i in range(len(class_names))}

inputs = gr.inputs.Image(label="Upload an image")
outputs = gr.outputs.Label(num_top_classes=2, label="Classification")

gr.Interface(fn=classify_bat, inputs=inputs, outputs=outputs).launch()