diff --git a/README.md b/README.md
index 667e157..70ae680 100644
--- a/README.md
+++ b/README.md
@@ -2,7 +2,7 @@
- 
+ 
diff --git a/README_RU.md b/README_RU.md
index 1e76995..4a9e121 100644
--- a/README_RU.md
+++ b/README_RU.md
@@ -2,7 +2,7 @@
- 
+ 
diff --git a/modules/reactor_mask.py b/modules/reactor_mask.py
new file mode 100644
index 0000000..5f939d1
--- /dev/null
+++ b/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
diff --git a/scripts/entities/face.py b/scripts/entities/face.py
index 21a6167..8ba1f7e 100644
--- a/scripts/entities/face.py
+++ b/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
diff --git a/scripts/inferencers/bisenet_mask_generator.py b/scripts/inferencers/bisenet_mask_generator.py
index 27eb2e7..3de09e8 100644
--- a/scripts/inferencers/bisenet_mask_generator.py
+++ b/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:
- logger.info("Use fallback mask generator")
- mask = self.fallback_mask_generator.generate_mask(
- original_face_image, face_area_on_image, use_minimal_area=True
- )"""
+ # """if MaskGenerator.calculate_mask_coverage(mask) < fallback_ratio:
+ # logger.info("Use fallback mask generator")
+ # mask = self.fallback_mask_generator.generate_mask(
+ # original_face_image, face_area_on_image, use_minimal_area=True
+ # )"""
return mask
diff --git a/scripts/inferencers/mask_generator.py b/scripts/inferencers/mask_generator.py
index 9359523..0c46b89 100644
--- a/scripts/inferencers/mask_generator.py
+++ b/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:
diff --git a/scripts/reactor_faceswap.py b/scripts/reactor_faceswap.py
index 0c54371..d2fbb9e 100644
--- a/scripts/reactor_faceswap.py
+++ b/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("
")
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
diff --git a/scripts/reactor_swapper.py b/scripts/reactor_swapper.py
index 061513d..c244724 100644
--- a/scripts/reactor_swapper.py
+++ b/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, ImageDraw
-from scripts.inferencers.bisenet_mask_generator import BiSeNetMaskGenerator
-from scripts.entities.face import Face
-from scripts.entities.rect import Rect
+from PIL import Image
+
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))
diff --git a/scripts/reactor_version.py b/scripts/reactor_version.py
index 90c70d6..3230c48 100644
--- a/scripts/reactor_version.py
+++ b/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