CalvinBackup/sd-webui-faceswaplab
1
0
Fork 0
Code Issues 56 Pull Requests 4 Actions Packages Projects Releases 8 Wiki Activity

add pre-commit hooks configuration

Browse Source
This commit is contained in:
Tran Xen
2023-07-28 18:25:28 +02:00
parent 8577d0186d
commit 5d4a29ff1e
33 changed files with 1674 additions and 820 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/faceswaplab_swapping/facemask.py
+52 -13
View File
@@ -7,21 +7,27 @@ from functools import lru_cache
from typing import Union, List
from torch import device as torch_device
@lru_cache
def get_parsing_model(device: torch_device) -> torch.nn.Module:
"""
Returns an instance of the parsing model.
Returns an instance of the parsing model.
The returned model is cached for faster subsequent access.
Args:
device: The torch device to use for computations.
Returns:
The parsing model.
"""
return init_parsing_model(device=device)
def convert_image_to_tensor(images: Union[np.ndarray, List[np.ndarray]], convert_bgr_to_rgb: bool = True, use_float32: bool = True) -> Union[torch.Tensor, List[torch.Tensor]]:
def convert_image_to_tensor(
images: Union[np.ndarray, List[np.ndarray]],
convert_bgr_to_rgb: bool = True,
use_float32: bool = True,
) -> Union[torch.Tensor, List[torch.Tensor]]:
"""
Converts an image or a list of images to PyTorch tensor.
@@ -33,10 +39,13 @@ def convert_image_to_tensor(images: Union[np.ndarray, List[np.ndarray]], convert
Returns:
PyTorch tensor or a list of PyTorch tensors.
"""
def _convert_single_image_to_tensor(image: np.ndarray, convert_bgr_to_rgb: bool, use_float32: bool) -> torch.Tensor:
def _convert_single_image_to_tensor(
image: np.ndarray, convert_bgr_to_rgb: bool, use_float32: bool
) -> torch.Tensor:
if image.shape[2] == 3 and convert_bgr_to_rgb:
if image.dtype == 'float64':
image = image.astype('float32')
if image.dtype == "float64":
image = image.astype("float32")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_tensor = torch.from_numpy(image.transpose(2, 0, 1))
if use_float32:
@@ -44,10 +53,14 @@ def convert_image_to_tensor(images: Union[np.ndarray, List[np.ndarray]], convert
return image_tensor
if isinstance(images, list):
return [_convert_single_image_to_tensor(image, convert_bgr_to_rgb, use_float32) for image in images]
return [
_convert_single_image_to_tensor(image, convert_bgr_to_rgb, use_float32)
for image in images
]
else:
return _convert_single_image_to_tensor(images, convert_bgr_to_rgb, use_float32)
def generate_face_mask(face_image: np.ndarray, device: torch.device) -> np.ndarray:
"""
Generates a face mask given a face image.
@@ -60,12 +73,18 @@ def generate_face_mask(face_image: np.ndarray, device: torch.device) -> np.ndarr
The face mask as a numpy.ndarray.
"""
# Resize the face image for the model
resized_face_image = cv2.resize(face_image, (512, 512), interpolation=cv2.INTER_LINEAR)
resized_face_image = cv2.resize(
face_image, (512, 512), interpolation=cv2.INTER_LINEAR
)
# Preprocess the image
face_input = convert_image_to_tensor((resized_face_image.astype('float32') / 255.0), convert_bgr_to_rgb=True, use_float32=True)
face_input = convert_image_to_tensor(
(resized_face_image.astype("float32") / 255.0),
convert_bgr_to_rgb=True,
use_float32=True,
)
normalize(face_input, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
assert isinstance(face_input,torch.Tensor)
assert isinstance(face_input, torch.Tensor)
face_input = torch.unsqueeze(face_input, 0).to(device)
# Pass the image through the model
@@ -75,7 +94,27 @@ def generate_face_mask(face_image: np.ndarray, device: torch.device) -> np.ndarr
# Generate the mask from the model output
parse_mask = np.zeros(model_output.shape)
MASK_COLOR_MAP = [0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 255, 0, 0, 0]
MASK_COLOR_MAP = [
0,
255,
255,
255,
255,
255,
255,
255,
255,
255,
255,
255,
255,
255,
0,
255,
0,
0,
0,
]
for idx, color in enumerate(MASK_COLOR_MAP):
parse_mask[model_output == idx] = color
scripts/faceswaplab_swapping/parsing/__init__.py
+28 -26
View File
@@ -5,36 +5,36 @@ S-Lab License 1.0
Copyright 2022 S-Lab
Redistribution and use for non-commercial purpose in source and
binary forms, with or without modification, are permitted provided
Redistribution and use for non-commercial purpose in source and
binary forms, with or without modification, are permitted provided
that the following conditions are met:
1. Redistributions of source code must retain the above copyright
1. Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in
the documentation and/or other materials provided with the
2. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in
the documentation and/or other materials provided with the
distribution.
3. Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived
3. Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived
from this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
In the event that redistribution and/or use for commercial purpose in
source or binary forms, with or without modification is required,
In the event that redistribution and/or use for commercial purpose in
source or binary forms, with or without modification is required,
please contact the contributor(s) of the work.
"""
@@ -50,12 +50,12 @@ from scripts.faceswaplab_globals import FACE_PARSER_DIR
ROOT_DIR = FACE_PARSER_DIR
def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
"""Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
"""
"""Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py"""
if model_dir is None:
hub_dir = get_dir()
model_dir = os.path.join(hub_dir, 'checkpoints')
model_dir = os.path.join(hub_dir, "checkpoints")
os.makedirs(os.path.join(ROOT_DIR, model_dir), exist_ok=True)
@@ -70,10 +70,12 @@ def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
return cached_file
def init_parsing_model(device='cuda'):
def init_parsing_model(device="cuda"):
model = ParseNet(in_size=512, out_size=512, parsing_ch=19)
model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'
model_path = load_file_from_url(url=model_url, model_dir='weights/facelib', progress=True, file_name=None)
model_url = "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth"
model_path = load_file_from_url(
url=model_url, model_dir="weights/facelib", progress=True, file_name=None
)
load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
model.load_state_dict(load_net, strict=True)
model.eval()
scripts/faceswaplab_swapping/parsing/parsenet.py
+89 -73
View File
@@ -5,36 +5,36 @@ S-Lab License 1.0
Copyright 2022 S-Lab
Redistribution and use for non-commercial purpose in source and
binary forms, with or without modification, are permitted provided
Redistribution and use for non-commercial purpose in source and
binary forms, with or without modification, are permitted provided
that the following conditions are met:
1. Redistributions of source code must retain the above copyright
1. Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in
the documentation and/or other materials provided with the
2. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in
the documentation and/or other materials provided with the
distribution.
3. Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived
3. Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived
from this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
In the event that redistribution and/or use for commercial purpose in
source or binary forms, with or without modification is required,
In the event that redistribution and/or use for commercial purpose in
source or binary forms, with or without modification is required,
please contact the contributor(s) of the work.
Modified from https://github.com/chaofengc/PSFRGAN
@@ -98,7 +98,7 @@ exhaustive, and do not form part of our licenses.
such as asking that all changes be marked or described.
Although not required by our licenses, you are encouraged to
respect those requests where reasonable. More_considerations
for the public:
for the public:
wiki.creativecommons.org/Considerations_for_licensees
=======================================================================
@@ -499,27 +499,27 @@ class NormLayer(nn.Module):
input_size: input shape without batch size, for layer norm.
"""
def __init__(self, channels, normalize_shape=None, norm_type='bn'):
def __init__(self, channels, normalize_shape=None, norm_type="bn"):
super(NormLayer, self).__init__()
norm_type = norm_type.lower()
self.norm_type = norm_type
if norm_type == 'bn':
if norm_type == "bn":
self.norm = nn.BatchNorm2d(channels, affine=True)
elif norm_type == 'in':
elif norm_type == "in":
self.norm = nn.InstanceNorm2d(channels, affine=False)
elif norm_type == 'gn':
elif norm_type == "gn":
self.norm = nn.GroupNorm(32, channels, affine=True)
elif norm_type == 'pixel':
elif norm_type == "pixel":
self.norm = lambda x: F.normalize(x, p=2, dim=1)
elif norm_type == 'layer':
elif norm_type == "layer":
self.norm = nn.LayerNorm(normalize_shape)
elif norm_type == 'none':
elif norm_type == "none":
self.norm = lambda x: x * 1.0
else:
assert 1 == 0, f'Norm type {norm_type} not support.'
assert 1 == 0, f"Norm type {norm_type} not support."
def forward(self, x, ref=None):
if self.norm_type == 'spade':
if self.norm_type == "spade":
return self.norm(x, ref)
else:
return self.norm(x)
@@ -537,51 +537,56 @@ class ReluLayer(nn.Module):
- none: direct pass
"""
def __init__(self, channels, relu_type='relu'):
def __init__(self, channels, relu_type="relu"):
super(ReluLayer, self).__init__()
relu_type = relu_type.lower()
if relu_type == 'relu':
if relu_type == "relu":
self.func = nn.ReLU(True)
elif relu_type == 'leakyrelu':
elif relu_type == "leakyrelu":
self.func = nn.LeakyReLU(0.2, inplace=True)
elif relu_type == 'prelu':
elif relu_type == "prelu":
self.func = nn.PReLU(channels)
elif relu_type == 'selu':
elif relu_type == "selu":
self.func = nn.SELU(True)
elif relu_type == 'none':
elif relu_type == "none":
self.func = lambda x: x * 1.0
else:
assert 1 == 0, f'Relu type {relu_type} not support.'
assert 1 == 0, f"Relu type {relu_type} not support."
def forward(self, x):
return self.func(x)
class ConvLayer(nn.Module):
def __init__(self,
in_channels,
out_channels,
kernel_size=3,
scale='none',
norm_type='none',
relu_type='none',
use_pad=True,
bias=True):
def __init__(
self,
in_channels,
out_channels,
kernel_size=3,
scale="none",
norm_type="none",
relu_type="none",
use_pad=True,
bias=True,
):
super(ConvLayer, self).__init__()
self.use_pad = use_pad
self.norm_type = norm_type
if norm_type in ['bn']:
if norm_type in ["bn"]:
bias = False
stride = 2 if scale == 'down' else 1
stride = 2 if scale == "down" else 1
self.scale_func = lambda x: x
if scale == 'up':
self.scale_func = lambda x: nn.functional.interpolate(x, scale_factor=2, mode='nearest')
if scale == "up":
self.scale_func = lambda x: nn.functional.interpolate(
x, scale_factor=2, mode="nearest"
)
self.reflection_pad = nn.ReflectionPad2d(int(np.ceil((kernel_size - 1.) / 2)))
self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, bias=bias)
self.reflection_pad = nn.ReflectionPad2d(int(np.ceil((kernel_size - 1.0) / 2)))
self.conv2d = nn.Conv2d(
in_channels, out_channels, kernel_size, stride, bias=bias
)
self.relu = ReluLayer(out_channels, relu_type)
self.norm = NormLayer(out_channels, norm_type=norm_type)
@@ -601,19 +606,27 @@ class ResidualBlock(nn.Module):
Residual block recommended in: http://torch.ch/blog/2016/02/04/resnets.html
"""
def __init__(self, c_in, c_out, relu_type='prelu', norm_type='bn', scale='none'):
def __init__(self, c_in, c_out, relu_type="prelu", norm_type="bn", scale="none"):
super(ResidualBlock, self).__init__()
if scale == 'none' and c_in == c_out:
if scale == "none" and c_in == c_out:
self.shortcut_func = lambda x: x
else:
self.shortcut_func = ConvLayer(c_in, c_out, 3, scale)
scale_config_dict = {'down': ['none', 'down'], 'up': ['up', 'none'], 'none': ['none', 'none']}
scale_config_dict = {
"down": ["none", "down"],
"up": ["up", "none"],
"none": ["none", "none"],
}
scale_conf = scale_config_dict[scale]
self.conv1 = ConvLayer(c_in, c_out, 3, scale_conf[0], norm_type=norm_type, relu_type=relu_type)
self.conv2 = ConvLayer(c_out, c_out, 3, scale_conf[1], norm_type=norm_type, relu_type='none')
self.conv1 = ConvLayer(
c_in, c_out, 3, scale_conf[0], norm_type=norm_type, relu_type=relu_type
)
self.conv2 = ConvLayer(
c_out, c_out, 3, scale_conf[1], norm_type=norm_type, relu_type="none"
)
def forward(self, x):
identity = self.shortcut_func(x)
@@ -624,20 +637,21 @@ class ResidualBlock(nn.Module):
class ParseNet(nn.Module):
def __init__(self,
in_size=128,
out_size=128,
min_feat_size=32,
base_ch=64,
parsing_ch=19,
res_depth=10,
relu_type='LeakyReLU',
norm_type='bn',
ch_range=[32, 256]):
def __init__(
self,
in_size=128,
out_size=128,
min_feat_size=32,
base_ch=64,
parsing_ch=19,
res_depth=10,
relu_type="LeakyReLU",
norm_type="bn",
ch_range=[32, 256],
):
super().__init__()
self.res_depth = res_depth
act_args = {'norm_type': norm_type, 'relu_type': relu_type}
act_args = {"norm_type": norm_type, "relu_type": relu_type}
min_ch, max_ch = ch_range
ch_clip = lambda x: max(min_ch, min(x, max_ch)) # noqa: E731
@@ -652,17 +666,19 @@ class ParseNet(nn.Module):
head_ch = base_ch
for i in range(down_steps):
cin, cout = ch_clip(head_ch), ch_clip(head_ch * 2)
self.encoder.append(ResidualBlock(cin, cout, scale='down', **act_args))
self.encoder.append(ResidualBlock(cin, cout, scale="down", **act_args))
head_ch = head_ch * 2
self.body = []
for i in range(res_depth):
self.body.append(ResidualBlock(ch_clip(head_ch), ch_clip(head_ch), **act_args))
self.body.append(
ResidualBlock(ch_clip(head_ch), ch_clip(head_ch), **act_args)
)
self.decoder = []
for i in range(up_steps):
cin, cout = ch_clip(head_ch), ch_clip(head_ch // 2)
self.decoder.append(ResidualBlock(cin, cout, scale='up', **act_args))
self.decoder.append(ResidualBlock(cin, cout, scale="up", **act_args))
head_ch = head_ch // 2
self.encoder = nn.Sequential(*self.encoder)
scripts/faceswaplab_swapping/swapper.py
+138 -64
View File
@@ -12,7 +12,11 @@ from PIL import Image
from sklearn.metrics.pairwise import cosine_similarity
from scripts.faceswaplab_swapping import upscaled_inswapper
from scripts.faceswaplab_utils.imgutils import cv2_to_pil, pil_to_cv2, check_against_nsfw
from scripts.faceswaplab_utils.imgutils import (
cv2_to_pil,
pil_to_cv2,
check_against_nsfw,
)
from scripts.faceswaplab_utils.faceswaplab_logging import logger, save_img_debug
from scripts import faceswaplab_globals
from modules.shared import opts
@@ -48,19 +52,20 @@ def cosine_similarity_face(face1, face2) -> float:
# Return the maximum of 0 and the calculated similarity as the final similarity score
return max(0, similarity[0, 0])
def compare_faces(img1: Image.Image, img2: Image.Image) -> float:
"""
Compares the similarity between two faces extracted from images using cosine similarity.
Args:
img1: The first image containing a face.
img2: The second image containing a face.
Returns:
A float value representing the similarity between the two faces (0 to 1).
A float value representing the similarity between the two faces (0 to 1).
Returns -1 if one or both of the images do not contain any faces.
"""
# Extract faces from the images
face1 = get_or_default(get_faces(pil_to_cv2(img1)), 0, None)
face2 = get_or_default(get_faces(pil_to_cv2(img2)), 0, None)
@@ -69,13 +74,14 @@ def compare_faces(img1: Image.Image, img2: Image.Image) -> float:
if face1 is not None and face2 is not None:
# Calculate the cosine similarity between the faces
return cosine_similarity_face(face1, face2)
# Return -1 if one or both of the images do not contain any faces
return -1
class FaceModelException(Exception):
"""Exception raised when an error is encountered in the face model."""
def __init__(self, message: str) -> None:
"""
Args:
@@ -84,15 +90,16 @@ class FaceModelException(Exception):
self.message = message
super().__init__(self.message)
@lru_cache(maxsize=1)
def getAnalysisModel():
"""
Retrieves the analysis model for face analysis.
Returns:
insightface.app.FaceAnalysis: The analysis model for face analysis.
"""
try :
try:
if not os.path.exists(faceswaplab_globals.ANALYZER_DIR):
os.makedirs(faceswaplab_globals.ANALYZER_DIR)
@@ -101,10 +108,13 @@ def getAnalysisModel():
return insightface.app.FaceAnalysis(
name="buffalo_l", providers=providers, root=faceswaplab_globals.ANALYZER_DIR
)
except Exception as e :
logger.error("Loading of swapping model failed, please check the requirements (On Windows, download and install Visual Studio. During the install, make sure to include the Python and C++ packages.)")
except Exception as e:
logger.error(
"Loading of swapping model failed, please check the requirements (On Windows, download and install Visual Studio. During the install, make sure to include the Python and C++ packages.)"
)
raise FaceModelException("Loading of swapping model failed")
@lru_cache(maxsize=1)
def getFaceSwapModel(model_path: str):
"""
@@ -116,14 +126,23 @@ def getFaceSwapModel(model_path: str):
Returns:
insightface.model_zoo.FaceModel: The face swap model.
"""
try :
try:
# Initializes the face swap model using the specified model path.
return upscaled_inswapper.UpscaledINSwapper(insightface.model_zoo.get_model(model_path, providers=providers))
except Exception as e :
logger.error("Loading of swapping model failed, please check the requirements (On Windows, download and install Visual Studio. During the install, make sure to include the Python and C++ packages.)")
return upscaled_inswapper.UpscaledINSwapper(
insightface.model_zoo.get_model(model_path, providers=providers)
)
except Exception as e:
logger.error(
"Loading of swapping model failed, please check the requirements (On Windows, download and install Visual Studio. During the install, make sure to include the Python and C++ packages.)"
)
def get_faces(img_data: np.ndarray, det_size=(640, 640), det_thresh : Optional[int]=None, sort_by_face_size = False) -> List[Face]:
def get_faces(
img_data: np.ndarray,
det_size=(640, 640),
det_thresh: Optional[int] = None,
sort_by_face_size=False,
) -> List[Face]:
"""
Detects and retrieves faces from an image using an analysis model.
@@ -136,7 +155,7 @@ def get_faces(img_data: np.ndarray, det_size=(640, 640), det_thresh : Optional[i
list: A list of detected faces, sorted by their x-coordinate of the bounding box.
"""
if det_thresh is None :
if det_thresh is None:
det_thresh = opts.data.get("faceswaplab_detection_threshold", 0.5)
# Create a deep copy of the analysis model (otherwise det_size is attached to the analysis model and can't be changed)
@@ -155,8 +174,12 @@ def get_faces(img_data: np.ndarray, det_size=(640, 640), det_thresh : Optional[i
return get_faces(img_data, det_size=det_size_half, det_thresh=det_thresh)
try:
if sort_by_face_size :
return sorted(face, reverse=True, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
if sort_by_face_size:
return sorted(
face,
reverse=True,
key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]),
)
# Sort the detected faces based on their x-coordinate of the bounding box
return sorted(face, key=lambda x: x.bbox[0])
@@ -164,7 +187,6 @@ def get_faces(img_data: np.ndarray, det_size=(640, 640), det_thresh : Optional[i
return []
@dataclass
class ImageResult:
"""
@@ -222,12 +244,15 @@ def get_faces_from_img_files(files):
if len(files) > 0:
for file in files:
img = Image.open(file.name) # Open the image file
face = get_or_default(get_faces(pil_to_cv2(img)), 0, None) # Extract faces from the image
face = get_or_default(
get_faces(pil_to_cv2(img)), 0, None
) # Extract faces from the image
if face is not None:
faces.append(face) # Add the detected face to the list of faces
return faces
def blend_faces(faces: List[Face]) -> Face:
"""
Blends the embeddings of multiple faces into a single face.
@@ -238,16 +263,16 @@ def blend_faces(faces: List[Face]) -> Face:
Returns:
Face: The blended Face object with the averaged embedding.
Returns None if the input list is empty.
Raises:
ValueError: If the embeddings have different shapes.
"""
embeddings = [face.embedding for face in faces]
if len(embeddings) > 0:
embedding_shape = embeddings[0].shape
# Check if all embeddings have the same shape
for embedding in embeddings:
if embedding.shape != embedding_shape:
@@ -255,15 +280,21 @@ def blend_faces(faces: List[Face]) -> Face:
# Compute the mean of all embeddings
blended_embedding = np.mean(embeddings, axis=0)
# Create a new Face object using the properties of the first face in the list
# Assign the blended embedding to the blended Face object
blended = Face(embedding=blended_embedding, gender=faces[0].gender, age=faces[0].age)
blended = Face(
embedding=blended_embedding, gender=faces[0].gender, age=faces[0].age
)
assert (
not np.array_equal(blended.embedding, faces[0].embedding)
if len(faces) > 1
else True
), "If len(faces)>0, the blended embedding should not be the same than the first image"
assert not np.array_equal(blended.embedding,faces[0].embedding) if len(faces) > 1 else True, "If len(faces)>0, the blended embedding should not be the same than the first image"
return blended
# Return None if the input list is empty
return None
@@ -275,9 +306,9 @@ def swap_face(
model: str,
faces_index: Set[int] = {0},
same_gender=True,
upscaled_swapper = False,
compute_similarity = True,
sort_by_face_size = False
upscaled_swapper=False,
compute_similarity=True,
sort_by_face_size=False,
) -> ImageResult:
"""
Swaps faces in the target image with the source face.
@@ -293,9 +324,9 @@ def swap_face(
Returns:
ImageResult: An object containing the swapped image and similarity scores.
"""
"""
return_result = ImageResult(target_img, {}, {})
try :
try:
target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
gender = source_face["gender"]
logger.info("Source Gender %s", gender)
@@ -313,19 +344,23 @@ def swap_face(
for i, swapped_face in enumerate(target_faces):
logger.info(f"swap face {i}")
if i in faces_index:
result = face_swapper.get(result, swapped_face, source_face, upscale = upscaled_swapper)
result = face_swapper.get(
result, swapped_face, source_face, upscale=upscaled_swapper
)
result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
return_result.image = result_image
if compute_similarity :
if compute_similarity:
try:
result_faces = get_faces(
cv2.cvtColor(np.array(result_image), cv2.COLOR_RGB2BGR), sort_by_face_size=sort_by_face_size
cv2.cvtColor(np.array(result_image), cv2.COLOR_RGB2BGR),
sort_by_face_size=sort_by_face_size,
)
if same_gender:
result_faces = [x for x in result_faces if x["gender"] == gender]
result_faces = [
x for x in result_faces if x["gender"] == gender
]
for i, swapped_face in enumerate(result_faces):
logger.info(f"compare face {i}")
@@ -343,13 +378,20 @@ def swap_face(
except Exception as e:
logger.error("Similarity processing failed %s", e)
raise e
except Exception as e :
except Exception as e:
logger.error("Conversion failed %s", e)
raise e
return return_result
def process_image_unit(model, unit : FaceSwapUnitSettings, image: Image.Image, info = None, upscaled_swapper = False, force_blend = False) -> List:
def process_image_unit(
model,
unit: FaceSwapUnitSettings,
image: Image.Image,
info=None,
upscaled_swapper=False,
force_blend=False,
) -> List:
"""Process one image and return a List of (image, info) (one if blended, many if not).
Args:
@@ -362,23 +404,28 @@ def process_image_unit(model, unit : FaceSwapUnitSettings, image: Image.Image, i
"""
results = []
if unit.enable :
if check_against_nsfw(image) :
if unit.enable:
if check_against_nsfw(image):
return [(image, info)]
if not unit.blend_faces and not force_blend :
if not unit.blend_faces and not force_blend:
src_faces = unit.faces
logger.info(f"will generate {len(src_faces)} images")
else :
else:
logger.info("blend all faces together")
src_faces = [unit.blended_faces]
assert(not np.array_equal(unit.reference_face.embedding,src_faces[0].embedding) if len(unit.faces)>1 else True), "Reference face cannot be the same as blended"
assert (
not np.array_equal(
unit.reference_face.embedding, src_faces[0].embedding
)
if len(unit.faces) > 1
else True
), "Reference face cannot be the same as blended"
for i,src_face in enumerate(src_faces):
for i, src_face in enumerate(src_faces):
logger.info(f"Process face {i}")
if unit.reference_face is not None :
if unit.reference_face is not None:
reference_face = unit.reference_face
else :
else:
logger.info("Use source face as reference face")
reference_face = src_face
@@ -392,14 +439,30 @@ def process_image_unit(model, unit : FaceSwapUnitSettings, image: Image.Image, i
same_gender=unit.same_gender,
upscaled_swapper=upscaled_swapper,
compute_similarity=unit.compute_similarity,
sort_by_face_size=unit.sort_by_size
sort_by_face_size=unit.sort_by_size,
)
save_img_debug(result.image, "After swap")
if result.image is None :
if result.image is None:
logger.error("Result image is None")
if (not unit.check_similarity) or result.similarity and all([result.similarity.values()!=0]+[x >= unit.min_sim for x in result.similarity.values()]) and all([result.ref_similarity.values()!=0]+[x >= unit.min_ref_sim for x in result.ref_similarity.values()]):
results.append((result.image, f"{info}, similarity = {result.similarity}, ref_similarity = {result.ref_similarity}"))
if (
(not unit.check_similarity)
or result.similarity
and all(
[result.similarity.values() != 0]
+ [x >= unit.min_sim for x in result.similarity.values()]
)
and all(
[result.ref_similarity.values() != 0]
+ [x >= unit.min_ref_sim for x in result.ref_similarity.values()]
)
):
results.append(
(
result.image,
f"{info}, similarity = {result.similarity}, ref_similarity = {result.ref_similarity}",
)
)
else:
logger.warning(
f"skip, similarity to low, sim = {result.similarity} (target {unit.min_sim}) ref sim = {result.ref_similarity} (target = {unit.min_ref_sim})"
@@ -407,22 +470,33 @@ def process_image_unit(model, unit : FaceSwapUnitSettings, image: Image.Image, i
logger.debug("process_image_unit : Unit produced %s results", len(results))
return results
def process_images_units(model, units : List[FaceSwapUnitSettings], images: List[Tuple[Optional[Image.Image], Optional[str]]], upscaled_swapper = False, force_blend = False) -> Union[List,None]:
if len(units) == 0 :
def process_images_units(
model,
units: List[FaceSwapUnitSettings],
images: List[Tuple[Optional[Image.Image], Optional[str]]],
upscaled_swapper=False,
force_blend=False,
) -> Union[List, None]:
if len(units) == 0:
logger.info("Finished processing image, return %s images", len(images))
return None
logger.debug("%s more units", len(units))
processed_images = []
for i,(image, info) in enumerate(images) :
for i, (image, info) in enumerate(images):
logger.debug("Processing image %s", i)
swapped = process_image_unit(model,units[0],image, info, upscaled_swapper, force_blend)
swapped = process_image_unit(
model, units[0], image, info, upscaled_swapper, force_blend
)
logger.debug("Image %s -> %s images", i, len(swapped))
nexts = process_images_units(model,units[1:],swapped, upscaled_swapper,force_blend)
if nexts :
nexts = process_images_units(
model, units[1:], swapped, upscaled_swapper, force_blend
)
if nexts:
processed_images.extend(nexts)
else :
else:
processed_images.extend(swapped)
return processed_images
return processed_images
scripts/faceswaplab_swapping/upscaled_inswapper.py
+123 -71
View File
@@ -1,4 +1,3 @@
import cv2
import numpy as np
import onnx
@@ -14,18 +13,22 @@ from PIL import Image
from scripts.faceswaplab_utils.faceswaplab_logging import logger
from scripts.faceswaplab_postprocessing import upscaling
from scripts.faceswaplab_postprocessing.postprocessing_options import \
PostProcessingOptions
from scripts.faceswaplab_postprocessing.postprocessing_options import (
PostProcessingOptions,
)
from scripts.faceswaplab_swapping.facemask import generate_face_mask
from scripts.faceswaplab_utils.imgutils import cv2_to_pil, pil_to_cv2
def get_upscaler() -> UpscalerData:
for upscaler in shared.sd_upscalers:
if upscaler.name == opts.data.get("faceswaplab_upscaled_swapper_upscaler", "LDSR"):
if upscaler.name == opts.data.get(
"faceswaplab_upscaled_swapper_upscaler", "LDSR"
):
return upscaler
return None
def merge_images_with_mask(image1, image2, mask):
if image1.shape != image2.shape or image1.shape[:2] != mask.shape:
raise ValueError("Img should have the same shape")
@@ -36,153 +39,202 @@ def merge_images_with_mask(image1, image2, mask):
merged_image = cv2.add(empty_region, masked_region)
return merged_image
def erode_mask(mask, kernel_size=3, iterations=1):
kernel = np.ones((kernel_size, kernel_size), np.uint8)
eroded_mask = cv2.erode(mask, kernel, iterations=iterations)
return eroded_mask
def apply_gaussian_blur(mask, kernel_size=(5, 5), sigma_x=0):
blurred_mask = cv2.GaussianBlur(mask, kernel_size, sigma_x)
return blurred_mask
def dilate_mask(mask, kernel_size=5, iterations=1):
kernel = np.ones((kernel_size, kernel_size), np.uint8)
dilated_mask = cv2.dilate(mask, kernel, iterations=iterations)
return dilated_mask
def get_face_mask(aimg,bgr_fake):
mask1 = generate_face_mask(aimg, device = shared.device)
mask2 = generate_face_mask(bgr_fake, device = shared.device)
mask = dilate_mask(cv2.bitwise_or(mask1,mask2))
def get_face_mask(aimg, bgr_fake):
mask1 = generate_face_mask(aimg, device=shared.device)
mask2 = generate_face_mask(bgr_fake, device=shared.device)
mask = dilate_mask(cv2.bitwise_or(mask1, mask2))
return mask
class UpscaledINSwapper():
def __init__(self, inswapper : INSwapper):
class UpscaledINSwapper:
def __init__(self, inswapper: INSwapper):
self.__dict__.update(inswapper.__dict__)
def forward(self, img, latent):
img = (img - self.input_mean) / self.input_std
pred = self.session.run(self.output_names, {self.input_names[0]: img, self.input_names[1]: latent})[0]
pred = self.session.run(
self.output_names, {self.input_names[0]: img, self.input_names[1]: latent}
)[0]
return pred
def super_resolution(self,img, k = 2) :
def super_resolution(self, img, k=2):
pil_img = cv2_to_pil(img)
options = PostProcessingOptions(
upscaler_name=opts.data.get('faceswaplab_upscaled_swapper_upscaler', 'LDSR'),
upscaler_name=opts.data.get(
"faceswaplab_upscaled_swapper_upscaler", "LDSR"
),
upscale_visibility=1,
scale=k,
face_restorer_name=opts.data.get('faceswaplab_upscaled_swapper_face_restorer', ""),
codeformer_weight= opts.data.get('faceswaplab_upscaled_swapper_face_restorer_weight', 1),
restorer_visibility=opts.data.get('faceswaplab_upscaled_swapper_face_restorer_visibility', 1))
face_restorer_name=opts.data.get(
"faceswaplab_upscaled_swapper_face_restorer", ""
),
codeformer_weight=opts.data.get(
"faceswaplab_upscaled_swapper_face_restorer_weight", 1
),
restorer_visibility=opts.data.get(
"faceswaplab_upscaled_swapper_face_restorer_visibility", 1
),
)
upscaled = upscaling.upscale_img(pil_img, options)
upscaled = upscaling.restore_face(upscaled, options)
return pil_to_cv2(upscaled)
def get(self, img, target_face, source_face, paste_back=True, upscale = True):
def get(self, img, target_face, source_face, paste_back=True, upscale=True):
aimg, M = face_align.norm_crop2(img, target_face.kps, self.input_size[0])
blob = cv2.dnn.blobFromImage(aimg, 1.0 / self.input_std, self.input_size,
(self.input_mean, self.input_mean, self.input_mean), swapRB=True)
latent = source_face.normed_embedding.reshape((1,-1))
blob = cv2.dnn.blobFromImage(
aimg,
1.0 / self.input_std,
self.input_size,
(self.input_mean, self.input_mean, self.input_mean),
swapRB=True,
)
latent = source_face.normed_embedding.reshape((1, -1))
latent = np.dot(latent, self.emap)
latent /= np.linalg.norm(latent)
pred = self.session.run(self.output_names, {self.input_names[0]: blob, self.input_names[1]: latent})[0]
#print(latent.shape, latent.dtype, pred.shape)
img_fake = pred.transpose((0,2,3,1))[0]
bgr_fake = np.clip(255 * img_fake, 0, 255).astype(np.uint8)[:,:,::-1]
try :
pred = self.session.run(
self.output_names, {self.input_names[0]: blob, self.input_names[1]: latent}
)[0]
# print(latent.shape, latent.dtype, pred.shape)
img_fake = pred.transpose((0, 2, 3, 1))[0]
bgr_fake = np.clip(255 * img_fake, 0, 255).astype(np.uint8)[:, :, ::-1]
try:
if not paste_back:
return bgr_fake, M
else:
target_img = img
def compute_diff(bgr_fake,aimg) :
def compute_diff(bgr_fake, aimg):
fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32)
fake_diff = np.abs(fake_diff).mean(axis=2)
fake_diff[:2,:] = 0
fake_diff[-2:,:] = 0
fake_diff[:,:2] = 0
fake_diff[:,-2:] = 0
fake_diff[:2, :] = 0
fake_diff[-2:, :] = 0
fake_diff[:, :2] = 0
fake_diff[:, -2:] = 0
return fake_diff
if upscale :
if upscale:
print("*" * 80)
print(
f"Upscaled inswapper using {opts.data.get('faceswaplab_upscaled_swapper_upscaler', 'LDSR')}"
)
print("*" * 80)
print("*"*80)
print(f"Upscaled inswapper using {opts.data.get('faceswaplab_upscaled_swapper_upscaler', 'LDSR')}")
print("*"*80)
k = 4
aimg, M = face_align.norm_crop2(img, target_face.kps, self.input_size[0]*k)
aimg, M = face_align.norm_crop2(
img, target_face.kps, self.input_size[0] * k
)
# upscale and restore face :
bgr_fake = self.super_resolution(bgr_fake, k)
if opts.data.get("faceswaplab_upscaled_improved_mask", True) :
mask = get_face_mask(aimg,bgr_fake)
bgr_fake = merge_images_with_mask(aimg, bgr_fake,mask)
if opts.data.get("faceswaplab_upscaled_improved_mask", True):
mask = get_face_mask(aimg, bgr_fake)
bgr_fake = merge_images_with_mask(aimg, bgr_fake, mask)
# compute fake_diff before sharpen and color correction (better result)
fake_diff = compute_diff(bgr_fake, aimg)
if opts.data.get("faceswaplab_upscaled_swapper_sharpen", True) :
if opts.data.get("faceswaplab_upscaled_swapper_sharpen", True):
print("sharpen")
# Add sharpness
blurred = cv2.GaussianBlur(bgr_fake, (0, 0), 3)
bgr_fake = cv2.addWeighted(bgr_fake, 1.5, blurred, -0.5, 0)
# Apply color corrections
if opts.data.get("faceswaplab_upscaled_swapper_fixcolor", True) :
if opts.data.get("faceswaplab_upscaled_swapper_fixcolor", True):
print("color correction")
correction = processing.setup_color_correction(cv2_to_pil(aimg))
bgr_fake_pil = processing.apply_color_correction(correction, cv2_to_pil(bgr_fake))
bgr_fake_pil = processing.apply_color_correction(
correction, cv2_to_pil(bgr_fake)
)
bgr_fake = pil_to_cv2(bgr_fake_pil)
else :
else:
fake_diff = compute_diff(bgr_fake, aimg)
IM = cv2.invertAffineTransform(M)
img_white = np.full((aimg.shape[0],aimg.shape[1]), 255, dtype=np.float32)
bgr_fake = cv2.warpAffine(bgr_fake, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0)
img_white = cv2.warpAffine(img_white, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0)
fake_diff = cv2.warpAffine(fake_diff, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0)
img_white[img_white>20] = 255
fthresh = opts.data.get('faceswaplab_upscaled_swapper_fthresh', 10)
img_white = np.full(
(aimg.shape[0], aimg.shape[1]), 255, dtype=np.float32
)
bgr_fake = cv2.warpAffine(
bgr_fake,
IM,
(target_img.shape[1], target_img.shape[0]),
borderValue=0.0,
)
img_white = cv2.warpAffine(
img_white,
IM,
(target_img.shape[1], target_img.shape[0]),
borderValue=0.0,
)
fake_diff = cv2.warpAffine(
fake_diff,
IM,
(target_img.shape[1], target_img.shape[0]),
borderValue=0.0,
)
img_white[img_white > 20] = 255
fthresh = opts.data.get("faceswaplab_upscaled_swapper_fthresh", 10)
print("fthresh", fthresh)
fake_diff[fake_diff<fthresh] = 0
fake_diff[fake_diff>=fthresh] = 255
fake_diff[fake_diff < fthresh] = 0
fake_diff[fake_diff >= fthresh] = 255
img_mask = img_white
mask_h_inds, mask_w_inds = np.where(img_mask==255)
mask_h_inds, mask_w_inds = np.where(img_mask == 255)
mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
mask_size = int(np.sqrt(mask_h*mask_w))
erosion_factor = opts.data.get('faceswaplab_upscaled_swapper_erosion', 1)
k = max(int(mask_size//10*erosion_factor), int(10*erosion_factor))
kernel = np.ones((k,k),np.uint8)
img_mask = cv2.erode(img_mask,kernel,iterations = 1)
kernel = np.ones((2,2),np.uint8)
fake_diff = cv2.dilate(fake_diff,kernel,iterations = 1)
k = max(int(mask_size//20*erosion_factor), int(5*erosion_factor))
mask_size = int(np.sqrt(mask_h * mask_w))
erosion_factor = opts.data.get(
"faceswaplab_upscaled_swapper_erosion", 1
)
k = max(int(mask_size // 10 * erosion_factor), int(10 * erosion_factor))
kernel = np.ones((k, k), np.uint8)
img_mask = cv2.erode(img_mask, kernel, iterations=1)
kernel = np.ones((2, 2), np.uint8)
fake_diff = cv2.dilate(fake_diff, kernel, iterations=1)
k = max(int(mask_size // 20 * erosion_factor), int(5 * erosion_factor))
kernel_size = (k, k)
blur_size = tuple(2*i+1 for i in kernel_size)
blur_size = tuple(2 * i + 1 for i in kernel_size)
img_mask = cv2.GaussianBlur(img_mask, blur_size, 0)
k = int(5*erosion_factor)
k = int(5 * erosion_factor)
kernel_size = (k, k)
blur_size = tuple(2*i+1 for i in kernel_size)
blur_size = tuple(2 * i + 1 for i in kernel_size)
fake_diff = cv2.GaussianBlur(fake_diff, blur_size, 0)
img_mask /= 255
fake_diff /= 255
img_mask = np.reshape(img_mask, [img_mask.shape[0],img_mask.shape[1],1])
fake_merged = img_mask * bgr_fake + (1-img_mask) * target_img.astype(np.float32)
img_mask = np.reshape(
img_mask, [img_mask.shape[0], img_mask.shape[1], 1]
)
fake_merged = img_mask * bgr_fake + (1 - img_mask) * target_img.astype(
np.float32
)
fake_merged = fake_merged.astype(np.uint8)
return fake_merged
except Exception as e :
except Exception as e:
import traceback
traceback.print_exc()
raise e