ignore bottom of face - draft

reactor_modules/reactor_mask.py

@@ -8,7 +8,9 @@
 from scripts.reactor_inferencers.bisenet_mask_generator import BiSeNetMaskGenerator
 from scripts.reactor_entities.face import FaceArea
 from scripts.reactor_entities.rect import Rect
+from insightface.app.common import Face
 
+from scripts.reactor_inferencers.mask_generator import MaskGenerator
 
 colors = [
     (255, 0, 0),
@@ -31,24 +33,98 @@ def color_generator(colors):
             yield color
 
 
+# def process_face_image(
+#         face: FaceArea,
+#         exclude_mouth: bool = False,  # New parameter to control mouth exclusion
+#         **kwargs,
+#     ) -> Image:
+#         image = np.array(face.image)
+#         overlay = image.copy()
+#         color_iter = color_generator(colors)
+
+#         # Draw a rectangle over the entire face
+#         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)
+
+#         print("checking landsmarks_on_image:",face.landmarks_on_image)
+#         if face.landmarks_on_image is not None:
+#             for landmark in face.landmarks_on_image:
+#                 # Check if the landmark is part of the mouth, if exclude_mouth is True
+#                 if exclude_mouth and is_mouth_landmark(landmark):
+#                     continue  # Skip mouth landmarks
+
+#                 # Draw a circle for each landmark
+#                 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 process_face_image(
         face: FaceArea,
+        exclude_bottom_half: bool = False,  # New parameter to control exclusion of bottom half
         **kwargs,
     ) -> Image:
         image = np.array(face.image)
         overlay = image.copy()
         color_iter = color_generator(colors)
+
+        # Draw a rectangle over the entire face
         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)
+
+        print("checking landmarks_on_image:", face.landmarks_on_image)
         if face.landmarks_on_image is not None:
+            # Determine the y-coordinate of the nose to define the exclusion boundary
+            nose_y = get_nose_y_coordinate(face.landmarks_on_image)
+
             for landmark in face.landmarks_on_image:
+                # Exclude everything below the nose if exclude_bottom_half is True
+                if exclude_bottom_half and int(landmark.y) >= nose_y:
+                    continue  # Skip landmarks in the bottom half
+
+                # Draw a circle for each landmark
                 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 get_nose_y_coordinate(landmarks):
+    """
+    Determine the y-coordinate of the nose landmark to define the boundary for exclusion.
+    This assumes that landmarks are provided and one of them represents the nose.
+    Adjust this function based on how your landmarks are structured and named.
+    """
+    nose_y = None
+    for landmark in landmarks:
+        if is_nose_landmark(landmark):  # Implement this function based on your landmarks
+            nose_y = int(landmark.y)
+            break
+    return nose_y if nose_y is not None else 0  # Default to 0 if nose isn't found
+
+def is_nose_landmark(landmark):
+    """
+    Determine if a given landmark represents the nose.
+    You'll need to adjust the logic here based on your specific landmark detection system.
+    """
+    # Placeholder condition; replace with your actual condition for identifying a nose landmark.
+    print("landmark",landmark)
+    return landmark.part == "Nose"  # Adjust based on your system
+
+def is_mouth_landmark(landmark):
+    # This function needs to be tailored to your specific landmark system
+    # Typically, you'd check if the landmark's index or name indicates it's part of the mouth
+    # For example:
+    # return landmark.part in ["Mouth_Lower_Lip", "Mouth_Upper_Lip"]  # Adjust based on your system
+    print("landmark",landmark)
+    return False  # Placeholder; implement your own logic here
+
 
 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")
@@ -174,3 +250,128 @@ def create_mask_from_bbox(
         mask_draw.rectangle(bbox, fill=255)
         masks.append(mask)
     return masks
+
+
+def mask_bottom_half_of_face(mask: np.ndarray, face_area: FaceArea) -> np.ndarray:
+    """
+    Modify the mask to cover only the bottom half of the face from the nose down.
+
+    Parameters:
+    - mask (np.ndarray): The original face mask.
+    - face_area (FaceArea): The FaceArea object containing the Rect with landmarks.
+
+    Returns:
+    - np.ndarray: The modified mask with only the bottom half of the face covered.
+    """
+    rect = face_area.face_area  # Extract the Rect from the FaceArea
+    if rect.landmarks and rect.landmarks.nose:
+        # Use the nose landmark to define the starting point of the bottom half
+        nose = rect.landmarks.nose
+
+        # Calculate the starting y-coordinate for the mask
+        y_start = int(nose.y * face_area.scale_factor)
+
+        # Calculate the height of the face to find the bottom
+        face_height = int(face_area.height * face_area.scale_factor)
+
+        # Set the top half of the face to 0 (unmasked) in the mask
+        mask[:y_start, :] = 0
+
+        # Optionally, if you want to ensure the mask only covers the area up to the chin:
+        # Determine the chin's y-coordinate (you might need a chin landmark or another method)
+        # y_chin = int(chin.y * face_area.scale_factor)
+        # mask[y_chin:, :] = 0  # Uncomment and adjust if you have a chin landmark or method
+
+    return mask
+
+def exclude_mouth_from_mask(mask: np.ndarray, face_area: FaceArea) -> np.ndarray:
+    """
+    Modify the mask to exclude the mouth region based on the provided landmarks in FaceArea.
+
+    Parameters:
+    - mask (np.ndarray): The original face mask.
+    - face_area (FaceArea): The FaceArea object containing the Rect with landmarks.
+
+    Returns:
+    - np.ndarray: The modified mask with the mouth area excluded.
+    """
+    rect = face_area.face_area  # Extract the Rect from the FaceArea
+    if rect.landmarks and rect.landmarks.mouth1 and rect.landmarks.mouth2:
+        # Use the mouth landmarks to define the exclusion area
+        mouth1, mouth2 = rect.landmarks.mouth1, rect.landmarks.mouth2
+
+        # Calculate the bounding box for the mouth
+        x_min, y_min = min(mouth1.x, mouth2.x), min(mouth1.y, mouth2.y)
+        x_max, y_max = max(mouth1.x, mouth2.x), max(mouth1.y, mouth2.y)
+
+        # Adjust for the scale and position of the face in the entire image
+        x_min, y_min, x_max, y_max = [int(val * face_area.scale_factor) for val in [x_min, y_min, x_max, y_max]]
+
+        # Set the mouth region to 0 (black) in the mask
+        mask[y_min:y_max, x_min:x_max] = 0
+
+    return mask
+
+
+
+
+def apply_face_mask_with_exclusion(swapped_image: np.ndarray, target_image: np.ndarray, target_face: Face, entire_mask_image: np.ndarray) -> np.ndarray:
+    """
+    Apply the face mask with an exclusion zone for the mouth.
+
+    Parameters:
+    - swapped_image (np.ndarray): The image with the swapped face.
+    - target_image (np.ndarray): The target image where the face will be placed.
+    - target_face: Face with bbox The bounding box of the target face.
+    - entire_mask_image (np.ndarray): The initial entire mask image.
+
+    Returns:
+    - np.ndarray: The result image with the face swapped, excluding the mouth region.
+    """
+
+    logger.status("mask_bottom_half_of_face Mouth Mask")
+
+
+    # Extract the bbox array from the target_face object
+    target_face_bbox = target_face.bbox if hasattr(target_face, 'bbox') else None
+
+    if target_face_bbox is None:
+        logger.error("No bounding box found in the target face object.")
+        return target_image  # or handle this scenario appropriately
+
+    # Now you can safely create a Rect object from the bbox array
+    rect = Rect.from_ndarray(np.array(target_face_bbox))
+
+    face = FaceArea(target_image, rect, 1.6, 512, "")
+    face_image = np.array(face.image)
+    process_face_image(face,exclude_bottom_half=True)
+    face_area_on_image = face.face_area_on_image
+# Then call the generate_mask method with all required arguments
+    mask_generator =  BiSeNetMaskGenerator()
+    mask = mask_generator.generate_mask(
+        face_image=face_image,  # make sure this is the first required positional argument
+        face_area_on_image=face_area_on_image,
+        affected_areas=["Face"],
+        mask_size=0,
+        use_minimal_area=True
+    )
+    mask = cv2.blur(mask, (12, 12))
+
+    # Modify the mask to exclude the mouth using the FaceArea object.
+    # mask = exclude_mouth_from_mask(mask, face)
+    mask = mask_bottom_half_of_face(mask, face)
+
+    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)
+
+
+
+
+
+# Use apply_face_mask_with_exclusion instead of apply_face_mask when performing the face swap

scripts/reactor_entities/face.py

@@ -1,41 +1,43 @@
 import traceback
-
 import cv2
 import numpy as np
 from modules import images
 from PIL import Image
-
-
 from scripts.reactor_entities.rect import Point, Rect
 
-
 class FaceArea:
     def __init__(self, entire_image: np.ndarray, face_area: Rect, face_margin: float, face_size: int, upscaler: str):
+        # Initialize the FaceArea with the entire image, a specified face area (Rect object), margin, size, and upscaler.
         self.face_area = face_area
-        self.center = face_area.center
-        left, top, right, bottom = face_area.to_square()
+        self.center = face_area.center  # Center point of the face area.
+        left, top, right, bottom = face_area.to_square()  # Convert face area to a square for uniformity.
 
+        # Ensure the face area has a margin around it for context.
         self.left, self.top, self.right, self.bottom = self.__ensure_margin(
             left, top, right, bottom, entire_image, face_margin
         )
 
+        # Dimensions of the face area.
         self.width = self.right - self.left
         self.height = self.bottom - self.top
 
+        # Crop and possibly upscale the face image from the entire image.
         self.image = self.__crop_face_image(entire_image, face_size, upscaler)
-        self.face_size = face_size
-        self.scale_factor = face_size / self.width
-        self.face_area_on_image = self.__get_face_area_on_image()
-        self.landmarks_on_image = self.__get_landmarks_on_image()
+        self.face_size = face_size  # Desired face size.
+        self.scale_factor = face_size / self.width  # Scaling factor for resizing.
+        self.face_area_on_image = self.__get_face_area_on_image()  # Actual face area on the cropped image.
+        self.landmarks_on_image = self.__get_landmarks_on_image()  # Facial landmarks on the image.
 
     def __get_face_area_on_image(self):
+        # Calculate the face area on the cropped and possibly upscaled image.
         left = int((self.face_area.left - self.left) * self.scale_factor)
         top = int((self.face_area.top - self.top) * self.scale_factor)
         right = int((self.face_area.right - self.left) * self.scale_factor)
         bottom = int((self.face_area.bottom - self.top) * self.scale_factor)
         return self.__clip_values(left, top, right, bottom)
 
     def __get_landmarks_on_image(self):
+        # Adjust the landmarks' positions based on the cropped and possibly upscaled image.
         landmarks = []
         if self.face_area.landmarks is not None:
             for landmark in self.face_area.landmarks:
@@ -48,24 +50,28 @@ def __get_landmarks_on_image(self):
         return landmarks
 
     def __crop_face_image(self, entire_image: np.ndarray, face_size: int, upscaler: str):
+        # Crop the face image from the entire image and resize it according to the desired face size.
         cropped = entire_image[self.top : self.bottom, self.left : self.right, :]
         if upscaler:
             return images.resize_image(0, Image.fromarray(cropped), face_size, face_size, upscaler)
         else:
             return Image.fromarray(cv2.resize(cropped, dsize=(face_size, face_size)))
 
     def __ensure_margin(self, left: int, top: int, right: int, bottom: int, entire_image: np.ndarray, margin: float):
+        # Ensure there's a margin around the face area by expanding it proportionally.
         entire_height, entire_width = entire_image.shape[:2]
 
         side_length = right - left
         margin = min(min(entire_height, entire_width) / side_length, margin)
         diff = int((side_length * margin - side_length) / 2)
 
+        # Adjust the face area with the margin and ensure it doesn't go out of image bounds.
         top = top - diff
         bottom = bottom + diff
         left = left - diff
         right = right + diff
 
+        # Correct positions if they go out of the image boundaries.
         if top < 0:
             bottom = bottom - top
             top = 0
@@ -83,6 +89,7 @@ def __ensure_margin(self, left: int, top: int, right: int, bottom: int, entire_i
         return left, top, right, bottom
 
     def get_angle(self) -> float:
+        # Calculate the angle of the face based on the eye positions.
         landmarks = getattr(self.face_area, "landmarks", None)
         if landmarks is None:
             return 0
@@ -93,12 +100,14 @@ def get_angle(self) -> float:
             return 0
 
         try:
+            # Calculate angle between the eyes.
             dx = eye2.x - eye1.x
             dy = eye2.y - eye1.y
             if dx == 0:
                 dx = 1
             angle = np.arctan(dy / dx) * 180 / np.pi
 
+            # Adjust angle based on the quadrant.
             if dx < 0:
                 angle = (angle + 180) % 360
             return angle
@@ -107,16 +116,19 @@ def get_angle(self) -> float:
             return 0
 
     def rotate_face_area_on_image(self, angle: float):
+        # Rotate the face area on the image based on the given angle.
         center = [
             (self.face_area_on_image[0] + self.face_area_on_image[2]) / 2,
             (self.face_area_on_image[1] + self.face_area_on_image[3]) / 2,
         ]
 
+        # Define points to represent the face area.
         points = [
             [self.face_area_on_image[0], self.face_area_on_image[1]],
             [self.face_area_on_image[2], self.face_area_on_image[3]],
         ]
 
+        # Calculate rotation matrix and apply it to the points.
         angle = np.radians(angle)
         rot_matrix = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]])
 
@@ -125,6 +137,7 @@ def rotate_face_area_on_image(self, angle: float):
         points += center
         left, top, right, bottom = (int(points[0][0]), int(points[0][1]), int(points[1][0]), int(points[1][1]))
 
+        # Adjust the face area based on the rotation.
         left, right = (right, left) if left > right else (left, right)
         top, bottom = (bottom, top) if top > bottom else (top, bottom)
 
@@ -136,6 +149,7 @@ def rotate_face_area_on_image(self, angle: float):
         return self.__clip_values(left, top, right, bottom)
 
     def __clip_values(self, *args):
+        # Ensure that the values don't go beyond the specified face size.
         result = []
         for val in args:
             if val < 0: