advancedretreival.py

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

Automatically generated by Colab.

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

!pip install --upgrade --quiet transformers accelerate timm open_clip_torch ftfy regex tqdm

import os
import torch
import pandas as pd
from PIL import Image
from tqdm import tqdm
import open_clip
import transformers

print(f"Transformers version: {transformers.__version__} (must be >=4.26)")

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

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

#tunable parameters
MODEL_NAME = "ViT-bigG-14"        # alternatives: "ViT-H-14", "ViT-L-14", "ViT-B-32"
PRETRAINED_WEIGHTS = "laion2b_s39b_b160k"  # alternatives: "laion2b_s34b_b88k", "laion400m_e32"
TOP_K = 10                        # how many top candidates before reranking
PROMPT_VERSION = "rich"           # options: "simple", "rich"
BLEND_WEIGHTS = (0.7, 0.3)        # (clip_weight, blip_weight)

# Load data
df = pd.read_csv("/content/drive/MyDrive/MediaEval - AFourP/newsarticles.csv")
IMAGE_FOLDER = "/content/drive/MyDrive/MediaEval - AFourP/newsimages"

# Load model
print(f"Loading {MODEL_NAME} with weights {PRETRAINED_WEIGHTS}...")
clip_model, _, preprocess = open_clip.create_model_and_transforms(
    MODEL_NAME, pretrained=PRETRAINED_WEIGHTS
)
clip_model = clip_model.to(device)
tokenizer = open_clip.get_tokenizer(MODEL_NAME)

# Compute embeddings
image_embeddings = {}
print("🔄 Computing OpenCLIP image embeddings...")
for file in tqdm(os.listdir(IMAGE_FOLDER)):
    if not file.endswith(".jpg"):
        continue
    image_id = file.split(".")[0]
    img_path = os.path.join(IMAGE_FOLDER, file)
    try:
        image = preprocess(Image.open(img_path).convert("RGB")).unsqueeze(0).to(device)
        with torch.no_grad():
            emb = clip_model.encode_image(image)
            emb /= emb.norm(dim=-1, keepdim=True)
        image_embeddings[image_id] = emb.cpu()
    except:
        pass

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

def build_prompt(row, version="simple"):
    title = str(row["article_title"])
    tags = ",".join(str(row["article_tags"]).split(";")[:5])

    if version == "simple":
        return f"News headline: {title}. Related topics: {tags}."
    elif version == "rich":
        return f"Headline: {title}. Topics: {tags}. Context: A photo illustrating people, events, or locations described."
    else:
        return title

# CLIP Reranker (fallback)
# ===============================
def rerank_with_clip(text, top_images):
    text_tokens = tokenizer([text])
    with torch.no_grad():
        text_emb = clip_model.encode_text(text_tokens.to(device))
        text_emb /= text_emb.norm(dim=-1, keepdim=True)

    best_img, best_score = None, -1
    for img_id, img_emb, _ in top_images:
        sim = torch.cosine_similarity(text_emb.cpu(), img_emb).item()
        if sim > best_score:
            best_score = sim
            best_img = img_id
    return best_img, best_score

!pip install --upgrade --quiet transformers accelerate timm

# Retrieval Loop
# ===============================
results = []
similarities = []

group_name = "ShivShiv2"
approach_name = "OPENCLIP_BLIPITM" if use_blip_itm else "OPENCLIP_CLIPRERANK"
out_folder = f"{group_name}/RET_{approach_name}_LARGE"
os.makedirs(out_folder, exist_ok=True)

print(f"Running retrieval with {approach_name} reranker...")

for _, row in tqdm(df.iterrows(), total=len(df)):
    text = build_prompt(row, version=PROMPT_VERSION)

    # Stage 1: Find Top-K similar images (CLIP)
    text_tokens = tokenizer([text])
    with torch.no_grad():
        text_emb = clip_model.encode_text(text_tokens.to(device))
        text_emb /= text_emb.norm(dim=-1, keepdim=True)

    scored_embeddings = [
        (img_id, img_emb, torch.cosine_similarity(text_emb.cpu(), img_emb).item())
        for img_id, img_emb in image_embeddings.items()
    ]
    top_k = sorted(scored_embeddings, key=lambda x: x[2], reverse=True)[:TOP_K]

    # Stage 2: Rerank with BLIP (if available)
    if use_blip_itm:
        best_img, best_score = None, -1
        for img_id, _, clip_score in top_k:
            img_path = os.path.join(IMAGE_FOLDER, img_id + ".jpg")
            image = Image.open(img_path).convert("RGB")
            inputs = blip_processor(images=image, text=text, return_tensors="pt").to(device)
            with torch.no_grad():
                outputs = blip_model(**inputs)
                probs = outputs.logits.softmax(dim=1)
                blip_score = probs[:, 1].item()

            # Blend CLIP + BLIP
            final_score = BLEND_WEIGHTS[0] * clip_score + BLEND_WEIGHTS[1] * blip_score

            if final_score > best_score:
                best_score = final_score
                best_img = img_id
        best_img_id, rerank_score = best_img, best_score
    else:
        best_img_id, rerank_score = rerank_with_clip(text, top_k)

    similarities.append(rerank_score)

    # Save output
    img = Image.open(os.path.join(IMAGE_FOLDER, best_img_id + ".jpg")).convert("RGB")
    img = img.resize((460, 260))
    img.save(f"{out_folder}/{row['article_id']}_{group_name}_{approach_name}.png")
    results.append((row['article_id'], best_img_id, rerank_score))

# Results + Metrics
# ===============================
avg_score = sum(similarities) / len(similarities)
print(f"\nAverage match score ({approach_name}): {avg_score:.4f}")

results_df = pd.DataFrame(results, columns=["article_id", "retrieved_image_id", "score"])
results_df.to_csv(f"{group_name}_retrieval_results.csv", index=False)
print(f"Results saved to {group_name}_retrieval_results.csv")
print(f"Submission images saved in {out_folder}")

# Display a visual preview of the first 3 results
from IPython.display import display
print("\n Previewing first 3 retrievals:")
for _, row in results_df.head(3).iterrows():
    article_row = df[df["article_id"] == row["article_id"]].iloc[0]
    prompt = build_prompt(article_row)
    print(f"\n {article_row['article_title']}")
    print(f"Prompt: {prompt}")
    print(f"Retrieved: {row['retrieved_image_id']} | Score: {row['score']:.4f}")
    img_path = os.path.join(IMAGE_FOLDER, row['retrieved_image_id'] + ".jpg")
    display(Image.open(img_path).resize((460, 260)))