Update app.py

app.py

@@ -23,15 +23,18 @@ print(predicted_labels)
 
 labels = []
 def predict(img):
-    img = PILImage.create(img)
+    image = PILImage.create(img)
 
-    inputs = feature_extractor(images=img, return_tensors="pt")
+    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 {"predicted_labels" : predicted_labels , "predictions": predictions}
+    return {"predicted_label" : predicted_labels[0] }
  
 title = "Geometric Shape Classifier"
 description = "A geometric shape setector."