From cfcf0ee2bdb9f57744ebb31f0a067451a49a446f Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Tue, 21 Jan 2025 23:36:06 +0530 Subject: [PATCH] add glob image pattern to config --- face_swapper/config.ini | 2 ++ face_swapper/src/data_loader.py | 10 ++++++---- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/face_swapper/config.ini b/face_swapper/config.ini index f9c40b3..116c6ec 100644 --- a/face_swapper/config.ini +++ b/face_swapper/config.ini @@ -1,5 +1,7 @@ [preparing.dataset] dataset_path = /assets/VGGface2_None_norm_512_true_bygfpgan +folder_pattern = {}/* +image_pattern = {}/*.*g [preparing.dataloader] same_person_probability = 0.2 diff --git a/face_swapper/src/data_loader.py b/face_swapper/src/data_loader.py index d203054..c4f8a9f 100644 --- a/face_swapper/src/data_loader.py +++ b/face_swapper/src/data_loader.py @@ -22,13 +22,15 @@ def read_image(image_path: str) -> VisionFrame: class DataLoaderVGG(TensorDataset): def __init__(self, dataset_path : str) -> None: self.same_person_probability = CONFIG.getfloat('preparing.dataloader', 'same_person_probability') - self.image_paths = glob.glob('{}/*/*.*g'.format(dataset_path)) # @todo globs belong to the config - self.folder_paths = glob.glob('{}/*'.format(dataset_path)) + image_pattern = CONFIG.get('preparing.dataset', 'image_pattern') + folder_pattern = CONFIG.get('preparing.dataset', 'folder_pattern') + self.folder_paths = glob.glob(folder_pattern.format(dataset_path)) + self.image_paths = [] self.image_path_set = {} - self._current_index = 0 for folder_path in self.folder_paths: - image_paths = glob.glob('{}/*'.format(folder_path)) # @todo not sure about alls this globs being used here :-) + image_paths = glob.glob(image_pattern.format(folder_path)) + self.image_paths.extend(image_paths) self.image_path_set[folder_path] = image_paths self.dataset_total = len(self.image_paths) self.transforms = transforms.Compose(