Update app.py

app.py

@@ -1,8 +1,7 @@
 import streamlit as st
 from PIL import Image
 import torch
-
-
+import json
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 
 extractor = AutoFeatureExtractor.from_pretrained("Amite5h/convnext-tiny-finetuned-eurosat")
@@ -11,8 +10,6 @@ model = AutoModelForImageClassification.from_pretrained("Amite5h/convnext-tiny-f
 
 
 
-#pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
-
 st.title("EuroSAT Detection")
 
 file_name = st.file_uploader("Upload a hot dog candidate image")
@@ -27,9 +24,21 @@ 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(predictions.logits, dim=1)
+
+    predicted_probabilities = torch.softmax(outputs.logits, dim=1)[0]
+    predicted_labels = model.config.id2label
+    
+    # Create a dictionary of labels and probabilities
+    label_probabilities = {
+        predicted_labels[i]: predicted_probabilities[i].item() for i in range(len(predicted_labels))
+    }
+    
+    # Convert the output to JSON string
+    json_output = json.dumps(label_probabilities)
+        
+    #predicted_class = torch.argmax(predictions.logits, dim=1)
     col2.header("Probabilities")
-    col2.subheader(predicted_class.item())
+    col2.subheader(json_output)
 
     # col2.header("Probabilities")
     # for p in predictions: