Update app.py

app.py

@@ -16,22 +16,24 @@ model_bytes = requests.get(model_url).content
 model = torch.load(BytesIO(model_bytes), map_location=device)
 model.eval()
 
-# Load CLIP for preprocessing
+# Load CLIP model
 model_clip, preprocess = clip.load("ViT-B/32", device=device)
 
-# Load full test split from Flickr30k
+# Load test dataset
 dataset = load_dataset("nlphuji/flickr30k", split="test")
 
-captions = []
-images = []
-image_embeddings = []
-text_embeddings = []
+captions, images = [], []
+image_embeddings, text_embeddings = [], []
+
+print("Extracting features...")
 
 for example in dataset:
     try:
         img = Image.open(requests.get(example["image"], stream=True).raw).convert("RGB")
+        caption = example["sentence"]
+
         images.append(img)
-        captions.append(example["sentence"])
+        captions.append(caption)
 
         img_tensor = preprocess(img).unsqueeze(0).to(device)
         with torch.no_grad():
@@ -39,71 +41,62 @@ for example in dataset:
             img_feat /= img_feat.norm(dim=-1, keepdim=True)
             image_embeddings.append(img_feat.cpu())
 
-            txt_token = clip.tokenize([example["sentence"]]).to(device)
+            txt_token = clip.tokenize([caption]).to(device)
             txt_feat = model_clip.encode_text(txt_token)
             txt_feat /= txt_feat.norm(dim=-1, keepdim=True)
             text_embeddings.append(txt_feat.cpu())
-    except:
+    except Exception as e:
+        print(f"Skipped one sample due to error: {e}")
         continue
 
 image_embeddings = torch.cat(image_embeddings, dim=0)
 text_embeddings = torch.cat(text_embeddings, dim=0)
 
-# Create FAISS indexes
+# Build FAISS indexes
 image_index = faiss.IndexFlatIP(image_embeddings.shape[1])
 image_index.add(image_embeddings.numpy())
 
 text_index = faiss.IndexFlatIP(text_embeddings.shape[1])
 text_index.add(text_embeddings.numpy())
 
-# Define functions
+# Inference Functions
 def image_to_text(image):
     image_input = preprocess(image).unsqueeze(0).to(device)
     with torch.no_grad():
         image_feature = model_clip.encode_image(image_input)
         image_feature /= image_feature.norm(dim=-1, keepdim=True)
+
     D, I = text_index.search(image_feature.cpu().numpy(), 1)
+    matched_caption = captions[I[0][0]]
     score = round(float(D[0][0]) * 100, 2)
-    return f"{captions[I[0][0]]}\n(Match Score: {score}%)"
-
-    image_input = preprocess(image).unsqueeze(0).to(device)
-    with torch.no_grad():
-        image_feature = model_clip.encode_image(image_input)
-        image_feature /= image_feature.norm(dim=-1, keepdim=True)
-    _, I = text_index.search(image_feature.cpu().numpy(), 1)
-    return captions[I[0][0]]
+    return f"{matched_caption}\n(Match Score: {score}%)"
 
 def text_to_image(text):
     text_input = clip.tokenize([text]).to(device)
     with torch.no_grad():
         text_feature = model_clip.encode_text(text_input)
         text_feature /= text_feature.norm(dim=-1, keepdim=True)
+
     D, I = image_index.search(text_feature.cpu().numpy(), 1)
+    matched_image = images[I[0][0]]
     score = round(float(D[0][0]) * 100, 2)
-    img = images[I[0][0]]
-    return img, f"Match Score: {score}%"
-
-    text_input = clip.tokenize([text]).to(device)
-    with torch.no_grad():
-        text_feature = model_clip.encode_text(text_input)
-        text_feature /= text_feature.norm(dim=-1, keepdim=True)
-    _, I = image_index.search(text_feature.cpu().numpy(), 1)
-    return images[I[0][0]]
+    return matched_image, f"Match Score: {score}%"
 
 # Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("## 🖼️📝 Cross-Modal Retriever on Flickr30k Test Split")
+    
     with gr.Tab("Image to Text"):
         img_input = gr.Image(type="pil")
         text_output = gr.Textbox(label="Retrieved Caption")
         btn1 = gr.Button("Search Caption")
-        btn1.click(image_to_text, inputs=img_input, outputs=text_output)
+        btn1.click(fn=image_to_text, inputs=img_input, outputs=text_output)
 
     with gr.Tab("Text to Image"):
         text_input = gr.Textbox(label="Enter Text Prompt")
         img_output = gr.Image(label="Most Similar Image")
         score_output = gr.Textbox(label="Similarity Score")
         btn2 = gr.Button("Search Image")
-        btn2.click(text_to_image, inputs=text_input, outputs=[img_output, score_output])
+        btn2.click(fn=text_to_image, inputs=text_input, outputs=[img_output, score_output])
 
-demo.launch()
+demo.launch()