Recrosoftware/sd-webui-reactor
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactoring after PR

Browse Source 
+VersionUP (0.5.1 beta1)
This commit is contained in:
Gourieff 2023-11-23 23:13:42 +07:00
parent 8b66464e6f
commit 2c2d40508a
9 changed files with 204 additions and 188 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@ -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>

2
README_RU.md
View File

@ -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>

176
modules/reactor_mask.py Normal file
View File

@ -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

2
scripts/entities/face.py
View File

@ -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

14
scripts/inferencers/bisenet_mask_generator.py
View File

@ -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

1
scripts/inferencers/mask_generator.py
View File

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

24
scripts/reactor_faceswap.py
View File

@ -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

169
scripts/reactor_swapper.py
View File

@ -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))

2
scripts/reactor_version.py
View File

@ -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