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

Time, Log Image Search Performance

Browse Source
This commit is contained in:
Debanjum Singh Solanky
2022-08-27 23:49:01 +03:00
parent c3ca99841b
commit 2eae32d743
1 changed files with 15 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/search_type/image_search.py
View File

@@ -4,6 +4,7 @@ import glob
import pathlib import pathlib
import copy import copy
import shutil import shutil
import time
# External Packages # External Packages
from sentence_transformers import SentenceTransformer, util from sentence_transformers import SentenceTransformer, util
@@ -16,6 +17,7 @@ from src.utils.helpers import get_absolute_path, resolve_absolute_path, load_mod
import src.utils.exiftool as exiftool import src.utils.exiftool as exiftool
from src.utils.config import ImageSearchModel from src.utils.config import ImageSearchModel
from src.utils.rawconfig import ImageContentConfig, ImageSearchConfig from src.utils.rawconfig import ImageContentConfig, ImageSearchConfig
from src.utils import state
def initialize_model(search_config: ImageSearchConfig): def initialize_model(search_config: ImageSearchConfig):
@@ -145,22 +147,34 @@ def query(raw_query, count, model: ImageSearchModel):
print(f"Find Images similar to Image at {query_imagepath}") print(f"Find Images similar to Image at {query_imagepath}")
else: else:
query = raw_query query = raw_query
if model.verbose > 0: if state.verbose > 0:
print(f"Find Images by Text: {query}") print(f"Find Images by Text: {query}")
# Now we encode the query (which can either be an image or a text string) # Now we encode the query (which can either be an image or a text string)
start = time.time()
query_embedding = model.image_encoder.encode([query], convert_to_tensor=True, show_progress_bar=False) query_embedding = model.image_encoder.encode([query], convert_to_tensor=True, show_progress_bar=False)
end = time.time()
if state.verbose > 1:
print(f"Query Encode Time: {end - start:.3f} seconds")
# Compute top_k ranked images based on cosine-similarity b/w query and all image embeddings. # Compute top_k ranked images based on cosine-similarity b/w query and all image embeddings.
start = time.time()
image_hits = {result['corpus_id']: {'image_score': result['score'], 'score': result['score']} image_hits = {result['corpus_id']: {'image_score': result['score'], 'score': result['score']}
for result for result
in util.semantic_search(query_embedding, model.image_embeddings, top_k=count)[0]} in util.semantic_search(query_embedding, model.image_embeddings, top_k=count)[0]}
end = time.time()
if state.verbose > 1:
print(f"Search Time: {end - start:.3f} seconds")
# Compute top_k ranked images based on cosine-similarity b/w query and all image metadata embeddings. # Compute top_k ranked images based on cosine-similarity b/w query and all image metadata embeddings.
if model.image_metadata_embeddings: if model.image_metadata_embeddings:
start = time.time()
metadata_hits = {result['corpus_id']: result['score'] metadata_hits = {result['corpus_id']: result['score']
for result for result
in util.semantic_search(query_embedding, model.image_metadata_embeddings, top_k=count)[0]} in util.semantic_search(query_embedding, model.image_metadata_embeddings, top_k=count)[0]}
end = time.time()
if state.verbose > 1:
print(f"Metadata Search Time: {end - start:.3f} seconds")
# Sum metadata, image scores of the highest ranked images # Sum metadata, image scores of the highest ranked images
for corpus_id, score in metadata_hits.items(): for corpus_id, score in metadata_hits.items():