import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array

# Load the trained model
MODEL_PATH = "best_model.weights.h5"
model = load_model(MODEL_PATH)

# Define the class names
class_names = [
    "Bear", "Bird", "Cat", "Cow", "Deer", 
    "Dog", "Dolphin", "Elephant", "Giraffe", 
    "Horse", "Kangaroo", "Lion", "Panda", 
    "Tiger", "Zebra"
]

def classify_image(image):
    img = image.resize((256, 256))
    img_array = img_to_array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    predictions = model.predict(img_array)
    predicted_class = class_names[np.argmax(predictions)]
    return f"Predicted Class: {predicted_class}"

def instruction():
    return (
        "**Important Note:**\n\n"
        "This model is specifically trained to classify images into the following **15 animal categories**:\n\n"
        "- Bear\n"
        "- Bird\n"
        "- Cat\n"
        "- Cow\n"
        "- Deer\n"
        "- Dog\n"
        "- Dolphin\n"
        "- Elephant\n"
        "- Giraffe\n"
        "- Horse\n"
        "- Kangaroo\n"
        "- Lion\n"
        "- Panda\n"
        "- Tiger\n"
        "- Zebra\n\n"
        "**Usage Limitation:**\n\n"
        "- The model will only recognize images containing these animals.\n"
        "- Uploading an image of an animal not listed above or a non-animal image may result in inaccurate or undefined predictions.\n\n"
        "Ensure the uploaded image is clear, contains a single animal, and resembles the categories listed for the best results."
    )

# Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("# Animal Classifier")
    gr.Markdown(instruction())

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload an Image", type="pil")
            predict_button = gr.Button("Classify Image")
        with gr.Column():
            result_output = gr.Textbox(label="Prediction Result", lines=3)

    predict_button.click(classify_image, inputs=image_input, outputs=result_output)

app.launch()