Merge 37486f03e7
into 28109e93bb
commit
615c9f05a9
|
@ -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.")
|
|
@ -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
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue