Spaces:
Sleeping
Sleeping
File size: 4,914 Bytes
8e790ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import gradio as gr
import torch
import torchvision.transforms as transforms
from PIL import Image
import os
class DogBreedClassifier:
def __init__(self, model_path="traced_models/model_tracing.pt"):
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Load the traced model
self.model = torch.jit.load(model_path)
self.model = self.model.to(self.device)
self.model.eval()
# Define the same transforms used during training/testing
self.transform = transforms.Compose([
transforms.Resize((160, 160)),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
])
# Class labels
self.labels = ['Beagle', 'Boxer', 'Bulldog', 'Dachshund', 'German_Shepherd', 'Golden_Retriever', 'Labrador_Retriever', 'Poodle', 'Rottweiler', 'Yorkshire_Terrier']
# Add dog breed facts dictionary
self.breed_facts = {
'Beagle': "Beagles have approximately 220 million scent receptors, compared to a human's mere 5 million!",
'Boxer': "Boxers were among the first dogs to be employed as police dogs and were used as messenger dogs during wartime.",
'Bulldog': "Despite their tough appearance, Bulldogs were bred to be companion dogs and are known for being gentle and patient.",
'Dachshund': "Dachshunds were originally bred to hunt badgers - their name literally means 'badger dog' in German!",
'German_Shepherd': "German Shepherds can learn a new command in as little as 5 repetitions and obey it 95% of the time.",
'Golden_Retriever': "Golden Retrievers were originally bred as hunting dogs to retrieve waterfowl without damaging them.",
'Labrador_Retriever': "Labs have a special water-resistant coat and a unique otter-like tail that helps them swim efficiently.",
'Poodle': "Despite their elegant appearance, Poodles were originally water retrievers, and their fancy haircut had a practical purpose!",
'Rottweiler': "Rottweilers are descendants of Roman drover dogs and were used to herd livestock and pull carts for butchers.",
'Yorkshire_Terrier': "Yorkies were originally bred to catch rats in clothing mills. Despite their small size, they're true working dogs!"
}
@torch.no_grad()
def predict(self, image):
if image is None:
return None, None
# Convert to PIL Image if needed
if not isinstance(image, Image.Image):
image = Image.fromarray(image).convert('RGB')
# Preprocess image
img_tensor = self.transform(image).unsqueeze(0).to(self.device)
# Get prediction
output = self.model(img_tensor)
probabilities = torch.nn.functional.softmax(output[0], dim=0)
# Get the breed with highest probability
max_prob_idx = torch.argmax(probabilities).item()
predicted_breed = self.labels[max_prob_idx]
breed_fact = self.breed_facts[predicted_breed]
# Create prediction dictionary
predictions = {
self.labels[idx]: float(prob)
for idx, prob in enumerate(probabilities)
}
return predictions, breed_fact
classifier = DogBreedClassifier()
demo = gr.Interface(
fn=classifier.predict,
inputs=gr.Image(type="pil", label="Upload a dog image"),
outputs=[
gr.Label(num_top_classes=5, label="Breed Predictions"),
gr.Textbox(label="Fun Fact About This Breed!")
],
title="🐕 Dog Breed Classifier",
description="""
## Identify Your Dog's Breed!
Upload a clear photo of a dog, and I'll tell you its breed and share an interesting fact about it!
This model can identify 10 popular dog breeds with high accuracy.
### Supported Breeds:
Beagle, Boxer, Bulldog, Dachshund, German Shepherd, Golden Retriever, Labrador Retriever, Poodle, Rottweiler, Yorkshire Terrier
""",
article="""
### Tips for best results:
- Use clear, well-lit photos
- Ensure the dog's face is visible
- Avoid blurry or dark images
Created with PyTorch and Gradio | [GitHub](your_github_link)
""",
examples=[
["examples/Beagle_56.jpg"],
["examples/Boxer_30.jpg"],
["examples/Bulldog_73.jpg"],
["examples/Dachshund_43.jpg"],
["examples/German Shepherd_57.jpg"],
["examples/Golden Retriever_78.jpg"],
["examples/Labrador Retriever_25.jpg"],
["examples/Poodle_85.jpg"],
["examples/Rottweiler_30.jpg"],
["examples/Yorkshire Terrier_92.jpg"]
],
theme=gr.themes.Citrus(),
css="footer {display: none !important;}"
)
if __name__ == "__main__":
demo.launch() |