diff --git a/embedding_converter/src/dataset.py b/embedding_converter/src/dataset.py index 3d61aea..f9e074c 100644 --- a/embedding_converter/src/dataset.py +++ b/embedding_converter/src/dataset.py @@ -15,8 +15,8 @@ class DynamicDataset(Dataset[Tensor]): self.transforms = self.compose_transforms() def __getitem__(self, index : int) -> Batch: - image_path = random.choice(self.file_paths) - vision_frame = cv2.imread(image_path) + file_path = random.choice(self.file_paths) + vision_frame = cv2.imread(file_path) return self.transforms(vision_frame) def __len__(self) -> int: diff --git a/face_swapper/src/dataset.py b/face_swapper/src/dataset.py index de7f344..61bc4a3 100644 --- a/face_swapper/src/dataset.py +++ b/face_swapper/src/dataset.py @@ -16,12 +16,12 @@ class DynamicDataset(Dataset[Tensor]): self.batch_ratio = batch_ratio def __getitem__(self, index : int) -> Batch: - source_image_path = self.file_paths[index] + file_path = self.file_paths[index] if random.random() < self.batch_ratio: - return self.prepare_equal_batch(source_image_path) + return self.prepare_equal_batch(file_path) - return self.prepare_different_batch(source_image_path) + return self.prepare_different_batch(file_path) def __len__(self) -> int: return len(self.file_paths)