Introduce batch mode via config for equal and same batches

face_swapper/src/dataset.py

@@ -7,21 +7,25 @@ from torch.utils.data import Dataset
 from torchvision import io, transforms
 
 from .helper import warp_tensor
-from .types import Batch, WarpTemplate
+from .types import Batch, BatchMode, WarpTemplate
 
 
 class DynamicDataset(Dataset[Tensor]):
-	def __init__(self, file_pattern : str, warp_template : WarpTemplate, batch_ratio : float) -> None:
+	def __init__(self, file_pattern : str, warp_template : WarpTemplate, batch_mode : BatchMode, batch_ratio : float) -> None:
 		self.file_paths = glob.glob(file_pattern)
-		self.transforms = self.compose_transforms()
 		self.warp_template = warp_template
+		self.batch_mode = batch_mode
 		self.batch_ratio = batch_ratio
+		self.transforms = self.compose_transforms()
 
 	def __getitem__(self, index : int) -> Batch:
 		file_path = self.file_paths[index]
 
 		if random.random() < self.batch_ratio:
-			return self.prepare_equal_batch(file_path)
+			if self.batch_mode == 'equal':
+				return self.prepare_equal_batch(file_path)
+			if self.batch_mode == 'same':
+				return self.prepare_same_batch(file_path)
 
 		return self.prepare_different_batch(file_path)
 
@@ -51,6 +55,11 @@ class DynamicDataset(Dataset[Tensor]):
 		target_tensor = self.transforms(target_tensor)
 		return source_tensor, target_tensor
 
+	def prepare_same_batch(self, source_path : str) -> Batch:
+		source_tensor = io.read_image(source_path)
+		source_tensor = self.transforms(source_tensor)
+		return source_tensor, source_tensor
+
 	def prepare_equal_batch(self, source_path : str) -> Batch:
 		target_directory_path = os.path.dirname(source_path)
 		target_file_name_and_extension = random.choice(os.listdir(target_directory_path))