sij/khoj
1
0
Fork 0
mirror of https://github.com/khoj-ai/khoj.git synced 2025-02-17 08:04:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Batch encode XMP metadata from images too for image_search

Browse source
This commit is contained in:
Debanjum Singh Solanky 2021-09-16 11:11:36 -07:00
parent 3afe054312
commit a4a23d7a72
1 changed files with 4 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
src/search_type/image_search.py
View file

@ -61,8 +61,10 @@ def compute_embeddings(image_names, model, embeddings_file, batch_size=50, regen
print(f"Saved computed embeddings to {embeddings_file}")
if image_metadata_embeddings is None:
image_metadata = [extract_metadata(image_name, verbose) for image_name in image_names],
image_metadata_embeddings = model.encode(image_metadata, batch_size=batch_size, convert_to_tensor=True, show_progress_bar=True)
image_metadata_embeddings = []
for index in trange(0, len(image_names), batch_size):
image_metadata = [extract_metadata(image_name, verbose) for image_name in image_names[index:index+batch_size]]
image_metadata_embeddings += model.encode(image_metadata, convert_to_tensor=True, batch_size=batch_size)
torch.save(image_metadata_embeddings, f"{embeddings_file}_metadata")
if verbose > 0:
print(f"Saved computed metadata embeddings to {embeddings_file}_metadata")