Update app.py

app.py

@@ -15,19 +15,20 @@ labels =  [
 feature_extractor = AutoImageProcessor.from_pretrained('0-ma/vit-geometric-shapes-tiny')
 model = AutoModelForImageClassification.from_pretrained('0-ma/vit-geometric-shapes-tiny')
 
-labels = []
+
 def predict(image):
     feature_extractor = AutoImageProcessor.from_pretrained('0-ma/vit-geometric-shapes-tiny')
     model = AutoModelForImageClassification.from_pretrained('0-ma/vit-geometric-shapes-tiny')
     inputs = feature_extractor(images=[image], return_tensors="pt")
     logits = model(**inputs)['logits'].cpu().detach().numpy()
-    predictions = np.argmax(logits, axis=1)    
+    #predictions = np.argmax(logits, axis=1)    
     #predicted_labels = [labels[prediction] for prediction in predictions]
     #print(predicted_labels[0],logits[0][predictions[0]])
-
-    return {"predictions" : predictions }
+    confidences = {labels[i]: float(logits[i]) for i in range(len(labels))}
+    return confidences
+    #return {"predictions" : predictions }
  
-    return {"predicted_label" : predicted_labels[0] }
+    #return {"predicted_label" : predicted_labels[0] }
  
 title = "Geometric Shape Classifier"
 description = "A geometric shape setector."