Update app.py

app.py

@@ -41,7 +41,7 @@ def calculate_score(image, text, model_name):
     model_type = MODELS[model_name][2]
 
     # Preprocess the image and text
-    inputs = processor(text=labels, images=[image], return_tensors="pt", padding=True)
+    inputs = processor(text=labels, images=[image], return_tensors="pt", padding="max_length")
     inputs = {k: v.to("cuda") for k, v in inputs.items()}
 
     # Calculate embeddings
@@ -58,14 +58,18 @@ def calculate_score(image, text, model_name):
     image_embeds = F.normalize(image_embeds, p=2, dim=1)
     text_embeds = F.normalize(text_embeds, p=2, dim=1)
 
-    # Calculate cosine similarity
-    cosine_similarities = torch.mm(text_embeds, image_embeds.t()).squeeze(1)
-
-    # Ensure values are between 0 and 1
-    cosine_similarities = torch.clamp(cosine_similarities, min=0, max=1)
+    # Calculate similarity
+    if model_type == "clip":
+        # For CLIP, use cosine similarity
+        similarities = torch.mm(text_embeds, image_embeds.t()).squeeze(1)
+        similarities = torch.clamp(similarities, min=0, max=1)
+    elif model_type == "siglip":
+        # For SigLIP, use sigmoid on dot product
+        logits = torch.mm(text_embeds, image_embeds.t()).squeeze(1)
+        similarities = torch.sigmoid(logits)
 
     # Convert to numpy array
-    similarities = cosine_similarities.cpu().numpy()
+    similarities = similarities.cpu().numpy()
 
     results_dict = {label: float(score) for label, score in zip(labels, similarities)}
     return results_dict