Update app.py

app.py

@@ -27,7 +27,7 @@ def inference(input_image):
         transforms.Resize(256),
         transforms.CenterCrop(224),
         transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
     ])
     input_tensor = preprocess(input_image)
     input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
@@ -42,7 +42,7 @@ def inference(input_image):
     with torch.no_grad():
         output = model(input_batch)
     # The output has unnormalized scores. To get probabilities, you can run a softmax on it.
-    probabilities = torch.nn.functional.softmax(output[0], dim=1)
+    probabilities = torch.nn.functional.softmax(output[0], dim=0)
 
     # Read the categories
     with open("artist_classes.txt", "r") as f:
@@ -61,14 +61,13 @@ def inference(input_image):
     top5_prob, top5_catid = torch.topk(probabilities, 6)
     result = {}
     for i in range(top5_prob.size(0)):
-        result[categories[top5_catid[i]]] = top5_prob[i].item()
-    print(result)
+        result[categories[top5_catid[i],item()]] = top5_prob[i].item()
     return result
 
 inputs = gr.Image(type='pil')
 outputs = gr.Label(type="confidences",num_top_classes=5)
 
-title = "MOBILENET V2"
+title = "Artist Classifier"
 description = "Gradio demo for MOBILENET V2, Efficient networks optimized for speed and memory, with residual blocks. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
 article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1801.04381'>MobileNetV2: Inverted Residuals and Linear Bottlenecks</a> | <a href='https://github.com/pytorch/vision/blob/master/torchvision/models/mobilenet.py'>Github Repo</a></p>"