sij/khoj
1
0
Fork 0
mirror of https://github.com/khoj-ai/khoj.git synced 2024-12-17 18:17:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

Search for images similar to query image provided by the user

Browse source 
Example user passes path to an image in query. e.g ~/Pictures/photo.jpg
The script should return images in images_embedding most similar to
the query image
This commit is contained in:
Debanjum Singh Solanky 2021-08-08 23:11:15 -07:00
parent 00d0065c5b
commit ca0a22f4dd
1 changed files with 12 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
image-search.py
View file

@ -48,8 +48,17 @@ def compute_embeddings(image_names, model, embeddings_file, verbose=False):
return image_embeddings
def search(query, image_embeddings, model, count=3):
# First, we encode the query (which can either be an image or a text string)
def search(query, image_embeddings, model, count=3, verbose=False):
# Set query to image content if query is a filepath
if pathlib.Path(query).expanduser().is_file():
query_imagepath = pathlib.Path(query).expanduser().resolve(strict=True)
query = copy.deepcopy(Image.open(query_imagepath))
if verbose:
print(f"Find Images similar to Image at {query_imagepath}")
else:
print(f"Find Images by Text: {query}")
# Now we encode the query (which can either be an image or a text string)
query_embedding = model.encode([query], convert_to_tensor=True, show_progress_bar=False)
# Then, we use the util.semantic_search function, which computes the cosine-similarity
@ -95,7 +104,7 @@ if __name__ == '__main__':
exit(0)
# query notes
hits = search(user_query, image_embeddings, model, args.results_count)
hits = search(user_query, image_embeddings, model, args.results_count, args.verbose)
# render results
render_results(hits, image_names, args.image_directory, count=args.results_count)