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

Use input filter in image search setup. Input filter wasn't used earlier

Browse Source
This commit is contained in:
Debanjum Singh Solanky
2022-08-05 01:24:21 +03:00
parent ca5a8bd113
commit d5b43eb836
2 changed files with 16 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/main.py
View File

@@ -293,14 +293,14 @@ def run():
global config_file global config_file
config_file = args.config_file config_file = args.config_file
# Store the verbose flag
global verbose
verbose = args.verbose
# Store the raw config data. # Store the raw config data.
global config global config
config = args.config config = args.config
# Store the verbose flag
global verbose
verbose = args.verbose
# Set device to GPU if available # Set device to GPU if available
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")

17
src/search_type/image_search.py
View File

@@ -1,5 +1,6 @@
# Standard Packages # Standard Packages
import argparse import argparse
import glob
import pathlib import pathlib
import copy import copy
import shutil import shutil
@@ -11,7 +12,7 @@ from tqdm import trange
import torch import torch
# Internal Packages # Internal Packages
from src.utils.helpers import resolve_absolute_path, load_model from src.utils.helpers import get_absolute_path, resolve_absolute_path, load_model
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
@@ -213,13 +214,19 @@ def setup(config: ImageContentConfig, search_config: ImageSearchConfig, regenera
encoder = initialize_model(search_config) encoder = initialize_model(search_config)
# Extract Entries # Extract Entries
image_directories = [resolve_absolute_path(directory, strict=True) for directory in config.input_directories] absolute_image_files, filtered_image_files = set(), set()
image_names = extract_entries(image_directories, verbose) if config.input_directories:
image_directories = [resolve_absolute_path(directory, strict=True) for directory in config.input_directories]
absolute_image_files = set(extract_entries(image_directories, verbose))
if config.input_filter:
filtered_image_files = set(glob.glob(get_absolute_path(config.input_filter)))
all_image_files = sorted(list(absolute_image_files | filtered_image_files))
# Compute or Load Embeddings # Compute or Load Embeddings
embeddings_file = resolve_absolute_path(config.embeddings_file) embeddings_file = resolve_absolute_path(config.embeddings_file)
image_embeddings, image_metadata_embeddings = compute_embeddings( image_embeddings, image_metadata_embeddings = compute_embeddings(
image_names, all_image_files,
encoder, encoder,
embeddings_file, embeddings_file,
batch_size=config.batch_size, batch_size=config.batch_size,
@@ -227,7 +234,7 @@ def setup(config: ImageContentConfig, search_config: ImageSearchConfig, regenera
use_xmp_metadata=config.use_xmp_metadata, use_xmp_metadata=config.use_xmp_metadata,
verbose=verbose) verbose=verbose)
return ImageSearchModel(image_names, return ImageSearchModel(all_image_files,
image_embeddings, image_embeddings,
image_metadata_embeddings, image_metadata_embeddings,
encoder, encoder,