Merge pull request #1 from jiveabillion/multi-facefile

Multi facefile
This commit is contained in:
jiveabillion 2023-11-28 14:45:31 -05:00 committed by GitHub
commit 6bd5c4250b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 347 additions and 78 deletions

View File

@ -4,16 +4,19 @@ import cv2
import modules.shared as shared import modules.shared as shared
import numpy as np import numpy as np
import torch import torch
from scripts.reactor_logger import logger
from facexlib.parsing import init_parsing_model from facexlib.parsing import init_parsing_model
from facexlib.utils.misc import img2tensor from facexlib.utils.misc import img2tensor
from torchvision.transforms.functional import normalize from torchvision.transforms.functional import normalize
from PIL import Image from PIL import Image
from scripts.inferencers.vignette_mask_generator import VignetteMaskGenerator
from scripts.inferencers.mask_generator import MaskGenerator from scripts.inferencers.mask_generator import MaskGenerator
from scripts.reactor_logger import logger
class BiSeNetMaskGenerator(MaskGenerator): class BiSeNetMaskGenerator(MaskGenerator):
def __init__(self) -> None: def __init__(self) -> None:
self.mask_model = init_parsing_model(device=shared.device) self.mask_model = init_parsing_model(device=shared.device)
self.fallback_mask_generator = VignetteMaskGenerator()
def name(self): def name(self):
return "BiSeNet" return "BiSeNet"
@ -25,7 +28,7 @@ class BiSeNetMaskGenerator(MaskGenerator):
affected_areas: List[str], affected_areas: List[str],
mask_size: int, mask_size: int,
use_minimal_area: bool, use_minimal_area: bool,
fallback_ratio: float = 0.25, fallback_ratio: float = 0.10,
**kwargs, **kwargs,
) -> np.ndarray: ) -> np.ndarray:
original_face_image = face_image original_face_image = face_image
@ -59,11 +62,11 @@ class BiSeNetMaskGenerator(MaskGenerator):
if w != 512 or h != 512: if w != 512 or h != 512:
mask = cv2.resize(mask, dsize=(w, h)) mask = cv2.resize(mask, dsize=(w, h))
"""if MaskGenerator.calculate_mask_coverage(mask) < fallback_ratio: if MaskGenerator.calculate_mask_coverage(mask) < fallback_ratio:
logger.info("Use fallback mask generator") logger.status(F"Mask coverage less than fallback ratio of {fallback_ratio}. Using vignette mask generator.")
mask = self.fallback_mask_generator.generate_mask( mask = self.fallback_mask_generator.generate_mask(
original_face_image, face_area_on_image, use_minimal_area=True original_face_image, face_area_on_image, use_minimal_area=True
)""" )
return mask return mask

View File

@ -0,0 +1,50 @@
from typing import Tuple
import cv2
import numpy as np
from scripts.inferencers.mask_generator import MaskGenerator
class VignetteMaskGenerator(MaskGenerator):
def name(self):
return "Vignette"
def generate_mask(
self,
face_image: np.ndarray,
face_area_on_image: Tuple[int, int, int, int],
use_minimal_area: bool,
sigma: float = -1,
keep_safe_area: bool = False,
**kwargs,
) -> np.ndarray:
(left, top, right, bottom) = face_area_on_image
w, h = right - left, bottom - top
mask = np.zeros((face_image.shape[0], face_image.shape[1]), dtype=np.uint8)
if use_minimal_area:
sigma = 120 if sigma == -1 else sigma
mask[top : top + h, left : left + w] = 255
else:
sigma = 180 if sigma == -1 else sigma
h, w = face_image.shape[0], face_image.shape[1]
mask[:, :] = 255
Y = np.linspace(0, h, h, endpoint=False)
X = np.linspace(0, w, w, endpoint=False)
Y, X = np.meshgrid(Y, X)
Y -= h / 2
X -= w / 2
gaussian = np.exp(-(X**2 + Y**2) / (2 * sigma**2))
gaussian_mask = np.uint8(255 * gaussian.T)
if use_minimal_area:
mask[top : top + h, left : left + w] = gaussian_mask
else:
mask[:, :] = gaussian_mask
if keep_safe_area:
mask = cv2.ellipse(mask, ((left + right) // 2, (top + bottom) // 2), (w // 2, h // 2), 0, 0, 360, 255, -1)
mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
return mask

View File

@ -1,5 +1,6 @@
import os, glob import os, glob
import gradio as gr import gradio as gr
import tempfile
from PIL import Image from PIL import Image
try: try:
import torch.cuda as cuda import torch.cuda as cuda
@ -8,7 +9,7 @@ except:
EP_is_visible = False EP_is_visible = False
from typing import List from typing import List
from PIL import Image
import modules.scripts as scripts import modules.scripts as scripts
from modules.upscaler import Upscaler, UpscalerData from modules.upscaler import Upscaler, UpscalerData
from modules import scripts, shared, images, scripts_postprocessing from modules import scripts, shared, images, scripts_postprocessing
@ -28,13 +29,14 @@ except:
model_path = os.path.abspath("models") model_path = os.path.abspath("models")
from scripts.reactor_logger import logger from scripts.reactor_logger import logger
from scripts.reactor_swapper import EnhancementOptions, swap_face, check_process_halt, reset_messaged from scripts.reactor_swapper import EnhancementOptions,MaskOptions,MaskOption, swap_face, check_process_halt, reset_messaged
from scripts.reactor_version import version_flag, app_title from scripts.reactor_version import version_flag, app_title
from scripts.console_log_patch import apply_logging_patch from scripts.console_log_patch import apply_logging_patch
from scripts.reactor_helpers import make_grid, get_image_path, set_Device from scripts.reactor_helpers import make_grid, get_image_path, set_Device
from scripts.reactor_globals import DEVICE, DEVICE_LIST from scripts.reactor_globals import DEVICE, DEVICE_LIST
MODELS_PATH = None MODELS_PATH = None
def get_models(): def get_models():
@ -61,13 +63,18 @@ class FaceSwapScript(scripts.Script):
def ui(self, is_img2img): def ui(self, is_img2img):
with gr.Accordion(f"{app_title}", open=False): with gr.Accordion(f"{app_title}", open=False):
enable = gr.Checkbox(False, label="Enable", info=f"The Fast and Simple FaceSwap Extension - {version_flag}")
gr.Markdown("<br>")
with gr.Tab("Main"): with gr.Tab("Main"):
with gr.Column():
img = gr.Image(type="pil")
enable = gr.Checkbox(False, label="Enable", info=f"The Fast and Simple FaceSwap Extension - {version_flag}")
save_original = gr.Checkbox(False, label="Save Original", info="Save the original image(s) made before swapping; If you use \"img2img\" - this option will affect with \"Swap in generated\" only")
mask_face = gr.Checkbox(False, label="Mask Faces", info="Attempt to mask only the faces and eliminate pixelation of the image around the contours.")
with gr.Column():
with gr.Tab("Single Source Image"):
img = gr.Image(type="pil")
with gr.Tab("Multiple Source Images"):
face_files = gr.File(label="Multiple Source Face Files",file_count="multiple",file_types=["image"],info="Upload multiple face files and each file will be processed in post processing")
save_original = gr.Checkbox(False, label="Save Original", info="Save the original image(s) made before swapping; If you use \"img2img\" - this option will affect with \"Swap in generated\" only")
mask_face = gr.Checkbox(False, label="Mask Faces", info="Attempt to mask only the faces and eliminate pixelation of the image around the contours. Additional settings in the Masking tab.")
gr.Markdown("<br>") gr.Markdown("<br>")
gr.Markdown("Source Image (above):") gr.Markdown("Source Image (above):")
with gr.Row(): with gr.Row():
@ -140,6 +147,26 @@ class FaceSwapScript(scripts.Script):
upscaler_visibility = gr.Slider( upscaler_visibility = gr.Slider(
0, 1, 1, step=0.1, label="Upscaler Visibility (if scale = 1)" 0, 1, 1, step=0.1, label="Upscaler Visibility (if scale = 1)"
) )
with gr.Tab("Masking"):
save_face_mask = gr.Checkbox(False, label="Save Face Mask", info="Save the face mask as a separate image with alpha transparency.")
use_minimal_area = gr.Checkbox(MaskOption.DEFAULT_USE_MINIMAL_AREA, label="Use Minimal Area", info="Use the least amount of area for the mask as possible. This is good for multiple faces that are close together or for preserving the most of the surrounding image.")
mask_areas = gr.CheckboxGroup(
label="Mask areas", choices=["Face", "Hair", "Hat", "Neck"], type="value", value= MaskOption.DEFAULT_FACE_AREAS
)
face_size = gr.Radio(
label = "Face Size", choices = [512,256,128],value=MaskOption.DEFAULT_FACE_SIZE,type="value", info="Size of the masked area. Use larger numbers if the face is expected to be large, smaller if small. Default is 512."
)
mask_blur = gr.Slider(label="Mask blur", minimum=0, maximum=64, step=1, value=12,info="The number of pixels from the outer edge of the mask to blur.")
mask_vignette_fallback_threshold = gr.Slider(
minimum=0.1,
maximum=1.0,
step=0.01,
value=MaskOption.DEFAULT_VIGNETTE_THRESHOLD,
label="Vignette fallback threshold",
info="Switch to a rectangular vignette mask when masked area is only this specified percentage of Face Size."
)
with gr.Tab("Settings"): with gr.Tab("Settings"):
models = get_models() models = get_models()
with gr.Row(visible=EP_is_visible): with gr.Row(visible=EP_is_visible):
@ -213,7 +240,14 @@ class FaceSwapScript(scripts.Script):
source_hash_check, source_hash_check,
target_hash_check, target_hash_check,
device, device,
mask_face mask_face,
save_face_mask,
mask_areas,
mask_blur,
use_minimal_area,
face_size,
mask_vignette_fallback_threshold,
face_files
] ]
@ -242,7 +276,17 @@ class FaceSwapScript(scripts.Script):
restorer_visibility=self.face_restorer_visibility, restorer_visibility=self.face_restorer_visibility,
codeformer_weight=self.codeformer_weight, codeformer_weight=self.codeformer_weight,
) )
@property
def mask_options(self) -> MaskOptions:
return MaskOptions(
mask_areas = self.mask_areas,
save_face_mask = self.save_face_mask,
mask_blur = self.mask_blur,
face_size = self.mask_face_size,
vignette_fallback_threshold = self.mask_vignette_fallback_threshold,
use_minimal_area = self.mask_use_minimal_area,
)
def process( def process(
self, self,
p: StableDiffusionProcessing, p: StableDiffusionProcessing,
@ -267,16 +311,24 @@ class FaceSwapScript(scripts.Script):
source_hash_check, source_hash_check,
target_hash_check, target_hash_check,
device, device,
mask_face mask_face,
save_face_mask:bool,
mask_areas,
mask_blur:int,
mask_use_minimal_area,
mask_face_size,
mask_vignette_fallback_threshold,
face_files
): ):
self.enable = enable self.enable = enable
if self.enable: if self.enable:
reset_messaged() reset_messaged()
if check_process_halt(): if check_process_halt():
return return
global MODELS_PATH global MODELS_PATH
self.source = img self.source = img
self.face_restorer_name = face_restorer_name self.face_restorer_name = face_restorer_name
self.upscaler_scale = upscaler_scale self.upscaler_scale = upscaler_scale
@ -296,6 +348,13 @@ class FaceSwapScript(scripts.Script):
self.target_hash_check = target_hash_check self.target_hash_check = target_hash_check
self.device = device self.device = device
self.mask_face = mask_face self.mask_face = mask_face
self.save_face_mask = save_face_mask
self.mask_blur = mask_blur
self.mask_areas = mask_areas
self.mask_face_size = mask_face_size
self.mask_vignette_fallback_threshold = mask_vignette_fallback_threshold
self.mask_use_minimal_area = mask_use_minimal_area
self.face_files = face_files
if self.gender_source is None or self.gender_source == "No": if self.gender_source is None or self.gender_source == "No":
self.gender_source = 0 self.gender_source = 0
if self.gender_target is None or self.gender_target == "No": if self.gender_target is None or self.gender_target == "No":
@ -318,9 +377,10 @@ class FaceSwapScript(scripts.Script):
self.target_hash_check = False self.target_hash_check = False
set_Device(self.device) set_Device(self.device)
apply_logging_patch(console_logging_level)
if self.source is not None: if self.source is not None:
apply_logging_patch(console_logging_level)
if isinstance(p, StableDiffusionProcessingImg2Img) and self.swap_in_source: if isinstance(p, StableDiffusionProcessingImg2Img) and self.swap_in_source:
logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index) logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index)
@ -339,7 +399,8 @@ class FaceSwapScript(scripts.Script):
source_hash_check=self.source_hash_check, source_hash_check=self.source_hash_check,
target_hash_check=self.target_hash_check, target_hash_check=self.target_hash_check,
device=self.device, device=self.device,
mask_face=mask_face mask_face=mask_face,
mask_options=self.mask_options
) )
p.init_images[i] = result p.init_images[i] = result
# result_path = get_image_path(p.init_images[i], p.outpath_samples, "", p.all_seeds[i], p.all_prompts[i], "txt", p=p, suffix="-swapped") # result_path = get_image_path(p.init_images[i], p.outpath_samples, "", p.all_seeds[i], p.all_prompts[i], "txt", p=p, suffix="-swapped")
@ -350,7 +411,7 @@ class FaceSwapScript(scripts.Script):
if shared.state.interrupted or shared.state.skipped: if shared.state.interrupted or shared.state.skipped:
return return
else: elif self.face_files is None or len(self.face_files) == 0:
logger.error("Please provide a source face") logger.error("Please provide a source face")
def postprocess(self, p: StableDiffusionProcessing, processed: Processed, *args): def postprocess(self, p: StableDiffusionProcessing, processed: Processed, *args):
@ -359,27 +420,29 @@ class FaceSwapScript(scripts.Script):
reset_messaged() reset_messaged()
if check_process_halt(): if check_process_halt():
return return
postprocess_run: bool = True
orig_images : List[Image.Image] = processed.images[processed.index_of_first_image:]
orig_infotexts : List[str] = processed.infotexts[processed.index_of_first_image:]
result_images:List[Image.Image] = []
if self.save_original: if self.save_original:
postprocess_run: bool = True
orig_images : List[Image.Image] = processed.images[processed.index_of_first_image:]
orig_infotexts : List[str] = processed.infotexts[processed.index_of_first_image:]
result_images: List = processed.images result_images: List = processed.images
# result_info: List = processed.infotexts # result_info: List = processed.infotexts
if self.swap_in_generated: if self.swap_in_generated:
logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index) logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index)
if self.source is not None: if self.source is not None:
for i,(img,info) in enumerate(zip(orig_images, orig_infotexts)): for i,(img,info) in enumerate(zip(orig_images, orig_infotexts)):
if check_process_halt(): if check_process_halt():
postprocess_run = False postprocess_run = False
break break
if len(orig_images) > 1: if len(orig_images) > 1:
logger.status("Swap in %s", i) logger.status("Swap in %s", i)
result, output, swapped = swap_face( result, output, swapped, masked_faces = swap_face(
self.source, self.source,
img, img,
source_faces_index=self.source_faces_index, source_faces_index=self.source_faces_index,
@ -391,7 +454,8 @@ class FaceSwapScript(scripts.Script):
source_hash_check=self.source_hash_check, source_hash_check=self.source_hash_check,
target_hash_check=self.target_hash_check, target_hash_check=self.target_hash_check,
device=self.device, device=self.device,
mask_face=self.mask_face mask_face=self.mask_face,
mask_options=self.mask_options
) )
if result is not None and swapped > 0: if result is not None and swapped > 0:
result_images.append(result) result_images.append(result)
@ -400,6 +464,14 @@ class FaceSwapScript(scripts.Script):
img_path = save_image(result, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix) img_path = save_image(result, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix)
except: except:
logger.error("Cannot save a result image - please, check SD WebUI Settings (Saving and Paths)") logger.error("Cannot save a result image - please, check SD WebUI Settings (Saving and Paths)")
if self.mask_face and self.save_face_mask and masked_faces is not None:
result_images.append(masked_faces)
suffix = "-mask"
try:
img_path = save_image(masked_faces, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix)
except:
logger.error("Cannot save a Masked Face image - please, check SD WebUI Settings (Saving and Paths)")
elif result is None: elif result is None:
logger.error("Cannot create a result image") logger.error("Cannot create a result image")
@ -408,7 +480,48 @@ class FaceSwapScript(scripts.Script):
# fullfn = split_fullfn[0] + ".txt" # fullfn = split_fullfn[0] + ".txt"
# with open(fullfn, 'w', encoding="utf8") as f: # with open(fullfn, 'w', encoding="utf8") as f:
# f.writelines(output) # f.writelines(output)
if self.face_files is not None and len(self.face_files) > 0:
for i,(img,info) in enumerate(zip(orig_images, orig_infotexts)):
for j,f_img in enumerate(self.face_files):
if check_process_halt():
postprocess_run = False
break
if len(self.face_files) > 1:
logger.status("Swap in face file #%s", j+1)
result, output, swapped, masked_faces = swap_face(
Image.open(os.path.abspath(f_img.name)),
img,
source_faces_index=self.source_faces_index,
faces_index=self.faces_index,
model=self.model,
enhancement_options=self.enhancement_options,
gender_source=self.gender_source,
gender_target=self.gender_target,
source_hash_check=self.source_hash_check,
target_hash_check=self.target_hash_check,
device=self.device,
mask_face=self.mask_face,
mask_options=self.mask_options
)
if result is not None and swapped > 0:
result_images.append(result)
suffix = f"-swapped-ff-{j+1}"
try:
img_path = save_image(result, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix)
except:
logger.error("Cannot save a result image - please, check SD WebUI Settings (Saving and Paths)")
if self.mask_face and self.save_face_mask and masked_faces is not None:
result_images.append(masked_faces)
suffix = f"-mask-ff-{j+1}"
try:
img_path = save_image(masked_faces, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix)
except:
logger.error("Cannot save a Masked Face image - please, check SD WebUI Settings (Saving and Paths)")
elif result is None:
logger.error("Cannot create a result image")
if shared.opts.return_grid and len(result_images) > 2 and postprocess_run: if shared.opts.return_grid and len(result_images) > 2 and postprocess_run:
grid = make_grid(result_images) grid = make_grid(result_images)
result_images.insert(0, grid) result_images.insert(0, grid)
@ -419,13 +532,62 @@ class FaceSwapScript(scripts.Script):
processed.images = result_images processed.images = result_images
# processed.infotexts = result_info # processed.infotexts = result_info
elif self.face_files is not None and len(self.face_files) > 0:
for i,(img,info) in enumerate(zip(orig_images, orig_infotexts)):
for j,f_img in enumerate(self.face_files):
if check_process_halt():
postprocess_run = False
break
if len(self.face_files) > 1:
logger.status("Swap in face file #%s", j+1)
result, output, swapped, masked_faces = swap_face(
Image.open(os.path.abspath(f_img.name)),
img,
source_faces_index=self.source_faces_index,
faces_index=self.faces_index,
model=self.model,
enhancement_options=self.enhancement_options,
gender_source=self.gender_source,
gender_target=self.gender_target,
source_hash_check=self.source_hash_check,
target_hash_check=self.target_hash_check,
device=self.device,
mask_face=self.mask_face,
mask_options=self.mask_options
)
if result is not None and swapped > 0:
result_images.append(result)
suffix = f"-swapped-ff-{j+1}"
try:
img_path = save_image(result, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix)
except:
logger.error("Cannot save a result image - please, check SD WebUI Settings (Saving and Paths)")
if self.mask_face and self.save_face_mask and masked_faces is not None:
result_images.append(masked_faces)
suffix = f"-mask-ff-{j+1}"
try:
img_path = save_image(masked_faces, p.outpath_samples, "", p.all_seeds[0], p.all_prompts[0], "png",info=info, p=p, suffix=suffix)
except:
logger.error("Cannot save a Masked Face image - please, check SD WebUI Settings (Saving and Paths)")
elif result is None:
logger.error("Cannot create a result image")
if shared.opts.return_grid and len(result_images) > 2 and postprocess_run:
grid = make_grid(result_images)
result_images.insert(0, grid)
try:
save_image(grid, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0], shared.opts.grid_format, info=info, short_filename=not shared.opts.grid_extended_filename, p=p, grid=True)
except:
logger.error("Cannot save a grid - please, check SD WebUI Settings (Saving and Paths)")
processed.images = result_images
def postprocess_batch(self, p, *args, **kwargs): def postprocess_batch(self, p, *args, **kwargs):
if self.enable and not self.save_original: if self.enable and not self.save_original:
images = kwargs["images"] images = kwargs["images"]
def postprocess_image(self, p, script_pp: scripts.PostprocessImageArgs, *args): def postprocess_image(self, p, script_pp: scripts.PostprocessImageArgs, *args):
if self.enable and self.swap_in_generated and not self.save_original: if self.enable and self.swap_in_generated and not ( self.save_original or ( self.face_files is not None and len(self.face_files) > 0)):
current_job_number = shared.state.job_no + 1 current_job_number = shared.state.job_no + 1
job_count = shared.state.job_count job_count = shared.state.job_count
@ -437,7 +599,7 @@ class FaceSwapScript(scripts.Script):
if self.source is not None: if self.source is not None:
logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index) logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index)
image: Image.Image = script_pp.image image: Image.Image = script_pp.image
result, output, swapped = swap_face( result, output, swapped, masked_faces = swap_face(
self.source, self.source,
image, image,
source_faces_index=self.source_faces_index, source_faces_index=self.source_faces_index,
@ -449,7 +611,8 @@ class FaceSwapScript(scripts.Script):
source_hash_check=self.source_hash_check, source_hash_check=self.source_hash_check,
target_hash_check=self.target_hash_check, target_hash_check=self.target_hash_check,
device=self.device, device=self.device,
mask_face=self.mask_face mask_face=self.mask_face,
mask_options=self.mask_options
) )
try: try:
pp = scripts_postprocessing.PostprocessedImage(result) pp = scripts_postprocessing.PostprocessedImage(result)
@ -464,7 +627,7 @@ class FaceSwapScript(scripts.Script):
# f.writelines(output) # f.writelines(output)
except: except:
logger.error("Cannot create a result image") logger.error("Cannot create a result image")
class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing): class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
name = 'ReActor' name = 'ReActor'
@ -475,9 +638,9 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
with gr.Tab("Main"): with gr.Tab("Main"):
with gr.Column(): with gr.Column():
img = gr.Image(type="pil") img = gr.Image(type="pil")
#face_files = gr.File(file_count="multiple",file_types=["image"],info="Upload multiple face files and each file will be processed in post processing")
enable = gr.Checkbox(False, label="Enable", info=f"The Fast and Simple FaceSwap Extension - {version_flag}") enable = gr.Checkbox(False, label="Enable", info=f"The Fast and Simple FaceSwap Extension - {version_flag}")
mask_face = gr.Checkbox(False, label="Mask Faces", info="Attempt to mask only the faces and eliminate pixelation of the image around the contours.") mask_face = gr.Checkbox(False, label="Mask Faces", info="Attempt to mask only the faces and eliminate pixelation of the image around the contours. Additional settings in the Masking tab.")
gr.Markdown("Source Image (above):") gr.Markdown("Source Image (above):")
with gr.Row(): with gr.Row():
source_faces_index = gr.Textbox( source_faces_index = gr.Textbox(
@ -536,6 +699,27 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
upscaler_visibility = gr.Slider( upscaler_visibility = gr.Slider(
0, 1, 1, step=0.1, label="Upscaler Visibility (if scale = 1)" 0, 1, 1, step=0.1, label="Upscaler Visibility (if scale = 1)"
) )
with gr.Tab("Masking"):
#save_face_mask = gr.Checkbox(False, label="Save Face Mask", info="Save the face mask as a separate image with alpha transparency.")
use_minimal_area = gr.Checkbox(MaskOption.DEFAULT_USE_MINIMAL_AREA, label="Use Minimal Area", info="Use the least amount of area for the mask as possible. This is good for multiple faces that are close together or for preserving the most of the surrounding image.")
mask_areas = gr.CheckboxGroup(
label="Mask areas", choices=["Face", "Hair", "Hat", "Neck"], type="value", value= MaskOption.DEFAULT_FACE_AREAS
)
face_size = gr.Radio(
label = "Face Size", choices = [512,256,128],value=MaskOption.DEFAULT_FACE_SIZE,type="value", info="Size of the masked area. Use larger numbers if the face is expected to be large, smaller if small. Default is 512."
)
mask_blur = gr.Slider(label="Mask blur", minimum=0, maximum=64, step=1, value=12,info="The number of pixels from the outer edge of the mask to blur.")
mask_vignette_fallback_threshold = gr.Slider(
minimum=0.1,
maximum=1.0,
step=0.01,
value=MaskOption.DEFAULT_VIGNETTE_THRESHOLD,
label="Vignette fallback threshold",
info="Switch to a rectangular vignette mask when masked area is only this specified percentage of Face Size."
)
with gr.Tab("Settings"): with gr.Tab("Settings"):
models = get_models() models = get_models()
with gr.Row(visible=EP_is_visible): with gr.Row(visible=EP_is_visible):
@ -574,7 +758,7 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
label="Console Log Level", label="Console Log Level",
type="index", type="index",
) )
args = { args = {
'img': img, 'img': img,
'enable': enable, 'enable': enable,
@ -592,7 +776,14 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
'gender_target': gender_target, 'gender_target': gender_target,
'codeformer_weight': codeformer_weight, 'codeformer_weight': codeformer_weight,
'device': device, 'device': device,
'mask_face':mask_face 'mask_face':mask_face,
'mask_areas':mask_areas,
'mask_blur':mask_blur,
'mask_vignette_fallback_threshold':mask_vignette_fallback_threshold,
'face_size':face_size,
'use_minimal_area':use_minimal_area,
} }
return args return args
@ -621,14 +812,23 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
restorer_visibility=self.face_restorer_visibility, restorer_visibility=self.face_restorer_visibility,
codeformer_weight=self.codeformer_weight, codeformer_weight=self.codeformer_weight,
) )
@property
def mask_options(self) -> MaskOptions:
return MaskOptions(
mask_areas = self.mask_areas,
save_face_mask = self.save_face_mask,
mask_blur = self.mask_blur,
face_size = self.face_size,
vignette_fallback_threshold = self.mask_vignette_fallback_threshold,
use_minimal_area = self.use_minimal_area,
)
def process(self, pp: scripts_postprocessing.PostprocessedImage, **args): def process(self, pp: scripts_postprocessing.PostprocessedImage, **args):
if args['enable']: if args['enable']:
reset_messaged() reset_messaged()
if check_process_halt(): if check_process_halt():
return return
global MODELS_PATH global MODELS_PATH
self.source = args['img'] self.source = args['img']
self.face_restorer_name = args['face_restorer_name'] self.face_restorer_name = args['face_restorer_name']
self.upscaler_scale = args['upscaler_scale'] self.upscaler_scale = args['upscaler_scale']
@ -643,6 +843,13 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
self.codeformer_weight = args['codeformer_weight'] self.codeformer_weight = args['codeformer_weight']
self.device = args['device'] self.device = args['device']
self.mask_face = args['mask_face'] self.mask_face = args['mask_face']
self.save_face_mask = None
self.mask_areas= args['mask_areas']
self.mask_blur= args['mask_blur']
self.mask_vignette_fallback_threshold= args['mask_vignette_fallback_threshold']
self.face_size= args['face_size']
self.use_minimal_area= args['use_minimal_area']
self.face_files = None
if self.gender_source is None or self.gender_source == "No": if self.gender_source is None or self.gender_source == "No":
self.gender_source = 0 self.gender_source = 0
if self.gender_target is None or self.gender_target == "No": if self.gender_target is None or self.gender_target == "No":
@ -664,12 +871,12 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
reset_messaged() reset_messaged()
set_Device(self.device) set_Device(self.device)
apply_logging_patch(self.console_logging_level)
if self.source is not None: if self.source is not None:
apply_logging_patch(self.console_logging_level)
logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index) logger.status("Working: source face index %s, target face index %s", self.source_faces_index, self.faces_index)
image: Image.Image = pp.image image: Image.Image = pp.image
result, output, swapped = swap_face( result, output, swapped,masked_faces = swap_face(
self.source, self.source,
image, image,
source_faces_index=self.source_faces_index, source_faces_index=self.source_faces_index,
@ -681,7 +888,8 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
source_hash_check=True, source_hash_check=True,
target_hash_check=True, target_hash_check=True,
device=self.device, device=self.device,
mask_face=self.mask_face mask_face=self.mask_face,
mask_options=self.mask_options
) )
try: try:
pp.info["ReActor"] = True pp.info["ReActor"] = True

View File

@ -1,7 +1,7 @@
import copy import copy
import os import os
from dataclasses import dataclass from dataclasses import dataclass
from typing import List, Union from typing import List, Tuple, Union
import cv2 import cv2
import numpy as np import numpy as np
@ -41,7 +41,12 @@ if DEVICE == "CUDA":
PROVIDERS = ["CUDAExecutionProvider"] PROVIDERS = ["CUDAExecutionProvider"]
else: else:
PROVIDERS = ["CPUExecutionProvider"] PROVIDERS = ["CPUExecutionProvider"]
class MaskOption:
DEFAULT_FACE_AREAS = ["Face"]
DEFAULT_FACE_SIZE:int = 512
DEFAULT_VIGNETTE_THRESHOLD:float = 0.1
DEFAULT_MASK_BLUR:int = 12,
DEFAULT_USE_MINIMAL_AREA:bool = True
@dataclass @dataclass
class EnhancementOptions: class EnhancementOptions:
@ -52,7 +57,16 @@ class EnhancementOptions:
face_restorer: FaceRestoration = None face_restorer: FaceRestoration = None
restorer_visibility: float = 0.5 restorer_visibility: float = 0.5
codeformer_weight: float = 0.5 codeformer_weight: float = 0.5
@dataclass
class MaskOptions:
mask_areas:List[str]
save_face_mask: bool = False
mask_blur:int = 12
face_size:int = 512
vignette_fallback_threshold:float =0.10
use_minimal_area:bool = True
MESSAGED_STOPPED = False MESSAGED_STOPPED = False
MESSAGED_SKIPPED = False MESSAGED_SKIPPED = False
@ -175,7 +189,7 @@ def enhance_image(image: Image, enhancement_options: EnhancementOptions):
result_image = restore_face(result_image, enhancement_options) result_image = restore_face(result_image, enhancement_options)
return result_image return result_image
def enhance_image_and_mask(image: Image.Image, enhancement_options: EnhancementOptions,target_img_orig:Image.Image,entire_mask_image:Image.Image)->Image.Image: def enhance_image_and_mask(image: Image.Image, enhancement_options: EnhancementOptions,target_img_orig:Image.Image,entire_mask_image:Image.Image)->Tuple[Image.Image,Image.Image]:
result_image = image result_image = image
if check_process_halt(msgforced=True): if check_process_halt(msgforced=True):
@ -183,7 +197,11 @@ def enhance_image_and_mask(image: Image.Image, enhancement_options: EnhancementO
if enhancement_options.do_restore_first: if enhancement_options.do_restore_first:
result_image = restore_face(result_image, enhancement_options) result_image = restore_face(result_image, enhancement_options)
transparent = Image.new("RGBA",result_image.size)
masked_faces = Image.composite(result_image.convert("RGBA"),transparent,entire_mask_image)
result_image = Image.composite(result_image,target_img_orig,entire_mask_image) result_image = Image.composite(result_image,target_img_orig,entire_mask_image)
result_image = upscale_image(result_image, enhancement_options) result_image = upscale_image(result_image, enhancement_options)
@ -191,10 +209,13 @@ def enhance_image_and_mask(image: Image.Image, enhancement_options: EnhancementO
result_image = upscale_image(result_image, enhancement_options) result_image = upscale_image(result_image, enhancement_options)
entire_mask_image = Image.fromarray(cv2.resize(np.array(entire_mask_image),result_image.size, interpolation=cv2.INTER_AREA)).convert("L") entire_mask_image = Image.fromarray(cv2.resize(np.array(entire_mask_image),result_image.size, interpolation=cv2.INTER_AREA)).convert("L")
result_image = Image.composite(result_image,target_img_orig,entire_mask_image) result_image = Image.composite(result_image,target_img_orig,entire_mask_image)
result_image = restore_face(result_image, enhancement_options) result_image = restore_face(result_image, enhancement_options)
transparent = Image.new("RGBA",result_image.size)
masked_faces = Image.composite(result_image.convert("RGBA"),transparent,entire_mask_image)
return result_image, masked_faces
return result_image
def get_gender(face, face_index): def get_gender(face, face_index):
@ -310,16 +331,19 @@ def swap_face(
source_hash_check: bool = True, source_hash_check: bool = True,
target_hash_check: bool = False, target_hash_check: bool = False,
device: str = "CPU", device: str = "CPU",
mask_face:bool = False mask_face:bool = False,
mask_options:Union[MaskOptions, None]= None
): ):
global SOURCE_FACES, SOURCE_IMAGE_HASH, TARGET_FACES, TARGET_IMAGE_HASH, PROVIDERS global SOURCE_FACES, SOURCE_IMAGE_HASH, TARGET_FACES, TARGET_IMAGE_HASH, PROVIDERS
result_image = target_img result_image = target_img
masked_faces = None
PROVIDERS = ["CUDAExecutionProvider"] if device == "CUDA" else ["CPUExecutionProvider"] PROVIDERS = ["CUDAExecutionProvider"] if device == "CUDA" else ["CPUExecutionProvider"]
if check_process_halt(): if check_process_halt():
return result_image, [], 0 return result_image, [], 0
if mask_options is None:
mask_options = MaskOptions()
if model is not None: if model is not None:
if isinstance(source_img, str): # source_img is a base64 string if isinstance(source_img, str): # source_img is a base64 string
@ -444,7 +468,7 @@ def swap_face(
swapped_image = face_swapper.get(result, target_face, source_face) swapped_image = face_swapper.get(result, target_face, source_face)
if mask_face: if mask_face:
result = apply_face_mask(swapped_image=swapped_image,target_image=result,target_face=target_face,entire_mask_image=entire_mask_image) result = apply_face_mask(swapped_image=swapped_image,target_image=result,target_face=target_face,entire_mask_image=entire_mask_image,mask_options=mask_options)
else: else:
result = swapped_image result = swapped_image
swapped += 1 swapped += 1
@ -480,8 +504,8 @@ def swap_face(
result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB)) result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
if enhancement_options is not None and swapped > 0: if enhancement_options is not None and swapped > 0:
if mask_face and entire_mask_image is not None: if mask_face and entire_mask_image is not None:
result_image = enhance_image_and_mask(result_image, enhancement_options,Image.fromarray(target_img_orig),Image.fromarray(entire_mask_image).convert("L")) result_image, masked_faces = enhance_image_and_mask(result_image, enhancement_options,Image.fromarray(target_img_orig),Image.fromarray(entire_mask_image).convert("L"))
else: else:
result_image = enhance_image(result_image, enhancement_options) result_image = enhance_image(result_image, enhancement_options)
elif mask_face and entire_mask_image is not None and swapped > 0: elif mask_face and entire_mask_image is not None and swapped > 0:
@ -492,20 +516,19 @@ def swap_face(
else: else:
logger.status("No source face(s) found") logger.status("No source face(s) found")
return result_image, output, swapped return result_image, output, swapped,masked_faces
def apply_face_mask(swapped_image:np.ndarray,target_image:np.ndarray,target_face,entire_mask_image:np.array)->np.ndarray: def apply_face_mask(swapped_image:np.ndarray,target_image:np.ndarray,target_face,entire_mask_image:np.array,mask_options:Union[MaskOptions,None] = None)->np.ndarray:
logger.status("Masking Face") logger.status("Masking Face")
mask_generator = BiSeNetMaskGenerator() mask_generator = BiSeNetMaskGenerator()
face = Face(target_image,Rect.from_ndarray(np.array(target_face.bbox)),1.6,512,"") face = Face(target_image,Rect.from_ndarray(np.array(target_face.bbox)),1.6,mask_options.face_size,"")
face_image = np.array(face.image) face_image = np.array(face.image)
process_face_image(face)
face_area_on_image = face.face_area_on_image face_area_on_image = face.face_area_on_image
mask = mask_generator.generate_mask(face_image,face_area_on_image=face_area_on_image,affected_areas=["Face"],mask_size=0,use_minimal_area=True)
mask = cv2.blur(mask, (12, 12)) mask = mask_generator.generate_mask(face_image,face_area_on_image=face_area_on_image,affected_areas=mask_options.mask_areas,mask_size=0,use_minimal_area=mask_options.use_minimal_area)
"""entire_mask_image = np.zeros_like(target_image)""" mask = cv2.blur(mask, (mask_options.mask_blur, mask_options.mask_blur))
larger_mask = cv2.resize(mask, dsize=(face.width, face.height)) larger_mask = cv2.resize(mask, dsize=(face.width, face.height))
entire_mask_image[ entire_mask_image[
face.top : face.bottom, face.top : face.bottom,
@ -554,22 +577,7 @@ def color_generator(colors):
color_iter = color_generator(colors) color_iter = color_generator(colors)
def process_face_image(
face: Face,
**kwargs,
) -> Image:
image = np.array(face.image)
overlay = image.copy()
cv2.rectangle(overlay, (0, 0), (image.shape[1], image.shape[0]), next(color_iter), -1)
l, t, r, b = face.face_area_on_image
cv2.rectangle(overlay, (l, t), (r, b), (0, 0, 0), 10)
if face.landmarks_on_image is not None:
for landmark in face.landmarks_on_image:
cv2.circle(overlay, (int(landmark.x), int(landmark.y)), 6, (0, 0, 0), 10)
alpha = 0.3
output = cv2.addWeighted(image, 1 - alpha, overlay, alpha, 0)
return Image.fromarray(output)
def dilate_erode(img: Image.Image, value: int) -> Image.Image: def dilate_erode(img: Image.Image, value: int) -> Image.Image:
""" """
The dilate_erode function takes an image and a value. The dilate_erode function takes an image and a value.