from fastai.vision.all import load_learner
import gradio as gr
import numpy as np

import pathlib
temp = pathlib.PosixPath
pathlib.PosixPath = pathlib.WindowsPath

# Function for recognizing species and behavior from an image
def recognize_image(input_image):
    # Add your model loading code here if not already loaded
    model = load_learner("final-v0.pkl")

    # Make predictions
    pred, idx, probabilities = model.predict(input_image)

    # Get all labels from the model's vocabulary
    all_labels = model.dls.vocab

    # List of behavior labels
    behavior_labels = ['F', 'H', 'M', 'M', 'P', 'R', 'T']

    # Get the top-k predicted labels
    top_k = 4
    top_indices = (-probabilities).argsort()[:top_k]
    top_labels = [all_labels[i] for i in top_indices]
    top_probabilities = [float(probabilities[i]) for i in top_indices]

    # Separate labels into behavior and species categories
    behavior_predictions = [label for label in top_labels if label in behavior_labels]
    species_predictions = [label for label in top_labels if label not in behavior_labels]

    # Create dictionaries for species and behavior predictions with their probabilities
    species_result = {
        'Species Predictions': dict(zip(species_predictions, [round(prob, 4) for prob in top_probabilities if prob not in behavior_labels])),
    }

    behavior_result = {
        'Behavior Predictions': dict(zip(behavior_predictions, [round(prob, 4) for prob in top_probabilities if prob in behavior_labels])),
    }

    return species_result, behavior_result

# Gradio interface
input_image = gr.inputs.Image(type='numpy', label='Upload Image')
output_species = gr.outputs.Label(label='Species Predictions')
output_behavior = gr.outputs.Label(label='Behavior Predictions')

gr.Interface(
    fn=recognize_image,
    inputs=input_image,
    outputs=[output_species, output_behavior],
    title='Species and Behavior Recognition',
    description='Upload an image and get predictions for species and behavior.'
).launch()