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

huge changes, inpainting in faces unit, change faces processing, change api, refactor, requires further testing

Browse Source
This commit is contained in:
Tran Xen
2023-08-01 20:17:43 +02:00
parent 1d9b3a64dc
commit ee7f7d09d2
24 changed files with 786 additions and 418 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/faceswaplab_swapping/facemask.py
+1 -1
View File
@@ -20,7 +20,7 @@ def get_parsing_model(device: torch_device) -> torch.nn.Module:
Returns:
The parsing model.
"""
return init_parsing_model(device=device)
return init_parsing_model(device=device) # type: ignore
def convert_image_to_tensor(
scripts/faceswaplab_swapping/parsing/__init__.py
+1 -1
View File
@@ -50,7 +50,7 @@ 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):
def load_file_from_url(url: str, model_dir=None, progress=True, file_name=None):
"""Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py"""
if model_dir is None:
hub_dir = get_dir()
scripts/faceswaplab_swapping/swapper.py
+130 -98
View File
@@ -7,7 +7,7 @@ import tempfile
import cv2
import insightface
import numpy as np
from insightface.app.common import Face
from insightface.app.common import Face as ISFace
from PIL import Image
from sklearn.metrics.pairwise import cosine_similarity
@@ -28,7 +28,8 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import (
)
from scripts.faceswaplab_utils.models_utils import get_current_model
import gradio as gr
from scripts.faceswaplab_utils.typing import CV2ImgU8, PILImage, Face
from scripts.faceswaplab_inpainting.i2i_pp import img2img_diffusion
providers = ["CPUExecutionProvider"]
@@ -60,7 +61,7 @@ def cosine_similarity_face(face1: Face, face2: Face) -> float:
return max(0, similarity[0, 0])
def compare_faces(img1: Image.Image, img2: Image.Image) -> float:
def compare_faces(img1: PILImage, img2: PILImage) -> float:
"""
Compares the similarity between two faces extracted from images using cosine similarity.
@@ -87,22 +88,22 @@ def compare_faces(img1: Image.Image, img2: Image.Image) -> float:
def batch_process(
src_images: List[Image.Image],
src_images: List[PILImage],
save_path: Optional[str],
units: List[FaceSwapUnitSettings],
postprocess_options: PostProcessingOptions,
) -> Optional[List[Image.Image]]:
) -> Optional[List[PILImage]]:
"""
Process a batch of images, apply face swapping according to the given settings, and optionally save the resulting images to a specified path.
Args:
src_images (List[Image.Image]): List of source PIL Images to process.
src_images (List[PILImage]): List of source PIL Images to process.
save_path (Optional[str]): Destination path where the processed images will be saved. If None, no images are saved.
units (List[FaceSwapUnitSettings]): List of FaceSwapUnitSettings to apply to the images.
postprocess_options (PostProcessingOptions): Post-processing settings to be applied to the images.
Returns:
Optional[List[Image.Image]]: List of processed images, or None in case of an exception.
Optional[List[PILImage]]: List of processed images, or None in case of an exception.
Raises:
Any exceptions raised by the underlying process will be logged and the function will return None.
@@ -149,7 +150,7 @@ def batch_process(
def extract_faces(
images: List[Image.Image],
images: List[PILImage],
extract_path: Optional[str],
postprocess_options: PostProcessingOptions,
) -> Optional[List[str]]:
@@ -206,7 +207,7 @@ def extract_faces(
return result_images
except Exception as e:
logger.info("Failed to extract : %s", e)
logger.error("Failed to extract : %s", e)
import traceback
traceback.print_exc()
@@ -273,16 +274,15 @@ def getFaceSwapModel(model_path: str) -> upscaled_inswapper.UpscaledINSwapper:
def get_faces(
img_data: np.ndarray, # type: ignore
img_data: CV2ImgU8,
det_size: Tuple[int, int] = (640, 640),
det_thresh: Optional[float] = None,
sort_by_face_size: bool = False,
) -> List[Face]:
"""
Detects and retrieves faces from an image using an analysis model.
Args:
img_data (np.ndarray): The image data as a NumPy array.
img_data (CV2ImgU8): The image data as a NumPy array.
det_size (tuple): The desired detection size (width, height). Defaults to (640, 640).
sort_by_face_size (bool) : Will sort the faces by their size from larger to smaller face
@@ -309,26 +309,55 @@ def get_faces(
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]),
)
# Sort the detected faces based on their x-coordinate of the bounding box
return sorted(face, key=lambda x: x.bbox[0])
except Exception as e:
return []
def filter_faces(
all_faces: List[Face],
faces_index: Set[int],
source_gender: int = None,
sort_by_face_size: bool = False,
) -> List[Face]:
"""
Sorts and filters a list of faces based on specified criteria.
This function takes a list of Face objects and can sort them by face size and filter them by gender.
Sorting by face size is performed if sort_by_face_size is set to True, and filtering by gender is
performed if source_gender is provided.
:param faces: A list of Face objects representing the faces to be sorted and filtered.
:param faces_index: A set of faces index
:param source_gender: An optional integer representing the gender by which to filter the faces.
If provided, only faces with the specified gender will be included in the result.
:param sort_by_face_size: A boolean indicating whether to sort the faces by size. If True, faces are
sorted in descending order by size, calculated as the area of the bounding box.
:return: A list of Face objects sorted and filtered according to the specified criteria.
"""
filtered_faces = copy.copy(all_faces)
if sort_by_face_size:
filtered_faces = sorted(
all_faces,
reverse=True,
key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]),
)
if source_gender is not None:
filtered_faces = [
face for face in filtered_faces if face["gender"] == source_gender
]
return [face for i, face in enumerate(filtered_faces) if i in faces_index]
@dataclass
class ImageResult:
"""
Represents the result of an image swap operation
"""
image: Image.Image
image: PILImage
"""
The image object with the swapped face
"""
@@ -362,7 +391,7 @@ def get_or_default(l: List[Any], index: int, default: Any) -> Any:
return l[index] if index < len(l) else default
def get_faces_from_img_files(files: List[gr.File]) -> List[Optional[np.ndarray]]: # type: ignore
def get_faces_from_img_files(files: List[gr.File]) -> List[Optional[CV2ImgU8]]:
"""
Extracts faces from a list of image files.
@@ -388,7 +417,7 @@ def get_faces_from_img_files(files: List[gr.File]) -> List[Optional[np.ndarray]]
return faces
def blend_faces(faces: List[Face]) -> Face:
def blend_faces(faces: List[Face]) -> Optional[Face]:
"""
Blends the embeddings of multiple faces into a single face.
@@ -418,16 +447,10 @@ def blend_faces(faces: List[Face]) -> Face:
# 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(
blended = ISFace(
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"
return blended
# Return None if the input list is empty
@@ -435,85 +458,80 @@ def blend_faces(faces: List[Face]) -> Face:
def swap_face(
reference_face: np.ndarray, # type: ignore
source_face: np.ndarray, # type: ignore
target_img: Image.Image,
reference_face: CV2ImgU8,
source_face: Face,
target_img: PILImage,
target_faces: List[Face],
model: str,
faces_index: Set[int] = {0},
same_gender: bool = True,
upscaled_swapper: bool = False,
compute_similarity: bool = True,
sort_by_face_size: bool = False,
) -> ImageResult:
"""
Swaps faces in the target image with the source face.
Args:
reference_face (np.ndarray): The reference face used for similarity comparison.
source_face (np.ndarray): The source face to be swapped.
target_img (Image.Image): The target image to swap faces in.
reference_face (CV2ImgU8): The reference face used for similarity comparison.
source_face (CV2ImgU8): The source face to be swapped.
target_img (PILImage): The target image to swap faces in.
model (str): Path to the face swap model.
faces_index (Set[int], optional): Set of indices specifying which faces to swap. Defaults to {0}.
same_gender (bool, optional): If True, only swap faces with the same gender as the source face. Defaults to True.
Returns:
ImageResult: An object containing the swapped image and similarity scores.
"""
return_result = ImageResult(target_img, {}, {})
target_img_cv2: CV2ImgU8 = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
try:
target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
gender = source_face["gender"]
logger.info("Source Gender %s", gender)
if source_face is not None:
result = target_img
result = target_img_cv2
model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model)
face_swapper = getFaceSwapModel(model_path)
target_faces = get_faces(target_img, sort_by_face_size=sort_by_face_size)
logger.info("Target faces count : %s", len(target_faces))
if same_gender:
target_faces = [x for x in target_faces if x["gender"] == gender]
logger.info("Target Gender Matches count %s", len(target_faces))
for i, swapped_face in enumerate(target_faces):
logger.info(f"swap face {i}")
if i in faces_index:
# type : ignore
result = face_swapper.get(
result, swapped_face, source_face, upscale=upscaled_swapper
)
result = face_swapper.get(
img=result,
target_face=swapped_face,
source_face=source_face,
upscale=upscaled_swapper,
) # type: ignore
result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
return_result.image = result_image
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,
)
if same_gender:
result_faces = [
x for x in result_faces if x["gender"] == gender
]
# FIXME : recompute similarity
for i, swapped_face in enumerate(result_faces):
logger.info(f"compare face {i}")
if i in faces_index and i < len(target_faces):
return_result.similarity[i] = cosine_similarity_face(
source_face, swapped_face
)
return_result.ref_similarity[i] = cosine_similarity_face(
reference_face, swapped_face
)
# 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,
# )
# if same_gender:
# result_faces = [
# x for x in result_faces if x["gender"] == gender
# ]
logger.info(f"similarity {return_result.similarity}")
logger.info(f"ref similarity {return_result.ref_similarity}")
# for i, swapped_face in enumerate(result_faces):
# logger.info(f"compare face {i}")
# if i in faces_index and i < len(target_faces):
# return_result.similarity[i] = cosine_similarity_face(
# source_face, swapped_face
# )
# return_result.ref_similarity[i] = cosine_similarity_face(
# reference_face, swapped_face
# )
except Exception as e:
logger.error("Similarity processing failed %s", e)
raise e
# logger.info(f"similarity {return_result.similarity}")
# logger.info(f"ref similarity {return_result.ref_similarity}")
# except Exception as e:
# logger.error("Similarity processing failed %s", e)
# raise e
except Exception as e:
logger.error("Conversion failed %s", e)
raise e
@@ -523,11 +541,11 @@ def swap_face(
def process_image_unit(
model: str,
unit: FaceSwapUnitSettings,
image: Image.Image,
image: PILImage,
info: str = None,
upscaled_swapper: bool = False,
force_blend: bool = False,
) -> List[Tuple[Image.Image, str]]:
) -> List[Tuple[PILImage, str]]:
"""Process one image and return a List of (image, info) (one if blended, many if not).
Args:
@@ -541,6 +559,8 @@ def process_image_unit(
results = []
if unit.enable:
faces = get_faces(pil_to_cv2(image))
if check_against_nsfw(image):
return [(image, info)]
if not unit.blend_faces and not force_blend:
@@ -549,15 +569,10 @@ def process_image_unit(
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"
for i, src_face in enumerate(src_faces):
current_image = image
logger.info(f"Process face {i}")
if unit.reference_face is not None:
reference_face = unit.reference_face
@@ -565,18 +580,35 @@ def process_image_unit(
logger.info("Use source face as reference face")
reference_face = src_face
save_img_debug(image, "Before swap")
result: ImageResult = swap_face(
reference_face,
src_face,
image,
target_faces = filter_faces(
faces,
faces_index=unit.faces_index,
model=model,
same_gender=unit.same_gender,
upscaled_swapper=upscaled_swapper,
compute_similarity=unit.compute_similarity,
source_gender=src_face["gender"] if unit.same_gender else None,
sort_by_face_size=unit.sort_by_size,
)
# Apply pre-inpainting to image
if unit.pre_inpainting.inpainting_denoising_strengh > 0:
current_image = img2img_diffusion(
img=current_image, faces=target_faces, options=unit.pre_inpainting
)
save_img_debug(image, "Before swap")
result: ImageResult = swap_face(
reference_face=reference_face,
source_face=src_face,
target_img=current_image,
target_faces=target_faces,
model=model,
upscaled_swapper=upscaled_swapper,
compute_similarity=unit.compute_similarity,
)
# Apply post-inpainting to image
if unit.post_inpainting.inpainting_denoising_strengh > 0:
result.image = img2img_diffusion(
img=result.image, faces=target_faces, options=unit.post_inpainting
)
save_img_debug(result.image, "After swap")
if result.image is None:
@@ -610,17 +642,17 @@ def process_image_unit(
def process_images_units(
model: str,
units: List[FaceSwapUnitSettings],
images: List[Tuple[Optional[Image.Image], Optional[str]]],
images: List[Tuple[Optional[PILImage], Optional[str]]],
upscaled_swapper: bool = False,
force_blend: bool = False,
) -> Optional[List[Tuple[Image.Image, str]]]:
) -> Optional[List[Tuple[PILImage, str]]]:
"""
Process a list of images using a specified model and unit settings for face swapping.
Args:
model (str): The name of the model to use for processing.
units (List[FaceSwapUnitSettings]): A list of settings for face swap units to apply on each image.
images (List[Tuple[Optional[Image.Image], Optional[str]]]): A list of tuples, each containing
images (List[Tuple[Optional[PILImage], Optional[str]]]): A list of tuples, each containing
an image and its associated info string. If an image or info string is not available,
its value can be None.
upscaled_swapper (bool, optional): If True, uses an upscaled version of the face swapper.
@@ -629,7 +661,7 @@ def process_images_units(
image. Defaults to False.
Returns:
Optional[List[Tuple[Image.Image, str]]]: A list of tuples, each containing a processed image
Optional[List[Tuple[PILImage, str]]]: A list of tuples, each containing a processed image
and its associated info string. If no units are provided for processing, returns None.
"""
scripts/faceswaplab_swapping/upscaled_inswapper.py
+81 -16
View File
@@ -1,3 +1,4 @@
from typing import Any, Tuple, Union
import cv2
import numpy as np
from insightface.model_zoo.inswapper import INSwapper
@@ -12,6 +13,7 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import (
)
from scripts.faceswaplab_swapping.facemask import generate_face_mask
from scripts.faceswaplab_utils.imgutils import cv2_to_pil, pil_to_cv2
from scripts.faceswaplab_utils.typing import CV2ImgU8, Face
def get_upscaler() -> UpscalerData:
@@ -23,7 +25,25 @@ def get_upscaler() -> UpscalerData:
return None
def merge_images_with_mask(image1, image2, mask):
def merge_images_with_mask(
image1: CV2ImgU8, image2: CV2ImgU8, mask: CV2ImgU8
) -> CV2ImgU8:
"""
Merges two images using a given mask. The regions where the mask is set will be replaced with the corresponding
areas of the second image.
Args:
image1 (CV2Img): The base image, which must have the same shape as image2.
image2 (CV2Img): The image to be merged, which must have the same shape as image1.
mask (CV2Img): A binary mask specifying the regions to be merged. The mask shape should match image1's first two dimensions.
Returns:
CV2Img: The merged image.
Raises:
ValueError: If the shapes of the images and mask do not match.
"""
if image1.shape != image2.shape or image1.shape[:2] != mask.shape:
raise ValueError("Img should have the same shape")
mask = mask.astype(np.uint8)
@@ -34,42 +54,80 @@ def merge_images_with_mask(image1, image2, mask):
return merged_image
def erode_mask(mask, kernel_size=3, iterations=1):
def erode_mask(mask: CV2ImgU8, kernel_size: int = 3, iterations: int = 1) -> CV2ImgU8:
"""
Erodes a binary mask using a given kernel size and number of iterations.
Args:
mask (CV2Img): The binary mask to erode.
kernel_size (int, optional): The size of the kernel. Default is 3.
iterations (int, optional): The number of erosion iterations. Default is 1.
Returns:
CV2Img: The eroded mask.
"""
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):
def apply_gaussian_blur(
mask: CV2ImgU8, kernel_size: Tuple[int, int] = (5, 5), sigma_x: int = 0
) -> CV2ImgU8:
"""
Applies a Gaussian blur to a mask.
Args:
mask (CV2Img): The mask to blur.
kernel_size (tuple, optional): The size of the kernel, e.g. (5, 5). Default is (5, 5).
sigma_x (int, optional): The standard deviation in the X direction. Default is 0.
Returns:
CV2Img: The blurred mask.
"""
blurred_mask = cv2.GaussianBlur(mask, kernel_size, sigma_x)
return blurred_mask
def dilate_mask(mask, kernel_size=5, iterations=1):
def dilate_mask(mask: CV2ImgU8, kernel_size: int = 5, iterations: int = 1) -> CV2ImgU8:
"""
Dilates a binary mask using a given kernel size and number of iterations.
Args:
mask (CV2Img): The binary mask to dilate.
kernel_size (int, optional): The size of the kernel. Default is 5.
iterations (int, optional): The number of dilation iterations. Default is 1.
Returns:
CV2Img: The dilated mask.
"""
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):
def get_face_mask(aimg: CV2ImgU8, bgr_fake: CV2ImgU8) -> CV2ImgU8:
"""
Generates a face mask by performing bitwise OR on two face masks and then dilating the result.
Args:
aimg (CV2Img): Input image for generating the first face mask.
bgr_fake (CV2Img): Input image for generating the second face mask.
Returns:
CV2Img: The combined and dilated face mask.
"""
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:
class UpscaledINSwapper(INSwapper):
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]
return pred
def super_resolution(self, img, k=2):
def super_resolution(self, img: CV2ImgU8, k: int = 2) -> CV2ImgU8:
pil_img = cv2_to_pil(img)
options = PostProcessingOptions(
upscaler_name=opts.data.get(
@@ -91,7 +149,14 @@ class UpscaledINSwapper:
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: CV2ImgU8,
target_face: Face,
source_face: Face,
paste_back: bool = True,
upscale: bool = True,
) -> Union[CV2ImgU8, Tuple[CV2ImgU8, Any]]:
aimg, M = face_align.norm_crop2(img, target_face.kps, self.input_size[0])
blob = cv2.dnn.blobFromImage(
aimg,
@@ -116,7 +181,7 @@ class UpscaledINSwapper:
else:
target_img = img
def compute_diff(bgr_fake, aimg):
def compute_diff(bgr_fake: CV2ImgU8, aimg: CV2ImgU8) -> CV2ImgU8:
fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32)
fake_diff = np.abs(fake_diff).mean(axis=2)
fake_diff[:2, :] = 0