rehanbgmi 2025-05-21 18:50:06 +00:00 committed by GitHub
commit 615c9f05a9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 449 additions and 75 deletions

View File

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

View File

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

View File

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

View File

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