Update app.py

app.py

@@ -1,22 +1,38 @@
 import torch
-from transformers import pipeline
+from transformers import pipeline, SiglipModel, AutoProcessor
+import numpy as np
 import gradio as gr
 
 
 siglip_checkpoint = "nielsr/siglip-base-patch16-224"
 clip_checkpoint = "openai/clip-vit-base-patch16"
-siglip_detector = pipeline(model=siglip_checkpoint, task="zero-shot-image-classification")
 clip_detector = pipeline(model=clip_checkpoint, task="zero-shot-image-classification")
+siglip_model = SiglipModel.from_pretrained("nielsr/siglip-base-patch16-224")
+siglip_processor = AutoProcessor.from_pretrained("nielsr/siglip-base-patch16-224")
+
 
 def postprocess(output):
   return {out["label"]: float(out["score"]) for out in output}
 
+def postprocess_siglip(output, labels):
+  return {labels[i]: float(np.array(output[0])[i]) for i in range(len(labels))}
+    
+def siglip_detector(image, texts):
+  inputs = siglip_processor(text=texts, images=image, return_tensors="pt",
+                     padding="max_length")
+
+  with torch.no_grad():
+    outputs = model(**inputs)
+    logits_per_image = outputs.logits_per_image
+    probs = torch.sigmoid(logits_per_image) 
+  return probs
+
 
 def infer(image, candidate_labels):
   candidate_labels = [label.lstrip(" ") for label in candidate_labels.split(",")]
-  siglip_out = siglip_detector(image, candidate_labels=candidate_labels)
+  siglip_out = siglip_detector(image, candidate_labels)
   clip_out = clip_detector(image, candidate_labels=candidate_labels)
-  return postprocess(clip_out), postprocess(siglip_out)
+  return postprocess(clip_out), postprocess_siglip(siglip_out, labels=candidate_labels)
 
 
 with gr.Blocks() as demo: