MediaEval2025/faiss.py at main · shivaaneesk/MediaEval2025 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
# -*- coding: utf-8 -*-
"""Faiss.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1xGFu3gCUOUmjy84eMUNQChCQoBFyB7NM
"""

!pip install open_clip_torch faiss-cpu pillow pandas tqdm

import os
import pandas as pd
import torch
import open_clip
import faiss
from PIL import Image
from torchvision import transforms
from tqdm import tqdm
import numpy as np

from google.colab import drive
drive.mount('/content/drive')

df = pd.read_csv("/content/drive/MyDrive/MediaEval - AFourP/newsarticles.csv")

IMAGE_FOLDER = "/content/drive/MyDrive/MediaEval - AFourP/newsimages"

print(df.head())

device = "cuda" if torch.cuda.is_available() else "cpu"

model, _, preprocess = open_clip.create_model_and_transforms(
    'ViT-B-32',
    pretrained='openai'
)
tokenizer = open_clip.get_tokenizer('ViT-B-32')
model = model.to(device)

# Preprocess pipeline for images
transform = preprocess

image_embeddings = []
image_ids = []

for fname in tqdm(os.listdir(IMAGE_FOLDER)):
    try:
        img_path = os.path.join(IMAGE_FOLDER, fname)
        img = Image.open(img_path).convert("RGB")
        img_tensor = transform(img).unsqueeze(0).to(device)

        with torch.no_grad():
            emb = model.encode_image(img_tensor).cpu().numpy()

        image_embeddings.append(emb)
        image_ids.append(fname)

    except Exception as e:
        print("Error with:", fname, e)

image_embeddings = np.vstack(image_embeddings).astype("float32")

# Build FAISS index
index = faiss.IndexFlatIP(image_embeddings.shape[1])  # cosine similarity
faiss.normalize_L2(image_embeddings)
index.add(image_embeddings)

def get_text_embedding(texts):
    """
    Encode text(s) into CLIP embeddings.
    texts: str or list of str
    """
    if isinstance(texts, str):
        texts = [texts]  # wrap single string as list

    tokens = tokenizer(texts).to(device)
    with torch.no_grad():
        text_emb = model.encode_text(tokens).cpu().numpy()
    faiss.normalize_L2(text_emb)
    return text_emb

print(df.columns)

# Example: retrieve for first article
example_text = str(df.iloc[0]["article_title"]) + " " + str(df.iloc[0]["article_tags"])
text_emb = get_text_embedding(example_text)

D, I = index.search(text_emb, k=3)  # top-3 results
print("Top retrieved images for:", df.iloc[0]["article_title"])
for idx in I[0]:
    print(image_ids[idx])

results = []

for i, row in df.iterrows():
    text = str(row["article_title"]) + " " + str(row["article_tags"])
    text_emb = get_text_embedding(text)
    D, I = index.search(text_emb, k=1)  # best image

    best_img = image_ids[I[0][0]]
    results.append({"article_id": row["article_id"], "retrieved_image": best_img})

retrieval_df = pd.DataFrame(results)
retrieval_df.to_csv("/content/drive/MyDrive/MediaEval - AFourP/retrieval_results.csv", index=False)
print("✅ Retrieval results saved!")

# Batch encode all article texts at once
texts = (df["article_title"].astype(str) + " " + df["article_tags"].astype(str)).tolist()

# Convert to embeddings in batches
batch_size = 64
all_embeddings = []
for i in range(0, len(texts), batch_size):
    batch = texts[i:i+batch_size]
    emb = get_text_embedding(batch)   # <-- modify function to handle list of strings
    all_embeddings.append(emb)

all_embeddings = np.vstack(all_embeddings)  # shape (8500, 512)