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

# Load the trained model
model = load_model("VGG.h5")

# Define class names (order from your dataset's subfolders)
class_names = ['cat', 'dog', 'wild']  # Change if your folder names differ

IMG_SIZE = 224

def predict(img):
    # Preprocess the image
    img = img.resize((IMG_SIZE, IMG_SIZE))
    img_array = image.img_to_array(img)
    img_array = img_array / 255.0  # Rescale
    img_array = np.expand_dims(img_array, axis=0)

    # Predict
    preds = model.predict(img_array)[0]
    result = {class_names[i]: float(preds[i]) for i in range(len(class_names))}
    return result

# Build Gradio Interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="Animal Face Classifier (VGG16)",
    description="Upload an image of an animal face (cat, dog, or wild) and get the predicted class probabilities."
)

if __name__ == "__main__":
    demo.launch()