Update app.py

app.py

@@ -37,6 +37,8 @@ def inference(input_image):
     if torch.cuda.is_available():
         input_batch = input_batch.to('cuda')
         model.to('cuda')
+    else:
+        model.to('cpu')
 
     with torch.no_grad():
         output = model(input_batch)
@@ -46,6 +48,16 @@ def inference(input_image):
     # Read the categories
     with open("artist_classes.txt", "r") as f:
         categories = [s.strip() for s in f.readlines()]
+    
+    categories = {
+        0:"vanGogh",
+        1:"Monet",
+        2:"Leonardo da Vinci",
+        3:"Rembrandt",
+        4:"Pablo Picasso",
+        5:"Salvador Dali"
+    }    
+    
     # Show top categories per image
     top5_prob, top5_catid = torch.topk(probabilities, 6)
     result = {}