import gradio as gr
import tensorflow as tf
import numpy as np
import json
import tensorflow_hub as hub
from PIL import Image, ImageDraw, ImageFont

# Register the custom layer
tf.keras.utils.get_custom_objects().update({'KerasLayer': hub.KerasLayer})

# Load the model
model = tf.keras.models.load_model("dog_breed_model.h5")

def load_breeds(file_path='unique_breeds.json'):
    with open(file_path, 'r') as file:
        data = json.load(file)
    return data["breeds"]

unique_breeds = load_breeds('unique_breeds.json')  # Load breed names into your array

def process_image(image, img_size=224):
    """
    Takes a PIL image and turns it into a preprocessed Tensor.
    """
    image = image.resize((img_size, img_size))  # Resize
    img_array = tf.keras.preprocessing.image.img_to_array(image)  # Convert to array
    img_array = img_array / 255.0  # Normalize
    return img_array

def predict_breed(image):
    img_array = process_image(image)
    img_array = tf.expand_dims(img_array, axis=0)  # Add batch dimension
    predictions = model.predict(img_array)
    predicted_class_index = np.argmax(predictions[0])
    predicted_class = unique_breeds[predicted_class_index]
    confidence = predictions[0][predicted_class_index] * 100  # Convert to percentage
    
    # Replace underscores with spaces in the class name and capitalize
    predicted_class = predicted_class.replace('_', ' ').upper()
   

    return f"{predicted_class} ({confidence:.2f}%)"

# Create Gradio interface
interface = gr.Interface(
    fn=predict_breed,
    inputs=gr.Image(type="pil"),
    outputs=gr.Text(),
    title="Dog Breed Identifier 🐶",
    description="Upload a dog image and the model will predict the breed along with confidence!"
)

# Launch app with public link
interface.launch(share=True)