Spaces:
Paused
Paused
| import json | |
| import torch | |
| import numpy as np | |
| from pathlib import Path | |
| from tqdm import tqdm | |
| from search import model, tokenizer | |
| # 1. Load filenames.json and image_embeddings.npy | |
| # build a lookup: filename → row number (0..999) | |
| EMB_DIR = Path("embeddings") | |
| image_embeddings = np.load(EMB_DIR / "image_embeddings.npy") # (1000, 512) | |
| filenames = json.loads((EMB_DIR / "filenames.json").read_text()) # 1000 names | |
| row_of = {name: i for i, name in enumerate(filenames)} | |
| # 2. Read Flickr8k.token.txt | |
| # for each line: split on "\t", strip the "#n" from the filename | |
| # keep it only if the filename is in the lookup | |
| # → captions: list of 5,000 strings | |
| # → true_rows: list of 5,000 row numbers (which image each caption belongs to) | |
| CAPTIONS_FILE = Path("data/Flickr8k.token.txt") | |
| captions = [] # the caption texts | |
| true_rows = [] # for each caption, the row of the image it describes | |
| for line in CAPTIONS_FILE.read_text().splitlines(): | |
| if not line.strip(): | |
| continue | |
| image_ref, caption = line.split("\t", 1) # "abc.jpg#0", "A dog runs ..." | |
| filename = image_ref.split("#")[0] # "abc.jpg" | |
| if filename in row_of: # only the 1,000 indexed images | |
| captions.append(caption.strip()) | |
| true_rows.append(row_of[filename]) | |
| true_rows = np.array(true_rows) | |
| # 3. Embed the 5,000 captions with CLIP's text encoder | |
| # in batches (e.g. 256), normalize, like build_index.py did for images | |
| # → text_embeddings: (5000, 512) | |
| BATCH_SIZE = 256 | |
| text_vectors = [] | |
| for start in tqdm(range(0, len(captions), BATCH_SIZE)): | |
| batch = captions[start:start + BATCH_SIZE] | |
| tokens = tokenizer(batch) | |
| with torch.no_grad(): | |
| vectors = model.encode_text(tokens) | |
| vectors = vectors / vectors.norm(dim=-1, keepdim=True) | |
| text_vectors.append(vectors.cpu().numpy()) | |
| text_embeddings = np.concatenate(text_vectors) | |
| # 4. Similarity matrix | |
| # sims = text_embeddings @ image_embeddings.T → (5000, 1000) | |
| sims = text_embeddings @ image_embeddings.T # (5000, 1000) | |
| # 5. Rank of the true image for each caption | |
| # true_scores = sims[i, true_rows[i]] for each i | |
| # rank = how many images scored HIGHER than the true one | |
| # (0 = it came first, 4 = it came fifth, …) | |
| true_scores = sims[np.arange(len(captions)), true_rows] # (5000,) | |
| ranks = (sims > true_scores[:, None]).sum(axis=1) # (5000,) 0 = first place | |
| # 6. Recall@K = fraction of captions with rank < K, for K = 1, 5, 10 | |
| # print a small table | |
| print() | |
| print("Text → image retrieval on Flickr8k test (1,000 images, 5,000 captions)") | |
| print("Model: CLIP ViT-B/32 (laion2b_s34b_b79k), zero-shot") | |
| print("-" * 40) | |
| for k in (1, 5, 10): | |
| recall = (ranks < k).mean() | |
| print(f"Recall@{k:<2} {recall:6.1%}") | |
| print("-" * 40) | |
| print(f"Median rank: {int(np.median(ranks)) + 1}") |