diff --git a/modules/hair_segmenter.py b/modules/hair_segmenter.py new file mode 100644 index 0000000..3f7daac --- /dev/null +++ b/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.") diff --git a/modules/processors/frame/face_swapper.py b/modules/processors/frame/face_swapper.py index 36b83d6..e101cb6 100644 --- a/modules/processors/frame/face_swapper.py +++ b/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( - swapped_frame, mouth_cutout, mouth_box, face_mask, lower_lip_polygon + # Apply to final_swapped_frame if hair blending happened, otherwise to swapped_frame + 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( - swapped_frame, target_face, mouth_mask_data + final_swapped_frame = draw_mouth_mask_visualization( + 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: - temp_frame = swap_face(source_face, target_face, temp_frame) + if source_face_obj and target_face: + 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: - temp_frame = swap_face(source_face, target_face, temp_frame) + if target_face and source_face_obj: + 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 - -def process_frame_v2(temp_frame: Frame, temp_frame_path: str = "") -> Frame: +# process_frame_v2 needs to accept source_frame_full as well +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() - for map in modules.globals.source_target_map: - target_face = map["target"]["face"] - temp_frame = swap_face(source_face, target_face, temp_frame) + source_face_obj = default_source_face() # This function needs to be checked if it needs source_frame_full + if source_face_obj: # Ensure default_source_face actually returns a face + for map_item in modules.globals.source_target_map: # Renamed map to map_item to avoid conflict + 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: - if "source" in map: - source_face = map["source"]["face"] - target_face = map["target"]["face"] - temp_frame = swap_face(source_face, target_face, temp_frame) + for map_item in modules.globals.source_target_map: # Renamed map to map_item + if "source" in map_item: + source_face_obj = map_item["source"]["face"] + target_face = map_item["target"]["face"] + 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() - for map in modules.globals.source_target_map: - target_frame = [ - f - 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 = [ + source_face_obj = default_source_face() # This function needs to be checked + if source_face_obj: + for map_item in modules.globals.source_target_map: # Renamed map to map_item + target_frames_data = [ # Renamed target_frame to target_frames_data 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 target_face in frame["faces"]: - temp_frame = swap_face(source_face, target_face, temp_frame) + 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: + 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() - for target_face in detected_faces: - temp_frame = swap_face(source_face, target_face, temp_frame) + source_face_obj = default_source_face() # This function needs to be checked + if source_face_obj: + 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"][ - closest_centroid_index - ], - detected_face, + source_face_obj_from_map, # This is source_face_obj + detected_face, # This is target_face + source_frame_full, # This is source_frame_full 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], - detected_faces[closest_centroid_index], + source_face_obj_from_map, # source_face_obj + 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)) - target_frame = cv2.imread(target_path) - result = process_frame(source_face, target_frame) + 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 + 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) - result = process_frame_v2(target_frame) + # Pass source_img (as source_frame_full) to process_frame_v2 + 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 diff --git a/modules/ui.py b/modules/ui.py index ce599d6..53eeef2 100644 --- a/modules/ui.py +++ b/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: - source_image = get_one_face(cv2.imread(modules.globals.source_path)) + # Case 1: map_faces is False + 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() diff --git a/requirements.txt b/requirements.txt index 6d9f8b8..7611804 100644 --- a/requirements.txt +++ b/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