From c1f39a73dd4eab8eece21142504bd5d1ed9d9496 Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Sun, 23 Mar 2025 18:18:10 +0530 Subject: [PATCH] change face-parser to face-masker --- face_swapper/README.md | 2 +- face_swapper/config.ini | 2 +- face_swapper/src/models/loss.py | 8 ++++---- face_swapper/src/training.py | 6 +++--- face_swapper/src/types.py | 2 +- 5 files changed, 10 insertions(+), 10 deletions(-) diff --git a/face_swapper/README.md b/face_swapper/README.md index 7a74819..4cdc960 100644 --- a/face_swapper/README.md +++ b/face_swapper/README.md @@ -46,7 +46,7 @@ split_ratio = 0.9995 embedder_path = .models/arcface.pt gazer_path = .models/gazer.pt motion_extractor_path = .models/motion_extractor.pt -face_parser_path = .models/face_parser.pt +face_masker_path = .models/face_masker.pt ``` ``` diff --git a/face_swapper/config.ini b/face_swapper/config.ini index 29443e6..7377124 100644 --- a/face_swapper/config.ini +++ b/face_swapper/config.ini @@ -14,7 +14,7 @@ split_ratio = embedder_path = gazer_path = motion_extractor_path = -face_parser_path = +face_masker_path = [training.model.generator] source_channels = diff --git a/face_swapper/src/models/loss.py b/face_swapper/src/models/loss.py index 2549af0..6f30ba3 100644 --- a/face_swapper/src/models/loss.py +++ b/face_swapper/src/models/loss.py @@ -7,7 +7,7 @@ from torch import Tensor, nn from torchvision import transforms from ..helper import calc_embedding -from ..types import EmbedderModule, FaceParserModule, Feature, GazerModule, Loss, Mask, MotionExtractorModule +from ..types import EmbedderModule, FaceMaskerModule, Feature, GazerModule, Loss, Mask, MotionExtractorModule class DiscriminatorLoss(nn.Module): @@ -180,11 +180,11 @@ class GazeLoss(nn.Module): class MaskLoss(nn.Module): - def __init__(self, config_parser : ConfigParser, face_parser : FaceParserModule) -> None: + def __init__(self, config_parser : ConfigParser, face_masker : FaceMaskerModule) -> None: super().__init__() self.config_mask_weight = config_parser.getfloat('training.losses', 'mask_weight') self.config_output_size = config_parser.getint('training.model.generator', 'output_size') - self.face_parser = face_parser + self.face_masker = face_masker self.mse_loss = nn.MSELoss() def forward(self, target_tensor : Tensor, output_mask : Mask) -> Tuple[Loss, Loss]: @@ -200,7 +200,7 @@ class MaskLoss(nn.Module): target_tensor = (target_tensor.clip(-1, 1) + 1) * 0.5 with torch.no_grad(): - output_tensor = self.face_parser(target_tensor) + output_tensor = self.face_masker(target_tensor) output_tensor = output_tensor.clamp(0, 1) output_tensor = torch.nn.functional.interpolate(output_tensor, (self.config_output_size, self.config_output_size), mode = 'bilinear') diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index b53feac..c105652 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -31,14 +31,14 @@ class FaceSwapperTrainer(LightningModule): self.config_embedder_path = config_parser.get('training.model', 'embedder_path') self.config_gazer_path = config_parser.get('training.model', 'gazer_path') self.config_motion_extractor_path = config_parser.get('training.model', 'motion_extractor_path') - self.config_face_parser_path = config_parser.get('training.model', 'face_parser_path') + self.config_face_masker_path = config_parser.get('training.model', 'face_masker_path') self.config_accumulate_size = config_parser.getfloat('training.trainer', 'accumulate_size') self.config_learning_rate = config_parser.getfloat('training.trainer', 'learning_rate') self.config_preview_frequency = config_parser.getint('training.trainer', 'preview_frequency') self.embedder = torch.jit.load(self.config_embedder_path, map_location = 'cpu').eval() self.gazer = torch.jit.load(self.config_gazer_path, map_location = 'cpu').eval() self.motion_extractor = torch.jit.load(self.config_motion_extractor_path, map_location = 'cpu').eval() - self.face_parser = torch.jit.load(self.config_face_parser_path, map_location ='cpu').eval() + self.face_masker = torch.jit.load(self.config_face_masker_path, map_location ='cpu').eval() self.generator = Generator(config_parser) self.discriminator = Discriminator(config_parser) self.discriminator_loss = DiscriminatorLoss() @@ -48,7 +48,7 @@ class FaceSwapperTrainer(LightningModule): self.identity_loss = IdentityLoss(config_parser, self.embedder) self.motion_loss = MotionLoss(config_parser, self.motion_extractor) self.gaze_loss = GazeLoss(config_parser, self.gazer) - self.mask_loss = MaskLoss(config_parser, self.face_parser) + self.mask_loss = MaskLoss(config_parser, self.face_masker) self.automatic_optimization = False def forward(self, source_embedding : Embedding, target_tensor : Tensor) -> Tuple[Tensor, Mask]: diff --git a/face_swapper/src/types.py b/face_swapper/src/types.py index 417f6b7..7c0c1c9 100644 --- a/face_swapper/src/types.py +++ b/face_swapper/src/types.py @@ -17,7 +17,7 @@ GeneratorModule : TypeAlias = Module EmbedderModule : TypeAlias = Module GazerModule : TypeAlias = Module MotionExtractorModule : TypeAlias = Module -FaceParserModule : TypeAlias = Module +FaceMaskerModule : TypeAlias = Module OptimizerSet : TypeAlias = Any