Update app.py

app.py

@@ -20,13 +20,19 @@ 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)    
-    #predicted_labels = [labels[prediction] for prediction in predictions]
-    #print(predicted_labels[0],logits[0][predictions[0]])
-    confidences = {labels[i]: float(logits[0][i]) for i in range(len(labels))}
+    logits = model(**inputs)['logits'].cpu().detach().numpy()[0]
+    logits_positive = logits
+    logits_positive[logits < 0] = 0
+    logits_positive = logits_positive/np.sum(logits_positive)
+    confidences = {}
+    for i in range(len(labels):
+        if logits[i]>0:
+            confidences[labels[i]] = float(logits_positive[i])
+            
+        
+    confidences = {labels[i]:  )}
     return confidences
-    #return {"predictions" : predictions }
+
  
     #return {"predicted_label" : predicted_labels[0] }