flava.py

import torch
import numpy as np
from tqdm import tqdm
import torch.nn as nn
from transformers import FlavaProcessor, FlavaForPreTraining, BertTokenizer, FlavaFeatureExtractor


class FlavaWrapper:
    def __init__(self, root_dir, device):
        self.model = FlavaForPreTraining.from_pretrained("facebook/flava-full", cache_dir=root_dir).eval().to(device)
        self.feature_extractor = FlavaFeatureExtractor.from_pretrained("facebook/flava-full", cache_dir=root_dir)
        self.tokenizer = BertTokenizer.from_pretrained("facebook/flava-full", cache_dir=root_dir)
        self.processor = FlavaProcessor.from_pretrained("facebook/flava-full", cache_dir=root_dir)
        self.device = device
    
    
    @torch.no_grad()
    def get_text_embeddings(self, texts, text_batch_size=64, normalize=False):
        num_text = len(texts)
        text_embeds = []
        for i in tqdm(range(0, num_text, text_batch_size)):
            text = texts[i: min(num_text, i+text_batch_size)]
            text_input = self.tokenizer(text=text, return_tensors="pt", padding="max_length", max_length=77).to(self.device)
            text_feats = self.model.flava.get_text_features(**text_input).cpu().numpy()[:, 0, :]
            if normalize:
                text_feats = text_feats / np.linalg.norm(text_feats, axis=1, keepdims=True)          
            text_embeds.append(text_feats)   
            
        return np.concatenate(text_embeds, axis=0)
    

    @torch.no_grad()
    def get_image_embeddings(self, image_loader, normalize=False):
        image_embeds = []
        for batch in tqdm(image_loader):
            images = batch["image"]
            inputs = self.feature_extractor(images=images, return_tensors="pt").to(self.device)
            image_feats = self.model.flava.get_image_features(**inputs).cpu().numpy()[:, 0, :]
            if normalize:
                image_feats = image_feats / np.linalg.norm(image_feats, axis=1, keepdims=True)
            image_embeds.append(image_feats)

        return np.concatenate(image_embeds, axis=0)
    
    
    def get_retrieval_scores_dataset(self, loader):
        texts = loader.dataset.text
        text_embeds = self.get_text_embeddings(texts, normalize=True)
        image_embeds = self.get_image_embeddings(loader, normalize=True)
        scores = image_embeds @ text_embeds.T
        return scores
    
    
    @torch.no_grad()
    def get_retrieval_scores_batched(self, joint_loader):
        """Computes the scores for each image_option / caption_option pair in the joint loader.

        Args:
            joint_loader (DataLoader): batches have "image_options" and "caption_options" fields.
            "image_options" is a list of images, and "caption_options" is a list of captions.

        Returns:
            all_scores: A numpy array containing the scores of the shape NxKxL,
            where N is the number of test cases, K is the number of image options per the test case,
            and L is the number of caption options per the test case.
        """
        scores = []
        for batch in tqdm(joint_loader):
            
            batch_scores = []
            for i_option in batch["image_options"]:
                im_scores = []
                for c_option in batch["caption_options"]:
                    inputs = self.processor(text=list(c_option), images=list(i_option), return_tensors="pt", padding="max_length", max_length=77, return_codebook_pixels=True, return_image_mask=True).to(self.device)
                    inputs["input_ids_masked"] = inputs["input_ids"].detach().clone() 
                    inputs["bool_masked_pos"] = torch.zeros_like(inputs["bool_masked_pos"])
                    outputs = self.model(**inputs)
                    score = nn.functional.softmax(outputs.itm_logits, dim=-1)[:, 1:2].cpu().numpy()
                    im_scores.append(np.expand_dims(score, -1))
                batch_scores.append(np.concatenate(im_scores, axis=-1))
                
            batch_scores = np.concatenate(batch_scores, axis=1)
            scores.append(batch_scores)
        
        all_scores = np.concatenate(scores, axis=0) # N x K x L
        return all_scores