Spaces:
Sleeping
Sleeping
import gradio as gr | |
from fastai.vision.all import * | |
def is_cat(x): return x[0].isupper() | |
def get_labels(path): | |
pet_type = 'cat' if is_cat(path.name) else 'dog' | |
breed = RegexLabeller(r'(.+)_\d+.jpg$')(path.name) | |
return [pet_type, breed] | |
learn = load_learner('dog_cat_multi_v3.pkl') | |
def predict(image): | |
# Transform the image | |
pil_image = PILImage.create(image) | |
# Predict | |
preds, mask, probs = learn.predict(pil_image) | |
# Apply the threshold | |
threshold = 0.570 | |
classes_with_probs = [(learn.dls.vocab[i], probs[i].item()) for i in range(len(mask)) if probs[i] > threshold] | |
confidences = {learn.dls.vocab[i]: float(probs[i].item()) for i in range(len(probs)) if probs[i] > threshold} | |
# Check if the prediction includes "dog" or "cat" | |
pet_type = None | |
breed = 'unknown' # Default to 'unknown' if no breed is identified | |
pet_prob = 0 | |
breed_prob = 0 | |
for class_name, prob in classes_with_probs: | |
if class_name == 'dog' or class_name == 'cat': | |
pet_type = class_name | |
pet_prob = prob | |
else: | |
breed = class_name | |
breed_prob = prob | |
# Check if pet_type is None (i.e., neither dog nor cat) | |
if pet_type is None: | |
return "This is not a cat, nor a dog." | |
result = f"This is a {pet_type}, the breed is {breed if breed else 'unknown'}.\n" \ | |
f"The probability for it being a {pet_type} is {pet_prob * 100:.2f}%, the probability of being the breed is {breed_prob * 100:.2f}%." | |
return result, confidences | |
# Define the Gradio interface | |
iface = gr.Interface( | |
fn=predict, | |
inputs=gr.inputs.Image(shape=(224, 224)), | |
outputs=[gr.outputs.Textbox(label="Prediction"), gr.outputs.Label(label='confidences',num_top_classes=2)], | |
theme=gr.themes.Soft(), | |
examples=["/images/cazou103_Generate_a_high-resolution_image_of_a_Beagle_showcasin_609a1bae-22ac-4158-9091-dbf3220b2765.PNG", | |
"/images/cazou103_Generate_a_high-resolution_image_of_a_Shiba_Inu_showca_998f6162-289f-450e-ab02-558c3a575f61.PNG", | |
"/images/cazou103_Generate_a_high-resolution_image_of_a_Siamese_cat_show_68b1d78f-c304-4164-9342-4a4b56d2c4c5.PNG" | |
], | |
title="Cat and Dog Image Classifier", | |
description="Upload an image of a cat or a dog, and the model will identify the type and breed.", | |
article="This model has been trained on the Oxford Pets dataset and might not recognize all types dog and cat breeds. For best results, use clear images of pets." | |
) | |
# Launch the interface | |
iface.launch() |