Update app.py

app.py

@@ -9,6 +9,8 @@ import torchvision.transforms as transforms
 # model = timm.create_model("hf_hub:nateraw/resnet18-random", pretrained=True)
 # model.train()
 
+id_label = {0:'airplane', 1:'automobile', 2:'bird', 3:'cat', 4:'deer', 5:'dog', 6:'frog', 7:'horse', 8:'ship', 9:'trunk'}
+
 import os 
 
 def print_bn():
@@ -45,7 +47,7 @@ def greet_backdoor(image):
         image = transform_nor(image).unsqueeze(0)
         print(image.shape)
         output = model(image).squeeze()
-        return 'classified as label: ' + str(int(torch.argmax(output)))
+        return 'classified as: ' + id_label[int(torch.argmax(output))]
 
 
 def greet(image):