import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
from tensorflow.keras.models import load_model
import gradio as gr

# List of class labels from the Stanford Dogs dataset
class_labels = [
    'Affenpinscher', 'African_hunting_dog', 'Airedale', 'American_Staffordshire_terrier', 'American_water_spaniel', 
    'Anatolian_shepherd_dog', 'Australian_terrier', 'Basenji', 'Basset_hound', 'Beagle', 'Bearded_collie', 
    'Beauceron', 'Bedlington_terrier', 'Belgian_malinois', 'Belgian_sheepdog', 'Bernese_mountain_dog', 
    'Biewer', 'Black-and-tan_coonhound', 'Black_russian_terrier', 'Border_collie', 'Border_terrier', 
    'Borzoi', 'Boston_bull', 'Bouvier_des_Flandres', 'Boxer', 'Boykin_spaniel', 'Briard', 'Brittany', 
    'Bull_mastiff', 'Cairn_terrier', 'Canaan_dog', 'Cavalier_king_charles_spaniel', 'Chihuahua', 
    'Chinese_crested', 'Chow', 'Clumber_spaniel', 'Cocker_spaniel', 'Collie', 'Curly-coated_retriever', 
    'Dachshund', 'Dalmatian', 'Dandie_Dinmont_terrier', 'Doberman', 'English_cocker_spaniel', 
    'English_setter', 'English_springer_spaniel', 'Entlebucher_mountain_dog', 'Field_spaniel', 'Finnish_spitz', 
    'Flat-coated_retriever', 'French_bulldog', 'German_pinscher', 'German_shepherd', 'German_short-haired_pointer', 
    'Giant_schnauzer', 'Glen_of_imaal_terrier', 'Golden_retriever', 'Goldendoodle', 'Great_dane', 'Great_pyrenees', 
    'Greater_swiss_mountain_dog', 'Havanese', 'Irish_setter', 'Irish_terrier', 'Irish_water_spaniel', 
    'Italian_greyhound', 'Japanese_chin', 'Keeshond', 'Kerry_blue_terrier', 'King_charles_spaniel', 
    'Klee_kai', 'Labrador_retriever', 'Lakeland_terrier', 'Lhasa', 'Maltese_dog', 'Manchester_terrier', 
    'Mastiff', 'Miniature_pinscher', 'Miniature_schnauzer', 'Newfoundland', 'Norfolk_terrier', 'Norwegian_elkhound', 
    'Norwich_terrier', 'Old_english_sheepdog', 'Otterhound', 'Papillon', 'Pekingese', 'Pembroke', 'Pharaoh_hound', 
    'Plott', 'Pointer', 'Pomeranian', 'Poodle', 'Portuguese_water_dog', 'Rottweiler', 'Saint_bernard', 
    'Saluki', 'Samoyed', 'Schipperke', 'Scotch_terrier', 'Shiba_inu', 'Shih-tzu', 'Siberian_husky', 'Silky_terrier', 
    'Staffordshire_bullterrier', 'Standard_schnauzer', 'Tibetan_mastiff', 'Tibetan_terrier', 'Weimaraner', 
    'Welsh_springer_spaniel', 'West_highland_white_terrier', 'Whippet', 'Yorkshire_terrier'
]

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

# Preprocessing function
def preprocess_image(image):
    # No need to use load_img, the image is already a PIL image from Gradio
    img = image.resize((224, 224))  # Resize image to 224x224
    img_array = img_to_array(img)
    img_array = preprocess_input(img_array)  # EfficientNetV2 preprocessing
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
    return img_array

# Prediction function
def predict_dog_breed(image):
    img_array = preprocess_image(image)
    predictions = model.predict(img_array)
    class_idx = np.argmax(predictions)  # Get index of the class with the highest probability
    
    # Ensure that class_idx is within the range of class_labels
    if class_idx < len(class_labels):
        breed = class_labels[class_idx]  # Map the index to the breed name
        confidence = predictions[0][class_idx]  # Confidence score
        return [breed, confidence]  # Return as a list
    else:
        return ["Error", "Invalid model output"]


        
def old_predict_dog_breed(image):
    img_array = preprocess_image(image)
    predictions = model.predict(img_array)
    class_idx = np.argmax(predictions)  # Get index of the class with the highest probability
    breed = class_labels[class_idx]  # Map the index to the breed name
    confidence = predictions[0][class_idx]  # Confidence score
    return [breed, confidence]  # Return as a list

# Create a Gradio interface
iface = gr.Interface(fn=predict_dog_breed, 
                     inputs=gr.Image(type="pil"),  # Gradio automatically handles PIL images
                     outputs=gr.JSON(),  # Outputs will be a JSON response (list of results)
                     live=True)

# Launch the Gradio app
iface.launch()