import os import cv2 import numpy as np from PIL import Image import json import gradio as gr import matplotlib.pyplot as plt import torch from torch import nn import torch.nn.functional as F import timm from transformers import DistilBertModel, DistilBertConfig, DistilBertTokenizer class CFG: image_path = './images' captions_path = './captions' batch_size = 64 num_workers = 4 head_lr = 1e-3 image_encoder_lr = 1e-4 text_encoder_lr = 1e-5 weight_decay = 1e-3 patience = 1 factor = 0.8 epochs = 2 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_name = 'resnet50' image_embedding = 2048 text_encoder_model = "distilbert-base-uncased" text_embedding = 768 text_tokenizer = "distilbert-base-uncased" max_length = 200 pretrained = True # for both image encoder and text encoder trainable = True # for both image encoder and text encoder temperature = 1.0 # image size size = 224 # for projection head; used for both image and text encoders num_projection_layers = 1 projection_dim = 256 dropout = 0.1 class ImageEncoder(nn.Module): """ Encode images to a fixed size vector """ def __init__( self, model_name=CFG.model_name, pretrained=CFG.pretrained, trainable=CFG.trainable ): super().__init__() self.model = timm.create_model( model_name, pretrained, num_classes=0, global_pool="avg" ) for p in self.model.parameters(): p.requires_grad = trainable def forward(self, x): return self.model(x) class TextEncoder(nn.Module): def __init__(self, model_name=CFG.text_encoder_model, pretrained=CFG.pretrained, trainable=CFG.trainable): super().__init__() if pretrained: self.model = DistilBertModel.from_pretrained(model_name) else: self.model = DistilBertModel(config=DistilBertConfig()) for p in self.model.parameters(): p.requires_grad = trainable # we are using the CLS token hidden representation as the sentence's embedding self.target_token_idx = 0 def forward(self, input_ids, attention_mask): output = self.model(input_ids=input_ids, attention_mask=attention_mask) last_hidden_state = output.last_hidden_state return last_hidden_state[:, self.target_token_idx, :] class ProjectionHead(nn.Module): def __init__( self, embedding_dim, projection_dim=CFG.projection_dim, dropout=CFG.dropout ): super().__init__() self.projection = nn.Linear(embedding_dim, projection_dim) self.gelu = nn.GELU() self.fc = nn.Linear(projection_dim, projection_dim) self.dropout = nn.Dropout(dropout) self.layer_norm = nn.LayerNorm(projection_dim) def forward(self, x): projected = self.projection(x) x = self.gelu(projected) x = self.fc(x) x = self.dropout(x) x = x + projected x = self.layer_norm(x) return x class CLIPModel(nn.Module): def __init__( self, temperature=CFG.temperature, image_embedding=CFG.image_embedding, text_embedding=CFG.text_embedding, ): super().__init__() self.image_encoder = ImageEncoder() self.text_encoder = TextEncoder() self.image_projection = ProjectionHead(embedding_dim=image_embedding) self.text_projection = ProjectionHead(embedding_dim=text_embedding) self.temperature = temperature def forward(self, batch): # Getting Image and Text Features image_features = self.image_encoder(batch["image"]) text_features = self.text_encoder( input_ids=batch["input_ids"], attention_mask=batch["attention_mask"] ) # Getting Image and Text Embeddings (with same dimension) image_embeddings = self.image_projection(image_features) text_embeddings = self.text_projection(text_features) # Calculating the Loss images_similarity = image_embeddings @ text_embeddings.T / self.temperature texts_similarity = images_similarity.T labels = torch.arange(batch["image"].shape[0]).long().to(CFG.device) total_loss = ( F.cross_entropy(images_similarity, labels) + F.cross_entropy(texts_similarity, labels) ) / 2 return total_loss def find_matches_cpu(model, image_embeddings, query, image_filenames, n=4): tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer) encoded_query = tokenizer([query]) batch = { key: torch.tensor(values).to('cpu') 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]] return matches def rle_decode(img_rle_array, img_name, img_size): encoded_image = img_rle_array # Initialize variables for decoding decoded_image = [] for i in range(0, len(encoded_image), 2): pixel_value = encoded_image[i] run_length = encoded_image[i + 1] decoded_image.extend([pixel_value] * run_length) # Convert the decoded image back to a NumPy array decoded_array = np.array(decoded_image, dtype=np.uint8) # Reshape the decoded array to the original image shape (224, 224) decoded_image = decoded_array.reshape(img_size) # Use original shape # Create a PIL Image from the decoded array decoded_image = Image.fromarray(decoded_image) decoded_image_save_path = './' + str(img_name) # Save or display the decoded image decoded_image.save(decoded_image_save_path) # Save the decoded image to a file return decoded_image_save_path def get_matched_image(matches, val_file_dict_loaded): img_size = (112, 112) match_img_list = [] for img_name in matches: img_rle_array = val_file_dict_loaded[img_name] decoded_image_path = rle_decode(img_rle_array, img_name, img_size) match_img_list.append(decoded_image_path) return match_img_list def get_grayscale_image(text_query): model_inf = CLIPModel().to('cpu') model_inf.load_state_dict(torch.load('best_clip_model_cpu.pt', map_location='cpu')) clip_image_embeddings_np_inf = np.load('clip_image_embeddings.npy') image_embeddings_inf = torch.tensor(clip_image_embeddings_np_inf) img_file_names = np.load('val_img_file_names.npy',allow_pickle=True) with open("val_imgs_rle_encode.json", "r") as json_file: val_file_dict_loaded = json.load(json_file) matches = find_matches_cpu(model_inf, image_embeddings_inf, query=text_query, image_filenames=img_file_names, n=1) matched_images = get_matched_image(matches, val_file_dict_loaded) return matched_images def gradio_fn(text): text_query = str(text) match_img_list = get_grayscale_image(text_query) pil_img = Image.open(match_img_list[0]) pil_img = pil_img.resize((224, 224)) np_img_array = np.array(pil_img) return np_img_array demo = gr.Interface(fn=gradio_fn, inputs=gr.Textbox(info="Enter the description of image you wish to search, CLIP will give the best image available in corpus that matches your search"), outputs=gr.Image(height=224, width=224), title="CLIP Image Search") demo.launch(share=True)