Update app.py

app.py

@@ -1,49 +1,54 @@
 import gradio as gr
 from transformers import AutoModel
 from PIL import Image
+import numpy as np
 import torch
 
 # Load JinaAI CLIP model
 model = AutoModel.from_pretrained("jinaai/jina-clip-v1", trust_remote_code=True)
 
+# Function to process input (convert to text or PIL image)
+def process_input(input_data, input_type):
+    if input_type == "Text":
+        return model.encode_text([input_data]) if input_data.strip() else None
+    elif input_type == "Image":
+        if isinstance(input_data, str):  # If it's a file path
+            image = Image.open(input_data).convert("RGB")
+        elif isinstance(input_data, np.ndarray):  # If it's a NumPy array (Gradio default)
+            image = Image.fromarray(input_data)
+        else:
+            return None  # Invalid input type
+        return model.encode_image([image])
+    return None
+
 # Function to compute similarity
 def compute_similarity(input1, input2, input1_type, input2_type):
-    # Ensure inputs are valid
-    if input1_type == "Text" and (not input1 or input1.strip() == ""):
+    # Validate inputs
+    if input1_type == "Text" and not input1.strip():
         return "Error: Input 1 is empty!"
-    if input2_type == "Text" and (not input2 or input2.strip() == ""):
+    if input2_type == "Text" and not input2.strip():
         return "Error: Input 2 is empty!"
     if input1_type == "Image" and input1 is None:
         return "Error: Image 1 is missing!"
     if input2_type == "Image" and input2 is None:
         return "Error: Image 2 is missing!"
 
-    # Encode inputs
-    inputs = []
-    
-    if input1_type == "Text":
-        text1_embedding = model.encode_text([input1])
-        inputs.append(text1_embedding)
-    elif input1_type == "Image":
-        image1_embedding = model.encode_image([Image.fromarray(input1)])
-        inputs.append(image1_embedding)
+    # Process inputs
+    embedding1 = process_input(input1, input1_type)
+    embedding2 = process_input(input2, input2_type)
 
-    if input2_type == "Text":
-        text2_embedding = model.encode_text([input2])
-        inputs.append(text2_embedding)
-    elif input2_type == "Image":
-        image2_embedding = model.encode_image([Image.fromarray(input2)])
-        inputs.append(image2_embedding)
+    if embedding1 is None or embedding2 is None:
+        return "Error: Failed to process input!"
 
     # Compute cosine similarity
-    similarity_score = (inputs[0] @ inputs[1].T).item()
+    similarity_score = (embedding1 @ embedding2.T).item()
     return f"Similarity Score: {similarity_score:.4f}"
 
 # Function to toggle input fields dynamically
 def update_visibility(input1_type, input2_type):
     return (
-        gr.update(visible=(input1_type == "Text")),  # Show text input if Text is selected
-        gr.update(visible=(input1_type == "Image")), # Show image input if Image is selected
+        gr.update(visible=(input1_type == "Text")),  
+        gr.update(visible=(input1_type == "Image")), 
         gr.update(visible=(input2_type == "Text")),  
         gr.update(visible=(input2_type == "Image"))
     )