clip_inferencing.py

import torch
import torch.nn.functional as F
from transformers import DistilBertTokenizer
from tqdm.autonotebook import tqdm
import pickle

from clip_model import CLIPModel
from configuration import CFG

import matplotlib.pyplot as plt
import cv2

def load_model(model_path):
 model = CLIPModel().to(CFG.device)
 model.load_state_dict(torch.load(model_path, map_location=CFG.device))
 model.eval()
 return model

def load_df():
 with open("pickles/valid_df.pkl", 'rb') as file:
 valid_df = pickle.load(file)
 return valid_df

def load_image_embeddings():
 with open("pickles/image_embeddings.pkl", 'rb') as file:
 image_embeddings = pickle.load(file)
 return image_embeddings

def find_matches(model, image_embeddings, query, image_filenames, n=9):
 tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
 encoded_query = tokenizer([query])
 batch = {
 key: torch.tensor(values).to(CFG.device)
 for key, values in encoded_query.items()
 }
 with torch.no_grad():
 text_features = model.text_encoder(
 input_ids=batch["input_ids"], attention_mask=batch["attention_mask"]
 )
 text_embeddings = model.text_projection(text_features)
 
 image_embeddings_n = F.normalize(image_embeddings, p=2, dim=-1)
 text_embeddings_n = F.normalize(text_embeddings, p=2, dim=-1)
 dot_similarity = text_embeddings_n @ image_embeddings_n.T
 
 values, indices = torch.topk(dot_similarity.squeeze(0), n * 5)
 matches = [image_filenames[idx] for idx in indices[::5]]
 
 _, axes = plt.subplots(3, 3, figsize=(10, 10))
 for match, ax in zip(matches, axes.flatten()):
 image = cv2.imread(f"Images/{match}")
 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 ax.imshow(image)
 ax.axis("off")
 
 plt.show()

def inference():
 valid_df = load_df()
 image_embeddings = load_image_embeddings()
 find_matches(load_model(model_path="model/best.pt"), 
 image_embeddings, 
 query="dogs on the grass", 
 image_filenames=valid_df['image'].values, n=9)