Spaces:
Sleeping
Sleeping
import numpy as np | |
import torch as th | |
import cv2 | |
from PIL import Image | |
import time | |
import copy | |
import numpy as np | |
import open_clip | |
import torchvision.transforms as T | |
def calculate_l2_distances(query, gallery): | |
return np.linalg.norm(gallery - query, axis=1) | |
def get_k_nearest_neighbors(distances, k): | |
indices = np.argsort(distances)[:k] | |
return indices | |
def get_similarity_l2(embeddings_gallery, embeddings_query, k): | |
print('Processing indices...') | |
s = time.time() | |
scores = [] | |
indices = [] | |
for query in embeddings_query: | |
distances = calculate_l2_distances(query, embeddings_gallery) | |
nearest_indices = get_k_nearest_neighbors(distances, k) | |
scores.append(distances[nearest_indices]) | |
indices.append(nearest_indices) | |
e = time.time() | |
print(f'Finished processing indices, took {e - s}s') | |
return np.array(scores), np.array(indices) | |
def get_transform(): | |
transform = T.Compose([ | |
T.Resize( | |
size=(224, 224), | |
interpolation=T.InterpolationMode.BICUBIC, | |
antialias=True), | |
T.ToTensor(), | |
T.Normalize( | |
mean=(0.48145466, 0.4578275, 0.40821073), | |
std=(0.26862954, 0.26130258, 0.27577711) | |
) | |
]) | |
return transform | |
def convert_indices_to_labels(indices, labels): | |
indices_copy = copy.deepcopy(indices) | |
for row in indices_copy: | |
for j in range(len(row)): | |
row[j] = labels[row[j]] | |
return indices_copy | |
def read_image(image_file): | |
img = cv2.imread( | |
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION | |
) | |
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
if img is None: | |
raise ValueError('Failed to read {}'.format(image_file)) | |
return img | |
def transform_img(image): | |
img = image | |
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
if isinstance(img, np.ndarray): | |
img = Image.fromarray(img) | |
transform = get_transform() | |
img = transform(img) | |
return img | |
def get_feature_vector_img(model, imgs, epoch=10, use_cuda=False): | |
features = [] | |
if use_cuda: | |
imgs = imgs.cuda() | |
x = (model(imgs)).detach().cpu().numpy().astype(np.float32) # .half() | |
features.append(x) | |
return np.concatenate(features, axis=0) | |
def Model(): | |
backbone = open_clip.create_model_and_transforms('ViT-H-14', None)[0].visual | |
backbone.load_state_dict(th.load('./model1.pt')) #https://huggingface.co/hca97/aicrowd-visual-recognition-models/blob/main/model1.pt | |
# backbone, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:Thenujan/ViT-H-14') | |
backbone.eval() # Dropping unecessary layers | |
return backbone | |
def find(image): | |
# img_path = '/kaggle/input/vprtestdata/public_dataset/queries/abiding-inchworm-of-ultimate-freedom.jpeg' | |
# image = read_image(img_path) | |
image = np.array(image) | |
image = transform_img(image) | |
image = image.unsqueeze(dim=0) | |
feature_vectors_gallery = np.load('./VPR image similarity search/gallary_embedding/feature_vectors_gallery.npy') | |
model = Model() | |
feature_vectors_query = get_feature_vector_img(model, image, 1) | |
_, indices = get_similarity_l2(feature_vectors_gallery, feature_vectors_query, 1000) | |
indices = indices.tolist() | |
return indices |