Add pose loss

hyperswap/README.md

@@ -48,6 +48,7 @@ split_ratio = 0.9995
 generator_embedder_path = .models/blendface.pt
 loss_embedder_path = .models/arcface.pt
 face_masker_path = .models/face_masker.pt
+face_aligner_path = .models/face_aligner.pt
 ```
 
 ```
@@ -80,6 +81,7 @@ feature_weight = 10.0
 reconstruction_weight = 10.0
 identity_weight = 20.0
 mask_weight = 5.0
+pose_weight = 25.0
 ```
 
 ```

hyperswap/config.ini

@@ -16,6 +16,7 @@ split_ratio =
 generator_embedder_path =
 loss_embedder_path =
 face_masker_path =
+face_aligner_path =
 
 [training.model.generator]
 source_channels =
@@ -40,6 +41,7 @@ feature_weight =
 reconstruction_weight =
 identity_weight =
 mask_weight =
+pose_weight =
 
 [training.trainer]
 accumulate_size =

hyperswap/src/helper.py

@@ -64,6 +64,19 @@ def overlay_mask(input_tensor : Tensor, input_mask : Mask) -> Tensor:
 	return output_tensor
 
 
+def blur_tensor(input_tensor : Tensor, kernel_size : int, sigma : float) -> Tensor:
+	coordinate_tensor = torch.arange(kernel_size, dtype = input_tensor.dtype, device = input_tensor.device) - (kernel_size - 1) * 0.5
+	kernel_tensor = torch.exp(-coordinate_tensor ** 2 / (2 * sigma ** 2))
+	kernel_tensor = kernel_tensor / kernel_tensor.sum()
+	horizontal_kernel = kernel_tensor.view(1, 1, 1, kernel_size).repeat(input_tensor.shape[1], 1, 1, 1)
+	vertical_kernel = kernel_tensor.view(1, 1, kernel_size, 1).repeat(input_tensor.shape[1], 1, 1, 1)
+	padding = kernel_size // 2
+	temp_tensor = nn.functional.pad(input_tensor, (padding, padding, padding, padding), mode = 'reflect')
+	temp_tensor = nn.functional.conv2d(temp_tensor, horizontal_kernel, groups = input_tensor.shape[1])
+	output_tensor = nn.functional.conv2d(temp_tensor, vertical_kernel, groups = input_tensor.shape[1])
+	return output_tensor
+
+
 def dilate_mask(input_tensor : Tensor, factor : float) -> Tensor:
 	padding = int(input_tensor.shape[2] * factor + 0.5)
 	kernel_size = 1 + 2 * padding

hyperswap/src/models/loss.py

@@ -5,8 +5,8 @@
 from pytorch_msssim import ssim
 from torch import Tensor, nn
 
-from ..helper import calculate_face_embedding, dilate_mask, oddify_size
-from ..types import EmbedderModule, FaceMaskerModule, Feature, Loss, Mask
+from ..helper import blur_tensor, calculate_face_embedding, dilate_mask, oddify_size
+from ..types import EmbedderModule, FaceAlignerModule, FaceMaskerModule, Feature, Loss, Mask
 
 
 class DiscriminatorLoss(nn.Module):
@@ -140,3 +140,30 @@ def calculate_mask(self, target_tensor : Tensor) -> Tensor:
 			output_tensor = torch.nn.functional.interpolate(output_tensor, (self.config_output_size, self.config_output_size), mode = 'bilinear')
 
 		return output_tensor
+
+
+class PoseLoss(nn.Module):
+	def __init__(self, config_parser : ConfigParser, face_aligner : FaceAlignerModule) -> None:
+		super().__init__()
+		self.config_pose_weight = config_parser.getfloat('training.losses', 'pose_weight')
+		self.face_aligner = face_aligner
+		self.mse_loss = nn.MSELoss()
+
+	def forward(self, target_tensor : Tensor, output_tensor : Tensor) -> Tuple[Loss, Loss]:
+		with torch.no_grad():
+			target_pose = self.calculate_pose(target_tensor)
+
+		output_pose = self.calculate_pose(output_tensor)
+		pose_loss = self.mse_loss(output_pose, target_pose)
+		weighted_pose_loss = pose_loss * self.config_pose_weight
+		return pose_loss, weighted_pose_loss
+
+	def calculate_pose(self, input_tensor : Tensor) -> Tensor:
+		temp_tensor = torch.nn.functional.interpolate(input_tensor, (256, 256), mode = 'area')
+		temp_tensor = (temp_tensor.clip(-1, 1) + 1) * 0.5
+		temp_tensor = blur_tensor(temp_tensor, 9, 1.7)
+
+		_, pose_tensor, _ = self.face_aligner(temp_tensor)
+		yaw_tensor = nn.functional.normalize(pose_tensor[:, 0:2], p = 2)
+		pitch_tensor = nn.functional.normalize(pose_tensor[:, 2:4], p = 2)
+		return torch.cat([ yaw_tensor, pitch_tensor ], dim = 1)

hyperswap/src/training.py

@@ -19,7 +19,7 @@
 from .helper import apply_noise, calculate_face_embedding, erode_mask, overlay_mask
 from .models.discriminator import Discriminator
 from .models.generator import Generator
-from .models.loss import AdversarialLoss, DiscriminatorLoss, FeatureLoss, IdentityLoss, MaskLoss, ReconstructionLoss
+from .models.loss import AdversarialLoss, DiscriminatorLoss, FeatureLoss, IdentityLoss, MaskLoss, PoseLoss, ReconstructionLoss
 from .types import Batch, Embedding, Mask, OptimizerSet, TrainerPrecision, TrainerStrategy
 
 warnings.filterwarnings('ignore', category = UserWarning, module = 'torch')
@@ -34,6 +34,7 @@ def __init__(self, config_parser : ConfigParser) -> None:
 		self.config_generator_embedder_path = config_parser.get('training.model', 'generator_embedder_path')
 		self.config_loss_embedder_path = config_parser.get('training.model', 'loss_embedder_path')
 		self.config_face_masker_path = config_parser.get('training.model', 'face_masker_path')
+		self.config_face_aligner_path = config_parser.get('training.model', 'face_aligner_path')
 		self.config_accumulate_size = config_parser.getfloat('training.trainer', 'accumulate_size')
 		self.config_discriminator_ratio = config_parser.getfloat('training.trainer', 'discriminator_ratio')
 		self.config_gradient_clip = config_parser.getfloat('training.trainer', 'gradient_clip')
@@ -51,6 +52,7 @@ def __init__(self, config_parser : ConfigParser) -> None:
 		self.generator_embedder = torch.jit.load(self.config_generator_embedder_path, map_location = 'cpu').eval()
 		self.loss_embedder = torch.jit.load(self.config_loss_embedder_path, map_location = 'cpu').eval()
 		self.face_masker = torch.jit.load(self.config_face_masker_path, map_location ='cpu').eval()
+		self.face_aligner = torch.jit.load(self.config_face_aligner_path, map_location = 'cpu').eval()
 		self.generator = Generator(config_parser)
 		self.discriminator = Discriminator(config_parser)
 		self.discriminator_loss = DiscriminatorLoss()
@@ -59,6 +61,7 @@ def __init__(self, config_parser : ConfigParser) -> None:
 		self.reconstruction_loss = ReconstructionLoss(config_parser, self.loss_embedder)
 		self.identity_loss = IdentityLoss(config_parser, self.loss_embedder)
 		self.mask_loss = MaskLoss(config_parser, self.face_masker)
+		self.pose_loss = PoseLoss(config_parser, self.face_aligner)
 		self.automatic_optimization = False
 
 	def forward(self, source_embedding : Embedding, target_tensor : Tensor) -> Tuple[Tensor, Mask]:
@@ -115,7 +118,8 @@ def training_step(self, batch : Batch, batch_index : int) -> Tensor:
 		reconstruction_loss, weighted_reconstruction_loss = self.reconstruction_loss(source_tensor, target_tensor, generator_output_tensor)
 		identity_loss, weighted_identity_loss = self.identity_loss(generator_output_tensor, source_tensor)
 		mask_loss, weighted_mask_loss = self.mask_loss(target_tensor, generator_output_mask)
-		generator_loss = weighted_adversarial_loss + weighted_feature_loss + weighted_reconstruction_loss + weighted_identity_loss + weighted_mask_loss
+		pose_loss, weighted_pose_loss = self.pose_loss(target_tensor, generator_output_tensor)
+		generator_loss = weighted_adversarial_loss + weighted_feature_loss + weighted_reconstruction_loss + weighted_identity_loss + weighted_mask_loss + weighted_pose_loss
 
 		if torch.randn(1).item() < self.config_discriminator_ratio:
 			discriminator_real_tensors = self.discriminator(source_tensor)
@@ -162,6 +166,7 @@ def training_step(self, batch : Batch, batch_index : int) -> Tensor:
 		self.log('reconstruction_loss', reconstruction_loss)
 		self.log('identity_loss', identity_loss)
 		self.log('mask_loss', mask_loss)
+		self.log('pose_loss', pose_loss)
 
 		if do_update:
 			generator_scheduler.step(generator_loss)

hyperswap/src/types.py

@@ -20,6 +20,7 @@
 GeneratorModule : TypeAlias = Module
 EmbedderModule : TypeAlias = Module
 FaceMaskerModule : TypeAlias = Module
+FaceAlignerModule : TypeAlias = Module
 
 OptimizerSet : TypeAlias = Any