Update app.py

app.py

@@ -4,7 +4,7 @@ from PIL import Image
 from torchvision import transforms
 import gradio as gr
 #https://huggingface.co/spaces/yuhe6/final_project/blob/main/CIFAR10_cnn.pth
-os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")
+#os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")
 
 #model = torch.hub.load('huawei-noah/ghostnet', 'ghostnet_1x', pretrained=True)
 model = torch.hub.load('/', 'CIFAR10_cnn', pretrained=True)
@@ -36,8 +36,10 @@ def inference(input_image):
     # Read the categories
     with open("imagenet_classes.txt", "r") as f:
         categories = [s.strip() for s in f.readlines()]
+    with open("dog_cat.txt", "r") as f:
+        categories = [s.strip() for s in f.readlines()]
     # Show top categories per image
-    top5_prob, top5_catid = torch.topk(probabilities, 5)
+    top5_prob, top5_catid = torch.topk(probabilities, 1)
     result = {}
     for i in range(top5_prob.size(0)):
         result[categories[top5_catid[i]]] = top5_prob[i].item()