Update measure_recall.py

measure_recall.py

@@ -1,89 +1,89 @@
-import torch
-import clip
-import deeplake
-from tqdm import tqdm
-from tabulate import tabulate
-
-# --- SETTINGS ---
-EMBEDDINGS_PATH = "flickr_embeddings.pt"
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-BATCH_SIZE = 32
-# ----------------
-
-def calculate_recall():
-    print(f"Loading CLIP on {DEVICE}...")
-    model, _ = clip.load("ViT-B/32", device=DEVICE)
-    
-    print(f"Loading Image Index from {EMBEDDINGS_PATH}...")
-    try:
-        # Load embeddings and force them to the correct device
-        image_embeds = torch.load(EMBEDDINGS_PATH, map_location=DEVICE)
-    except FileNotFoundError:
-        print(f"ERROR: Could not find '{EMBEDDINGS_PATH}'. Run indexer.py first.")
-        return
-
-    num_images = len(image_embeds)
-    print(f"Loaded {num_images} image fingerprints.")
-
-    print("Connecting to Deep Lake to fetch captions...")
-    # Ensure this matches the dataset used in indexer.py
-    ds = deeplake.load('hub://activeloop/flickr30k', access_method="stream")
-    
-    text_embeds = []
-    
-    print("Encoding Captions (this acts as the 'Search Query')...")
-    with torch.no_grad():
-        for i in tqdm(range(0, num_images, BATCH_SIZE)):
-            batch_captions = []
-            end_idx = min(i + BATCH_SIZE, num_images)
-            
-            # Fetch batch of captions
-            for idx in range(i, end_idx):
-                try:
-                    # Robust text extraction
-                    txt = ds[idx].caption_0.text().numpy().item()
-                except AttributeError:
-                    txt = ds[idx].caption_0.numpy().item()
-                except Exception:
-                    txt = "unknown"
-                
-                batch_captions.append(str(txt))
-            
-            # Tokenize and Encode
-            text_tokens = clip.tokenize(batch_captions, truncate=True).to(DEVICE)
-            batch_features = model.encode_text(text_tokens)
-            
-            # Normalize
-            batch_features /= batch_features.norm(dim=-1, keepdim=True)
-            text_embeds.append(batch_features)
-
-    text_embeds = torch.cat(text_embeds)
-
-    print("\nCalculating Similarity Matrix...")
-    
-    # --- FIX: Force both to float32 to avoid mismatched types (Half vs Float) ---
-    image_embeds = image_embeds.float()
-    text_embeds = text_embeds.float()
-    
-    # Cosine similarity
-    similarity = text_embeds @ image_embeds.t()
-    
-    print("Computing Recall Metrics...")
-    results = {}
-    k_values = [1, 5, 10]
-    
-    _, top_indices = similarity.topk(max(k_values), dim=1)
-    
-    for k in k_values:
-        correct_count = 0
-        for i in range(num_images):
-            # Check if the correct index 'i' is in the top 'k' predictions
-            if i in top_indices[i, :k]:
-                correct_count += 1
-        results[f"R@{k}"] = (correct_count / num_images) * 100
-
-    print("\n--- RETRIEVAL RESULTS (Zero-Shot) ---")
-    print(tabulate([[k, f"{v:.2f}%"] for k, v in results.items()], headers=["Metric", "Score"]))
-
-if __name__ == "__main__":
+import torch
+import clip
+import deeplake
+from tqdm import tqdm
+from tabulate import tabulate
+
+# --- SETTINGS ---
+EMBEDDINGS_PATH = "flickr_embeddings.pt"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+BATCH_SIZE = 32
+# ----------------
+
+def calculate_recall():
+    print(f"Loading CLIP on {DEVICE}...")
+    model, _ = clip.load("ViT-B/32", device=DEVICE)
+    
+    print(f"Loading Image Index from {EMBEDDINGS_PATH}...")
+    try:
+        # Load embeddings and force them to the correct device
+        image_embeds = torch.load(EMBEDDINGS_PATH, map_location=DEVICE)
+    except FileNotFoundError:
+        print(f"ERROR: Could not find '{EMBEDDINGS_PATH}'. Run indexer.py first.")
+        return
+
+    num_images = len(image_embeds)
+    print(f"Loaded {num_images} image fingerprints.")
+
+    print("Connecting to Deep Lake to fetch captions...")
+    # Ensure this matches the dataset used in indexer.py
+    ds = deeplake.load('hub://activeloop/flickr30k', access_method="stream")
+    
+    text_embeds = []
+    
+    print("Encoding Captions (this acts as the 'Search Query')...")
+    with torch.no_grad():
+        for i in tqdm(range(0, num_images, BATCH_SIZE)):
+            batch_captions = []
+            end_idx = min(i + BATCH_SIZE, num_images)
+            
+            # Fetch batch of captions
+            for idx in range(i, end_idx):
+                try:
+                    # Robust text extraction
+                    txt = ds[idx].caption_0.text().numpy().item()
+                except AttributeError:
+                    txt = ds[idx].caption_0.numpy().item()
+                except Exception:
+                    txt = "unknown"
+                
+                batch_captions.append(str(txt))
+            
+            # Tokenize and Encode
+            text_tokens = clip.tokenize(batch_captions, truncate=True).to(DEVICE)
+            batch_features = model.encode_text(text_tokens)
+            
+            # Normalize
+            batch_features /= batch_features.norm(dim=-1, keepdim=True)
+            text_embeds.append(batch_features)
+
+    text_embeds = torch.cat(text_embeds)
+
+    print("\nCalculating Similarity Matrix...")
+    
+    # Force both to float32 to avoid mismatched types
+    image_embeds = image_embeds.float()
+    text_embeds = text_embeds.float()
+    
+    # Cosine similarity
+    similarity = text_embeds @ image_embeds.t()
+    
+    print("Computing Recall Metrics...")
+    results = {}
+    k_values = [1, 5, 10]
+    
+    _, top_indices = similarity.topk(max(k_values), dim=1)
+    
+    for k in k_values:
+        correct_count = 0
+        for i in range(num_images):
+            # Check if the correct index 'i' is in the top 'k' predictions
+            if i in top_indices[i, :k]:
+                correct_count += 1
+        results[f"R@{k}"] = (correct_count / num_images) * 100
+
+    print("\n--- RETRIEVAL RESULTS (Zero-Shot) ---")
+    print(tabulate([[k, f"{v:.2f}%"] for k, v in results.items()], headers=["Metric", "Score"]))
+
+if __name__ == "__main__":
     calculate_recall()