ma4389 commited on
Commit
08ffcf3
·
verified ·
1 Parent(s): d3b764f

Upload 3 files

Browse files
Files changed (3) hide show
  1. app.py +89 -0
  2. best_model.pth +3 -0
  3. 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