From 35894e2cf94e4804d95c9eaf3eda540d4f3cd1a3 Mon Sep 17 00:00:00 2001 From: shuxun <1594510003@qq.com> Date: Tue, 30 Jan 2024 16:44:43 +0800 Subject: [PATCH] Update reader.py --- clip_retrieval/clip_inference/reader.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/clip_retrieval/clip_inference/reader.py b/clip_retrieval/clip_inference/reader.py index e49aa4d..7299748 100644 --- a/clip_retrieval/clip_inference/reader.py +++ b/clip_retrieval/clip_inference/reader.py @@ -15,7 +15,7 @@ def folder_to_keys(folder, enable_text=True, enable_image=True, enable_metadata= image_files = None if enable_text: text_files = [*path.glob("**/*.txt")] - text_files = {text_file.relative_to(path).as_posix(): text_file for text_file in text_files} + text_files = {text_file.relative_to(path).with_suffix('').as_posix(): text_file for text_file in text_files} if enable_image: image_files = [ *path.glob("**/*.png"), @@ -29,10 +29,10 @@ def folder_to_keys(folder, enable_text=True, enable_image=True, enable_metadata= *path.glob("**/*.BMP"), *path.glob("**/*.WEBP"), ] - image_files = {image_file.relative_to(path).as_posix(): image_file for image_file in image_files} + image_files = {image_file.relative_to(path).with_suffix('').as_posix(): image_file for image_file in image_files} if enable_metadata: metadata_files = [*path.glob("**/*.json")] - metadata_files = {metadata_file.relative_to(path).as_posix(): metadata_file for metadata_file in metadata_files} + metadata_files = {metadata_file.relative_to(path).with_suffix('').as_posix(): metadata_file for metadata_file in metadata_files} keys = None @@ -41,9 +41,9 @@ def join(new_set): if enable_text: keys = join(text_files.keys()) - elif enable_image: + if enable_image: keys = join(image_files.keys()) - elif enable_metadata: + if enable_metadata: keys = join(metadata_files.keys()) keys = list(sorted(keys))