Update app.py

app.py

@@ -69,11 +69,16 @@ img_encoder = SonarImageEnc().to(device).eval()
 img_encoder.load_state_dict(torch.load(model_path, map_location=device))
 
 # -------- Similarity Scoring --------
-def compute_similarity(image, image_url, option_a, option_b, option_c, option_d, lang):
+def compute_similarity(
+    image, image_url, 
+    option_a, option_b, option_c, option_d,
+    lang_opt_a, lang_opt_b, lang_opt_c, lang_opt_d
+):
+    
     if not image:
         try:
             headers = {
-                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
+                "User-Agent": "Mozilla/5.0"
             }
             response = requests.get(image_url, headers=headers)
             response.raise_for_status()
@@ -88,16 +93,38 @@ def compute_similarity(image, image_url, option_a, option_b, option_c, option_d,
         image_emb, _ = img_encoder(inputs.pixel_values)
         image_emb = image_emb.to(device, torch.float16)
 
-    # Process text
-    lang_code = language_mapping[lang]
+    # Map languages
+    lang_codes = [
+        language_mapping[lang_opt_a],
+        language_mapping[lang_opt_b],
+        language_mapping[lang_opt_c],
+        language_mapping[lang_opt_d],
+    ]
     texts = [option_a, option_b, option_c, option_d]
-    text_embeddings = t2t_model_emb.predict(texts, source_lang=lang_code)
+
+    # Get embeddings per option with corresponding language
+    text_embeddings = []
+    for text, lang in zip(texts, lang_codes):
+        emb = t2t_model_emb.predict([text], source_lang=lang)
+        text_embeddings.append(emb)
+
+    text_embeddings = torch.cat(text_embeddings, dim=0).to(device)
 
     scores = cos(image_emb, text_embeddings)
-    results = {f"Option {chr(65+i)}": round(score.item(), 3) for i, score in enumerate(scores)}
+
+    results = {
+        f"Option {chr(65+i)}": round(score.item(), 3)
+        for i, score in enumerate(scores)
+    }
+
+    results = {
+        k: f"{round(v * 100, 2)}%"
+        for k, v in sorted(results.items(), key=lambda item: item[1], reverse=True)
+    }
 
     return image, results
 
+
 # -------- Gradio UI --------
 with gr.Blocks() as demo:
     gr.Markdown("## 🔍 SONAR: Image-Text Similarity Scorer")
@@ -106,10 +133,20 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             image_url = gr.Textbox(label="Image URL", value="http://images.cocodataset.org/val2017/000000039769.jpg")
-            option_a = gr.Textbox(label="Option A", value="A cat with two remotes.")
-            option_b = gr.Textbox(label="Option B", value="Two cat with two remotes.")
-            option_c = gr.Textbox(label="Option C", value="Two remotes.")
-            option_d = gr.Textbox(label="Option D", value="Two cats.")
+
+            with gr.Row():
+                option_a = gr.Textbox(label="Option A", value="A cat with two remotes.")
+                lang_opt_a = gr.Dropdown(choices=list(language_mapping.keys()), value="English", label="Language")
+            
+                option_b = gr.Textbox(label="Option B", value="Two cat with two remotes.")
+                lang_opt_b = gr.Dropdown(choices=list(language_mapping.keys()), value="English", label="Language")
+                
+                option_c = gr.Textbox(label="Option C", value="Two remotes.")
+                lang_opt_c = gr.Dropdown(choices=list(language_mapping.keys()), value="English", label="Language")
+                
+                option_d = gr.Textbox(label="Option D", value="Two cats.")
+                lang_opt_d = gr.Dropdown(choices=list(language_mapping.keys()), value="English", label="Language")
+            
             language = gr.Dropdown(choices=list(language_mapping.keys()), value="English", label="Select Language")
 
         with gr.Column():
@@ -120,8 +157,15 @@ with gr.Blocks() as demo:
                 img_output = gr.Image(label="Input Image", type="pil", width=300, height=300)
                 result_output = gr.JSON(label="Similarity Scores")
 
-    btn.click(fn=compute_similarity,
-              inputs=[image_input, image_url, option_a, option_b, option_c, option_d, language],
-              outputs=[img_output, result_output])
+    btn.click(
+        fn=compute_similarity,
+        inputs=[
+            image_input, image_url,
+            option_a, option_b, option_c, option_d,
+            lang_opt_a, lang_opt_b, lang_opt_c, lang_opt_d
+        ],
+        outputs=[img_output, result_output]
+    )
+
 
 demo.launch()