CalvinBackup/facefusion-labs
1
0
Fork 0
mirror of https://github.com/facefusion/facefusion-labs.git synced 2026-07-28 16:08:50 +02:00
Code Issues Packages Projects Releases Wiki Activity

Move magic methods up

Browse Source
This commit is contained in:
henryruhs
2025-03-11 14:43:09 +01:00
parent 4d2038d4ce
commit 58c81cd646
1 changed files with 11 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
face_swapper/src/data_loader.py
+11 -11
View File
@@ -19,6 +19,17 @@ class DataLoaderVGG(TensorDataset):
self.dataset_total = len(self.image_paths)
self.transforms = self.compose_transforms()
def __getitem__(self, index : int) -> Batch:
source_image_path = self.image_paths[index]
if random.random() > self.same_person_probability:
return self.prepare_same_person(source_image_path)
return self.prepare_different_person(source_image_path)
def __len__(self) -> int:
return self.dataset_total
def prepare_image_paths(self, dataset_image_pattern : str) -> Tuple[ImagePathList, ImagePathSet]:
image_paths = []
image_path_set = {}
@@ -41,14 +52,6 @@ class DataLoaderVGG(TensorDataset):
])
return transform
def __getitem__(self, index : int) -> Batch:
source_image_path = self.image_paths[index]
if random.random() > self.same_person_probability:
return self.prepare_same_person(source_image_path)
return self.prepare_different_person(source_image_path)
def prepare_different_person(self, source_image_path : str) -> Batch:
is_same_person = torch.tensor(0)
target_image_path = random.choice(self.image_paths)
@@ -66,6 +69,3 @@ class DataLoaderVGG(TensorDataset):
source_tensor = self.transforms(source_vision_frame)
target_tensor = self.transforms(target_vision_frame)
return source_tensor, target_tensor, is_same_person
def __len__(self) -> int:
return self.dataset_total