single_pass.py

# -*- coding: utf-8 -*-
"""Single_Pass.ipynb

Automatically generated by Colab.

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

!pip install git+https://github.com/openai/CLIP.git
!pip install ftfy regex tqdm

import os
import pandas as pd
import torch
import clip
from PIL import Image
from tqdm import tqdm
from torchvision import transforms

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

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

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

device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)

resize_transform = transforms.Compose([
    transforms.Resize((260, 460)),
    transforms.ToTensor()
])

from torch.utils.data import DataLoader, Dataset

class NewsImageDataset(Dataset):
    def __init__(self, image_folder, preprocess):
        self.files = [f for f in os.listdir(image_folder) if f.endswith(".jpg")]
        self.folder = image_folder
        self.preprocess = preprocess

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        f = self.files[idx]
        img_path = os.path.join(self.folder, f)
        image_id = f.split(".")[0]
        image = self.preprocess(Image.open(img_path).convert("RGB"))
        return image_id, image

# Build dataset + loader
dataset = NewsImageDataset(IMAGE_FOLDER, preprocess)
loader = DataLoader(dataset, batch_size=64, shuffle=False, num_workers=2)

image_embeddings = {}
print("🔄 Computing embeddings in batches...")

with torch.no_grad():
    for batch_ids, batch_imgs in loader:
        batch_imgs = batch_imgs.to(device)
        emb = model.encode_image(batch_imgs)
        emb /= emb.norm(dim=-1, keepdim=True)
        for i, img_id in enumerate(batch_ids):
            image_embeddings[img_id] = emb[i].cpu()

print(f"✅ Stored embeddings for {len(image_embeddings)} images.")

# Pick first 3 articles
sample = df.head(3)

for _, row in sample.iterrows():
    article_id = row["article_id"]

    # Use a natural prompt for CLIP instead of raw concatenation
    text = f"A news photo about {row['article_title']}. Tags: {row['article_tags']}."

    # Encode text
    text_tokens = safe_tokenize(text, device)
    with torch.no_grad():
        text_emb = model.encode_text(text_tokens)
        text_emb /= text_emb.norm(dim=-1, keepdim=True)

    # Compare with ALL precomputed image embeddings
    best_img_id, best_score = None, -1
    for img_id, img_emb in image_embeddings.items():
        sim = torch.cosine_similarity(text_emb.cpu(), img_emb).item()
        if sim > best_score:
            best_img_id, best_score = img_id, sim

    print(f"\n📰 Article ID: {article_id}")
    print("Title:", row["article_title"])
    print("Tags :", row["article_tags"])
    print(f"Best Match Image: {best_img_id}, CLIP Similarity: {best_score:.4f}")

    # Display retrieved image
    img_path = os.path.join(IMAGE_FOLDER, best_img_id + ".jpg")
    display(Image.open(img_path).resize((460, 260)))

def safe_tokenize(text, device, max_len=77):
    # Tokenize and truncate automatically
    tokens = clip.tokenize([text], truncate=True).to(device)
    return tokens

similarities = []

for _, row in df.iterrows():
    text = str(row["article_title"]) + " " + str(row["article_tags"])

    # Use safe_tokenize here
    text_tokens = safe_tokenize(text, device)

    with torch.no_grad():
        text_emb = model.encode_text(text_tokens)
        text_emb /= text_emb.norm(dim=-1, keepdim=True)

    # Compute best similarity with all image embeddings
    best_score = max(
        torch.cosine_similarity(text_emb.cpu(), img_emb).item()
        for img_emb in image_embeddings.values()
    )
    similarities.append(best_score)

average_similarity = sum(similarities) / len(similarities)
print(f"\n📊 Average CLIP similarity: {average_similarity:.4f}")