klbr/khoj
1
0
Fork 0
mirror of https://github.com/khoaliber/khoj.git synced 2026-03-09 05:39:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

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
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}") print(f"Saved computed embeddings to {embeddings_file}")
if image_metadata_embeddings is None: if image_metadata_embeddings is None:
image_metadata = [extract_metadata(image_name, verbose) for image_name in image_names], image_metadata_embeddings = []
image_metadata_embeddings = model.encode(image_metadata, batch_size=batch_size, convert_to_tensor=True, show_progress_bar=True) 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") torch.save(image_metadata_embeddings, f"{embeddings_file}_metadata")
if verbose > 0: if verbose > 0:
print(f"Saved computed metadata embeddings to {embeddings_file}_metadata") print(f"Saved computed metadata embeddings to {embeddings_file}_metadata")