Replace naming warp_matrix with warp_template

face_swapper/README.md

@@ -28,7 +28,7 @@ This `config.ini` utilizes the MegaFace dataset to train the Face Swapper model.
 ```
 [training.dataset]
 file_pattern = .datasets/vggface2/**/*.jpg
-warp_matrix = vgg_face_hq_to_arcface_128_v2
+warp_template = vgg_face_hq_to_arcface_128_v2
 batch_ratio = 0.2
 ```
 

face_swapper/config.ini

@@ -1,6 +1,6 @@
 [training.dataset]
 file_pattern =
-warp_matrix =
+warp_template =
 batch_ratio =
 
 [training.loader]

face_swapper/src/dataset.py

@@ -7,14 +7,14 @@ from torch.utils.data import Dataset
 from torchvision import io, transforms
 
 from .helper import warp_tensor
-from .types import Batch, WarpMatrix
+from .types import Batch, WarpTemplate
 
 
 class DynamicDataset(Dataset[Tensor]):
-	def __init__(self, file_pattern : str, warp_matrix : WarpMatrix, batch_ratio : float) -> None:
+	def __init__(self, file_pattern : str, warp_template : WarpTemplate, batch_ratio : float) -> None:
 		self.file_paths = glob.glob(file_pattern)
 		self.transforms = self.compose_transforms()
-		self.warp_matrix = warp_matrix
+		self.warp_template = warp_template
 		self.batch_ratio = batch_ratio
 
 	def __getitem__(self, index : int) -> Batch:
@@ -41,7 +41,7 @@ class DynamicDataset(Dataset[Tensor]):
 		])
 
 	def warp_tensor(self, temp_tensor : Tensor) -> Tensor:
-		return warp_tensor(temp_tensor.unsqueeze(0), self.warp_matrix).squeeze(0)
+		return warp_tensor(temp_tensor.unsqueeze(0), self.warp_template).squeeze(0)
 
 	def prepare_different_batch(self, source_path : str) -> Batch:
 		target_path = random.choice(self.file_paths)

face_swapper/src/helper.py

@@ -1,9 +1,9 @@
 import torch
 from torch import Tensor, nn
 
-from .types import EmbedderModule, Embedding, Padding, WarpMatrix, WarpMatrixSet
+from .types import EmbedderModule, Embedding, Padding, WarpTemplate, WarpTemplateSet
 
-WARP_MATRIX_SET : WarpMatrixSet =\
+WARP_TEMPLATE_SET : WarpTemplateSet =\
 {
 	'vgg_face_hq_to_arcface_128_v2': torch.tensor(
 	[
@@ -18,10 +18,10 @@ WARP_MATRIX_SET : WarpMatrixSet =\
 }
 
 
-def warp_tensor(input_tensor : Tensor, warp_matrix : WarpMatrix) -> Tensor:
-	warp_matrix = WARP_MATRIX_SET.get(warp_matrix).repeat(input_tensor.shape[0], 1, 1)
-	grid = nn.functional.affine_grid(warp_matrix.to(input_tensor.device), list(input_tensor.shape))
-	output_tensor = nn.functional.grid_sample(input_tensor, grid, align_corners = False, padding_mode = 'reflection')
+def warp_tensor(input_tensor : Tensor, warp_template : WarpTemplate) -> Tensor:
+	normed_warp_template = WARP_TEMPLATE_SET.get(warp_template).repeat(input_tensor.shape[0], 1, 1)
+	affine_grid = nn.functional.affine_grid(normed_warp_template.to(input_tensor.device), list(input_tensor.shape))
+	output_tensor = nn.functional.grid_sample(input_tensor, affine_grid, align_corners = False, padding_mode = 'reflection')
 	return output_tensor
 
 

face_swapper/src/training.py

@@ -17,7 +17,7 @@ from .helper import calc_embedding
 from .models.discriminator import Discriminator
 from .models.generator import Generator
 from .models.loss import AdversarialLoss, AttributeLoss, DiscriminatorLoss, GazeLoss, IdentityLoss, PoseLoss, ReconstructionLoss
-from .types import Batch, Embedding, OptimizerConfig, WarpMatrix
+from .types import Batch, Embedding, OptimizerConfig, WarpTemplate
 
 CONFIG = configparser.ConfigParser()
 CONFIG.read('config.ini')
@@ -194,14 +194,14 @@ def create_trainer() -> Trainer:
 
 def train() -> None:
 	dataset_file_pattern = CONFIG.get('training.dataset', 'file_pattern')
-	dataset_warp_matrix = cast(WarpMatrix, CONFIG.get('training.dataset', 'warp_matrix'))
+	dataset_warp_template = cast(WarpTemplate, CONFIG.get('training.dataset', 'warp_template'))
 	dataset_batch_ratio = CONFIG.getfloat('training.dataset', 'batch_ratio')
 	output_resume_path = CONFIG.get('training.output', 'resume_path')
 
 	if torch.cuda.is_available():
 		torch.set_float32_matmul_precision('high')
 
-	dataset = DynamicDataset(dataset_file_pattern, dataset_warp_matrix, dataset_batch_ratio)
+	dataset = DynamicDataset(dataset_file_pattern, dataset_warp_template, dataset_batch_ratio)
 	training_loader, validation_loader = create_loaders(dataset)
 	face_swapper_trainer = FaceSwapperTrainer()
 	trainer = create_trainer()

face_swapper/src/types.py

@@ -18,5 +18,5 @@ MotionExtractorModule : TypeAlias = Module
 
 OptimizerConfig : TypeAlias = Any
 
-WarpMatrix = Literal['vgg_face_hq_to_arcface_128_v2', 'arcface_128_v2_to_arcface_112_v2']
-WarpMatrixSet : TypeAlias = Dict[WarpMatrix, Tensor]
+WarpTemplate = Literal['vgg_face_hq_to_arcface_128_v2', 'arcface_128_v2_to_arcface_112_v2']
+WarpTemplateSet : TypeAlias = Dict[WarpTemplate, Tensor]