Spaces:
Sleeping
Sleeping
| 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() | |