Update app.py

app.py

@@ -5,8 +5,6 @@ import torchvision.models as models
 import os
 import gradio as gr
 
-
-
 # Load the pre-trained model that was used in the training script. In this case it was ResNet18 model
 model = models.resnet18()
 
@@ -15,7 +13,7 @@ num_ftrs = model.fc.in_features
 model.fc = torch.nn.Linear(num_ftrs, 73)
 
 # Load the saved state dictionary (Device agnostic; inference for image classification is decent even on CPU)
-model.load_state_dict(torch.load('Dogrun2.pth', map_location=torch.device('cpu')))
+model.load_state_dict(torch.load('Dogrun2.pth', map_location=torch.device('cpu'), weights_only=True))
 
 # Set the model to evaluation mode
 model.eval()
@@ -33,8 +31,8 @@ transforms_test = v2.Compose([
     v2.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-
 labels = ['Afghan-Hound', 'Airedale-Terrier', 'Akita', 'Alaskan-Malamute', 'American-Foxhound', 'American-Hairless-Terrier', 'American-Water-Spaniel', 'Basenji', 'Basset-Hound', 'Beagle', 'Bearded-Collie', 'Belgian-Malinois', 'Belgian-Sheepdog', 'Bernese-Mountain-Dog', 'Bichon-Frise', 'Bloodhound', 'Bluetick-Coonhound', 'Border-Collie', 'Borzoi', 'Boston-Terrier', 'Boxer', 'Bull-Terrier', 'Bulldog', 'Bullmastiff', 'Cairn-Terrier', 'Cane-Corso', 'Cavalier-King-Charles-Spaniel', 'Chihuahua', 'Chinese-Crested', 'Chinese-Shar-Pei', 'Chow-Chow', 'Clumber-Spaniel', 'Cockapoo', 'Cocker-Spaniel', 'Collie', 'Dachshund', 'Dalmatian', 'Doberman-Pinscher', 'French-Bulldog', 'German-Shepherd', 'German-Shorthaired-Pointer', 'Golden-Retriever', 'Great-Dane', 'Great-Pyrenees', 'Greyhound', 'Irish-Water-Spaniel', 'Irish-Wolfhound', 'Japanese-Chin', 'Komondor', 'Labradoodle', 'Labrador-Retriever', 'Lhasa-Apso', 'Maltese', 'Miniature-Schnauzer', 'Newfoundland', 'Norwegian-Elkhound', 'Pekingese', 'Pembroke-Welsh-Corgi', 'Pomeranian', 'Poodle', 'Pug', 'Rhodesian-Ridgeback', 'Rottweiler', 'Saint-Bernard', 'Samoyed', 'Scottish-Terrier', 'Shiba-Inu', 'Shih-Tzu', 'Siberian-Husky', 'Staffordshire-Bull-Terrier', 'Vizsla', 'Xoloitzcuintli', 'Yorkshire-Terrier']
+
 #I have the breed_nicknames dictionary set because some breeds arent recognized that much by the official breed, such as the ones below.
 breed_nicknames = {
     'Xoloitzcuintli': ' (Mexican Hairless)',
@@ -49,20 +47,21 @@ def predict(input_img):
     topk_values, topk_indices = torch.topk(output_softmax, 3)
     topk_indices = topk_indices.tolist()[0]
     topk_labels = [labels[index] for index in topk_indices]
-    topk_probs = topk_values[0].tolist()
-    result = "\n".join([f"{label:<20}: {prob*100:.2f}%" for label, prob in zip(topk_labels, topk_probs)])
-    return result
-    
+    return topk_labels
+
+def get_user_selection(input_img, selected_breed):
+    topk_labels = predict(input_img)
+    return f"You selected {selected_breed}, one of the top predicted breeds: {', '.join(topk_labels)}. I love {selected_breed}!"
+
 gradio_app = gr.Interface(
-    predict,
-    inputs=gr.Image(label="Please select a clear image of your good dog to upload, or use your camera to take a picture.", sources=['upload', 'webcam'], type="pil"),
-    outputs=[gr.Label(label="Result")],
+    fn=get_user_selection,
+    inputs=[
+        gr.Image(label="Please select a clear image of your good dog to upload, or use your camera to take a picture.", sources=['upload', 'webcam'], type="pil"),
+        gr.Dropdown(label="Select the breed you think is correct", choices=labels, type="value")
+    ],
+    outputs="text",
     title="What is Your Dog Breed?",
 )
 
 if __name__ == "__main__":
     gradio_app.launch()
-
-
-
-