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