Spaces:
Build error
Build error
import gradio as gr | |
from transformers import AutoFeatureExtractor, AutoModelForImageClassification | |
from PIL import Image | |
import torch | |
# Load model and processor | |
model_name = "facebook/deit-base-distilled-patch16-224" | |
extractor = AutoFeatureExtractor.from_pretrained(model_name) | |
model = AutoModelForImageClassification.from_pretrained(model_name) | |
cat_keywords = ["cat", "kitten", "feline", "tabby", "siamese", "persian", "egyptian cat"] | |
def detect_cat(img): | |
inputs = extractor(images=img, return_tensors="pt") | |
with torch.no_grad(): | |
outputs = model(**inputs).logits | |
probs = torch.softmax(outputs, dim=-1) | |
predicted_class = torch.argmax(probs).item() | |
confidence = probs[0][predicted_class].item() | |
label = model.config.id2label[predicted_class].lower() | |
if confidence < 0.60: | |
return f"🤔 Not confident it's a cat. ({label}, {confidence:.2f})" | |
elif any(k in label for k in cat_keywords): | |
return f"😺 Yes, it's a cat! ({label}, confidence: {confidence:.2f})" | |
else: | |
return f"🐶 Nope, not a cat. ({label}, confidence: {confidence:.2f})" | |
gr.Interface( | |
fn=detect_cat, | |
inputs=gr.Image(type="pil", label="Upload your image 🖼️", height=300), | |
outputs=gr.Textbox(label="🐾 Result"), | |
title="😼 Is It a Cat?", | |
description="Upload an image to check if there's a cat in it.", | |
theme="default" | |
).launch()ß |