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

# Load your model and tokenizer
labels = {
    'class': ['amphibia', 'aves', 'invertebrates', 'lacertilia', 'mammalia', 'serpentes', 'testudines'],
    'serpentes': ["Butler's Gartersnake", "Dekay's Brownsnake", 'Eastern Gartersnake', 'Eastern Hog-nosed snake', 'Eastern Massasauga', 'Eastern Milksnake', 'Eastern Racer Snake', 'Eastern Ribbonsnake', 'Gray Ratsnake', "Kirtland's Snake", 'Northern Watersnake', 'Plains Gartersnake', 'Red-bellied Snake', 'Smooth Greensnake'],
    'mammalia': ['American Mink', 'Brown Rat', 'Eastern Chipmunk', 'Eastern Cottontail', 'Long-tailed Weasel', 'Masked Shrew', 'Meadow Jumping Mouse', 'Meadow Vole', 'N. Short-tailed Shrew', 'Raccoon', 'Star-nosed mole', 'Striped Skunk', 'Virginia Opossum', 'White-footed Mouse', 'Woodchuck', 'Woodland Jumping Mouse'],
    'aves': ['Common Yellowthroat', 'Gray Catbird', 'Indigo Bunting', 'Northern House Wren', 'Song Sparrow', 'Sora'],
    'amphibia': ['American Bullfrog', 'American Toad', 'Green Frog', 'Northern Leopard Frog']
}

hierarchical_models = {}

model_path = r"inceptionv3_class.h5"
print(model_path)
hierarchical_models['class'] = load_model(model_path)
print(hierarchical_models)
    
def load_and_preprocess_image(image, target_size=(299, 299)):
    img_array = img_to_array(image)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = tf.keras.applications.inception_v3.preprocess_input(img_array)
    return img_array

def predict(image):
    results = {}
    image_array = load_and_preprocess_image(image)
    
    # Predict class level
    class_preds = hierarchical_models['class'].predict(image_array)
    class_idx = np.argmax(class_preds)
    class_label = labels['class'][class_idx]
    results['class'] = class_label
    
    # Predict species level
    species_preds = hierarchical_models[class_label].predict(image_array)
    species_idx = np.argmax(species_preds)
    species_label = labels[class_label][species_idx]
    results['species'] = species_label
    print(results)
    return results

# Create Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Label(label="Class"), gr.Label(label="Species")],
    title="Image Classification",
    description="Upload an image to classify it into species and class level."
)

# Launch the interface
iface.launch()