Spaces:
Running
Running
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() | |