diff --git a/scripts/reactor_swapper.py b/scripts/reactor_swapper.py index 9cb24a8..fee61b1 100644 --- a/scripts/reactor_swapper.py +++ b/scripts/reactor_swapper.py @@ -28,6 +28,16 @@ except: import warnings +PILImage = Image.Image +CV2ImgU8 = np.ndarray[int, np.dtype[uint8]] +Face = IFace +BoxCoords = Tuple[int, int, int, int] + + +class Gender(Enum): + AUTO = -1 + FEMALE = 0 + MALE = 1 np.warnings = warnings np.warnings.filterwarnings('ignore') @@ -453,3 +463,101 @@ def swap_face( logger.status("No source face(s) found") return result_image, output, swapped + + +def merge_images_with_mask( + image1: CV2ImgU8, image2: CV2ImgU8, mask: CV2ImgU8 +) -> CV2ImgU8: + """ + Merges two images using a given mask. The regions where the mask is set will be replaced with the corresponding + areas of the second image. + + Args: + image1 (CV2Img): The base image, which must have the same shape as image2. + image2 (CV2Img): The image to be merged, which must have the same shape as image1. + mask (CV2Img): A binary mask specifying the regions to be merged. The mask shape should match image1's first two dimensions. + + Returns: + CV2Img: The merged image. + + Raises: + ValueError: If the shapes of the images and mask do not match. + """ + + if image1.shape != image2.shape or image1.shape[:2] != mask.shape: + raise ValueError("Img should have the same shape") + mask = mask.astype(np.uint8) + masked_region = cv2.bitwise_and(image2, image2, mask=mask) + inverse_mask = cv2.bitwise_not(mask) + empty_region = cv2.bitwise_and(image1, image1, mask=inverse_mask) + merged_image = cv2.add(empty_region, masked_region) + return merged_image + + +def erode_mask(mask: CV2ImgU8, kernel_size: int = 3, iterations: int = 1) -> CV2ImgU8: + """ + Erodes a binary mask using a given kernel size and number of iterations. + + Args: + mask (CV2Img): The binary mask to erode. + kernel_size (int, optional): The size of the kernel. Default is 3. + iterations (int, optional): The number of erosion iterations. Default is 1. + + Returns: + CV2Img: The eroded mask. + """ + kernel = np.ones((kernel_size, kernel_size), np.uint8) + eroded_mask = cv2.erode(mask, kernel, iterations=iterations) + return eroded_mask + + +def apply_gaussian_blur( + mask: CV2ImgU8, kernel_size: Tuple[int, int] = (5, 5), sigma_x: int = 0 +) -> CV2ImgU8: + """ + Applies a Gaussian blur to a mask. + + Args: + mask (CV2Img): The mask to blur. + kernel_size (tuple, optional): The size of the kernel, e.g. (5, 5). Default is (5, 5). + sigma_x (int, optional): The standard deviation in the X direction. Default is 0. + + Returns: + CV2Img: The blurred mask. + """ + blurred_mask = cv2.GaussianBlur(mask, kernel_size, sigma_x) + return blurred_mask + + +def dilate_mask(mask: CV2ImgU8, kernel_size: int = 5, iterations: int = 1) -> CV2ImgU8: + """ + Dilates a binary mask using a given kernel size and number of iterations. + + Args: + mask (CV2Img): The binary mask to dilate. + kernel_size (int, optional): The size of the kernel. Default is 5. + iterations (int, optional): The number of dilation iterations. Default is 1. + + Returns: + CV2Img: The dilated mask. + """ + kernel = np.ones((kernel_size, kernel_size), np.uint8) + dilated_mask = cv2.dilate(mask, kernel, iterations=iterations) + return dilated_mask + + +def get_face_mask(aimg: CV2ImgU8, bgr_fake: CV2ImgU8) -> CV2ImgU8: + """ + Generates a face mask by performing bitwise OR on two face masks and then dilating the result. + + Args: + aimg (CV2Img): Input image for generating the first face mask. + bgr_fake (CV2Img): Input image for generating the second face mask. + + Returns: + CV2Img: The combined and dilated face mask. + """ + mask1 = generate_face_mask(aimg, device=shared.device) + mask2 = generate_face_mask(bgr_fake, device=shared.device) + mask = dilate_mask(cv2.bitwise_or(mask1, mask2)) + return mask