Update app.py

app.py

@@ -22,10 +22,11 @@ def predict(image):
     inputs = feature_extractor(images=[image], return_tensors="pt")
     logits = model(**inputs)['logits'].cpu().detach().numpy()
     predictions = np.argmax(logits, axis=1)    
-    predicted_labels = [labels[prediction] for prediction in predictions]
-    print(predicted_labels[0],logits[0][predictions[0]])
-
+    #predicted_labels = [labels[prediction] for prediction in predictions]
+    #print(predicted_labels[0],logits[0][predictions[0]])
 
+    return {"predictions" : predictions }
+ 
     return {"predicted_label" : predicted_labels[0] }
  
 title = "Geometric Shape Classifier"