Spaces:
Sleeping
Sleeping
File size: 4,421 Bytes
08ffcf3 |
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 |
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import gradio as gr
# ====== 1. Class names ======
classes = [
"affenpinscher", "afghan_hound", "african", "airedale", "akita", "american_terrier",
"appenzeller", "australian_cattledog", "australian_terrier", "basenji", "basset_hound",
"beagle", "bedlington_terrier", "bernese_mountain", "bichon_frise", "blenheim_spaniel",
"blood_hound", "bluetick", "border_collie", "border_terrier", "borzoi", "boston_bulldog",
"bouvier", "boxer", "brabancon", "briard", "brittany_spaniel", "bull_mastiff",
"cairn_terrier", "cardigan_corgi", "caucasian_ovcharka", "cavapoo", "chesapeake_retriever",
"chihuahua", "chow", "clumber", "cockapoo", "cocker_spaniel", "coonhound", "cotondetulear",
"curly_retriever", "dachshund", "dalmatian", "dandie_terrier", "dhole", "dingo", "doberman",
"english_bulldog", "english_hound", "english_mastiff", "english_setter", "english_sheepdog",
"english_springer", "entlebucher", "eskimo", "flatcoated_retriever", "fox_terrier",
"french_bulldog", "german_pointer", "germanlonghair_pointer", "germanshepherd",
"golden_retriever", "gordon_setter", "great_dane", "groenendael", "havanese", "husky",
"ibizan_hound", "indian_bakharwal", "indian_chippiparai", "indian_gaddi", "indian_greyhound",
"indian_mastiff", "indian_mudhol", "indian_pariah", "indian_sheepdog", "indian_spitz",
"irish_setter", "irish_spaniel", "irish_terrier", "irish_wolfhound", "italian_greyhound",
"japanese_spaniel", "japanese_spitz", "keeshond", "kelpie", "kelpie_australian",
"kerryblue_terrier", "kombai", "komondor", "kuvasz", "labradoodle", "labrador",
"lakeland_terrier", "lapphund_finnish", "leonberg", "lhasa", "malamute", "malinois",
"maltese", "medium_poodle", "mexicanhairless", "miniature_pinscher", "miniature_poodle",
"mix", "newfoundland", "norfolk_terrier", "norwegian_buhund", "norwegian_elkhound",
"norwich_terrier", "otterhound", "papillon", "patterdale_terrier", "pekinese", "pembroke",
"pitbull", "plott_hound", "pomeranian", "pug", "puggle", "pyrenees", "redbone",
"rottweiler", "russell_terrier", "saluki", "samoyed", "schipperke", "scottish_deerhound",
"scottish_terrier", "sealyham_terrier", "sharpei", "shepherd_australian", "shetland_sheepdog",
"shiba", "shihtzu", "silky_terrier", "spanish_waterdog", "staffordshire_bullterrier",
"standard_poodle", "stbernard", "sussex_spaniel", "swedish_danish", "swiss_mountain",
"tervuren", "tibetan_mastiff", "tibetan_terrier", "toy_poodle", "toy_terrier", "vizsla",
"walker_hound", "weimaraner", "welsh_spaniel", "welsh_terrier", "westhighland_terrier",
"wheaten_terrier", "whippet", "yorkshire_terrier"
]
# ====== 2. Transform (same as training) ======
transform = transforms.Compose([
transforms.Lambda(lambda image: image.convert('RGB')),
transforms.Resize((224, 224)), # fixed size
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
])
# ====== 3. Load trained model ======
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
in_features = model.fc.in_features
model.fc = nn.Sequential(
nn.Linear(in_features, 1024),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(1024, len(classes))
)
model.load_state_dict(torch.load("best_model.pth", map_location=device))
model.to(device)
model.eval()
# ====== 4. Prediction function ======
def predict_breed(image):
image = image.convert('RGB')
img_tensor = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
outputs = model(img_tensor)
probs = torch.softmax(outputs, dim=1)[0]
top_probs, top_idxs = torch.topk(probs, 3)
results = {classes[idx]: float(prob.item()) for prob, idx in zip(top_probs, top_idxs)}
return results
# ====== 5. Gradio Interface ======
title = "🐶 Dog Breed Classifier"
description = "Upload a dog image to predict its breed from 157 possible classes."
demo = gr.Interface(
fn=predict_breed,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=3),
title=title,
description=description
)
# ====== 6. Launch ======
demo.launch()
|