diff --git a/face_swapper/README.md b/face_swapper/README.md index 0962162..48d0c51 100644 --- a/face_swapper/README.md +++ b/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 ``` diff --git a/face_swapper/config.ini b/face_swapper/config.ini index 63aa355..72233e8 100644 --- a/face_swapper/config.ini +++ b/face_swapper/config.ini @@ -1,6 +1,6 @@ [training.dataset] file_pattern = -warp_matrix = +warp_template = batch_ratio = [training.loader] diff --git a/face_swapper/src/dataset.py b/face_swapper/src/dataset.py index f1ad8b6..e289ac9 100644 --- a/face_swapper/src/dataset.py +++ b/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) diff --git a/face_swapper/src/helper.py b/face_swapper/src/helper.py index eb8f059..557fcbe 100644 --- a/face_swapper/src/helper.py +++ b/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 diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index d4cef33..590606c 100644 --- a/face_swapper/src/training.py +++ b/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() diff --git a/face_swapper/src/types.py b/face_swapper/src/types.py index ec28757..7742764 100644 --- a/face_swapper/src/types.py +++ b/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]