cleanup

2026-06-25 07:59:55 +02:00 · 2025-02-10 22:43:15 +05:30
parent b7e2d3ccd7
commit 2ed558a873
10 changed files with 310 additions and 286 deletions
@@ -0,0 +1,95 @@
+import configparser
+from itertools import chain
+from typing import List
+
+import numpy
+import torch.nn
+import torch.nn as nn
+from torch import Tensor
+
+from face_swapper.src.types import DiscriminatorOutputs
+
+CONFIG = configparser.ConfigParser()
+CONFIG.read('config.ini')
+
+
+class MultiscaleDiscriminator(nn.Module):
+	def __init__(self) -> None:
+		super(MultiscaleDiscriminator, self).__init__()
+		self.input_channels = CONFIG.getint('training.model.discriminator', 'input_channels')
+		self.num_filters = CONFIG.getint('training.model.discriminator', 'num_filters')
+		self.kernel_size = CONFIG.getint('training.model.discriminator', 'kernel_size')
+		self.num_layers = CONFIG.getint('training.model.discriminator', 'num_layers')
+		self.num_discriminators = CONFIG.getint('training.model.discriminator', 'num_discriminators')
+
+		self.downsample = nn.AvgPool2d(kernel_size = 3, stride = 2, padding = [ 1, 1 ], count_include_pad = False) # type:ignore[arg-type]
+		self.prepare_discriminators()
+
+	def prepare_discriminators(self) -> None:
+		for discriminator_index in range(self.num_discriminators):
+			single_discriminator = NLayerDiscriminator(self.input_channels, self.num_filters, self.num_layers, self.kernel_size)
+			setattr(self, 'discriminator_layer_{}'.format(discriminator_index), single_discriminator.model)
+
+	def forward(self, input_tensor : Tensor) -> DiscriminatorOutputs:
+		discriminator_outputs = []
+		temp_tensor = input_tensor
+
+		for discriminator_index in range(self.num_discriminators):
+			model_layers = getattr(self, 'discriminator_layer_{}'.format(self.num_discriminators - 1 - discriminator_index))
+			discriminator_outputs.append([ model_layers(temp_tensor) ])
+
+			if discriminator_index < (self.num_discriminators - 1):
+				temp_tensor = self.downsample(temp_tensor)
+
+		return discriminator_outputs
+
+
+class NLayerDiscriminator(nn.Module):
+	def __init__(self, input_channels : int, num_filters : int, num_layers : int, kernel_size : int) -> None:
+		super(NLayerDiscriminator, self).__init__()
+		self.num_layers = num_layers
+		model_layers = self.prepare_model_layers(input_channels, num_filters, num_layers, kernel_size)
+		self.model = nn.Sequential(*list(chain(*model_layers)))
+
+	def prepare_model_layers(self, input_channels : int, num_filters : int, num_layers : int, kernel_size : int) -> List[List[torch.nn.Module]]:
+		padding_size = int(numpy.ceil((kernel_size - 1.0) / 2))
+
+		model_layers =\
+		[
+			[
+				nn.Conv2d(input_channels, num_filters, kernel_size = kernel_size, stride = 2, padding = padding_size),
+			 	nn.LeakyReLU(0.2, True)
+			]
+		]
+		current_filters = num_filters
+
+		for layer_index in range(1, num_layers):
+			previous_filters = current_filters
+			current_filters = min(current_filters * 2, 512)
+			model_layers +=\
+			[
+				[
+					nn.Conv2d(previous_filters, current_filters, kernel_size = kernel_size, stride = 2, padding = padding_size),
+					nn.InstanceNorm2d(current_filters), nn.LeakyReLU(0.2, True)
+				]
+			]
+		previous_filters = current_filters
+		current_filters = min(current_filters * 2, 512)
+		model_layers +=\
+		[
+			[
+				nn.Conv2d(previous_filters, current_filters, kernel_size = kernel_size, padding = padding_size),
+				nn.InstanceNorm2d(current_filters),
+				nn.LeakyReLU(0.2, True)
+			]
+		]
+		model_layers +=\
+		[
+			[
+				nn.Conv2d(current_filters, 1, kernel_size = kernel_size, padding = padding_size)
+			]
+		]
+		return model_layers
+
+	def forward(self, input_tensor : Tensor) -> Tensor:
+		return self.model(input_tensor)
@@ -0,0 +1,43 @@
+import configparser
+from typing import Tuple
+
+import torch.nn as nn
+
+from face_swapper.src.networks.attribute_modulator import AADGenerator
+from face_swapper.src.networks.encoder import UNet
+from face_swapper.src.types import SourceEmbedding, TargetAttributes, VisionTensor
+
+CONFIG = configparser.ConfigParser()
+CONFIG.read('config.ini')
+
+
+class AdaptiveEmbeddingIntegrationNetwork(nn.Module):
+	def __init__(self) -> None:
+		super(AdaptiveEmbeddingIntegrationNetwork, self).__init__()
+		id_channels = CONFIG.getint('training.model.generator', 'id_channels')
+		num_blocks = CONFIG.getint('training.model.generator', 'num_blocks')
+
+		self.encoder = UNet()
+		self.generator = AADGenerator(id_channels, num_blocks)
+		self.encoder.apply(init_weight)
+		self.generator.apply(init_weight)
+
+	def forward(self, target : VisionTensor, source_embedding : SourceEmbedding) -> Tuple[VisionTensor, TargetAttributes]:
+		target_attributes = self.get_attributes(target)
+		swap_tensor = self.generator(target_attributes, source_embedding)
+		return swap_tensor, target_attributes
+
+	def get_attributes(self, target : VisionTensor) -> TargetAttributes:
+		return self.encoder(target)
+
+
+def init_weight(module : nn.Module) -> None:
+	if isinstance(module, nn.Linear):
+		module.weight.data.normal_(std = 0.001)
+		module.bias.data.zero_()
+
+	if isinstance(module, nn.Conv2d):
+		nn.init.xavier_normal_(module.weight.data)
+
+	if isinstance(module, nn.ConvTranspose2d):
+		nn.init.xavier_normal_(module.weight.data)
@@ -0,0 +1,137 @@
+import configparser
+from typing import Tuple
+
+import torch
+from pytorch_msssim import ssim
+from torch import Tensor
+
+from face_swapper.src.helper import calc_id_embedding, hinge_fake_loss, hinge_real_loss
+from face_swapper.src.types import Batch, DiscriminatorLossSet, DiscriminatorOutputs, FaceLandmark203, GeneratorLossSet, LossTensor, SwapAttributes, TargetAttributes, VisionTensor
+
+CONFIG = configparser.ConfigParser()
+CONFIG.read('config.ini')
+
+
+class FaceSwapperLoss:
+	def __init__(self) -> None:
+		id_embedder_path = CONFIG.get('training.model', 'id_embedder_path')
+		landmarker_path = CONFIG.get('training.model', 'landmarker_path')
+		motion_extractor_path = CONFIG.get('training.model', 'motion_extractor_path')
+		self.batch_size = CONFIG.getint('training.loader', 'batch_size')
+		self.mse_loss = torch.nn.MSELoss()
+		self.id_embedder = torch.jit.load(id_embedder_path, map_location = 'cpu')  # type:ignore[no-untyped-call]
+		self.landmarker = torch.jit.load(landmarker_path, map_location = 'cpu')  # type:ignore[no-untyped-call]
+		self.motion_extractor = torch.jit.load(motion_extractor_path, map_location = 'cpu')  # type:ignore[no-untyped-call]
+		self.id_embedder.eval()
+		self.landmarker.eval()
+		self.motion_extractor.eval()
+
+	def calc_generator_loss(self, swap_tensor : VisionTensor, target_attributes : TargetAttributes, swap_attributes : SwapAttributes, discriminator_outputs : DiscriminatorOutputs, batch : Batch) -> GeneratorLossSet:
+		source_tensor, target_tensor, is_same_person = batch
+		weight_adversarial = CONFIG.getfloat('training.losses', 'weight_adversarial')
+		weight_id = CONFIG.getfloat('training.losses', 'weight_id')
+		weight_attribute = CONFIG.getfloat('training.losses', 'weight_attribute')
+		weight_reconstruction = CONFIG.getfloat('training.losses', 'weight_reconstruction')
+		weight_pose = CONFIG.getfloat('training.losses', 'weight_pose')
+		weight_gaze = CONFIG.getfloat('training.losses', 'weight_gaze')
+		generator_loss_set = {}
+
+		generator_loss_set['loss_adversarial'] = self.calc_adversarial_loss(discriminator_outputs)
+		generator_loss_set['loss_id'] = self.calc_id_loss(source_tensor, swap_tensor)
+		generator_loss_set['loss_attribute'] = self.calc_attribute_loss(target_attributes, swap_attributes)
+		generator_loss_set['loss_reconstruction'] = self.calc_reconstruction_loss(swap_tensor, target_tensor, is_same_person)
+
+		if weight_pose > 0:
+			generator_loss_set['loss_pose'] = self.calc_pose_loss(swap_tensor, target_tensor)
+		else:
+			generator_loss_set['loss_pose'] = torch.tensor(0).to(swap_tensor.device).to(swap_tensor.dtype)
+
+		if weight_gaze > 0:
+			generator_loss_set['loss_gaze'] = self.calc_gaze_loss(swap_tensor, target_tensor)
+		else:
+			generator_loss_set['loss_gaze'] = torch.tensor(0).to(swap_tensor.device).to(swap_tensor.dtype)
+
+		generator_loss_set['loss_generator'] = generator_loss_set.get('loss_adversarial') * weight_adversarial
+		generator_loss_set['loss_generator'] += generator_loss_set.get('loss_id') * weight_id
+		generator_loss_set['loss_generator'] += generator_loss_set.get('loss_attribute') * weight_attribute
+		generator_loss_set['loss_generator'] += generator_loss_set.get('loss_reconstruction') * weight_reconstruction
+		generator_loss_set['loss_generator'] += generator_loss_set.get('loss_pose') * weight_pose
+		generator_loss_set['loss_generator'] += generator_loss_set.get('loss_gaze') * weight_gaze
+		return generator_loss_set
+
+	def calc_discriminator_loss(self, real_discriminator_outputs : DiscriminatorOutputs, fake_discriminator_outputs : DiscriminatorOutputs) -> DiscriminatorLossSet:
+		discriminator_loss_set = {}
+		loss_fake = torch.Tensor(0)
+
+		for fake_discriminator_output in fake_discriminator_outputs:
+			loss_fake += hinge_fake_loss(fake_discriminator_output[0]).mean()
+
+		loss_true = torch.Tensor(0)
+
+		for true_discriminator_output in real_discriminator_outputs:
+			loss_true += hinge_real_loss(true_discriminator_output[0]).mean()
+
+		discriminator_loss_set['loss_discriminator'] = (loss_true.mean() + loss_fake.mean()) * 0.5
+		return discriminator_loss_set
+
+	def calc_adversarial_loss(self, discriminator_outputs : DiscriminatorOutputs) -> LossTensor:
+		loss_adversarial = torch.Tensor(0)
+
+		for discriminator_output in discriminator_outputs:
+			loss_adversarial += hinge_real_loss(discriminator_output[0])
+
+		loss_adversarial = torch.mean(loss_adversarial)
+		return loss_adversarial
+
+	def calc_attribute_loss(self, target_attributes : TargetAttributes, swap_attributes : SwapAttributes) -> LossTensor:
+		loss_attribute = torch.Tensor(0)
+
+		for swap_attribute, target_attribute in zip(swap_attributes, target_attributes):
+			loss_attribute += torch.mean(torch.pow(swap_attribute - target_attribute, 2).reshape(self.batch_size, -1), dim = 1).mean()
+
+		loss_attribute *= 0.5
+		return loss_attribute
+
+	def calc_reconstruction_loss(self, swap_tensor : VisionTensor, target_tensor : VisionTensor, is_same_person : Tensor) -> LossTensor:
+		loss_reconstruction = torch.pow(swap_tensor - target_tensor, 2).reshape(self.batch_size, -1)
+		loss_reconstruction = torch.mean(loss_reconstruction, dim = 1) * 0.5
+		loss_reconstruction = torch.sum(loss_reconstruction * is_same_person) / (is_same_person.sum() + 1e-4)
+		loss_ssim = 1 - ssim(swap_tensor, target_tensor, data_range = float(torch.max(swap_tensor) - torch.min(swap_tensor))).mean()
+		loss_reconstruction = (loss_reconstruction + loss_ssim) * 0.5
+		return loss_reconstruction
+
+	def calc_id_loss(self, source_tensor : VisionTensor, swap_tensor : VisionTensor) -> LossTensor:
+		swap_embedding = calc_id_embedding(self.id_embedder, swap_tensor, (30, 0, 10, 10))
+		source_embedding = calc_id_embedding(self.id_embedder, source_tensor, (30, 0, 10, 10))
+		loss_id = (1 - torch.cosine_similarity(source_embedding, swap_embedding)).mean()
+		return loss_id
+
+	def calc_pose_loss(self, swap_tensor : VisionTensor, target_tensor : VisionTensor) -> LossTensor:
+		swap_motion_features = self.get_pose_features(swap_tensor)
+		target_motion_features = self.get_pose_features(target_tensor)
+		loss_pose = torch.tensor(0).to(swap_tensor.device).to(swap_tensor.dtype)
+
+		for swap_motion_feature, target_motion_feature in zip(swap_motion_features, target_motion_features):
+			loss_pose += self.mse_loss(swap_motion_feature, target_motion_feature)
+
+		return loss_pose
+
+	def calc_gaze_loss(self, swap_tensor : VisionTensor, target_tensor : VisionTensor) -> LossTensor:
+		swap_landmark = self.get_face_landmarks(swap_tensor)
+		target_landmark = self.get_face_landmarks(target_tensor)
+		left_gaze_loss = self.mse_loss(swap_landmark[:, 198], target_landmark[:, 198])
+		right_gaze_loss = self.mse_loss(swap_landmark[:, 197], target_landmark[:, 197])
+		gaze_loss = left_gaze_loss + right_gaze_loss
+		return gaze_loss
+
+	def get_face_landmarks(self, vision_tensor : VisionTensor) -> FaceLandmark203:
+		vision_tensor_norm = (vision_tensor + 1) * 0.5
+		vision_tensor_norm = torch.nn.functional.interpolate(vision_tensor_norm, size = (224, 224), mode = 'bilinear')
+		landmarks = self.landmarker(vision_tensor_norm)[2].view(-1, 203, 2)
+		return landmarks
+
+	def get_pose_features(self, vision_tensor : Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+		vision_tensor_norm = (vision_tensor + 1) * 0.5
+		pitch, yaw, roll, translation, expression, scale, _ = self.motion_extractor(vision_tensor_norm)
+		rotation = torch.cat([ pitch, yaw, roll ], dim = 1)
+		return translation, scale, rotation