Spaces:
Sleeping
Sleeping
Upload 3 files
Browse files- app.py +89 -0
- best_model.pth +3 -0
- requirements.txt +4 -0
app.py
ADDED
|
@@ -0,0 +1,89 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
import torch.nn as nn
|
| 3 |
+
from torchvision import models, transforms
|
| 4 |
+
from PIL import Image
|
| 5 |
+
import gradio as gr
|
| 6 |
+
|
| 7 |
+
# ====== 1. Class names ======
|
| 8 |
+
classes = [
|
| 9 |
+
"affenpinscher", "afghan_hound", "african", "airedale", "akita", "american_terrier",
|
| 10 |
+
"appenzeller", "australian_cattledog", "australian_terrier", "basenji", "basset_hound",
|
| 11 |
+
"beagle", "bedlington_terrier", "bernese_mountain", "bichon_frise", "blenheim_spaniel",
|
| 12 |
+
"blood_hound", "bluetick", "border_collie", "border_terrier", "borzoi", "boston_bulldog",
|
| 13 |
+
"bouvier", "boxer", "brabancon", "briard", "brittany_spaniel", "bull_mastiff",
|
| 14 |
+
"cairn_terrier", "cardigan_corgi", "caucasian_ovcharka", "cavapoo", "chesapeake_retriever",
|
| 15 |
+
"chihuahua", "chow", "clumber", "cockapoo", "cocker_spaniel", "coonhound", "cotondetulear",
|
| 16 |
+
"curly_retriever", "dachshund", "dalmatian", "dandie_terrier", "dhole", "dingo", "doberman",
|
| 17 |
+
"english_bulldog", "english_hound", "english_mastiff", "english_setter", "english_sheepdog",
|
| 18 |
+
"english_springer", "entlebucher", "eskimo", "flatcoated_retriever", "fox_terrier",
|
| 19 |
+
"french_bulldog", "german_pointer", "germanlonghair_pointer", "germanshepherd",
|
| 20 |
+
"golden_retriever", "gordon_setter", "great_dane", "groenendael", "havanese", "husky",
|
| 21 |
+
"ibizan_hound", "indian_bakharwal", "indian_chippiparai", "indian_gaddi", "indian_greyhound",
|
| 22 |
+
"indian_mastiff", "indian_mudhol", "indian_pariah", "indian_sheepdog", "indian_spitz",
|
| 23 |
+
"irish_setter", "irish_spaniel", "irish_terrier", "irish_wolfhound", "italian_greyhound",
|
| 24 |
+
"japanese_spaniel", "japanese_spitz", "keeshond", "kelpie", "kelpie_australian",
|
| 25 |
+
"kerryblue_terrier", "kombai", "komondor", "kuvasz", "labradoodle", "labrador",
|
| 26 |
+
"lakeland_terrier", "lapphund_finnish", "leonberg", "lhasa", "malamute", "malinois",
|
| 27 |
+
"maltese", "medium_poodle", "mexicanhairless", "miniature_pinscher", "miniature_poodle",
|
| 28 |
+
"mix", "newfoundland", "norfolk_terrier", "norwegian_buhund", "norwegian_elkhound",
|
| 29 |
+
"norwich_terrier", "otterhound", "papillon", "patterdale_terrier", "pekinese", "pembroke",
|
| 30 |
+
"pitbull", "plott_hound", "pomeranian", "pug", "puggle", "pyrenees", "redbone",
|
| 31 |
+
"rottweiler", "russell_terrier", "saluki", "samoyed", "schipperke", "scottish_deerhound",
|
| 32 |
+
"scottish_terrier", "sealyham_terrier", "sharpei", "shepherd_australian", "shetland_sheepdog",
|
| 33 |
+
"shiba", "shihtzu", "silky_terrier", "spanish_waterdog", "staffordshire_bullterrier",
|
| 34 |
+
"standard_poodle", "stbernard", "sussex_spaniel", "swedish_danish", "swiss_mountain",
|
| 35 |
+
"tervuren", "tibetan_mastiff", "tibetan_terrier", "toy_poodle", "toy_terrier", "vizsla",
|
| 36 |
+
"walker_hound", "weimaraner", "welsh_spaniel", "welsh_terrier", "westhighland_terrier",
|
| 37 |
+
"wheaten_terrier", "whippet", "yorkshire_terrier"
|
| 38 |
+
]
|
| 39 |
+
|
| 40 |
+
# ====== 2. Transform (same as training) ======
|
| 41 |
+
transform = transforms.Compose([
|
| 42 |
+
transforms.Lambda(lambda image: image.convert('RGB')),
|
| 43 |
+
transforms.Resize((224, 224)), # fixed size
|
| 44 |
+
transforms.ToTensor(),
|
| 45 |
+
transforms.Normalize(mean=[0.5, 0.5, 0.5],
|
| 46 |
+
std=[0.5, 0.5, 0.5])
|
| 47 |
+
])
|
| 48 |
+
|
| 49 |
+
# ====== 3. Load trained model ======
|
| 50 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 51 |
+
|
| 52 |
+
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
|
| 53 |
+
in_features = model.fc.in_features
|
| 54 |
+
model.fc = nn.Sequential(
|
| 55 |
+
nn.Linear(in_features, 1024),
|
| 56 |
+
nn.ReLU(),
|
| 57 |
+
nn.Dropout(0.4),
|
| 58 |
+
nn.Linear(1024, len(classes))
|
| 59 |
+
)
|
| 60 |
+
model.load_state_dict(torch.load("best_model.pth", map_location=device))
|
| 61 |
+
model.to(device)
|
| 62 |
+
model.eval()
|
| 63 |
+
|
| 64 |
+
# ====== 4. Prediction function ======
|
| 65 |
+
def predict_breed(image):
|
| 66 |
+
image = image.convert('RGB')
|
| 67 |
+
img_tensor = transform(image).unsqueeze(0).to(device)
|
| 68 |
+
with torch.no_grad():
|
| 69 |
+
outputs = model(img_tensor)
|
| 70 |
+
probs = torch.softmax(outputs, dim=1)[0]
|
| 71 |
+
top_probs, top_idxs = torch.topk(probs, 3)
|
| 72 |
+
|
| 73 |
+
results = {classes[idx]: float(prob.item()) for prob, idx in zip(top_probs, top_idxs)}
|
| 74 |
+
return results
|
| 75 |
+
|
| 76 |
+
# ====== 5. Gradio Interface ======
|
| 77 |
+
title = "🐶 Dog Breed Classifier"
|
| 78 |
+
description = "Upload a dog image to predict its breed from 157 possible classes."
|
| 79 |
+
|
| 80 |
+
demo = gr.Interface(
|
| 81 |
+
fn=predict_breed,
|
| 82 |
+
inputs=gr.Image(type="pil"),
|
| 83 |
+
outputs=gr.Label(num_top_classes=3),
|
| 84 |
+
title=title,
|
| 85 |
+
description=description
|
| 86 |
+
)
|
| 87 |
+
|
| 88 |
+
# ====== 6. Launch ======
|
| 89 |
+
demo.launch()
|
best_model.pth
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:083d4dc5c1d48aa49956e37700f17f60a070e2b8c302cdb29f708520ec25b9d5
|
| 3 |
+
size 103387776
|
requirements.txt
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
torch>=2.0.0
|
| 2 |
+
torchvision>=0.15.0
|
| 3 |
+
pillow
|
| 4 |
+
gradio>=4.0.0
|