interface.py

import torch
from transformers import CLIPProcessor, CLIPModel
import numpy as np
from PIL import Image
import gradio as gr
from sklearn.metrics.pairwise import cosine_similarity
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
from explain_image import generate_visual_explanation
from explain_text import word_occlusion_attribution


# Paths
IMAGE_EMB_FILE = "./data/image_embeds.pt"
TEXT_EMB_FILE = "./data/text_embeds.pt"
MAP_FILE = "./data/image_paths.txt"
CAPTION_FILE = "./data/captions.txt"

# Setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"[INFO] Using device: {device}")

print("[INFO] Loading CLIP model and processor...")
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
clip_model.eval()

# Load embeddings and data
print("[INFO] Loading cached embeddings...")
image_embeds = torch.load(IMAGE_EMB_FILE).to(device)
text_embeds = torch.load(TEXT_EMB_FILE).to(device)

with open(MAP_FILE, "r", encoding="utf-8") as f:
    caption_image_map = [line.strip() for line in f.readlines()]

with open(CAPTION_FILE, "r", encoding="utf-8") as f:
    captions = [line.strip().split("|", 1)[1] for line in f]

print(f"[INFO] Loaded {len(captions)} captions and {len(caption_image_map)} images.")

clip_transform = Compose([
    Resize(224, interpolation=Image.BICUBIC),
    CenterCrop(224),
    ToTensor(),
    Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
])


# Retrieval
def retrieve(text, image, explain):
    print(f"[INFO] Query received. Text: {bool(text)}, Image: {bool(image)}, Explain: {explain}")
    results = {"images": [], "captions": [], "explain": None}

    if text:
        print("[INFO] Processing text query...")
        inputs = clip_processor(text=[text], return_tensors="pt", truncation=True).to(device)
        with torch.no_grad():
            q_text_embed = clip_model.get_text_features(**inputs)
        q_text_embed = q_text_embed / q_text_embed.norm(p=2, dim=-1, keepdim=True)

        sim_img = cosine_similarity(q_text_embed.cpu(), image_embeds.cpu())[0]
        sim_txt = cosine_similarity(q_text_embed.cpu(), text_embeds.cpu())[0]

        top_imgs_raw = np.argsort(sim_img)[::-1]
        top_txts = np.argsort(sim_txt)[::-1][:3]
        seen = set()
        top_imgs = []
        for idx in top_imgs_raw:
            img_path = caption_image_map[idx]
            if img_path not in seen:
                seen.add(img_path)
                top_imgs.append(idx)
            if len(top_imgs) == 3:
                break

        results["images"] = [caption_image_map[i] for i in top_imgs]
        results["captions"] = [captions[i] for i in top_txts]

        
        if explain:
            print("[INFO] Generating text explanation...")
            # 1. Explain top retrieved texts
            caption_explanations = []
            for i in top_txts:
                retrieved_caption = captions[i]
                html = word_occlusion_attribution(retrieved_caption, q_text_embed, clip_model, clip_processor, device)
                caption_explanations.append(html)

            # 2. Explain top retrieved images (heatmaps showing influence of text)
            top_images = [caption_image_map[i] for i in top_imgs]
            print("==== top images,", top_images)
            image_heatmaps = []
            for img_path in top_images:
                retrieved_image = Image.open(img_path).convert("RGB")
                
                text_input = clip_processor(text=[text], return_tensors="pt").to(device)
                with torch.no_grad():
                    text_feat = clip_model.text_model(**text_input).last_hidden_state[:, 0, :]
                    reference_embed = clip_model.text_projection(text_feat)
                heatmap = generate_visual_explanation(target_image=retrieved_image, reference_embed=reference_embed[0], model=clip_model, clip_transform=clip_transform, device=device)
                image_heatmaps.append(heatmap)

            results["explain"] = ("<br><br>".join(caption_explanations), image_heatmaps)

        

    elif image:
        print("[INFO] Processing image query...")
        image = image.convert("RGB")
        inputs = clip_processor(images=image, return_tensors="pt").to(device)
        with torch.no_grad():
            q_img_embed = clip_model.get_image_features(**inputs)
        q_img_embed = q_img_embed / q_img_embed.norm(p=2, dim=-1, keepdim=True)

        sim_img = cosine_similarity(q_img_embed.cpu(), image_embeds.cpu())[0]
        sim_txt = cosine_similarity(q_img_embed.cpu(), text_embeds.cpu())[0]

        top_imgs_raw = np.argsort(sim_img)[::-1]
        top_txts = np.argsort(sim_txt)[::-1][:3]
        seen = set()
        top_imgs = []
        for idx in top_imgs_raw:
            img_path = caption_image_map[idx]
            if img_path not in seen:
                seen.add(img_path)
                top_imgs.append(idx)
            if len(top_imgs) == 3:
                break

        results["images"] = [caption_image_map[i] for i in top_imgs]
        results["captions"] = [captions[i] for i in top_txts]

        if explain:
            cap_expls = [word_occlusion_attribution(captions[i], q_img_embed, clip_model, clip_processor, device) for i in top_txts]
            
            query_input = clip_transform(image).unsqueeze(0).to(device)
            with torch.no_grad():
                reference_embed = clip_model.get_image_features(pixel_values=query_input) 
            img_img_expls = [generate_visual_explanation(target_image=Image.open(caption_image_map[i]), reference_embed=reference_embed[0], model=clip_model, clip_transform=clip_transform, device=device) for i in top_imgs]
            results["explain"] = ("<br><br>".join(cap_expls), img_img_expls)

    return results["images"], "\n\n".join(results["captions"]), results["explain"]


def ui(text, image, explainability):
    if (not text and not image) or (text and image):
        return None, "", "<span style='color:red'>Please provide either a text or an image query, not both.</span>", None
    
    imgs, caps, expl = retrieve(text, image, explainability)
    text_html = ""
    image_expls = []

    if isinstance(expl, tuple):
        text_html, image_expls = expl
        
    elif isinstance(expl, str):
        text_html = expl

    return imgs, caps, image_expls, text_html

with gr.Blocks() as demo:
    gr.Markdown("## Visual-Text Retrieval with Explainability")
    gr.Markdown("Use either a caption **or** an image to search. Turn on explainability to visualize what influenced the match.")

    with gr.Row():
        with gr.Column(scale=1):
            text_input = gr.Textbox(label="Text Query", placeholder="Describe an image...", lines=1)
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="Image Query")
        with gr.Column(scale=0.5):
            explain_checkbox = gr.Checkbox(label="Show Explainability")
            submit_btn = gr.Button("Submit", variant="primary")

    with gr.Row(visible=False) as output_row:
        with gr.Column():
            gallery_output = gr.Gallery(label="Top Retrieved Images", columns=3, height=300)
        with gr.Column():
            caption_output = gr.Textbox(label="Top Retrieved Captions", lines=4)

    with gr.Row(visible=False) as explain_row:
        with gr.Column():
            explain_images_output = gr.Gallery(label="Explanation Images", columns=3, height=300)
        with gr.Column():
            html_output = gr.HTML()
        
    # Reveal outputs only after valid input
    def update_ui(text, image, explainability):
        outputs = ui(text, image, explainability)
        if explainability:
            return (*outputs, gr.update(visible=True), gr.update(visible=True))
        else:
            return (*outputs, gr.update(visible=True), gr.update(visible=False))
            

    submit_btn.click(fn=update_ui,
                     inputs=[text_input, image_input, explain_checkbox],
                     outputs=[gallery_output, caption_output, explain_images_output, html_output, output_row, explain_row],
                     show_progress=True)

if __name__ == "__main__":
    print("[INFO] Starting Gradio interface...")
    demo.launch()