Update app.py

app.py

@@ -26,7 +26,10 @@ if file_name is not None:
     # Convert grayscale image to RGB format
     image_tensor = extractor(images=image, return_tensors="pt")["pixel_values"]
     predictions = model(image_tensor)
-
+    predicted_class = torch.argmax(outputs.logits, dim=1)
     col2.header("Probabilities")
-    for p in predictions:
-        col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")
+    col2.subheader(predicted_class.item())
+
+    # col2.header("Probabilities")
+    # for p in predictions:
+    #     col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")