Update app.py

app.py

@@ -42,8 +42,31 @@ label2id = {v: k for k, v in id2label.items()}
 
 model = NewheadDinov2ForImageClassification.from_pretrained(checkpoint_name)
 
-
+def sigmoid(_outputs):
+    return 1.0 / (1.0 + np.exp(-_outputs))
+    
 def predict(input_image):
+    image_processor = AutoImageProcessor.from_pretrained(checkpoint_name)
+    # predict
+    inputs = image_processor(input_image, return_tensors="pt")
+    inputs = inputs
+    with torch.no_grad():
+        model_outputs = model(**inputs)
+    outputs = model_outputs["logits"][0]
+    scores = sigmoid(outputs)
+    result = {}
+    i = 0
+    for score in scores:
+        label = id2label[i]
+        result[label] = float(score)
+        i += 1
+    result = {key: result[key] for key in result if result[key] > 0.5}
+    return result
+
+
+
+
+'''    
     image_processor = AutoImageProcessor.from_pretrained(checkpoint_name)
     # predict
     inputs = image_processor(input_image, return_tensors="pt")
@@ -52,7 +75,7 @@ def predict(input_image):
     pred = pred[0]
     confidences = {classes_names[i]: round(float(pred[i]), 2) for i in range(50)}
     return confidences
-
+'''
 
 gr.Interface(
     fn=predict,