Update app.py

app.py

@@ -23,6 +23,7 @@ model.load_state_dict(
     )
 )
 
+
 def predict(img) -> Tuple[Dict, float]:
     start_time = timer()
 
@@ -40,17 +41,12 @@ def predict(img) -> Tuple[Dict, float]:
         outputs = model(image).logits
         predicted_probs = torch.softmax(outputs, dim=1)
 
-    # Get the class name and its associated probability
-    predicted_class_idx = torch.argmax(predicted_probs).item()
-    predicted_class = class_names[predicted_class_idx]
-    predicted_probability = predicted_probs[0][predicted_class_idx].item()
+    # Create a prediction label and prediction probability dictionary for each prediction class
+    pred_labels_and_probs = {class_names[i]: float(predicted_probs[0][i]) for i in range(len(class_names))}
 
     # Calculate the prediction time
     pred_time = round(timer() - start_time, 5)
 
-    # Create a prediction label and prediction probability dictionary for the predicted class
-    pred_labels_and_probs = {predicted_class: predicted_probability}
-
     return pred_labels_and_probs, pred_time