import gradio as gr
from fastai.vision.all import *

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

# Load your multi-target model
model = load_learner("final-v0.pkl")

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

all_labels = model.dls.vocab

def recognize_image(input_image, top_k=4):
    # Make predictions
    _, _, probabilities = model.predict(input_image)
    
    # Get the indices of the top-k predicted labels based on probabilities
    top_indices = (-probabilities).argsort()[:top_k]
    
    # Get the labels corresponding to the top indices from the vocabulary
    top_labels = [all_labels[i] for i in top_indices]
    
    # Get the probabilities for the top-k labels
    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 = {
        'top_species_labels': species_predictions,
        'top_species_probabilities': [prob for label, prob in zip(top_labels, top_probabilities) if label not in behavior_labels],
    }
    
    behavior_result = {
        'top_behavior_labels': behavior_predictions,
        'top_behavior_probabilities': [prob for label, prob in zip(top_labels, top_probabilities) if label in behavior_labels],
    }
    
    # Format the results into a more readable format
    formatted_species_result = {
        'Species Predictions': dict(zip(species_predictions, [round(prob, 4) for prob in species_result['top_species_probabilities']])),
    }
    
    formatted_behavior_result = {
        'Behavior Predictions': dict(zip(behavior_predictions, [round(prob, 4) for prob in behavior_result['top_behavior_probabilities']])),
    }
    
    return formatted_species_result, formatted_behavior_result

def get_outputs(input_image):
    species_predictions, behavior_predictions = recognize_image(input_image)
    print("Species Predictions:")
    for species, probability in species_predictions['Species Predictions'].items():
        print(f"{species}: {probability}")

    print("Behavior Predictions:")
    for behavior, probability in behavior_predictions['Behavior Predictions'].items():
        print(f"{behavior}: {probability}")

    return species_predictions['Species Predictions'], behavior_predictions['Behavior Predictions']

path = 'test images/'

image_path = []

for i in os.listdir(path):
  image_path.append(path+i) 

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

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