diff --git a/src/search_type/image_search.py b/src/search_type/image_search.py index 08b4ec5b..323556e7 100644 --- a/src/search_type/image_search.py +++ b/src/search_type/image_search.py @@ -136,12 +136,12 @@ def extract_metadata(image_name): def query(raw_query, count, model: ImageSearchModel): - # Set query to image content if query is a filepath - if pathlib.Path(raw_query).is_file(): - query_imagepath = resolve_absolute_path(pathlib.Path(raw_query), strict=True) + # Set query to image content if query is of form file:/path/to/file.png + if raw_query.startswith("file:") and pathlib.Path(raw_query[5:]).is_file(): + query_imagepath = resolve_absolute_path(pathlib.Path(raw_query[5:]), strict=True) query = copy.deepcopy(Image.open(query_imagepath)) query.thumbnail((640, query.height)) # scale down image for faster processing - logger.info(f"Find Images similar to Image at {query_imagepath}") + logger.info(f"Find Images by Image: {query_imagepath}") else: query = raw_query logger.info(f"Find Images by Text: {query}") diff --git a/tests/test_image_search.py b/tests/test_image_search.py index 374168f5..2a396f54 100644 --- a/tests/test_image_search.py +++ b/tests/test_image_search.py @@ -1,4 +1,5 @@ # Standard Modules +import logging from pathlib import Path from PIL import Image @@ -77,3 +78,42 @@ def test_image_search(content_config: ContentConfig, search_config: SearchConfig # Cleanup # Delete the image files copied to results directory actual_image_path.unlink() + + +# ---------------------------------------------------------------------------------------------------- +def test_image_search_by_filepath(content_config: ContentConfig, search_config: SearchConfig, caplog): + # Arrange + output_directory = resolve_absolute_path(web_directory) + model.image_search = image_search.setup(content_config.image, search_config.image, regenerate=False) + image_directory = content_config.image.input_directories[0] + + query = f"file:{image_directory.joinpath('kitten_park.jpg')}" + expected_image_path = f"{image_directory.joinpath('kitten_park.jpg')}" + + # Act + with caplog.at_level(logging.INFO, logger="src.search_type.image_search"): + hits = image_search.query( + query, + count = 1, + model = model.image_search) + + results = image_search.collate_results( + hits, + model.image_search.image_names, + output_directory=output_directory, + image_files_url='/static/images', + count=1) + + actual_image_path = output_directory.joinpath(Path(results[0].entry).name) + actual_image = Image.open(actual_image_path) + expected_image = Image.open(expected_image_path) + + # Assert + # Ensure file search triggered instead of query with file path as string + assert f"Find Images by Image: {resolve_absolute_path(expected_image_path)}" in caplog.text, "File search not triggered" + # Ensure the correct image is returned + assert expected_image == actual_image, "Incorrect image returned by file search" + + # Cleanup + # Delete the image files copied to results directory + actual_image_path.unlink()