From ca0a22f4dd77811fe6e266b3208aea4121a12afc Mon Sep 17 00:00:00 2001 From: Debanjum Singh Solanky Date: Sun, 8 Aug 2021 23:11:15 -0700 Subject: [PATCH] Search for images similar to query image provided by the user 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 --- image-search.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/image-search.py b/image-search.py index f4bee31f..bec36783 100644 --- a/image-search.py +++ b/image-search.py @@ -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)