sij/khoj
1
0
Fork 0
mirror of https://github.com/khoj-ai/khoj.git synced 2024-11-23 23:48:56 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 16

Make image batch size to encode configurable via config.yml

Browse source
This commit is contained in:
Debanjum Singh Solanky 2021-09-16 10:51:39 -07:00
parent 41c328dae0
commit 3afe054312
4 changed files with 9 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
sample_config.yml
View file

@ -11,7 +11,7 @@ content-type:
image: image:
embeddings-file: '.image_embeddings.pt' embeddings-file: '.image_embeddings.pt'
batch-size: 50
search-type: search-type:
asymmetric: asymmetric:

5
src/main.py
View file

@ -187,8 +187,9 @@ if __name__ == '__main__':
image_names, image_embeddings, image_metadata_embeddings, image_encoder = image_search.setup( image_names, image_embeddings, image_metadata_embeddings, image_encoder = image_search.setup(
pathlib.Path(image_config['input-directory']), pathlib.Path(image_config['input-directory']),
pathlib.Path(image_config['embeddings-file']), pathlib.Path(image_config['embeddings-file']),
args.regenerate, batch_size=image_config['batch-size'],
args.verbose) regenerate=args.regenerate,
verbose=args.verbose)
# Start Application Server # Start Application Server
uvicorn.run(app) uvicorn.run(app)

7
src/search_type/image_search.py
View file

@ -30,7 +30,7 @@ def extract_entries(image_directory, verbose=0):
return image_names return image_names
def compute_embeddings(image_names, model, embeddings_file, regenerate=False, verbose=0): def compute_embeddings(image_names, model, embeddings_file, batch_size=50, regenerate=False, verbose=0):
"Compute (and Save) Embeddings or Load Pre-Computed Embeddings" "Compute (and Save) Embeddings or Load Pre-Computed Embeddings"
image_embeddings = None image_embeddings = None
image_metadata_embeddings = None image_metadata_embeddings = None
@ -51,7 +51,6 @@ def compute_embeddings(image_names, model, embeddings_file, regenerate=False, ve
if verbose > 0: if verbose > 0:
print(f"Loading the {len(image_names)} images into memory") print(f"Loading the {len(image_names)} images into memory")
batch_size = 50
if image_embeddings is None: if image_embeddings is None:
image_embeddings = [] image_embeddings = []
for index in trange(0, len(image_names), batch_size): for index in trange(0, len(image_names), batch_size):
@ -137,7 +136,7 @@ def collate_results(hits, image_names, image_directory, count=5):
in hits[0:count]] in hits[0:count]]
def setup(image_directory, embeddings_file, regenerate=False, verbose=0): def setup(image_directory, embeddings_file, batch_size=50, regenerate=False, verbose=0):
# Initialize Model # Initialize Model
model = initialize_model() model = initialize_model()
@ -147,7 +146,7 @@ def setup(image_directory, embeddings_file, regenerate=False, verbose=0):
# Compute or Load Embeddings # Compute or Load Embeddings
embeddings_file = resolve_absolute_path(embeddings_file) embeddings_file = resolve_absolute_path(embeddings_file)
image_embeddings, image_metadata_embeddings = compute_embeddings(image_names, model, embeddings_file, regenerate=regenerate, verbose=verbose) image_embeddings, image_metadata_embeddings = compute_embeddings(image_names, model, embeddings_file, batch_size=batch_size, regenerate=regenerate, verbose=verbose)
return image_names, image_embeddings, image_metadata_embeddings, model return image_names, image_embeddings, image_metadata_embeddings, model

3
src/utils/cli.py
View file

@ -56,7 +56,8 @@ default_config = {
}, },
'image': 'image':
{ {
'embeddings-file': '.image_embeddings.pt' 'embeddings-file': '.image_embeddings.pt',
'batch-size': 50
}, },
'music': 'music':
{ {