Refactoring after PR

+VersionUP (0.5.1 beta1)

README.md

@@ -2,7 +2,7 @@
 
   <img src="https://github.com/Gourieff/Assets/raw/main/sd-webui-reactor/ReActor_logo_red.png?raw=true" alt="logo" width="180px"/>
     
-  ![Version](https://img.shields.io/badge/version-0.5.0-brightgreen?style=for-the-badge&labelColor=darkgreen)
+  ![Version](https://img.shields.io/badge/version-0.5.1_beta1-green?style=for-the-badge&labelColor=darkgreen)
   
   <a href='https://ko-fi.com/gourieff' target='_blank'><img height='33' src='https://storage.ko-fi.com/cdn/kofi3.png?v=3' border='0' alt='Buy Me a Coffee at ko-fi.com' /></a>
   

README_RU.md

@@ -2,7 +2,7 @@
 
   <img src="https://github.com/Gourieff/Assets/raw/main/sd-webui-reactor/ReActor_logo_red.png?raw=true" alt="logo" width="180px"/>
     
-  ![Version](https://img.shields.io/badge/версия-0.5.0-brightgreen?style=for-the-badge&labelColor=darkgreen)
+  ![Version](https://img.shields.io/badge/версия-0.5.1_beta1-green?style=for-the-badge&labelColor=darkgreen)
   
   <a href='https://ko-fi.com/gourieff' target='_blank'><img height='33' src='https://storage.ko-fi.com/cdn/kofi3.png?v=3' border='0' alt='Buy Me a Coffee at ko-fi.com' /></a>
   

modules/reactor_mask.py

@@ -0,0 +1,176 @@
+import cv2
+import numpy as np
+from PIL import Image, ImageDraw
+
+from torchvision.transforms.functional import to_pil_image
+
+from scripts.reactor_logger import logger
+from scripts.inferencers.bisenet_mask_generator import BiSeNetMaskGenerator
+from scripts.entities.face import FaceArea
+from scripts.entities.rect import Rect
+
+
+colors = [
+    (255, 0, 0),
+    (0, 255, 0),
+    (0, 0, 255),
+    (255, 255, 0),
+    (255, 0, 255),
+    (0, 255, 255),
+    (255, 255, 255),
+    (128, 0, 0),
+    (0, 128, 0),
+    (128, 128, 0),
+    (0, 0, 128),
+    (0, 128, 128),
+]
+
+def color_generator(colors):
+    while True:
+        for color in colors:
+            yield color
+
+
+def process_face_image(
+        face: FaceArea,
+        **kwargs,
+    ) -> Image:
+        image = np.array(face.image)
+        overlay = image.copy()
+        color_iter = color_generator(colors)
+        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 apply_face_mask(swapped_image:np.ndarray,target_image:np.ndarray,target_face,entire_mask_image:np.array)->np.ndarray:
+    logger.status("Correcting Face Mask")
+    mask_generator =  BiSeNetMaskGenerator()
+    face = FaceArea(target_image,Rect.from_ndarray(np.array(target_face.bbox)),1.6,512,"")
+    face_image = np.array(face.image)
+    process_face_image(face)
+    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))
+    # """entire_mask_image = np.zeros_like(target_image)"""
+    larger_mask = cv2.resize(mask, dsize=(face.width, face.height))
+    entire_mask_image[
+        face.top : face.bottom,
+        face.left : face.right,
+    ] = larger_mask
+   
+    result = Image.composite(Image.fromarray(swapped_image),Image.fromarray(target_image), Image.fromarray(entire_mask_image).convert("L"))
+    return np.array(result)
+
+
+def rotate_array(image: np.ndarray, angle: float) -> np.ndarray:
+    if angle == 0:
+        return image
+
+    h, w = image.shape[:2]
+    center = (w // 2, h // 2)
+
+    M = cv2.getRotationMatrix2D(center, angle, 1.0)
+    return cv2.warpAffine(image, M, (w, h))
+
+
+def rotate_image(image: Image, angle: float) -> Image:
+    if angle == 0:
+        return image
+    return Image.fromarray(rotate_array(np.array(image), angle))
+
+
+def correct_face_tilt(angle: float) -> bool:
+    angle = abs(angle)
+    if angle > 180:
+        angle = 360 - angle
+    return angle > 40
+
+
+def _dilate(arr: np.ndarray, value: int) -> np.ndarray:
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (value, value))
+    return cv2.dilate(arr, kernel, iterations=1)
+
+
+def _erode(arr: np.ndarray, value: int) -> np.ndarray:
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (value, value))
+    return cv2.erode(arr, kernel, iterations=1)
+
+
+def dilate_erode(img: Image.Image, value: int) -> Image.Image:
+    """
+    The dilate_erode function takes an image and a value.
+    If the value is positive, it dilates the image by that amount.
+    If the value is negative, it erodes the image by that amount.
+
+    Parameters
+    ----------
+        img: PIL.Image.Image
+            the image to be processed
+        value: int
+            kernel size of dilation or erosion
+
+    Returns
+    -------
+        PIL.Image.Image
+            The image that has been dilated or eroded
+    """
+    if value == 0:
+        return img
+
+    arr = np.array(img)
+    arr = _dilate(arr, value) if value > 0 else _erode(arr, -value)
+
+    return Image.fromarray(arr)
+
+def mask_to_pil(masks, shape: tuple[int, int]) -> list[Image.Image]:
+    """
+    Parameters
+    ----------
+    masks: torch.Tensor, dtype=torch.float32, shape=(N, H, W).
+        The device can be CUDA, but `to_pil_image` takes care of that.
+
+    shape: tuple[int, int]
+        (width, height) of the original image
+    """
+    n = masks.shape[0]
+    return [to_pil_image(masks[i], mode="L").resize(shape) for i in range(n)]
+
+def create_mask_from_bbox(
+    bboxes: list[list[float]], shape: tuple[int, int]
+) -> list[Image.Image]:
+    """
+    Parameters
+    ----------
+        bboxes: list[list[float]]
+            list of [x1, y1, x2, y2]
+            bounding boxes
+        shape: tuple[int, int]
+            shape of the image (width, height)
+
+    Returns
+    -------
+        masks: list[Image.Image]
+        A list of masks
+
+    """
+    masks = []
+    for bbox in bboxes:
+        mask = Image.new("L", shape, 0)
+        mask_draw = ImageDraw.Draw(mask)
+        mask_draw.rectangle(bbox, fill=255)
+        masks.append(mask)
+    return masks

scripts/entities/face.py

@@ -9,7 +9,7 @@ from PIL import Image
 from scripts.entities.rect import Point, Rect
 
 
-class Face:
+class FaceArea:
     def __init__(self, entire_image: np.ndarray, face_area: Rect, face_margin: float, face_size: int, upscaler: str):
         self.face_area = face_area
         self.center = face_area.center

scripts/inferencers/bisenet_mask_generator.py

@@ -7,9 +7,7 @@ import torch
 from facexlib.parsing import init_parsing_model
 from facexlib.utils.misc import img2tensor
 from torchvision.transforms.functional import normalize
-from PIL import Image
 from scripts.inferencers.mask_generator import MaskGenerator
-from scripts.reactor_logger import logger
 
 class BiSeNetMaskGenerator(MaskGenerator):
     def __init__(self) -> None:
@@ -28,7 +26,7 @@ class BiSeNetMaskGenerator(MaskGenerator):
         fallback_ratio: float = 0.25,
         **kwargs,
     ) -> np.ndarray:
-        original_face_image = face_image
+        # original_face_image = face_image
         face_image = face_image.copy()
         face_image = face_image[:, :, ::-1]
 
@@ -59,11 +57,11 @@ class BiSeNetMaskGenerator(MaskGenerator):
         if w != 512 or h != 512:
             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.info("Use fallback 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
 

scripts/inferencers/mask_generator.py

@@ -4,7 +4,6 @@ from typing import Tuple
 import cv2
 import numpy as np
 
-
 class MaskGenerator(ABC):
     @abstractmethod
     def name(self) -> str:

scripts/reactor_faceswap.py

@@ -66,8 +66,7 @@ class FaceSwapScript(scripts.Script):
                     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.")
+                    mask_face = gr.Checkbox(False, label="Face Mask Correction", info="Apply this option if you see some pixelation around face contours")
-
                     gr.Markdown("<br>")
                     gr.Markdown("Source Image (above):")
                     with gr.Row():
@@ -213,7 +212,7 @@ class FaceSwapScript(scripts.Script):
             source_hash_check,
             target_hash_check,
             device,
-            mask_face
+            mask_face,
         ]
 
 
@@ -267,7 +266,7 @@ class FaceSwapScript(scripts.Script):
         source_hash_check,
         target_hash_check,
         device,
-        mask_face
+        mask_face,
     ):
         self.enable = enable
         if self.enable:
@@ -316,6 +315,8 @@ class FaceSwapScript(scripts.Script):
                 self.source_hash_check = True
             if self.target_hash_check is None:
                 self.target_hash_check = False
+            if self.mask_face is None:
+                self.mask_face = False
 
             set_Device(self.device)
             
@@ -339,7 +340,7 @@ class FaceSwapScript(scripts.Script):
                             source_hash_check=self.source_hash_check,
                             target_hash_check=self.target_hash_check,
                             device=self.device,
-                            mask_face=mask_face
+                            mask_face=self.mask_face,
                         )
                         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")
@@ -391,7 +392,7 @@ class FaceSwapScript(scripts.Script):
                                 source_hash_check=self.source_hash_check,
                                 target_hash_check=self.target_hash_check,
                                 device=self.device,
-                                mask_face=self.mask_face
+                                mask_face=self.mask_face,
                             )
                             if result is not None and swapped > 0:
                                 result_images.append(result)
@@ -449,7 +450,7 @@ class FaceSwapScript(scripts.Script):
                     source_hash_check=self.source_hash_check,
                     target_hash_check=self.target_hash_check,
                     device=self.device,
-                    mask_face=self.mask_face
+                    mask_face=self.mask_face,
                 )
                 try:
                     pp = scripts_postprocessing.PostprocessedImage(result)
@@ -476,8 +477,7 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
                 with gr.Column():
                     img = gr.Image(type="pil")
                     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="Face Mask Correction", info="Apply this option if you see some pixelation around face contours")
-                    
                     gr.Markdown("Source Image (above):")
                     with gr.Row():
                         source_faces_index = gr.Textbox(
@@ -592,7 +592,7 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
             'gender_target': gender_target,
             'codeformer_weight': codeformer_weight,
             'device': device,
-            'mask_face':mask_face
+            'mask_face': mask_face,
         }
         return args
 
@@ -657,6 +657,8 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
                 self.source_faces_index = [0]
             if len(self.faces_index) == 0:
                 self.faces_index = [0]
+            if self.mask_face is None:
+                self.mask_face = False
 
             current_job_number = shared.state.job_no + 1
             job_count = shared.state.job_count
@@ -681,7 +683,7 @@ class FaceSwapScriptExtras(scripts_postprocessing.ScriptPostprocessing):
                     source_hash_check=True,
                     target_hash_check=True,
                     device=self.device,
-                    mask_face=self.mask_face
+                    mask_face=self.mask_face,
                 )
                 try:
                     pp.info["ReActor"] = True

scripts/reactor_swapper.py

@@ -5,13 +5,10 @@ from typing import List, Union
 
 import cv2
 import numpy as np
-from numpy import uint8
+from PIL import Image
-from PIL import Image, ImageDraw
+
-from scripts.inferencers.bisenet_mask_generator import BiSeNetMaskGenerator
-from scripts.entities.face import Face
-from scripts.entities.rect import Rect
 import insightface
-from torchvision.transforms.functional import to_pil_image
+
 from scripts.reactor_helpers import get_image_md5hash, get_Device
 from modules.face_restoration import FaceRestoration
 try: # A1111
@@ -21,6 +18,7 @@ except: # SD.Next
 from modules.upscaler import UpscalerData
 from modules.shared import state
 from scripts.reactor_logger import logger
+from modules.reactor_mask import apply_face_mask
 
 try:
     from modules.paths_internal import models_path
@@ -310,7 +308,7 @@ def swap_face(
     source_hash_check: bool = True,
     target_hash_check: bool = False,
     device: str = "CPU",
-    mask_face:bool = False
+    mask_face: bool = False,
 ):
     global SOURCE_FACES, SOURCE_IMAGE_HASH, TARGET_FACES, TARGET_IMAGE_HASH, PROVIDERS
     result_image = target_img
@@ -493,160 +491,3 @@ def swap_face(
             logger.status("No source face(s) found")
     
     return result_image, output, swapped
-
-
-
-def apply_face_mask(swapped_image:np.ndarray,target_image:np.ndarray,target_face,entire_mask_image:np.array)->np.ndarray:
-    logger.status("Masking Face")
-    mask_generator =  BiSeNetMaskGenerator()
-    face = Face(target_image,Rect.from_ndarray(np.array(target_face.bbox)),1.6,512,"")
-    face_image = np.array(face.image)
-    process_face_image(face)
-    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))
-    """entire_mask_image = np.zeros_like(target_image)"""
-    larger_mask = cv2.resize(mask, dsize=(face.width, face.height))
-    entire_mask_image[
-        face.top : face.bottom,
-        face.left : face.right,
-    ] = larger_mask
-   
-   
-    result = Image.composite(Image.fromarray(swapped_image),Image.fromarray(target_image), Image.fromarray(entire_mask_image).convert("L"))
-    return np.array(result)
-    
-
-def correct_face_tilt(angle: float) -> bool:
-        
-        angle = abs(angle)
-        if angle > 180:
-            angle = 360 - angle
-        return angle > 40
-def _dilate(arr: np.ndarray, value: int) -> np.ndarray:
-    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (value, value))
-    return cv2.dilate(arr, kernel, iterations=1)
-
-
-def _erode(arr: np.ndarray, value: int) -> np.ndarray:
-    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (value, value))
-    return cv2.erode(arr, kernel, iterations=1)
-colors = [
-    (255, 0, 0),
-    (0, 255, 0),
-    (0, 0, 255),
-    (255, 255, 0),
-    (255, 0, 255),
-    (0, 255, 255),
-    (255, 255, 255),
-    (128, 0, 0),
-    (0, 128, 0),
-    (128, 128, 0),
-    (0, 0, 128),
-    (0, 128, 128),
-]
-
-
-def color_generator(colors):
-    while True:
-        for color in colors:
-            yield color
-
-
-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:
-    """
-    The dilate_erode function takes an image and a value.
-    If the value is positive, it dilates the image by that amount.
-    If the value is negative, it erodes the image by that amount.
-
-    Parameters
-    ----------
-        img: PIL.Image.Image
-            the image to be processed
-        value: int
-            kernel size of dilation or erosion
-
-    Returns
-    -------
-        PIL.Image.Image
-            The image that has been dilated or eroded
-    """
-    if value == 0:
-        return img
-
-    arr = np.array(img)
-    arr = _dilate(arr, value) if value > 0 else _erode(arr, -value)
-
-    return Image.fromarray(arr)
-
-def mask_to_pil(masks, shape: tuple[int, int]) -> list[Image.Image]:
-    """
-    Parameters
-    ----------
-    masks: torch.Tensor, dtype=torch.float32, shape=(N, H, W).
-        The device can be CUDA, but `to_pil_image` takes care of that.
-
-    shape: tuple[int, int]
-        (width, height) of the original image
-    """
-    n = masks.shape[0]
-    return [to_pil_image(masks[i], mode="L").resize(shape) for i in range(n)]
-
-def create_mask_from_bbox(
-    bboxes: list[list[float]], shape: tuple[int, int]
-) -> list[Image.Image]:
-    """
-    Parameters
-    ----------
-        bboxes: list[list[float]]
-            list of [x1, y1, x2, y2]
-            bounding boxes
-        shape: tuple[int, int]
-            shape of the image (width, height)
-
-    Returns
-    -------
-        masks: list[Image.Image]
-        A list of masks
-
-    """
-    masks = []
-    for bbox in bboxes:
-        mask = Image.new("L", shape, 0)
-        mask_draw = ImageDraw.Draw(mask)
-        mask_draw.rectangle(bbox, fill=255)
-        masks.append(mask)
-    return masks
-
-def rotate_image(image: Image, angle: float) -> Image:
-    if angle == 0:
-        return image
-    return Image.fromarray(rotate_array(np.array(image), angle))
-
-
-def rotate_array(image: np.ndarray, angle: float) -> np.ndarray:
-    if angle == 0:
-        return image
-
-    h, w = image.shape[:2]
-    center = (w // 2, h // 2)
-
-    M = cv2.getRotationMatrix2D(center, angle, 1.0)
-    return cv2.warpAffine(image, M, (w, h))

scripts/reactor_version.py

@@ -1,5 +1,5 @@
 app_title = "ReActor"
-version_flag = "v0.5.0"
+version_flag = "v0.5.1-b1"
 
 from scripts.reactor_logger import logger, get_Run, set_Run