From d755be132538d2f31f62b940cdb67f0a36784b6c Mon Sep 17 00:00:00 2001 From: Tran Xen <137925069+glucauze@users.noreply.github.com> Date: Sat, 29 Jul 2023 00:24:51 +0200 Subject: [PATCH] fix bugs, fix api, clear code --- example/api/client_utils.py | 29 +- example/api/roop_api_example.py | 10 +- scripts/faceswaplab.py | 200 +++++------ scripts/faceswaplab_api/faceswaplab_api.py | 92 +++-- .../faceswaplab_api/faceswaplab_api_types.py | 19 +- scripts/faceswaplab_globals.py | 4 + scripts/faceswaplab_postprocessing/i2i_pp.py | 3 +- .../postprocessing.py | 27 +- .../postprocessing_options.py | 2 +- scripts/faceswaplab_ui/faceswaplab_tab.py | 326 +++++++++++------- .../faceswaplab_unit_settings.py | 20 +- 11 files changed, 464 insertions(+), 268 deletions(-) diff --git a/example/api/client_utils.py b/example/api/client_utils.py index 98a4d16..a31f7c2 100644 --- a/example/api/client_utils.py +++ b/example/api/client_utils.py @@ -5,6 +5,7 @@ from enum import Enum import base64, io from io import BytesIO from typing import List, Tuple, Optional +import numpy as np class InpaintingWhen(Enum): @@ -41,7 +42,10 @@ class FaceSwapUnit(BaseModel): blend_faces: bool = Field(description="Will blend faces if True", default=True) # Use same gender filtering - same_gender: bool = Field(description="Use same gender filtering", default=True) + same_gender: bool = Field(description="Use same gender filtering", default=False) + + # Use same gender filtering + sort_by_size: bool = Field(description="Sort Faces by size", default=False) # If True, discard images with low similarity check_similarity: bool = Field( @@ -70,6 +74,18 @@ class FaceSwapUnit(BaseModel): default=(0,), ) + reference_face_index: int = Field( + description="The face index to use to extract face from reference", + default=0, + ) + + def get_batch_images(self) -> List[Image.Image]: + images = [] + if self.batch_images: + for img in self.batch_images: + images.append(base64_to_pil(img)) + return images + class PostProcessingOptions(BaseModel): face_restorer_name: str = Field(description="face restorer name", default=None) @@ -82,7 +98,7 @@ class PostProcessingOptions(BaseModel): upscaler_name: str = Field(description="upscaler name", default=None) scale: float = Field(description="upscaling scale", default=1, le=10, ge=0) - upscale_visibility: float = Field( + upscaler_visibility: float = Field( description="upscaler visibility", default=1, le=1, ge=0 ) @@ -116,6 +132,9 @@ class PostProcessingOptions(BaseModel): examples=[e.value for e in InpaintingWhen.__members__.values()], default=InpaintingWhen.NEVER, ) + inpainting_model: str = Field( + description="Inpainting model", examples=["Current"], default="Current" + ) class FaceSwapRequest(BaseModel): @@ -125,7 +144,7 @@ class FaceSwapRequest(BaseModel): default=None, ) units: List[FaceSwapUnit] - postprocessing: PostProcessingOptions + postprocessing: Optional[PostProcessingOptions] class FaceSwapResponse(BaseModel): @@ -133,11 +152,11 @@ class FaceSwapResponse(BaseModel): infos: List[str] @property - def pil_images(self): + def pil_images(self) -> Image.Image: return [base64_to_pil(img) for img in self.images] -def pil_to_base64(img): +def pil_to_base64(img: Image.Image) -> np.array: # type:ignore if isinstance(img, str): img = Image.open(img) diff --git a/example/api/roop_api_example.py b/example/api/roop_api_example.py index d262940..918d995 100644 --- a/example/api/roop_api_example.py +++ b/example/api/roop_api_example.py @@ -5,6 +5,7 @@ from client_utils import ( PostProcessingOptions, FaceSwapResponse, pil_to_base64, + InpaintingWhen, ) address = "http://127.0.0.1:7860" @@ -24,7 +25,14 @@ unit2 = FaceSwapUnit( # Post-processing config : pp = PostProcessingOptions( - face_restorer_name="CodeFormer", codeformer_weight=0.5, restorer_visibility=1 + face_restorer_name="CodeFormer", + codeformer_weight=0.5, + restorer_visibility=1, + upscaler_name="Lanczos", + scale=4, + inpainting_steps=30, + inpainting_denoising_strengh=0.1, + inpainting_when=InpaintingWhen.BEFORE_RESTORE_FACE, ) # Prepare the request diff --git a/scripts/faceswaplab.py b/scripts/faceswaplab.py index f058c3b..4967fc8 100644 --- a/scripts/faceswaplab.py +++ b/scripts/faceswaplab.py @@ -146,105 +146,111 @@ class FaceSwapScript(scripts.Script): def process( self, p: StableDiffusionProcessing, *components: List[gr.components.Component] ) -> None: - self.read_config(p, *components) - - # If is instance of img2img, we check if face swapping in source is required. - if isinstance(p, StableDiffusionProcessingImg2Img): - if self.enabled and len(self.swap_in_source_units) > 0: - init_images: List[Tuple[Optional[Image.Image], Optional[str]]] = [ - (img, None) for img in p.init_images - ] - new_inits = swapper.process_images_units( - get_current_model(), - self.swap_in_source_units, - images=init_images, - upscaled_swapper=self.upscaled_swapper_in_source, - force_blend=True, - ) - logger.info(f"processed init images: {len(init_images)}") - if new_inits is not None: - p.init_images = [img[0] for img in new_inits] + try: + self.read_config(p, *components) + + # If is instance of img2img, we check if face swapping in source is required. + if isinstance(p, StableDiffusionProcessingImg2Img): + if self.enabled and len(self.swap_in_source_units) > 0: + init_images: List[Tuple[Optional[Image.Image], Optional[str]]] = [ + (img, None) for img in p.init_images + ] + new_inits = swapper.process_images_units( + get_current_model(), + self.swap_in_source_units, + images=init_images, + upscaled_swapper=self.upscaled_swapper_in_source, + force_blend=True, + ) + logger.info(f"processed init images: {len(init_images)}") + if new_inits is not None: + p.init_images = [img[0] for img in new_inits] + except Exception as e: + logger.info("Failed to process : %s", e) def postprocess( self, p: StableDiffusionProcessing, processed: Processed, *args: List[Any] ) -> None: - if self.enabled: - # Get the original images without the grid - orig_images: List[Image.Image] = processed.images[ - processed.index_of_first_image : - ] - orig_infotexts: List[str] = processed.infotexts[ - processed.index_of_first_image : - ] - - keep_original = self.keep_original_images - - # These are were images and infos of swapped images will be stored - images = [] - infotexts = [] - if (len(self.swap_in_generated_units)) > 0: - for i, (img, info) in enumerate(zip(orig_images, orig_infotexts)): - batch_index = i % p.batch_size - swapped_images = swapper.process_images_units( - get_current_model(), - self.swap_in_generated_units, - images=[(img, info)], - upscaled_swapper=self.upscaled_swapper_in_generated, - ) - if swapped_images is None: - continue - - logger.info(f"{len(swapped_images)} images swapped") - for swp_img, new_info in swapped_images: - img = swp_img # Will only swap the last image in the batch in next units (FIXME : hard to fix properly but not really critical) - - if swp_img is not None: - save_img_debug(swp_img, "Before apply mask") - swp_img = imgutils.apply_mask(swp_img, p, batch_index) - save_img_debug(swp_img, "After apply mask") - - try: - if self.postprocess_options is not None: - swp_img = enhance_image( - swp_img, self.postprocess_options + try: + if self.enabled: + # Get the original images without the grid + orig_images: List[Image.Image] = processed.images[ + processed.index_of_first_image : + ] + orig_infotexts: List[str] = processed.infotexts[ + processed.index_of_first_image : + ] + + keep_original = self.keep_original_images + + # These are were images and infos of swapped images will be stored + images = [] + infotexts = [] + if (len(self.swap_in_generated_units)) > 0: + for i, (img, info) in enumerate(zip(orig_images, orig_infotexts)): + batch_index = i % p.batch_size + swapped_images = swapper.process_images_units( + get_current_model(), + self.swap_in_generated_units, + images=[(img, info)], + upscaled_swapper=self.upscaled_swapper_in_generated, + ) + if swapped_images is None: + continue + + logger.info(f"{len(swapped_images)} images swapped") + for swp_img, new_info in swapped_images: + img = swp_img # Will only swap the last image in the batch in next units (FIXME : hard to fix properly but not really critical) + + if swp_img is not None: + save_img_debug(swp_img, "Before apply mask") + swp_img = imgutils.apply_mask(swp_img, p, batch_index) + save_img_debug(swp_img, "After apply mask") + + try: + if self.postprocess_options is not None: + swp_img = enhance_image( + swp_img, self.postprocess_options + ) + except Exception as e: + logger.error("Failed to upscale : %s", e) + + logger.info("Add swp image to processed") + images.append(swp_img) + infotexts.append(new_info) + if p.outpath_samples and opts.samples_save: + save_image( + swp_img, + p.outpath_samples, + "", + p.all_seeds[batch_index], + p.all_prompts[batch_index], + opts.samples_format, + info=new_info, + p=p, + suffix="-swapped", ) - except Exception as e: - logger.error("Failed to upscale : %s", e) - - logger.info("Add swp image to processed") - images.append(swp_img) - infotexts.append(new_info) - if p.outpath_samples and opts.samples_save: - save_image( - swp_img, - p.outpath_samples, - "", - p.all_seeds[batch_index], - p.all_prompts[batch_index], - opts.samples_format, - info=new_info, - p=p, - suffix="-swapped", - ) - else: - logger.error("swp image is None") - else: - keep_original = True - - # Generate grid : - if opts.return_grid and len(images) > 1: - # FIXME :Use sd method, not that if blended is not active, the result will be a bit messy. - grid = imgutils.create_square_image(images) - text = processed.infotexts[0] - infotexts.insert(0, text) - if opts.enable_pnginfo: - grid.info["parameters"] = text - images.insert(0, grid) - - if keep_original: - # If we want to keep original images, we add all existing (including grid this time) - images += processed.images - infotexts += processed.infotexts - - processed.images = images - processed.infotexts = infotexts + else: + logger.error("swp image is None") + else: + keep_original = True + + # Generate grid : + if opts.return_grid and len(images) > 1: + # FIXME :Use sd method, not that if blended is not active, the result will be a bit messy. + grid = imgutils.create_square_image(images) + text = processed.infotexts[0] + infotexts.insert(0, text) + if opts.enable_pnginfo: + grid.info["parameters"] = text + images.insert(0, grid) + + if keep_original: + # If we want to keep original images, we add all existing (including grid this time) + images += processed.images + infotexts += processed.infotexts + + processed.images = images + processed.infotexts = infotexts + except Exception as e: + logger.error("Failed to swap face %s in postprocess method", e) diff --git a/scripts/faceswaplab_api/faceswaplab_api.py b/scripts/faceswaplab_api/faceswaplab_api.py index 6decff3..e896523 100644 --- a/scripts/faceswaplab_api/faceswaplab_api.py +++ b/scripts/faceswaplab_api/faceswaplab_api.py @@ -3,7 +3,6 @@ import numpy as np from fastapi import FastAPI from modules.api import api from scripts.faceswaplab_api.faceswaplab_api_types import ( - FaceSwapRequest, FaceSwapResponse, ) from scripts.faceswaplab_globals import VERSION_FLAG @@ -16,9 +15,15 @@ from scripts.faceswaplab_utils.imgutils import ( ) from scripts.faceswaplab_utils.models_utils import get_current_model from modules.shared import opts +from scripts.faceswaplab_postprocessing.postprocessing import enhance_image +from scripts.faceswaplab_postprocessing.postprocessing_options import ( + PostProcessingOptions, +) +from scripts.faceswaplab_api import faceswaplab_api_types +from scripts.faceswaplab_postprocessing.postprocessing_options import InpaintingWhen -def encode_to_base64(image: Union[str, Image.Image, np.ndarray]) -> str: +def encode_to_base64(image: Union[str, Image.Image, np.ndarray]) -> str: # type: ignore """ Encode an image to a base64 string. @@ -40,7 +45,7 @@ def encode_to_base64(image: Union[str, Image.Image, np.ndarray]) -> str: return "" -def encode_np_to_base64(image: np.ndarray) -> str: +def encode_np_to_base64(image: np.ndarray) -> str: # type: ignore """ Encode a NumPy array to a base64 string. @@ -56,6 +61,59 @@ def encode_np_to_base64(image: np.ndarray) -> str: return api.encode_pil_to_base64(pil) +def get_postprocessing_options( + options: faceswaplab_api_types.PostProcessingOptions, +) -> PostProcessingOptions: + pp_options = PostProcessingOptions( + face_restorer_name=options.face_restorer_name, + restorer_visibility=options.restorer_visibility, + codeformer_weight=options.codeformer_weight, + upscaler_name=options.upscaler_name, + scale=options.scale, + upscale_visibility=options.upscaler_visibility, + inpainting_denoising_strengh=options.inpainting_denoising_strengh, + inpainting_prompt=options.inpainting_prompt, + inpainting_negative_prompt=options.inpainting_negative_prompt, + inpainting_steps=options.inpainting_steps, + inpainting_sampler=options.inpainting_sampler, + inpainting_when=options.inpainting_when, + inpainting_model=options.inpainting_model, + ) + + assert isinstance( + pp_options.inpainting_when, InpaintingWhen + ), "Value is not a valid InpaintingWhen enum" + + return pp_options + + +def get_faceswap_units_settings( + api_units: List[faceswaplab_api_types.FaceSwapUnit], +) -> List[FaceSwapUnitSettings]: + units = [] + for u in api_units: + units.append( + FaceSwapUnitSettings( + source_img=base64_to_pil(u.source_img), + source_face=u.source_face, + _batch_files=u.get_batch_images(), + blend_faces=u.blend_faces, + enable=True, + same_gender=u.same_gender, + sort_by_size=u.sort_by_size, + check_similarity=u.check_similarity, + _compute_similarity=u.compute_similarity, + min_ref_sim=u.min_ref_sim, + min_sim=u.min_sim, + _faces_index=",".join([str(i) for i in (u.faces_index)]), + reference_face_index=u.reference_face_index, + swap_in_generated=True, + swap_in_source=False, + ) + ) + return units + + def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None: @app.get( "/faceswaplab/version", @@ -71,29 +129,17 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None: tags=["faceswaplab"], description="Swap a face in an image using units", ) - async def swap_face(request: FaceSwapRequest) -> FaceSwapResponse: + async def swap_face( + request: faceswaplab_api_types.FaceSwapRequest, + ) -> faceswaplab_api_types.FaceSwapResponse: units: List[FaceSwapUnitSettings] = [] src_image: Optional[Image.Image] = base64_to_pil(request.image) response = FaceSwapResponse(images=[], infos=[]) + if request.postprocessing: + pp_options = get_postprocessing_options(request.postprocessing) + if src_image is not None: - for u in request.units: - units.append( - FaceSwapUnitSettings( - source_img=base64_to_pil(u.source_img), - source_face=u.source_face, - _batch_files=u.get_batch_images(), - blend_faces=u.blend_faces, - enable=True, - same_gender=u.same_gender, - check_similarity=u.check_similarity, - _compute_similarity=u.compute_similarity, - min_ref_sim=u.min_ref_sim, - min_sim=u.min_sim, - _faces_index=",".join([str(i) for i in (u.faces_index)]), - swap_in_generated=True, - swap_in_source=False, - ) - ) + units = get_faceswap_units_settings(request.units) swapped_images = swapper.process_images_units( get_current_model(), @@ -102,6 +148,8 @@ def faceswaplab_api(_: gr.Blocks, app: FastAPI) -> None: upscaled_swapper=opts.data.get("faceswaplab_upscaled_swapper", False), ) for img, info in swapped_images: + if pp_options: + img = enhance_image(img, pp_options) response.images.append(encode_to_base64(img)) response.infos.append(info) diff --git a/scripts/faceswaplab_api/faceswaplab_api_types.py b/scripts/faceswaplab_api/faceswaplab_api_types.py index 383a1d9..8a08bbd 100644 --- a/scripts/faceswaplab_api/faceswaplab_api_types.py +++ b/scripts/faceswaplab_api/faceswaplab_api_types.py @@ -1,4 +1,4 @@ -from typing import List, Tuple +from typing import List, Optional, Tuple from PIL import Image from scripts.faceswaplab_utils.imgutils import ( base64_to_pil, @@ -34,7 +34,10 @@ class FaceSwapUnit(BaseModel): blend_faces: bool = Field(description="Will blend faces if True", default=True) # Use same gender filtering - same_gender: bool = Field(description="Use same gender filtering", default=True) + same_gender: bool = Field(description="Use same gender filtering", default=False) + + # Use same gender filtering + sort_by_size: bool = Field(description="Sort Faces by size", default=False) # If True, discard images with low similarity check_similarity: bool = Field( @@ -63,6 +66,11 @@ class FaceSwapUnit(BaseModel): default=(0,), ) + reference_face_index: int = Field( + description="The face index to use to extract face from reference", + default=0, + ) + def get_batch_images(self) -> List[Image.Image]: images = [] if self.batch_images: @@ -82,7 +90,7 @@ class PostProcessingOptions(BaseModel): upscaler_name: str = Field(description="upscaler name", default=None) scale: float = Field(description="upscaling scale", default=1, le=10, ge=0) - upscale_visibility: float = Field( + upscaler_visibility: float = Field( description="upscaler visibility", default=1, le=1, ge=0 ) @@ -116,6 +124,9 @@ class PostProcessingOptions(BaseModel): examples=[e.value for e in InpaintingWhen.__members__.values()], default=InpaintingWhen.NEVER, ) + inpainting_model: str = Field( + description="Inpainting model", examples=["Current"], default="Current" + ) class FaceSwapRequest(BaseModel): @@ -125,7 +136,7 @@ class FaceSwapRequest(BaseModel): default=None, ) units: List[FaceSwapUnit] - postprocessing: PostProcessingOptions + postprocessing: Optional[PostProcessingOptions] class FaceSwapResponse(BaseModel): diff --git a/scripts/faceswaplab_globals.py b/scripts/faceswaplab_globals.py index fc14c41..cdc4c66 100644 --- a/scripts/faceswaplab_globals.py +++ b/scripts/faceswaplab_globals.py @@ -1,8 +1,12 @@ import os +from modules import scripts MODELS_DIR = os.path.abspath(os.path.join("models", "faceswaplab")) ANALYZER_DIR = os.path.abspath(os.path.join(MODELS_DIR, "analysers")) FACE_PARSER_DIR = os.path.abspath(os.path.join(MODELS_DIR, "parser")) +REFERENCE_PATH = os.path.join( + scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references" +) VERSION_FLAG: str = "v1.1.0" EXTENSION_PATH = os.path.join("extensions", "sd-webui-faceswaplab") diff --git a/scripts/faceswaplab_postprocessing/i2i_pp.py b/scripts/faceswaplab_postprocessing/i2i_pp.py index 97952e4..09ebce2 100644 --- a/scripts/faceswaplab_postprocessing/i2i_pp.py +++ b/scripts/faceswaplab_postprocessing/i2i_pp.py @@ -14,6 +14,7 @@ from scripts.faceswaplab_swapping import swapper def img2img_diffusion(img: Image.Image, pp: PostProcessingOptions) -> Image.Image: if pp.inpainting_denoising_strengh == 0: + logger.info("Discard inpainting denoising strength is 0") return img try: @@ -25,7 +26,7 @@ inpainting_steps : {pp.inpainting_steps} """ ) if not isinstance(pp.inpainting_sampler, str): - pass + pp.inpainting_sampler = "Euler" logger.info("send faces to image to image") img = img.copy() diff --git a/scripts/faceswaplab_postprocessing/postprocessing.py b/scripts/faceswaplab_postprocessing/postprocessing.py index ddc8599..813a5dc 100644 --- a/scripts/faceswaplab_postprocessing/postprocessing.py +++ b/scripts/faceswaplab_postprocessing/postprocessing.py @@ -11,17 +11,32 @@ from scripts.faceswaplab_postprocessing.upscaling import upscale_img, restore_fa def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Image.Image: result_image = image try: - if pp_options.inpainting_when == InpaintingWhen.BEFORE_UPSCALING.value: - result_image = img2img_diffusion(image, pp_options) + logger.debug("enhance_image, inpainting : %s", pp_options.inpainting_when) + result_image = image + + if ( + pp_options.inpainting_when == InpaintingWhen.BEFORE_UPSCALING.value + or pp_options.inpainting_when == InpaintingWhen.BEFORE_UPSCALING + ): + logger.debug("Inpaint before upscale") + result_image = img2img_diffusion(result_image, pp_options) result_image = upscale_img(result_image, pp_options) - if pp_options.inpainting_when == InpaintingWhen.BEFORE_RESTORE_FACE.value: - result_image = img2img_diffusion(image, pp_options) + if ( + pp_options.inpainting_when == InpaintingWhen.BEFORE_RESTORE_FACE.value + or pp_options.inpainting_when == InpaintingWhen.BEFORE_RESTORE_FACE + ): + logger.debug("Inpaint before restore") + result_image = img2img_diffusion(result_image, pp_options) result_image = restore_face(result_image, pp_options) - if pp_options.inpainting_when == InpaintingWhen.AFTER_ALL.value: - result_image = img2img_diffusion(image, pp_options) + if ( + pp_options.inpainting_when == InpaintingWhen.AFTER_ALL.value + or pp_options.inpainting_when == InpaintingWhen.AFTER_ALL + ): + logger.debug("Inpaint after all") + result_image = img2img_diffusion(result_image, pp_options) except Exception as e: logger.error("Failed to upscale %s", e) diff --git a/scripts/faceswaplab_postprocessing/postprocessing_options.py b/scripts/faceswaplab_postprocessing/postprocessing_options.py index 6ed6096..e2f6f39 100644 --- a/scripts/faceswaplab_postprocessing/postprocessing_options.py +++ b/scripts/faceswaplab_postprocessing/postprocessing_options.py @@ -19,7 +19,7 @@ class PostProcessingOptions: codeformer_weight: float = 1 upscaler_name: str = "" - scale: int = 1 + scale: float = 1 upscale_visibility: float = 0.5 inpainting_denoising_strengh: float = 0 diff --git a/scripts/faceswaplab_ui/faceswaplab_tab.py b/scripts/faceswaplab_ui/faceswaplab_tab.py index bc3b1d5..6df0e3e 100644 --- a/scripts/faceswaplab_ui/faceswaplab_tab.py +++ b/scripts/faceswaplab_ui/faceswaplab_tab.py @@ -24,61 +24,132 @@ from scripts.faceswaplab_postprocessing.postprocessing_options import ( ) from scripts.faceswaplab_postprocessing.postprocessing import enhance_image from dataclasses import fields -from typing import Any, List, Optional, Union +from typing import Any, Dict, List, Optional from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings from scripts.faceswaplab_utils.models_utils import get_current_model import re +from scripts.faceswaplab_globals import REFERENCE_PATH -def compare(img1: Image.Image, img2: Image.Image) -> Union[float, str]: - if img1 is not None and img2 is not None: - return swapper.compare_faces(img1, img2) +def compare(img1: Image.Image, img2: Image.Image) -> str: + """ + 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 str of a float value representing the similarity between the two faces (0 to 1). + Returns"You need 2 images to compare" if one or both of the images do not contain any faces. + """ + try: + if img1 is not None and img2 is not None: + return str(swapper.compare_faces(img1, img2)) + except Exception as e: + logger.error("Fail to compare", e) return "You need 2 images to compare" def extract_faces( - files, - extract_path, + files: Optional[List[str]], + extract_path: Optional[str], *components: List[gr.components.Component], -): - postprocess_options = PostProcessingOptions(*components) # type: ignore - - if not extract_path: - tempfile.mkdtemp() - if files is not None: - images = [] - for file in files: - img = Image.open(file.name).convert("RGB") - faces = swapper.get_faces(pil_to_cv2(img)) - if faces: - face_images = [] - for face in faces: - bbox = face.bbox.astype(int) - x_min, y_min, x_max, y_max = bbox - face_image = img.crop((x_min, y_min, x_max, y_max)) - if ( - postprocess_options.face_restorer_name - or postprocess_options.restorer_visibility - ): - postprocess_options.scale = ( - 1 if face_image.width > 512 else 512 // face_image.width - ) - face_image = enhance_image( - face_image, - postprocess_options, - ) - path = tempfile.NamedTemporaryFile( - delete=False, suffix=".png", dir=extract_path - ).name - face_image.save(path) - face_images.append(path) - images += face_images - return images +) -> Optional[List[str]]: + """ + Extracts faces from a list of image files. + + Given a list of image file paths, this function opens each image, extracts the faces, + and saves them in a specified directory. Post-processing is applied to each extracted face, + and the processed faces are saved as separate PNG files. + + Parameters: + files (Optional[List[str]]): List of file paths to the images to extract faces from. + extract_path (Optional[str]): Path where the extracted faces will be saved. + If no path is provided, a temporary directory will be created. + components (List[gr.components.Component]): List of components for post-processing. + + Returns: + Optional[List[str]]: List of file paths to the saved images of the extracted faces. + If no faces are found, None is returned. + """ + + try: + postprocess_options = PostProcessingOptions(*components) # type: ignore + + if not extract_path: + extract_path = tempfile.mkdtemp() + + if files: + images = [] + for file in files: + img = Image.open(file).convert("RGB") + faces = swapper.get_faces(pil_to_cv2(img)) + + if faces: + face_images = [] + for face in faces: + bbox = face.bbox.astype(int) + x_min, y_min, x_max, y_max = bbox + face_image = img.crop((x_min, y_min, x_max, y_max)) + + if ( + postprocess_options.face_restorer_name + or postprocess_options.restorer_visibility + ): + postprocess_options.scale = ( + 1 if face_image.width > 512 else 512 // face_image.width + ) + face_image = enhance_image(face_image, postprocess_options) + + path = tempfile.NamedTemporaryFile( + delete=False, suffix=".png", dir=extract_path + ).name + face_image.save(path) + face_images.append(path) + + images += face_images + + return images + except Exception as e: + logger.info("Failed to extract : %s", e) + return None -def analyse_faces(image: Image.Image, det_threshold: float = 0.5) -> str: +def analyse_faces(image: Image.Image, det_threshold: float = 0.5) -> Optional[str]: + """ + Function to analyze the faces in an image and provide a detailed report. + + Parameters + ---------- + image : PIL.Image.Image + The input image where faces will be detected. The image must be a PIL Image object. + + det_threshold : float, optional + The detection threshold for the face detection process, by default 0.5. It determines + the confidence level at which the function will consider a detected object as a face. + Value should be in the range [0, 1], with higher values indicating greater certainty. + + Returns + ------- + str or None + Returns a formatted string providing details about each face detected in the image. + For each face, the string will include an index and a set of facial details. + In the event of an exception (e.g., analysis failure), the function will log the error + and return None. + + Raises + ------ + This function handles exceptions internally and does not raise. + + Examples + -------- + >>> image = Image.open("test.jpg") + >>> print(analyse_faces(image, 0.7)) + """ + try: faces = swapper.get_faces(imgutils.pil_to_cv2(image), det_thresh=det_threshold) result = "" @@ -86,11 +157,12 @@ def analyse_faces(image: Image.Image, det_threshold: float = 0.5) -> str: result += f"\nFace {i} \n" + "=" * 40 + "\n" result += pformat(face) + "\n" result += "=" * 40 - return result + return result if result else None except Exception as e: logger.error("Analysis Failed : %s", e) - return "Analysis Failed" + + return None def sanitize_name(name: str) -> str: @@ -116,96 +188,104 @@ def build_face_checkpoint_and_save( Returns: PIL.Image.Image or None: The resulting swapped face image if the process is successful; None otherwise. """ - name = sanitize_name(name) - batch_files = batch_files or [] - logger.info("Build %s %s", name, [x.name for x in batch_files]) - faces = swapper.get_faces_from_img_files(batch_files) - blended_face = swapper.blend_faces(faces) - preview_path = os.path.join( - scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references" - ) - faces_path = os.path.join(scripts.basedir(), "models", "faceswaplab", "faces") - - os.makedirs(faces_path, exist_ok=True) - - target_img = None - if blended_face: - if blended_face["gender"] == 0: - target_img = Image.open(os.path.join(preview_path, "woman.png")) - else: - target_img = Image.open(os.path.join(preview_path, "man.png")) - - if name == "": - name = "default_name" - pprint(blended_face) - result = swapper.swap_face( - blended_face, blended_face, target_img, get_models()[0] - ) - result_image = enhance_image( - result.image, - PostProcessingOptions( - face_restorer_name="CodeFormer", restorer_visibility=1 - ), - ) - file_path = os.path.join(faces_path, f"{name}.pkl") - file_number = 1 - while os.path.exists(file_path): - file_path = os.path.join(faces_path, f"{name}_{file_number}.pkl") - file_number += 1 - result_image.save(file_path + ".png") - with open(file_path, "wb") as file: - pickle.dump( - { - "embedding": blended_face.embedding, - "gender": blended_face.gender, - "age": blended_face.age, - }, - file, + try: + name = sanitize_name(name) + batch_files = batch_files or [] + logger.info("Build %s %s", name, [x.name for x in batch_files]) + faces = swapper.get_faces_from_img_files(batch_files) + blended_face = swapper.blend_faces(faces) + preview_path = REFERENCE_PATH + + faces_path = os.path.join(scripts.basedir(), "models", "faceswaplab", "faces") + + os.makedirs(faces_path, exist_ok=True) + + target_img = None + if blended_face: + if blended_face["gender"] == 0: + target_img = Image.open(os.path.join(preview_path, "woman.png")) + else: + target_img = Image.open(os.path.join(preview_path, "man.png")) + + if name == "": + name = "default_name" + pprint(blended_face) + result = swapper.swap_face( + blended_face, blended_face, target_img, get_models()[0] + ) + result_image = enhance_image( + result.image, + PostProcessingOptions( + face_restorer_name="CodeFormer", restorer_visibility=1 + ), ) - try: - with open(file_path, "rb") as file: - data = Face(pickle.load(file)) - print(data) - except Exception as e: - print(e) - return result_image - print("No face found") + file_path = os.path.join(faces_path, f"{name}.pkl") + file_number = 1 + while os.path.exists(file_path): + file_path = os.path.join(faces_path, f"{name}_{file_number}.pkl") + file_number += 1 + result_image.save(file_path + ".png") + with open(file_path, "wb") as file: + pickle.dump( + { + "embedding": blended_face.embedding, + "gender": blended_face.gender, + "age": blended_face.age, + }, + file, + ) + try: + with open(file_path, "rb") as file: + data = Face(pickle.load(file)) + print(data) + except Exception as e: + print(e) + return result_image + + print("No face found") + except Exception as e: + logger.error("Failed to build checkpoint %s", e) + return None return target_img -def explore_onnx_faceswap_model(model_path): - data = { - "Node Name": [], - "Op Type": [], - "Inputs": [], - "Outputs": [], - "Attributes": [], - } - if model_path: - model = onnx.load(model_path) - for node in model.graph.node: - data["Node Name"].append(pformat(node.name)) - data["Op Type"].append(pformat(node.op_type)) - data["Inputs"].append(pformat(node.input)) - data["Outputs"].append(pformat(node.output)) - attributes = [] - for attr in node.attribute: - attr_name = attr.name - attr_value = attr.t - attributes.append( - "{} = {}".format(pformat(attr_name), pformat(attr_value)) - ) - data["Attributes"].append(attributes) - - df = pd.DataFrame(data) +def explore_onnx_faceswap_model(model_path: str) -> pd.DataFrame: + try: + data: Dict[str, Any] = { + "Node Name": [], + "Op Type": [], + "Inputs": [], + "Outputs": [], + "Attributes": [], + } + if model_path: + model = onnx.load(model_path) + for node in model.graph.node: + data["Node Name"].append(pformat(node.name)) + data["Op Type"].append(pformat(node.op_type)) + data["Inputs"].append(pformat(node.input)) + data["Outputs"].append(pformat(node.output)) + attributes = [] + for attr in node.attribute: + attr_name = attr.name + attr_value = attr.t + attributes.append( + "{} = {}".format(pformat(attr_name), pformat(attr_value)) + ) + data["Attributes"].append(attributes) + + df = pd.DataFrame(data) + except Exception as e: + logger.info("Failed to explore model %s", e) + return None return df def batch_process( - files, save_path, *components: List[gr.components.Component] + files: List[gr.File], save_path: str, *components: List[gr.components.Component] ) -> Optional[List[Image.Image]]: try: if save_path is not None: @@ -216,7 +296,7 @@ def batch_process( # Parse and convert units flat components into FaceSwapUnitSettings for i in range(0, units_count): - units += [FaceSwapUnitSettings.get_unit_configuration(i, components)] + units += [FaceSwapUnitSettings.get_unit_configuration(i, components)] # type: ignore for i, u in enumerate(units): logger.debug("%s, %s", pformat(i), pformat(u)) diff --git a/scripts/faceswaplab_ui/faceswaplab_unit_settings.py b/scripts/faceswaplab_ui/faceswaplab_unit_settings.py index 7e9348f..cd40a11 100644 --- a/scripts/faceswaplab_ui/faceswaplab_unit_settings.py +++ b/scripts/faceswaplab_ui/faceswaplab_unit_settings.py @@ -3,7 +3,7 @@ import numpy as np import base64 import io from dataclasses import dataclass, fields -from typing import List, Union +from typing import Any, List, Optional, Set, Union import dill as pickle import gradio as gr from insightface.app.common import Face @@ -50,14 +50,16 @@ class FaceSwapUnitSettings: swap_in_generated: bool @staticmethod - def get_unit_configuration(unit: int, components): + def get_unit_configuration( + unit: int, components: List[gr.components.Component] + ) -> Any: fields_count = len(fields(FaceSwapUnitSettings)) return FaceSwapUnitSettings( *components[unit * fields_count : unit * fields_count + fields_count] ) @property - def faces_index(self): + def faces_index(self) -> Set[int]: """ Convert _faces_index from str to int """ @@ -72,18 +74,18 @@ class FaceSwapUnitSettings: return faces_index @property - def compute_similarity(self): + def compute_similarity(self) -> bool: return self._compute_similarity or self.check_similarity @property - def batch_files(self): + def batch_files(self) -> List[gr.File]: """ Return empty array instead of None for batch files """ return self._batch_files or [] @property - def reference_face(self): + def reference_face(self) -> Optional[Face]: """ Extract reference face (only once and store it for the rest of processing). Reference face is the checkpoint or the source image or the first image in the batch in that order. @@ -97,6 +99,7 @@ class FaceSwapUnitSettings: self._reference_face = face except Exception as e: logger.error("Failed to load checkpoint : %s", e) + raise e elif self.source_img is not None: if isinstance(self.source_img, str): # source_img is a base64 string if ( @@ -119,11 +122,12 @@ class FaceSwapUnitSettings: if self._reference_face is None: logger.error("You need at least one reference face") + raise Exception("No reference face found") return self._reference_face @property - def faces(self): + def faces(self) -> List[Face]: """_summary_ Extract all faces (including reference face) to provide an array of faces Only processed once. @@ -146,7 +150,7 @@ class FaceSwapUnitSettings: return self._faces @property - def blended_faces(self): + def blended_faces(self) -> Face: """ Blend the faces using the mean of all embeddings """