Update app.py

app.py

@@ -1,50 +1,51 @@
+import gradio as gr
 import numpy as np
 import clip
 import torch
-import gradio as gr
-import base64
 from PIL import Image
-from io import BytesIO
+import base64
 
 # Load the CLIP model
 model, preprocess = clip.load("ViT-B/32")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device).eval()
 
-def find_similarity(image_base64, text_input):
-    # Decode the base64 image string to bytes
-    image_bytes = base64.b64decode(image_base64)
+# Define a function to find similarity
+def find_similarity(base64_image, text_input):
+    # Decode the base64 image to bytes
+    image_bytes = base64.b64decode(base64_image)
+
+    # Convert the bytes to a PIL image
     image = Image.open(BytesIO(image_bytes))
 
     # Preprocess the image
     image = preprocess(image).unsqueeze(0).to(device)
 
-    # Prepare input text
+    # Tokenize the text input
     text_tokens = clip.tokenize([text_input]).to(device)
 
     # Encode image and text features
     with torch.no_grad():
-        image_features = model.encode_image(image).float()
-        text_features = model.encode_text(text_tokens).float()
+        image_features = model.encode_image(image)
+        text_features = model.encode_text(text_tokens)
 
-    # Normalize features and calculate similarity
-    image_features /= image_features.norm(dim=-1, keepdim=True)
-    text_features /= text_features.norm(dim=-1, keepdim=True)
-    similarity = (text_features @ image_features.T).cpu().numpy()
+    # Calculate cosine similarity
+    similarity = (image_features @ text_features.T).squeeze(0).cpu().numpy()
 
-    return similarity[0, 0]
+    return similarity
 
+# Create a Gradio interface
 iface = gr.Interface(
     fn=find_similarity,
     inputs=[
-        gr.inputs.Textbox(lines=3, label="Enter Base64 Image"),
-        gr.inputs.Textbox(lines=3, label="Enter Text"),
+        gr.inputs.Textbox(label="Base64 Image", lines=8),
+        "text"
     ],
     outputs="number",
     live=True,
     interpretation="default",
     title="CLIP Model Image-Text Cosine Similarity",
-    description="Enter a base64-encoded image and text to find their cosine similarity.",
+    description="Upload a base64 image and enter text to find their cosine similarity.",
 )
 
 iface.launch()