import numpy as np
import torch
import torch.nn.functional as F
from typing import Iterable, Tuple, Union
from pathlib import Path
from torchvision import transforms
import kornia
from omegaconf import DictConfig
from src.FaceDetector.face_detector import Detection
from src.FaceAlign.face_align import align_face, inverse_transform_batch
from src.PostProcess.utils import SoftErosion
from src.model_loader import get_model
from src.Misc.types import CheckpointType, FaceAlignmentType
from src.Misc.utils import tensor2img
class SimSwap:
def __init__(
self,
config: DictConfig,
id_image: Union[np.ndarray, None] = None,
specific_image: Union[np.ndarray, None] = None,
):
self.id_image: Union[np.ndarray, None] = id_image
self.id_latent: Union[torch.Tensor, None] = None
self.specific_id_image: Union[np.ndarray, None] = specific_image
self.specific_latent: Union[torch.Tensor, None] = None
self.use_mask: Union[bool, None] = True
self.crop_size: Union[int, None] = None
self.checkpoint_type: Union[CheckpointType, None] = None
self.face_alignment_type: Union[FaceAlignmentType, None] = None
self.smooth_mask_iter: Union[int, None] = None
self.smooth_mask_kernel_size: Union[int, None] = None
self.smooth_mask_threshold: Union[float, None] = None
self.face_detector_threshold: Union[float, None] = None
self.specific_latent_match_threshold: Union[float, None] = None
self.device = torch.device(config.device)
self.set_parameters(config)
# For BiSeNet and for official_224 SimSwap
self.to_tensor_normalize = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
# For SimSwap models trained with the updated code
self.to_tensor = transforms.ToTensor()
self.face_detector = get_model(
"face_detector",
device=self.device,
load_state_dice=False,
model_path=Path(config.face_detector_weights),
det_thresh=self.face_detector_threshold,
det_size=(640, 640),
mode="ffhq",
)
self.face_id_net = get_model(
"arcface",
device=self.device,
load_state_dice=False,
model_path=Path(config.face_id_weights),
)
self.bise_net = get_model(
"parsing_model",
device=self.device,
load_state_dice=True,
model_path=Path(config.parsing_model_weights),
n_classes=19,
)
gen_model = "generator_512" if self.crop_size == 512 else "generator_224"
self.simswap_net = get_model(
gen_model,
device=self.device,
load_state_dice=True,
model_path=Path(config.simswap_weights),
input_nc=3,
output_nc=3,
latent_size=512,
n_blocks=9,
deep=True if self.crop_size == 512 else False,
use_last_act=True
if self.checkpoint_type == CheckpointType.OFFICIAL_224
else False,
)
self.blend = get_model(
"blend_module",
device=self.device,
load_state_dice=False,
model_path=Path(config.blend_module_weights)
)
self.enhance_output = config.enhance_output
if config.enhance_output:
self.gfpgan_net = get_model(
"gfpgan",
device=self.device,
load_state_dice=True,
model_path=Path(config.gfpgan_weights)
)
def set_parameters(self, config) -> None:
self.set_crop_size(config.crop_size)
self.set_checkpoint_type(config.checkpoint_type)
self.set_face_alignment_type(config.face_alignment_type)
self.set_face_detector_threshold(config.face_detector_threshold)
self.set_specific_latent_match_threshold(config.specific_latent_match_threshold)
self.set_smooth_mask_kernel_size(config.smooth_mask_kernel_size)
self.set_smooth_mask_threshold(config.smooth_mask_threshold)
self.set_smooth_mask_iter(config.smooth_mask_iter)
def set_crop_size(self, crop_size: int) -> None:
if crop_size < 0:
raise "Invalid crop_size! Must be a positive value."
self.crop_size = crop_size
def set_checkpoint_type(self, checkpoint_type: str) -> None:
type = CheckpointType(checkpoint_type)
if type not in (CheckpointType.OFFICIAL_224, CheckpointType.UNOFFICIAL):
raise "Invalid checkpoint_type! Must be one of the predefined values."
self.checkpoint_type = type
def set_face_alignment_type(self, face_alignment_type: str) -> None:
type = FaceAlignmentType(face_alignment_type)
if type not in (
FaceAlignmentType.FFHQ,
FaceAlignmentType.DEFAULT,
):
raise "Invalid face_alignment_type! Must be one of the predefined values."
self.face_alignment_type = type
def set_face_detector_threshold(self, face_detector_threshold: float) -> None:
if face_detector_threshold < 0.0 or face_detector_threshold > 1.0:
raise "Invalid face_detector_threshold! Must be a positive value in range [0.0...1.0]."
self.face_detector_threshold = face_detector_threshold
def set_specific_latent_match_threshold(
self, specific_latent_match_threshold: float
) -> None:
if specific_latent_match_threshold < 0.0:
raise "Invalid specific_latent_match_th! Must be a positive value."
self.specific_latent_match_threshold = specific_latent_match_threshold
def re_initialize_soft_mask(self):
self.smooth_mask = SoftErosion(kernel_size=self.smooth_mask_kernel_size,
threshold=self.smooth_mask_threshold,
iterations=self.smooth_mask_iter).to(self.device)
def set_smooth_mask_kernel_size(self, smooth_mask_kernel_size: int) -> None:
if smooth_mask_kernel_size < 0:
raise "Invalid smooth_mask_kernel_size! Must be a positive value."
smooth_mask_kernel_size += 1 if smooth_mask_kernel_size % 2 == 0 else 0
self.smooth_mask_kernel_size = smooth_mask_kernel_size
self.re_initialize_soft_mask()
def set_smooth_mask_threshold(self, smooth_mask_threshold: int) -> None:
if smooth_mask_threshold < 0 or smooth_mask_threshold > 1.0:
raise "Invalid smooth_mask_threshold! Must be within 0...1 range."
self.smooth_mask_threshold = smooth_mask_threshold
self.re_initialize_soft_mask()
def set_smooth_mask_iter(self, smooth_mask_iter: float) -> None:
if smooth_mask_iter < 0:
raise "Invalid smooth_mask_iter! Must be a positive value.."
self.smooth_mask_iter = smooth_mask_iter
self.re_initialize_soft_mask()
def run_detect_align(self, image: np.ndarray, for_id: bool = False) -> Tuple[Union[Iterable[np.ndarray], None],
Union[Iterable[np.ndarray], None],
np.ndarray]:
detection: Detection = self.face_detector(image)
if detection.bbox is None:
if for_id:
raise "Can't detect a face! Please change the ID image!"
return None, None, detection.score
kps = detection.key_points
if for_id:
max_score_ind = np.argmax(detection.score, axis=0)
kps = detection.key_points[max_score_ind]
kps = kps[None, ...]
align_imgs, transforms = align_face(
image,
kps,
crop_size=self.crop_size,
mode="ffhq"
if self.face_alignment_type == FaceAlignmentType.FFHQ
else "none",
)
return align_imgs, transforms, detection.score
def __call__(self, att_image: np.ndarray) -> np.ndarray:
if self.id_latent is None:
align_id_imgs, id_transforms, _ = self.run_detect_align(
self.id_image, for_id=True
)
# normalize=True, because official SimSwap model trained with normalized id_lattent
self.id_latent: torch.Tensor = self.face_id_net(
align_id_imgs, normalize=True
)
if self.specific_id_image is not None and self.specific_latent is None:
align_specific_imgs, specific_transforms, _ = self.run_detect_align(
self.specific_id_image, for_id=True
)
self.specific_latent: torch.Tensor = self.face_id_net(
align_specific_imgs, normalize=False
)
# for_id=False, because we want to get all faces
align_att_imgs, att_transforms, att_detection_score = self.run_detect_align(
att_image, for_id=False
)
if align_att_imgs is None and att_transforms is None:
return att_image
# Select specific crop from the target image
if self.specific_latent is not None:
att_latent: torch.Tensor = self.face_id_net(align_att_imgs, normalize=False)
latent_dist = torch.mean(
F.mse_loss(
att_latent,
self.specific_latent.repeat(att_latent.shape[0], 1),
reduction="none",
),
dim=-1,
)
att_detection_score = torch.tensor(
att_detection_score, device=latent_dist.device
)
min_index = torch.argmin(latent_dist * att_detection_score)
min_value = latent_dist[min_index]
if min_value < self.specific_latent_match_threshold:
align_att_imgs = [align_att_imgs[min_index]]
att_transforms = [att_transforms[min_index]]
else:
return att_image
swapped_img: torch.Tensor = self.simswap_net(align_att_imgs, self.id_latent)
if self.enhance_output:
swapped_img = self.gfpgan_net.enhance(swapped_img, weight=0.5)
# Put all crops/transformations into a batch
align_att_img_batch_for_parsing_model: torch.Tensor = torch.stack(
[self.to_tensor_normalize(x) for x in align_att_imgs], dim=0
)
align_att_img_batch_for_parsing_model = (
align_att_img_batch_for_parsing_model.to(self.device)
)
att_transforms: torch.Tensor = torch.stack(
[torch.tensor(x).float() for x in att_transforms], dim=0
)
att_transforms = att_transforms.to(self.device, non_blocking=True)
align_att_img_batch: torch.Tensor = torch.stack(
[self.to_tensor(x) for x in align_att_imgs], dim=0
)
align_att_img_batch = align_att_img_batch.to(self.device, non_blocking=True)
# Get face masks for the attribute image
face_mask, ignore_mask_ids = self.bise_net.get_mask(
align_att_img_batch_for_parsing_model, self.crop_size
)
inv_att_transforms: torch.Tensor = inverse_transform_batch(att_transforms)
soft_face_mask, _ = self.smooth_mask(face_mask)
swapped_img[ignore_mask_ids, ...] = align_att_img_batch[ignore_mask_ids, ...]
frame_size = (att_image.shape[0], att_image.shape[1])
att_image = self.to_tensor(att_image).to(self.device, non_blocking=True).unsqueeze(0)
target_image = kornia.geometry.transform.warp_affine(
swapped_img,
inv_att_transforms,
frame_size,
mode="bilinear",
padding_mode="border",
align_corners=True,
fill_value=torch.zeros(3),
)
soft_face_mask = kornia.geometry.transform.warp_affine(
soft_face_mask,
inv_att_transforms,
frame_size,
mode="bilinear",
padding_mode="zeros",
align_corners=True,
fill_value=torch.zeros(3),
)
result = self.blend(target_image, soft_face_mask, att_image)
return tensor2img(result)