text2image / funcs /fiass_similaruty.py
konstantinG's picture
Upload 18 files
77fd092
import torch
import numpy as np
import clip
import torch.nn.functional as F
import faiss
device = 'cpu'
model_path = "weights/ViT-B-32.pt"
model, preprocess = clip.load('ViT-B/32', device)
def load_embeddings(path_to_emb_file):
features = np.load(path_to_emb_file)
features = torch.from_numpy(features)
features = features.squeeze(1)
features = F.normalize(features, p=2, dim=-1)
return features
def encode_text(query):
text = clip.tokenize([query]).to(device)
text_features = model.encode_text(text).to("cpu")
text_features= F.normalize(text_features, p=2, dim=-1)
text_features = text_features.to("cpu").detach().numpy()
return text_features
def find_matches_fiass(image_embeddings, query, image_filenames, n=5):
features = image_embeddings
index = faiss.IndexFlatL2(features.shape[1])
index.add(features)
text_features = encode_text(query)
_, I = index.search(text_features, n)
matches = [image_filenames[idx] for idx in I.squeeze(0)]
return matches