import tensorflow as tf
import numpy as np
from PIL import Image
import gradio as gr
import pickle
import os

# Define image dimensions (must match training dimensions)
img_height = 224
img_width = 224

# Load the trained Keras model from export.pkl
# Ensure 'export.pkl' is in the same directory as app.py
try:
    with open('export.pkl', 'rb') as file:
        model = pickle.load(file)
    print("Model loaded from 'export.pkl' successfully.")
except FileNotFoundError:
    print("Error: 'export.pkl' not found. Make sure the model file exists.")
    exit()
except Exception as e:
    print(f"Error loading model: {e}")
    exit()

# Define the class indices mapping explicitly
# This needs to match the order used by ImageDataGenerator during training
class_indices_map = {'Non Venomous': 0, 'Venomous': 1}
# Reverse the map to get class names from indices for prediction output
label_map = {v: k for k, v in class_indices_map.items()}

def predict_snake_class(img):
    # 1. Convert to PIL Image (if not already) and resize
    if not isinstance(img, Image.Image):
        img = Image.fromarray(img)
    img = img.resize((img_width, img_height))

    # 2. Convert to NumPy array
    img_array = np.array(img)

    # 3. Normalize the image
    img_array = img_array / 255.0

    # 4. Expand dimensions to create a batch
    img_array = np.expand_dims(img_array, axis=0)

    # 5. Make a prediction
    predictions = model.predict(img_array)

    # 6. Determine the predicted class and confidence
    predicted_class_index = np.argmax(predictions[0])
    confidence_score = predictions[0][predicted_class_index]

    predicted_class_name = label_map[predicted_class_index]

    # Format output for Gradio Label component
    return {label_map[i]: float(predictions[0][i]) for i in range(len(label_map))}


# Set up the Gradio interface
iface = gr.Interface(
    fn=predict_snake_class,  # Our prediction function
    inputs=gr.Image(type="pil", label="Upload Snake Image"),  # Input is an image
    outputs=gr.Label(num_top_classes=2),  # Output is a label with top 2 classes
    title="Snake Classification Model",
    description="Upload an image of a snake to predict if it's Non Venomous or Venomous.",
    live=True
)

# Launch the interface
iface.launch(share=True)