Merge 37486f03e7 into 28109e93bb

modules/hair_segmenter.py

@@ -0,0 +1,81 @@
+import torch
+import numpy as np
+from PIL import Image
+from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+import cv2 # Imported for BGR to RGB conversion, though PIL can also do it.
+
+def segment_hair(image_np: np.ndarray) -> np.ndarray:
+    """
+    Segments hair from an image.
+
+    Args:
+        image_np: NumPy array representing the image (BGR format from OpenCV).
+
+    Returns:
+        NumPy array representing the binary hair mask.
+    """
+    processor = SegformerImageProcessor.from_pretrained("isjackwild/segformer-b0-finetuned-segments-skin-hair-clothing")
+    model = SegformerForSemanticSegmentation.from_pretrained("isjackwild/segformer-b0-finetuned-segments-skin-hair-clothing")
+
+    # Convert BGR (OpenCV) to RGB (PIL)
+    image_rgb = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
+    image_pil = Image.fromarray(image_rgb)
+
+    inputs = processor(images=image_pil, return_tensors="pt")
+    outputs = model(**inputs)
+    logits = outputs.logits  # Shape: batch_size, num_labels, height, width
+
+    # Upsample logits to original image size
+    upsampled_logits = torch.nn.functional.interpolate(
+        logits,
+        size=(image_np.shape[0], image_np.shape[1]), # H, W
+        mode='bilinear',
+        align_corners=False
+    )
+
+    segmentation_map = upsampled_logits.argmax(dim=1).squeeze().cpu().numpy()
+
+    # Label 2 is for hair in this model
+    hair_mask = np.where(segmentation_map == 2, 255, 0).astype(np.uint8)
+
+    return hair_mask
+
+if __name__ == '__main__':
+    # This is a conceptual test.
+    # In a real scenario, you would load an image using OpenCV or Pillow.
+    # For example:
+    # sample_image_np = cv2.imread("path/to/your/image.jpg")
+    # if sample_image_np is not None:
+    #     hair_mask_output = segment_hair(sample_image_np)
+    #     cv2.imwrite("hair_mask_output.png", hair_mask_output)
+    #     print("Hair mask saved to hair_mask_output.png")
+    # else:
+    #     print("Failed to load sample image.")
+
+    print("Conceptual test: Hair segmenter module created.")
+    # Create a dummy image for a basic test run if no image is available.
+    dummy_image_np = np.zeros((100, 100, 3), dtype=np.uint8) # 100x100 BGR image
+    dummy_image_np[:, :, 1] = 255 # Make it green to distinguish from black mask
+    
+    try:
+        print("Running segment_hair with a dummy image...")
+        hair_mask_output = segment_hair(dummy_image_np)
+        print(f"segment_hair returned a mask of shape: {hair_mask_output.shape}")
+        # Check if the output is a 2D array (mask) and has the same H, W as input
+        assert hair_mask_output.shape == (dummy_image_np.shape[0], dummy_image_np.shape[1])
+        # Check if the mask is binary (0 or 255)
+        assert np.all(np.isin(hair_mask_output, [0, 255]))
+        print("Dummy image test successful. Hair mask seems to be generated correctly.")
+        
+        # Attempt to save the dummy mask (optional, just for visual confirmation if needed)
+        # cv2.imwrite("dummy_hair_mask_output.png", hair_mask_output)
+        # print("Dummy hair mask saved to dummy_hair_mask_output.png")
+
+    except ImportError as e:
+        print(f"An ImportError occurred: {e}. This might be due to missing dependencies like transformers, torch, or Pillow.")
+        print("Please ensure all required packages are installed by updating requirements.txt and installing them.")
+    except Exception as e:
+        print(f"An error occurred during the dummy image test: {e}")
+        print("This could be due to issues with model loading, processing, or other runtime errors.")
+
+    print("To perform a full test, replace the dummy image with a real image path.")

modules/processors/frame/face_swapper.py

@@ -9,6 +9,7 @@ import modules.processors.frame.core
 from modules.core import update_status
 from modules.face_analyser import get_one_face, get_many_faces, default_source_face
 from modules.typing import Face, Frame
+from modules.hair_segmenter import segment_hair
 from modules.utilities import (
     conditional_download,
     is_image,
@@ -67,14 +68,93 @@ def get_face_swapper() -> Any:
     return FACE_SWAPPER
 
 
-def swap_face(source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
+def swap_face(source_face_obj: Face, target_face: Face, source_frame_full: Frame, temp_frame: Frame) -> Frame:
     face_swapper = get_face_swapper()
 
     # Apply the face swap
     swapped_frame = face_swapper.get(
-        temp_frame, target_face, source_face, paste_back=True
+        temp_frame, target_face, source_face_obj, paste_back=True
     )
 
+    final_swapped_frame = swapped_frame.copy() # Initialize final_swapped_frame
+
+    # START of Hair Blending Logic
+    if source_face_obj.kps is not None and target_face.kps is not None and source_face_obj.kps.shape[0] >=2 and target_face.kps.shape[0] >=2 : # kps are 5x2 landmarks
+        hair_only_mask_source = segment_hair(source_frame_full)
+
+        # Ensure kps are float32 for estimateAffinePartial2D
+        source_kps_float = source_face_obj.kps.astype(np.float32)
+        target_kps_float = target_face.kps.astype(np.float32)
+
+        # b. Estimate Transformation Matrix
+        # Using LMEDS for robustness
+        matrix, _ = cv2.estimateAffinePartial2D(source_kps_float, target_kps_float, method=cv2.LMEDS)
+
+        if matrix is not None:
+            # c. Warp Source Hair and its Mask
+            dsize = (temp_frame.shape[1], temp_frame.shape[0]) # width, height
+
+            # Ensure hair_only_mask_source is 8-bit single channel
+            if hair_only_mask_source.ndim == 3 and hair_only_mask_source.shape[2] == 3:
+                hair_only_mask_source_gray = cv2.cvtColor(hair_only_mask_source, cv2.COLOR_BGR2GRAY)
+            else:
+                hair_only_mask_source_gray = hair_only_mask_source
+            
+            # Threshold to ensure binary mask for warping
+            _, hair_only_mask_source_binary = cv2.threshold(hair_only_mask_source_gray, 127, 255, cv2.THRESH_BINARY)
+
+            warped_hair_mask = cv2.warpAffine(hair_only_mask_source_binary, matrix, dsize)
+            warped_source_hair_image = cv2.warpAffine(source_frame_full, matrix, dsize)
+
+            # d. Color Correct Warped Source Hair
+            # Using swapped_frame (face-swapped output) as the target for color correction
+            color_corrected_warped_hair = apply_color_transfer(warped_source_hair_image, swapped_frame)
+
+            # e. Blend Hair onto Swapped Frame
+            # Ensure warped_hair_mask is binary (0 or 255) after warping
+            _, warped_hair_mask_binary = cv2.threshold(warped_hair_mask, 127, 255, cv2.THRESH_BINARY)
+
+            # Preferred: cv2.seamlessClone
+            x, y, w, h = cv2.boundingRect(warped_hair_mask_binary)
+            if w > 0 and h > 0:
+                center = (x + w // 2, y + h // 2)
+                # seamlessClone expects target image, source image, mask, center, method
+                # The mask should be single channel 8-bit.
+                # The source (color_corrected_warped_hair) and target (swapped_frame) should be 8-bit 3-channel.
+                
+                # Check if swapped_frame is suitable for seamlessClone (it should be the base)
+                # Ensure color_corrected_warped_hair is also 8UC3
+                if color_corrected_warped_hair.shape == swapped_frame.shape and \
+                   color_corrected_warped_hair.dtype == swapped_frame.dtype and \
+                   warped_hair_mask_binary.dtype == np.uint8:
+                    try:
+                        final_swapped_frame = cv2.seamlessClone(color_corrected_warped_hair, swapped_frame, warped_hair_mask_binary, center, cv2.NORMAL_CLONE)
+                    except cv2.error as e:
+                        logging.warning(f"cv2.seamlessClone failed: {e}. Falling back to simple blending.")
+                        # Fallback: Simple Blending (if seamlessClone fails)
+                        warped_hair_mask_3ch = cv2.cvtColor(warped_hair_mask_binary, cv2.COLOR_GRAY2BGR) > 0 # boolean mask
+                        final_swapped_frame[warped_hair_mask_3ch] = color_corrected_warped_hair[warped_hair_mask_3ch]
+                else:
+                    logging.warning("Mismatch in shape/type for seamlessClone. Falling back to simple blending.")
+                    # Fallback: Simple Blending
+                    warped_hair_mask_3ch = cv2.cvtColor(warped_hair_mask_binary, cv2.COLOR_GRAY2BGR) > 0
+                    final_swapped_frame[warped_hair_mask_3ch] = color_corrected_warped_hair[warped_hair_mask_3ch]
+            else:
+                # Mask is empty, no hair to blend, final_swapped_frame remains as is (copy of swapped_frame)
+                logging.info("Warped hair mask is empty. Skipping hair blending.")
+                # final_swapped_frame is already a copy of swapped_frame
+        else:
+            logging.warning("Failed to estimate affine transformation matrix for hair. Skipping hair blending.")
+            # final_swapped_frame is already a copy of swapped_frame
+    else:
+        if source_face_obj.kps is None or target_face.kps is None:
+            logging.warning("Source or target keypoints (kps) are None. Skipping hair blending.")
+        else:
+            logging.warning(f"Not enough keypoints for hair transformation. Source kps: {source_face_obj.kps.shape if source_face_obj.kps is not None else 'None'}, Target kps: {target_face.kps.shape if target_face.kps is not None else 'None'}. Skipping hair blending.")
+        # final_swapped_frame is already a copy of swapped_frame
+    # END of Hair Blending Logic
+
+    # f. Mouth Mask Logic
     if modules.globals.mouth_mask:
         # Create a mask for the target face
         face_mask = create_face_mask(target_face, temp_frame)
@@ -85,20 +165,21 @@ def swap_face(source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
         )
 
         # Apply the mouth area
-        swapped_frame = apply_mouth_area(
+        # Apply to final_swapped_frame if hair blending happened, otherwise to swapped_frame
-            swapped_frame, mouth_cutout, mouth_box, face_mask, lower_lip_polygon
+        final_swapped_frame = apply_mouth_area(
+            final_swapped_frame, mouth_cutout, mouth_box, face_mask, lower_lip_polygon
         )
 
         if modules.globals.show_mouth_mask_box:
             mouth_mask_data = (mouth_mask, mouth_cutout, mouth_box, lower_lip_polygon)
-            swapped_frame = draw_mouth_mask_visualization(
+            final_swapped_frame = draw_mouth_mask_visualization(
-                swapped_frame, target_face, mouth_mask_data
+                final_swapped_frame, target_face, mouth_mask_data
             )
 
-    return swapped_frame
+    return final_swapped_frame
 
 
-def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
+def process_frame(source_face_obj: Face, source_frame_full: Frame, temp_frame: Frame) -> Frame:
     if modules.globals.color_correction:
         temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)
 
@@ -106,70 +187,73 @@ def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
         many_faces = get_many_faces(temp_frame)
         if many_faces:
             for target_face in many_faces:
-                if source_face and target_face:
+                if source_face_obj and target_face:
-                    temp_frame = swap_face(source_face, target_face, temp_frame)
+                    temp_frame = swap_face(source_face_obj, target_face, source_frame_full, temp_frame)
                 else:
                     print("Face detection failed for target/source.")
     else:
         target_face = get_one_face(temp_frame)
-        if target_face and source_face:
+        if target_face and source_face_obj:
-            temp_frame = swap_face(source_face, target_face, temp_frame)
+            temp_frame = swap_face(source_face_obj, target_face, source_frame_full, temp_frame)
         else:
             logging.error("Face detection failed for target or source.")
     return temp_frame
 
 
-
+# process_frame_v2 needs to accept source_frame_full as well
-def process_frame_v2(temp_frame: Frame, temp_frame_path: str = "") -> Frame:
+def process_frame_v2(source_frame_full: Frame, temp_frame: Frame, temp_frame_path: str = "") -> Frame:
     if is_image(modules.globals.target_path):
         if modules.globals.many_faces:
-            source_face = default_source_face()
+            source_face_obj = default_source_face() # This function needs to be checked if it needs source_frame_full
-            for map in modules.globals.source_target_map:
+            if source_face_obj: # Ensure default_source_face actually returns a face
-                target_face = map["target"]["face"]
+                 for map_item in modules.globals.source_target_map: # Renamed map to map_item to avoid conflict
-                temp_frame = swap_face(source_face, target_face, temp_frame)
+                    target_face = map_item["target"]["face"]
+                    temp_frame = swap_face(source_face_obj, target_face, source_frame_full, temp_frame)
 
         elif not modules.globals.many_faces:
-            for map in modules.globals.source_target_map:
+            for map_item in modules.globals.source_target_map: # Renamed map to map_item
-                if "source" in map:
+                if "source" in map_item:
-                    source_face = map["source"]["face"]
+                    source_face_obj = map_item["source"]["face"]
-                    target_face = map["target"]["face"]
+                    target_face = map_item["target"]["face"]
-                    temp_frame = swap_face(source_face, target_face, temp_frame)
+                    temp_frame = swap_face(source_face_obj, target_face, source_frame_full, temp_frame)
 
     elif is_video(modules.globals.target_path):
         if modules.globals.many_faces:
-            source_face = default_source_face()
+            source_face_obj = default_source_face() # This function needs to be checked
-            for map in modules.globals.source_target_map:
+            if source_face_obj:
-                target_frame = [
+                for map_item in modules.globals.source_target_map: # Renamed map to map_item
-                    f
+                    target_frames_data = [ # Renamed target_frame to target_frames_data
-                    for f in map["target_faces_in_frame"]
-                    if f["location"] == temp_frame_path
-                ]
-
-                for frame in target_frame:
-                    for target_face in frame["faces"]:
-                        temp_frame = swap_face(source_face, target_face, temp_frame)
-
-        elif not modules.globals.many_faces:
-            for map in modules.globals.source_target_map:
-                if "source" in map:
-                    target_frame = [
                         f
-                        for f in map["target_faces_in_frame"]
+                        for f in map_item["target_faces_in_frame"]
                         if f["location"] == temp_frame_path
                     ]
-                    source_face = map["source"]["face"]
 
-                    for frame in target_frame:
+                    for frame_data in target_frames_data: # Renamed frame to frame_data
-                        for target_face in frame["faces"]:
+                        for target_face in frame_data["faces"]:
-                            temp_frame = swap_face(source_face, target_face, temp_frame)
+                            temp_frame = swap_face(source_face_obj, target_face, source_frame_full, temp_frame)
 
-    else:
+        elif not modules.globals.many_faces:
+            for map_item in modules.globals.source_target_map: # Renamed map to map_item
+                if "source" in map_item:
+                    target_frames_data = [ # Renamed target_frame to target_frames_data
+                        f
+                        for f in map_item["target_faces_in_frame"]
+                        if f["location"] == temp_frame_path
+                    ]
+                    source_face_obj = map_item["source"]["face"]
+
+                    for frame_data in target_frames_data: # Renamed frame to frame_data
+                        for target_face in frame_data["faces"]:
+                            temp_frame = swap_face(source_face_obj, target_face, source_frame_full, temp_frame)
+
+    else: # This is the live cam / generic case
         detected_faces = get_many_faces(temp_frame)
         if modules.globals.many_faces:
             if detected_faces:
-                source_face = default_source_face()
+                source_face_obj = default_source_face() # This function needs to be checked
-                for target_face in detected_faces:
+                if source_face_obj:
-                    temp_frame = swap_face(source_face, target_face, temp_frame)
+                    for target_face in detected_faces:
+                        temp_frame = swap_face(source_face_obj, target_face, source_frame_full, temp_frame)
 
         elif not modules.globals.many_faces:
             if detected_faces:
@@ -181,12 +265,13 @@ def process_frame_v2(temp_frame: Frame, temp_frame_path: str = "") -> Frame:
                             modules.globals.simple_map["target_embeddings"],
                             detected_face.normed_embedding,
                         )
-
+                        # Assuming simple_map["source_faces"] are Face objects
+                        # And default_source_face() logic might need to be more complex if source_frame_full is always from a single source_path
+                        source_face_obj_from_map = modules.globals.simple_map["source_faces"][closest_centroid_index]
                         temp_frame = swap_face(
-                            modules.globals.simple_map["source_faces"][
+                            source_face_obj_from_map, # This is source_face_obj
-                                closest_centroid_index
+                            detected_face,           # This is target_face
-                            ],
+                            source_frame_full,       # This is source_frame_full
-                            detected_face,
                             temp_frame,
                         )
                 else:
@@ -200,10 +285,11 @@ def process_frame_v2(temp_frame: Frame, temp_frame_path: str = "") -> Frame:
                         closest_centroid_index, _ = find_closest_centroid(
                             detected_faces_centroids, target_embedding
                         )
-
+                        source_face_obj_from_map = modules.globals.simple_map["source_faces"][i]
                         temp_frame = swap_face(
-                            modules.globals.simple_map["source_faces"][i],
+                            source_face_obj_from_map, # source_face_obj
-                            detected_faces[closest_centroid_index],
+                            detected_faces[closest_centroid_index], # target_face
+                            source_frame_full, # source_frame_full
                             temp_frame,
                         )
                         i += 1
@@ -213,44 +299,83 @@ def process_frame_v2(temp_frame: Frame, temp_frame_path: str = "") -> Frame:
 def process_frames(
     source_path: str, temp_frame_paths: List[str], progress: Any = None
 ) -> None:
+    source_img = cv2.imread(source_path)
+    if source_img is None:
+        logging.error(f"Failed to read source image from {source_path}")
+        return
+
     if not modules.globals.map_faces:
-        source_face = get_one_face(cv2.imread(source_path))
+        source_face_obj = get_one_face(source_img) # Use source_img here
+        if not source_face_obj:
+            logging.error(f"No face detected in source image {source_path}")
+            return
         for temp_frame_path in temp_frame_paths:
             temp_frame = cv2.imread(temp_frame_path)
+            if temp_frame is None:
+                logging.warning(f"Failed to read temp_frame from {temp_frame_path}, skipping.")
+                continue
             try:
-                result = process_frame(source_face, temp_frame)
+                result = process_frame(source_face_obj, source_img, temp_frame)
                 cv2.imwrite(temp_frame_path, result)
             except Exception as exception:
-                print(exception)
+                logging.error(f"Error processing frame {temp_frame_path}: {exception}", exc_info=True)
                 pass
             if progress:
                 progress.update(1)
-    else:
+    else: # This is for map_faces == True
+        # In map_faces=True, source_face is determined per mapping.
+        # process_frame_v2 will need source_frame_full for hair,
+        # which should be the original source_path image.
         for temp_frame_path in temp_frame_paths:
             temp_frame = cv2.imread(temp_frame_path)
+            if temp_frame is None:
+                logging.warning(f"Failed to read temp_frame from {temp_frame_path}, skipping.")
+                continue
             try:
-                result = process_frame_v2(temp_frame, temp_frame_path)
+                # Pass source_img (as source_frame_full) to process_frame_v2
+                result = process_frame_v2(source_img, temp_frame, temp_frame_path)
                 cv2.imwrite(temp_frame_path, result)
             except Exception as exception:
-                print(exception)
+                logging.error(f"Error processing frame {temp_frame_path} with map_faces: {exception}", exc_info=True)
                 pass
             if progress:
                 progress.update(1)
 
 
 def process_image(source_path: str, target_path: str, output_path: str) -> None:
+    source_img = cv2.imread(source_path)
+    if source_img is None:
+        logging.error(f"Failed to read source image from {source_path}")
+        return
+
+    target_frame = cv2.imread(target_path)
+    if target_frame is None:
+        logging.error(f"Failed to read target image from {target_path}")
+        return
+
     if not modules.globals.map_faces:
-        source_face = get_one_face(cv2.imread(source_path))
+        source_face_obj = get_one_face(source_img) # Use source_img here
-        target_frame = cv2.imread(target_path)
+        if not source_face_obj:
-        result = process_frame(source_face, target_frame)
+            logging.error(f"No face detected in source image {source_path}")
+            return
+        result = process_frame(source_face_obj, source_img, target_frame)
         cv2.imwrite(output_path, result)
     else:
+        # map_faces == True for process_image
+        # process_frame_v2 expects source_frame_full as its first argument.
+        # The output_path is often the same as target_path initially for images.
+        # We read the target_frame (which will be modified)
+        target_frame_for_v2 = cv2.imread(output_path) # Or target_path, depending on desired workflow
+        if target_frame_for_v2 is None:
+            logging.error(f"Failed to read image for process_frame_v2 from {output_path}")
+            return
+
         if modules.globals.many_faces:
             update_status(
                 "Many faces enabled. Using first source image. Progressing...", NAME
             )
-        target_frame = cv2.imread(output_path)
+        # Pass source_img (as source_frame_full) to process_frame_v2
-        result = process_frame_v2(target_frame)
+        result = process_frame_v2(source_img, target_frame_for_v2, target_path) # target_path as temp_frame_path hint
         cv2.imwrite(output_path, result)
 
 
@@ -620,3 +745,113 @@ def apply_color_transfer(source, target):
     source = (source - source_mean) * (target_std / source_std) + target_mean
 
     return cv2.cvtColor(np.clip(source, 0, 255).astype("uint8"), cv2.COLOR_LAB2BGR)
+
+
+def create_face_and_hair_mask(source_face: Face, source_frame: Frame) -> np.ndarray:
+    """
+    Creates a combined mask for the face and hair from the source image.
+    """
+    # 1. Generate the basic face mask (adapted from create_face_mask)
+    face_only_mask = np.zeros(source_frame.shape[:2], dtype=np.uint8)
+    landmarks = source_face.landmark_2d_106
+    if landmarks is not None:
+        landmarks = landmarks.astype(np.int32)
+
+        # Extract facial features (same logic as create_face_mask)
+        right_side_face = landmarks[0:16]
+        left_side_face = landmarks[17:32]
+        # right_eye = landmarks[33:42] # Not directly used for outline
+        right_eye_brow = landmarks[43:51]
+        # left_eye = landmarks[87:96] # Not directly used for outline
+        left_eye_brow = landmarks[97:105]
+
+        # Calculate forehead extension (same logic as create_face_mask)
+        right_eyebrow_top = np.min(right_eye_brow[:, 1])
+        left_eyebrow_top = np.min(left_eye_brow[:, 1])
+        eyebrow_top = min(right_eyebrow_top, left_eyebrow_top)
+
+        face_top = np.min([right_side_face[0, 1], left_side_face[-1, 1]])
+        # Ensure forehead_height is not negative if eyebrows are above the topmost landmark of face sides
+        forehead_height = max(0, face_top - eyebrow_top)
+        extended_forehead_height = int(forehead_height * 5.0)
+
+        forehead_left = right_side_face[0].copy()
+        forehead_right = left_side_face[-1].copy()
+        
+        # Ensure extended forehead points do not go into negative y values
+        forehead_left[1] = max(0, forehead_left[1] - extended_forehead_height)
+        forehead_right[1] = max(0, forehead_right[1] - extended_forehead_height)
+
+        face_outline = np.vstack(
+            [
+                [forehead_left],
+                right_side_face,
+                left_side_face[::-1],
+                [forehead_right],
+            ]
+        )
+
+        # Calculate padding (same logic as create_face_mask)
+        # Ensure face_outline has at least one point before calculating norm
+        if face_outline.shape[0] > 1:
+            padding = int(
+                np.linalg.norm(right_side_face[0] - left_side_face[-1]) * 0.05
+            )
+        else:
+            padding = 5 # Default padding if not enough points
+
+        hull = cv2.convexHull(face_outline)
+        hull_padded = []
+        center = np.mean(face_outline, axis=0).squeeze() # Squeeze to handle potential extra dim
+
+        # Ensure center is a 1D array for subtraction
+        if center.ndim > 1:
+            center = np.mean(center, axis=0)
+
+
+        for point_contour in hull:
+            point = point_contour[0] # cv2.convexHull returns points wrapped in an extra array
+            direction = point - center
+            norm_direction = np.linalg.norm(direction)
+            if norm_direction == 0: # Avoid division by zero if point is the center
+                unit_direction = np.array([0,0])
+            else:
+                unit_direction = direction / norm_direction
+            
+            padded_point = point + unit_direction * padding
+            hull_padded.append(padded_point)
+
+        if hull_padded: # Ensure hull_padded is not empty
+            hull_padded = np.array(hull_padded, dtype=np.int32)
+            cv2.fillConvexPoly(face_only_mask, hull_padded, 255)
+        else: # Fallback if hull_padded is empty (e.g. very few landmarks)
+            cv2.fillConvexPoly(face_only_mask, hull, 255) # Use unpadded hull
+
+
+        # Initial blur for face_only_mask is not strictly in the old one before combining,
+        # but can be applied here or after combining. Let's keep it like original for now.
+        # face_only_mask = cv2.GaussianBlur(face_only_mask, (5, 5), 3) # Original blur from create_face_mask
+
+    # 2. Generate the hair mask
+    # Ensure source_frame is contiguous, as some cv2 functions might require it.
+    source_frame_contiguous = np.ascontiguousarray(source_frame, dtype=np.uint8)
+    hair_mask_on_source = segment_hair(source_frame_contiguous)
+
+    # 3. Combine the masks
+    # Ensure masks are binary and of the same type for bitwise operations
+    _, face_only_mask_binary = cv2.threshold(face_only_mask, 127, 255, cv2.THRESH_BINARY)
+    _, hair_mask_on_source_binary = cv2.threshold(hair_mask_on_source, 127, 255, cv2.THRESH_BINARY)
+    
+    # Ensure shapes match. If not, hair_mask might be different. Resize if necessary.
+    # This should ideally not happen if segment_hair preserves dimensions.
+    if face_only_mask_binary.shape != hair_mask_on_source_binary.shape:
+        hair_mask_on_source_binary = cv2.resize(hair_mask_on_source_binary, 
+                                                (face_only_mask_binary.shape[1], face_only_mask_binary.shape[0]), 
+                                                interpolation=cv2.INTER_NEAREST)
+
+    combined_mask = cv2.bitwise_or(face_only_mask_binary, hair_mask_on_source_binary)
+
+    # 4. Apply Gaussian blur to the combined mask
+    combined_mask = cv2.GaussianBlur(combined_mask, (5, 5), 3)
+
+    return combined_mask

modules/ui.py

@@ -880,7 +880,7 @@ def create_webcam_preview(camera_index: int):
     PREVIEW.deiconify()
 
     frame_processors = get_frame_processors_modules(modules.globals.frame_processors)
-    source_image = None
+    # source_image = None # Replaced by source_face_obj_for_cam
     prev_time = time.time()
     fps_update_interval = 0.5
     frame_count = 0
@@ -907,23 +907,80 @@ def create_webcam_preview(camera_index: int):
             )
 
         if not modules.globals.map_faces:
-            if source_image is None and modules.globals.source_path:
+            # Case 1: map_faces is False
-                source_image = get_one_face(cv2.imread(modules.globals.source_path))
+            source_face_obj_for_cam = None
+            source_frame_full_for_cam = None
+            if modules.globals.source_path and os.path.exists(modules.globals.source_path):
+                source_frame_full_for_cam = cv2.imread(modules.globals.source_path)
+                if source_frame_full_for_cam is not None:
+                    source_face_obj_for_cam = get_one_face(source_frame_full_for_cam)
+                    if source_face_obj_for_cam is None:
+                        update_status(f"Error: No face detected in source image at {modules.globals.source_path}")
+                        # Optional: could return here or allow running without a source face if some processors handle it
+                else:
+                    update_status(f"Error: Could not read source image at {modules.globals.source_path}")
+                    cap.release()
+                    PREVIEW.withdraw()
+                    return 
+            elif modules.globals.source_path: 
+                update_status(f"Error: Source image not found at {modules.globals.source_path}")
+                cap.release()
+                PREVIEW.withdraw()
+                return
+            else: 
+                update_status("Error: No source image selected for webcam mode.")
+                cap.release()
+                PREVIEW.withdraw()
+                return
 
             for frame_processor in frame_processors:
                 if frame_processor.NAME == "DLC.FACE-ENHANCER":
                     if modules.globals.fp_ui["face_enhancer"]:
-                        temp_frame = frame_processor.process_frame(None, temp_frame)
+                        # Assuming face_enhancer's process_frame doesn't need source_face or source_frame_full
+                        temp_frame = frame_processor.process_frame(None, temp_frame) 
                 else:
-                    temp_frame = frame_processor.process_frame(source_image, temp_frame)
+                    if source_face_obj_for_cam and source_frame_full_for_cam is not None:
+                        temp_frame = frame_processor.process_frame(source_face_obj_for_cam, source_frame_full_for_cam, temp_frame)
+                    # else: temp_frame remains unchanged if source isn't ready
         else:
-            modules.globals.target_path = None
+            # Case 2: map_faces is True
+            source_frame_full_for_cam_map_faces = None
+            if modules.globals.source_path and os.path.exists(modules.globals.source_path):
+                source_frame_full_for_cam_map_faces = cv2.imread(modules.globals.source_path)
+                if source_frame_full_for_cam_map_faces is None:
+                    update_status(f"Error: Could not read source image (for hair/background) at {modules.globals.source_path}")
+                    cap.release()
+                    PREVIEW.withdraw()
+                    return
+            elif modules.globals.source_path:
+                 update_status(f"Error: Source image (for hair/background) not found at {modules.globals.source_path}")
+                 cap.release()
+                 PREVIEW.withdraw()
+                 return
+            else:
+                update_status("Error: No global source image selected (for hair/background in map_faces mode).")
+                cap.release()
+                PREVIEW.withdraw()
+                return
+
+            # Also check if map is defined, though process_frame_v2 handles specific face mapping internally
+            if not modules.globals.source_target_map and not modules.globals.simple_map: # Check both map types
+                update_status("Error: No face map defined for map_faces mode.")
+                # This might not need a return if some processors can run without map
+                # but for face_swapper, it's likely needed.
+                # For now, we proceed and let process_frame_v2 handle it.
+
+            modules.globals.target_path = None # Standard for live mode
             for frame_processor in frame_processors:
                 if frame_processor.NAME == "DLC.FACE-ENHANCER":
                     if modules.globals.fp_ui["face_enhancer"]:
-                        temp_frame = frame_processor.process_frame_v2(temp_frame)
+                        # Pass source_frame_full_for_cam_map_faces for signature consistency
+                        # The enhancer can choose to ignore it if not needed.
+                        temp_frame = frame_processor.process_frame_v2(source_frame_full_for_cam_map_faces, temp_frame)
                 else:
-                    temp_frame = frame_processor.process_frame_v2(temp_frame)
+                    if source_frame_full_for_cam_map_faces is not None:
+                        temp_frame = frame_processor.process_frame_v2(source_frame_full_for_cam_map_faces, temp_frame)
+                    # else: temp_frame remains unchanged if global source for map_faces isn't ready
 
         # Calculate and display FPS
         current_time = time.time()

requirements.txt

@@ -19,3 +19,4 @@ onnxruntime-gpu==1.17; sys_platform != 'darwin'
 tensorflow; sys_platform != 'darwin'
 opennsfw2==0.10.2
 protobuf==4.23.2
+transformers>=4.0.0