Update app.py

app.py

@@ -31,11 +31,11 @@ def predict_image(image):
     with torch.no_grad():
         outputs = model(processed_image)
         probabilities = torch.nn.functional.softmax(outputs, dim=1)
-        top_probs, top_indices = torch.topk(probabilities, 3)
+        top_probs, top_indices = torch.topk(probabilities, 2)
         top_classes = [class_names[idx] for idx in top_indices[0]]
 
     # Create a result dictionary with class names and probabilities
-    result = {top_classes[i]: top_probs[0][i].item() for i in range(3)}
+    result = {top_classes[i]: top_probs[0][i].item() for i in range(2)}
 
     return result