Spaces:
Runtime error
Runtime error
File size: 1,863 Bytes
77fd092 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import streamlit as st
import torch
import torchvision.transforms as transforms
from PIL import Image
import os
import clip
import numpy as np
import torch.nn.functional as F
import matplotlib.pyplot as plt
device = 'cpu'
model_path = "weights/ViT-B-32.pt"
model, preprocess = clip.load('ViT-B/32', device)
def get_similarity_score(text_query, image_features):
text_tokens = clip.tokenize([text_query]).to(device)
with torch.no_grad():
text_features = model.encode_text(text_tokens).squeeze(0)
text_features= F.normalize(text_features, p=2, dim=-1)
similarity_score = text_features @ image_features.T * 100.0
similarity_score = similarity_score.squeeze(0)
return similarity_score
def create_filelist(path_to_imagefolder):
image_folder = path_to_imagefolder
image_paths = []
for filename in os.listdir(image_folder):
if filename.endswith(".jpg") or filename.endswith(".jpeg") or filename.endswith(".png"):
image_path = os.path.join(image_folder, filename)
image_paths.append(image_path)
file_paths = image_paths
return file_paths
def load_embeddings(path_to_emb_file):
features = np.load(path_to_emb_file)
features = torch.from_numpy(features)
return features
def find_matches(image_embeddings, query, image_filenames, n=6):
text_query = query
features = image_embeddings
similarity_scores = []
for emb in features:
emb /= emb.norm(dim=-1, keepdim=True)
similarity_score = get_similarity_score(text_query, emb)
similarity_scores.append(similarity_score)
similarity_scores = torch.stack(similarity_scores)
values, indices = torch.topk(similarity_scores.squeeze(0), 6)
matches = [image_filenames[idx] for idx in indices]
return matches |