From 611618e41345653f5c99f6c5d06a7269c217a0ad Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 12 Feb 2025 10:31:53 +0100 Subject: [PATCH] Make Embedding great again --- arcface_converter/src/preparing.py | 12 ++++++------ arcface_converter/src/types.py | 2 +- face_swapper/src/helper.py | 4 ++-- face_swapper/src/inferencing.py | 4 ++-- face_swapper/src/models/generator.py | 4 ++-- face_swapper/src/networks/attribute_modulator.py | 10 +++++----- face_swapper/src/training.py | 4 ++-- face_swapper/src/types.py | 5 ++--- 8 files changed, 22 insertions(+), 23 deletions(-) diff --git a/arcface_converter/src/preparing.py b/arcface_converter/src/preparing.py index 92676d7..6ba75fe 100644 --- a/arcface_converter/src/preparing.py +++ b/arcface_converter/src/preparing.py @@ -9,7 +9,7 @@ from mxnet.io import ImageRecordIter from onnxruntime import InferenceSession from tqdm import tqdm -from .types import Embedding, EmbeddingPairs, VisionFrame +from .types import Embedding, VisionFrame CONFIG = configparser.ConfigParser() CONFIG.read('config.ini') @@ -35,9 +35,9 @@ def forward(inference_session : InferenceSession, crop_vision_frame : VisionFram return embedding -def process_embeddings(dataset_reader : ImageRecordIter, source_inference_session : InferenceSession, target_inference_session : InferenceSession) -> EmbeddingPairs: +def process_embeddings(dataset_reader : ImageRecordIter, source_inference_session : InferenceSession, target_inference_session : InferenceSession) -> Embedding: dataset_process_limit = CONFIG.getint('preparing.dataset', 'process_limit') - embedding_pairs = [] + embeddings = [] with tqdm(total = dataset_process_limit) as progress: for batch in dataset_reader: @@ -45,13 +45,13 @@ def process_embeddings(dataset_reader : ImageRecordIter, source_inference_sessio crop_vision_frame = prepare_crop_vision_frame(crop_vision_frame) source_embedding = forward(source_inference_session, crop_vision_frame) target_embedding = forward(target_inference_session, crop_vision_frame) - embedding_pairs.append([ source_embedding, target_embedding ]) + embeddings.append([ source_embedding, target_embedding ]) progress.update() if progress.n == dataset_process_limit: - return numpy.concatenate(embedding_pairs, axis = 1).T + return numpy.concatenate(embeddings, axis = 1).T - return numpy.concatenate(embedding_pairs, axis = 1).T + return numpy.concatenate(embeddings, axis = 1).T def prepare() -> None: diff --git a/arcface_converter/src/types.py b/arcface_converter/src/types.py index faeeec2..16e699a 100644 --- a/arcface_converter/src/types.py +++ b/arcface_converter/src/types.py @@ -8,6 +8,6 @@ Batch = Tuple[Tensor, Tensor] Loader = DataLoader[Tuple[Tensor, ...]] Embedding = NDArray[Any] -EmbeddingPairs = NDArray[Any] FaceLandmark5 = NDArray[Any] + VisionFrame = NDArray[Any] diff --git a/face_swapper/src/helper.py b/face_swapper/src/helper.py index c9fce5f..05ee770 100644 --- a/face_swapper/src/helper.py +++ b/face_swapper/src/helper.py @@ -4,7 +4,7 @@ import cv2 import numpy import torch -from .types import IdEmbedder, IdEmbedding, Padding, Tensor, VisionFrame, VisionTensor +from .types import Embedder, Embedding, Padding, Tensor, VisionFrame, VisionTensor def is_windows() -> bool: @@ -47,7 +47,7 @@ def hinge_fake_loss(tensor : Tensor) -> Tensor: return fake_loss -def calc_id_embedding(id_embedder : IdEmbedder, vision_tensor : VisionTensor, padding : Padding) -> IdEmbedding: +def calc_id_embedding(id_embedder : Embedder, vision_tensor : VisionTensor, padding : Padding) -> Embedding: crop_vision_tensor = vision_tensor[:, :, 15 : 241, 15 : 241] crop_vision_tensor = torch.nn.functional.interpolate(crop_vision_tensor, size = (112, 112), mode = 'area') crop_vision_tensor[:, :, :padding[0], :] = 0 diff --git a/face_swapper/src/inferencing.py b/face_swapper/src/inferencing.py index 3da1573..18451bc 100644 --- a/face_swapper/src/inferencing.py +++ b/face_swapper/src/inferencing.py @@ -5,13 +5,13 @@ import torch from .helper import calc_id_embedding, convert_to_vision_frame, convert_to_vision_tensor, read_image from .models.generator import AdaptiveEmbeddingIntegrationNetwork -from .types import Generator, IdEmbedder, VisionFrame +from .types import Generator, Embedder, VisionFrame CONFIG = configparser.ConfigParser() CONFIG.read('config.ini') -def run_swap(generator : Generator, id_embedder : IdEmbedder, source_vision_frame : VisionFrame, target_vision_frame : VisionFrame) -> VisionFrame: +def run_swap(generator : Generator, id_embedder : Embedder, source_vision_frame : VisionFrame, target_vision_frame : VisionFrame) -> VisionFrame: source_vision_tensor = convert_to_vision_tensor(source_vision_frame) target_vision_tensor = convert_to_vision_tensor(target_vision_frame) source_embedding = calc_id_embedding(id_embedder, source_vision_tensor, (0, 0, 0, 0)) diff --git a/face_swapper/src/models/generator.py b/face_swapper/src/models/generator.py index 12a4f76..1b38dac 100644 --- a/face_swapper/src/models/generator.py +++ b/face_swapper/src/models/generator.py @@ -5,7 +5,7 @@ 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 +from face_swapper.src.types import Embedding, TargetAttributes, VisionTensor CONFIG = configparser.ConfigParser() CONFIG.read('config.ini') @@ -22,7 +22,7 @@ class AdaptiveEmbeddingIntegrationNetwork(nn.Module): self.encoder.apply(init_weight) self.generator.apply(init_weight) - def forward(self, target : VisionTensor, source_embedding : SourceEmbedding) -> Tuple[VisionTensor, TargetAttributes]: + def forward(self, target : VisionTensor, source_embedding : Embedding) -> Tuple[VisionTensor, TargetAttributes]: target_attributes = self.get_attributes(target) swap_tensor = self.generator(target_attributes, source_embedding) return swap_tensor, target_attributes diff --git a/face_swapper/src/networks/attribute_modulator.py b/face_swapper/src/networks/attribute_modulator.py index cc34463..69d71e1 100644 --- a/face_swapper/src/networks/attribute_modulator.py +++ b/face_swapper/src/networks/attribute_modulator.py @@ -1,7 +1,7 @@ import torch from torch import Tensor, nn as nn -from face_swapper.src.types import SourceEmbedding, TargetAttributes +from face_swapper.src.types import Embedding, TargetAttributes class AADGenerator(nn.Module): @@ -17,7 +17,7 @@ class AADGenerator(nn.Module): self.res_block_7 = AADResBlock(128, 64, 64, id_channels, num_blocks) self.res_block_8 = AADResBlock(64, 3, 64, id_channels, num_blocks) - def forward(self, target_attributes : TargetAttributes, source_embedding : SourceEmbedding) -> Tensor: + def forward(self, target_attributes : TargetAttributes, source_embedding : Embedding) -> Tensor: feature_map = self.upsample(source_embedding) feature_map_1 = torch.nn.functional.interpolate(self.res_block_1(feature_map, target_attributes[0], source_embedding), scale_factor = 2, mode = 'bilinear', align_corners = False) feature_map_2 = torch.nn.functional.interpolate(self.res_block_2(feature_map_1, target_attributes[1], source_embedding), scale_factor = 2, mode = 'bilinear', align_corners = False) @@ -41,7 +41,7 @@ class AADLayer(nn.Module): self.instance_norm = nn.InstanceNorm2d(input_channels) self.conv_mask = nn.Conv2d(input_channels, 1, kernel_size = 1) - def forward(self, feature_map : Tensor, attribute_embedding : Tensor, id_embedding : SourceEmbedding) -> Tensor: + def forward(self, feature_map : Tensor, attribute_embedding : Tensor, id_embedding : Embedding) -> Tensor: feature_map = self.instance_norm(feature_map) gamma_attribute = self.conv_gamma(attribute_embedding) beta_attribute = self.conv_beta(attribute_embedding) @@ -59,7 +59,7 @@ class AADSequential(nn.Module): super(AADSequential, self).__init__() self.layers = nn.ModuleList(args) - def forward(self, feature_map: Tensor, attribute_embedding: Tensor, id_embedding: SourceEmbedding) -> Tensor: + def forward(self, feature_map: Tensor, attribute_embedding: Tensor, id_embedding: Embedding) -> Tensor: for layer in self.layers: if isinstance(layer, AADLayer): feature_map = layer(feature_map, attribute_embedding, id_embedding) @@ -99,7 +99,7 @@ class AADResBlock(nn.Module): ) self.auxiliary_add_blocks = auxiliary_add_blocks - def forward(self, feature_map : Tensor, attribute_embedding : Tensor, id_embedding : SourceEmbedding) -> Tensor: + def forward(self, feature_map : Tensor, attribute_embedding : Tensor, id_embedding : Embedding) -> Tensor: primary_feature = self.primary_add_blocks(feature_map, attribute_embedding, id_embedding) if self.input_channels > self.output_channels: diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index ef094f0..d023085 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -16,7 +16,7 @@ from .helper import calc_id_embedding from .models.discriminator import MultiscaleDiscriminator from .models.generator import AdaptiveEmbeddingIntegrationNetwork from .models.loss import FaceSwapperLoss -from .types import Batch, SourceEmbedding, TargetAttributes, VisionTensor +from .types import Batch, Embedding, TargetAttributes, VisionTensor CONFIG = configparser.ConfigParser() CONFIG.read('config.ini') @@ -29,7 +29,7 @@ class FaceSwapperTrain(pytorch_lightning.LightningModule, FaceSwapperLoss): self.discriminator = MultiscaleDiscriminator() self.automatic_optimization = CONFIG.getboolean('training.trainer', 'automatic_optimization') - def forward(self, target_tensor : VisionTensor, source_embedding : SourceEmbedding) -> Tuple[VisionTensor, TargetAttributes]: + def forward(self, target_tensor : VisionTensor, source_embedding : Embedding) -> Tuple[VisionTensor, TargetAttributes]: output = self.generator(target_tensor, source_embedding) return output diff --git a/face_swapper/src/types.py b/face_swapper/src/types.py index 96de434..e37001d 100644 --- a/face_swapper/src/types.py +++ b/face_swapper/src/types.py @@ -15,8 +15,7 @@ SwapAttributes = Tuple[Tensor, ...] TargetAttributes = Tuple[Tensor, ...] DiscriminatorOutputs = List[List[Tensor]] -IdEmbedding = Tensor -SourceEmbedding = IdEmbedding +Embedding = Tensor FaceLandmark203 = Tensor StateSet = OrderedDict[str, Any] @@ -30,4 +29,4 @@ GeneratorLossSet = Dict[str, Tensor] DiscriminatorLossSet = Dict[str, Tensor] Generator = torch.nn.Module -IdEmbedder = torch.nn.Module +Embedder = torch.nn.Module