diff --git a/nodes.py b/nodes.py index b51f6e9..e2e5d54 100644 --- a/nodes.py +++ b/nodes.py @@ -167,6 +167,8 @@ class ApplyFaceProcessor: "face_crop_scale": ("FLOAT", {"default": 1.5, "min": 1.0, "max": 3.0, "step": 0.1}), "confidence_threshold": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.05}), "with_neck": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), + # NEW: Face Only Mode + "face_only_mode": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), } } @@ -175,7 +177,7 @@ class ApplyFaceProcessor: FUNCTION = "apply_processing" CATEGORY = "Stand-In" - def apply_processing(self, face_processor, image, resize_to, border_thresh, face_crop_scale, confidence_threshold, with_neck): + def apply_processing(self, face_processor, image, resize_to, border_thresh, face_crop_scale, confidence_threshold, with_neck, face_only_mode): detection_model, parsing_model, device = face_processor frame = tensor_to_cv2_img(image) @@ -218,17 +220,27 @@ class ApplyFaceProcessor: face_tensor = _img2tensor(image_resized, bgr2rgb=True).unsqueeze(0).to(device) with torch.no_grad(): - normalized_face = normalize(face_tensor, [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) + # The BiSeNet model in facexlib was trained on a different normalization + normalized_face = normalize(face_tensor, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) parsing_out = parsing_model(normalized_face)[0] parsing_map_tensor = parsing_out.argmax(dim=1, keepdim=True) parsing_map_np = parsing_map_tensor.squeeze().cpu().numpy().astype(np.uint8) - if with_neck: + # --- UPDATED: MASK SELECTION LOGIC --- + if face_only_mode: + # If face_only_mode is ON, create a precise mask for core facial features + # This will ignore the with_neck setting + parts_to_exclude = [0, 14, 15, 16, 17, 18] # bg, neck, cloth, hair etc. + final_mask_np = np.isin(parsing_map_np, parts_to_exclude, invert=True).astype(np.uint8) + elif with_neck: + # Standard mode: include neck and hair final_mask_np = (parsing_map_np != 0).astype(np.uint8) else: - parts_to_remove = [0, 14, 16] # Typically background, neck, clothes - final_mask_np = np.isin(parsing_map_np, parts_to_remove, invert=True).astype(np.uint8) + # Standard mode: exclude neck (original behavior) + parts_to_exclude = [0, 14, 15, 16, 18] + final_mask_np = np.isin(parsing_map_np, parts_to_exclude, invert=True).astype(np.uint8) + # --- END OF MASK SELECTION --- # --- Create the primary output (masked face on white background) --- white_background = np.ones_like(image_resized, dtype=np.uint8) * 255 @@ -237,18 +249,13 @@ class ApplyFaceProcessor: result_tensor = cv2_img_to_tensor(result_img_bgr) # --- Create the secondary RGBA face output --- - # 1. Get the RGB channels from the resized BGR image image_resized_rgb = cv2.cvtColor(image_resized, cv2.COLOR_BGR2RGB) - # 2. Create the alpha channel from the final mask alpha_channel = final_mask_np * 255 - # 3. Stack RGB and Alpha channels to create a 4-channel RGBA image face_rgba_np = np.dstack((image_resized_rgb, alpha_channel)) - # 4. Convert the RGBA numpy array to a tensor face_rgba_tensor = torch.from_numpy(face_rgba_np.astype(np.float32) / 255.0).unsqueeze(0) return (result_tensor, face_rgba_tensor,) - # --- MAPPINGS --- NODE_CLASS_MAPPINGS = { "FaceProcessorLoader": FaceProcessorLoader, diff --git a/videocombine.py b/videocombine.py index c0df75e..1a8a948 100644 --- a/videocombine.py +++ b/videocombine.py @@ -22,7 +22,7 @@ class VideoBackgroundRestorer: """ Analyzes a synthesized video to create a face mask and then uses this mask to composite the synthesized face onto the background of an original video. - Includes edge dilation and feathering for seamless blending. + Includes edge dilation, feathering, color matching, and a face-only mode. """ @classmethod def INPUT_TYPES(cls): @@ -32,10 +32,14 @@ class VideoBackgroundRestorer: "synth_images": ("IMAGE",), "orig_images": ("IMAGE",), "confidence_threshold": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.05}), - "face_crop_scale": ("FLOAT", {"default": 1.8, "min": 1.0, "max": 3.0, "step": 0.1}), - "dilation_kernel_size": ("INT", {"default": 10, "min": 0, "max": 50, "step": 1}), + "face_crop_scale": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 10.0, "step": 0.1}), + "dilation_kernel_size": ("INT", {"default": 25, "min": 0, "max": 50, "step": 1}), "feather_amount": ("INT", {"default": 50, "min": 0, "max": 151, "step": 2, "display": "slider"}), - "with_neck": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), + "with_neck": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), + "color_match_enabled": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), + "color_match_strength": ("FLOAT", {"default": 0.75, "min": 0.0, "max": 1.0, "step": 0.05}), + # NEW: Face Only Pasting Mode + "face_only_mode": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), } } @@ -44,7 +48,7 @@ class VideoBackgroundRestorer: FUNCTION = "restore_background" CATEGORY = "Stand-In" - def restore_background(self, face_processor, synth_images: torch.Tensor, orig_images: torch.Tensor, confidence_threshold: float, face_crop_scale: float, dilation_kernel_size: int, feather_amount: int, with_neck: bool): + def restore_background(self, face_processor, synth_images: torch.Tensor, orig_images: torch.Tensor, confidence_threshold: float, face_crop_scale: float, dilation_kernel_size: int, feather_amount: int, with_neck: bool, color_match_enabled: bool, color_match_strength: float, face_only_mode: bool): detection_model, parsing_model, device = face_processor if synth_images.shape != orig_images.shape: @@ -52,15 +56,19 @@ class VideoBackgroundRestorer: total_frames, h, w = synth_images.shape[0], synth_images.shape[1], synth_images.shape[2] - print(f"Processing {total_frames} frames ({w}x{h}) to restore background with edge feathering.") + print(f"Processing {total_frames} frames ({w}x{h}) to restore background.") processed_frames_tensors = [] + # Define parts to exclude for face-only skin analysis and pasting + parts_to_exclude_for_face_only = [0, 14, 15, 16, 17, 18] # bg, neck, cloth, hair etc. + for i in tqdm(range(total_frames), desc="Restoring video background"): synth_frame_tensor = synth_images[i] orig_frame_tensor = orig_images[i] synth_frame_bgr = tensor_to_cv2_img(synth_frame_tensor) + orig_frame_bgr = tensor_to_cv2_img(orig_frame_tensor) results = detection_model(synth_frame_bgr, verbose=False) confident_boxes = results[0].boxes.xyxy[results[0].boxes.conf > confidence_threshold] @@ -90,24 +98,59 @@ class VideoBackgroundRestorer: parsing_map = parsing_model(normalized_face)[0].argmax(dim=1, keepdim=True) parsing_map_np = parsing_map.squeeze().cpu().numpy().astype(np.uint8) - if with_neck: + + if color_match_enabled and color_match_strength > 0: + face_skin_mask_512 = np.isin(parsing_map_np, parts_to_exclude_for_face_only, invert=True).astype(np.uint8) + + if np.sum(face_skin_mask_512) > 0: + face_skin_mask_crop = cv2.resize(face_skin_mask_512, (face_crop_bgr.shape[1], face_crop_bgr.shape[0]), interpolation=cv2.INTER_NEAREST) + + orig_face_crop_bgr = orig_frame_bgr[crop_y1:crop_y2, crop_x1:crop_x2] + source_lab = cv2.cvtColor(orig_face_crop_bgr, cv2.COLOR_BGR2LAB) + target_lab = cv2.cvtColor(face_crop_bgr, cv2.COLOR_BGR2LAB) + + source_mean, source_std = cv2.meanStdDev(source_lab, mask=face_skin_mask_crop) + target_mean, target_std = cv2.meanStdDev(target_lab, mask=face_skin_mask_crop) + + l, a, b = cv2.split(target_lab) + + eps = 1e-6 + l = (l - target_mean[0][0]) * (source_std[0][0] / (target_std[0][0] + eps)) + source_mean[0][0] + a = (a - target_mean[1][0]) * (source_std[1][0] / (target_std[1][0] + eps)) + source_mean[1][0] + b = (b - target_mean[2][0]) * (source_std[2][0] / (target_std[2][0] + eps)) + source_mean[2][0] + + corrected_lab = cv2.merge([l, a, b]) + corrected_lab = np.clip(corrected_lab, 0, 255).astype(np.uint8) + + corrected_face_crop_bgr = cv2.cvtColor(corrected_lab, cv2.COLOR_LAB2BGR) + + face_crop_bgr = cv2.addWeighted(corrected_face_crop_bgr, color_match_strength, face_crop_bgr, 1 - color_match_strength, 0) + + synth_frame_bgr[crop_y1:crop_y2, crop_x1:crop_x2] = face_crop_bgr + + corrected_synth_rgb = cv2.cvtColor(synth_frame_bgr, cv2.COLOR_BGR2RGB) + synth_frame_tensor = torch.from_numpy(corrected_synth_rgb.astype(np.float32) / 255.0) + + # --- NEW: MASK SELECTION LOGIC --- + if face_only_mode: + # If face_only_mode is ON, use the precise skin mask and ignore with_neck setting + final_mask_512 = np.isin(parsing_map_np, parts_to_exclude_for_face_only, invert=True).astype(np.uint8) * 255 + elif with_neck: + # Standard mode: include neck and hair final_mask_512 = (parsing_map_np != 0).astype(np.uint8) * 255 else: - parts_to_exclude = [0, 14, 15, 16, 18] - final_mask_512 = np.isin(parsing_map_np, parts_to_exclude, invert=True).astype(np.uint8) * 255 + # Standard mode: exclude neck + parts_to_exclude_neck = [0, 14, 15, 16, 18] + final_mask_512 = np.isin(parsing_map_np, parts_to_exclude_neck, invert=True).astype(np.uint8) * 255 + # --- END OF MASK SELECTION --- - # --- EDGE PROCESSING LOGIC --- - # 1. DILATION (Expansion) if dilation_kernel_size > 0: kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilation_kernel_size, dilation_kernel_size)) final_mask_512 = cv2.dilate(final_mask_512, kernel, iterations=1) - # 2. FEATHERING (Softening) if feather_amount > 0: - # Ensure the kernel size is odd if feather_amount % 2 == 0: feather_amount += 1 - # Apply Gaussian Blur to soften the edges of the mask final_mask_512 = cv2.GaussianBlur(final_mask_512, (feather_amount, feather_amount), 0) mask_resized_to_crop = cv2.resize(final_mask_512, (face_crop_bgr.shape[1], face_crop_bgr.shape[0]), interpolation=cv2.INTER_LINEAR) @@ -121,12 +164,4 @@ class VideoBackgroundRestorer: output_image_batch = torch.stack(processed_frames_tensors).cpu() - return (output_image_batch,) - -NODE_CLASS_MAPPINGS = { - "VideoBackgroundRestorer": VideoBackgroundRestorer, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "VideoBackgroundRestorer": "Stand-In Background Restorer", -} \ No newline at end of file + return (output_image_batch,) \ No newline at end of file diff --git a/videoinput.py b/videoinput.py index 32d1d56..4ac67c6 100644 --- a/videoinput.py +++ b/videoinput.py @@ -4,6 +4,7 @@ import numpy as np from torchvision.transforms.functional import normalize from tqdm import tqdm from PIL import Image +import random # Helper function to convert a ComfyUI IMAGE tensor to an OpenCV BGR image. def tensor_to_cv2_img(tensor_frame: torch.Tensor) -> np.ndarray: @@ -18,8 +19,10 @@ def tensor_to_pil(tensor_frame: torch.Tensor, mode='RGB') -> Image.Image: class VideoInputPreprocessor: """ - Processes a batch of image frames to generate a face mask, then uses that mask's - bounding box to paste a provided RGBA image onto each frame. + Processes a batch of image frames to generate a face mask. It can either paste a + provided RGBA image into the face's bounding box (default) or use a generated, + feathered mask for a precise, non-rectangular composite ("face only" mode). + It also supports random horizontal flipping of the input face on a per-frame basis. """ @classmethod def INPUT_TYPES(cls): @@ -33,6 +36,9 @@ class VideoInputPreprocessor: "face_crop_scale": ("FLOAT", {"default": 1.5, "min": 1.0, "max": 3.0, "step": 0.1}), "dilation_kernel_size": ("INT", {"default": 10, "min": 0, "max": 50, "step": 1}), "with_neck": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), + "face_only_mode": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), + "feather_amount": ("INT", {"default": 21, "min": 0, "max": 151, "step": 2, "display": "slider"}), + "random_horizontal_flip_chance": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.05}), } } @@ -41,17 +47,14 @@ class VideoInputPreprocessor: FUNCTION = "generate_mask_and_paste" CATEGORY = "Stand-In" - def generate_mask_and_paste(self, face_processor, images: torch.Tensor, face_rgba: torch.Tensor, denoise_strength: float, confidence_threshold: float, face_crop_scale: float, dilation_kernel_size: int, with_neck: bool): + def generate_mask_and_paste(self, face_processor, images: torch.Tensor, face_rgba: torch.Tensor, denoise_strength: float, confidence_threshold: float, face_crop_scale: float, dilation_kernel_size: int, with_neck: bool, face_only_mode: bool, feather_amount: int, random_horizontal_flip_chance: float): detection_model, parsing_model, device = face_processor total_frames, h, w = images.shape[0], images.shape[1], images.shape[2] print(f"Processing {total_frames} frames ({w}x{h}) to paste new face.") - # --- Prepare the face_to_paste image --- - # Ensure it is RGBA and convert to PIL for easy pasting with transparency if face_rgba.shape[3] != 4: raise ValueError("The 'face_to_paste' image must be an RGBA image.") - # We only need the first image from the batch face_to_paste_pil = tensor_to_pil(face_rgba[0], mode='RGBA') processed_frames_tensors = [] @@ -63,7 +66,7 @@ class VideoInputPreprocessor: results = detection_model(frame_bgr, verbose=False) confident_boxes = results[0].boxes.xyxy[results[0].boxes.conf > confidence_threshold] - full_mask_np = np.zeros((h, w), dtype=np.uint8) + target_frame_pil = tensor_to_pil(frame_tensor).copy() if confident_boxes.shape[0] > 0: areas = (confident_boxes[:, 2] - confident_boxes[:, 0]) * (confident_boxes[:, 3] - confident_boxes[:, 1]) @@ -76,44 +79,50 @@ class VideoInputPreprocessor: crop_y1, crop_x1 = max(center_y - half_side, 0), max(center_x - half_side, 0) crop_y2, crop_x2 = min(center_y + half_side, h), min(center_x + half_side, w) - face_crop_bgr = frame_bgr[crop_y1:crop_y2, crop_x1:crop_x2] - - if face_crop_bgr.size > 0: - face_resized = cv2.resize(face_crop_bgr, (512, 512), interpolation=cv2.INTER_AREA) - face_rgb = cv2.cvtColor(face_resized, cv2.COLOR_BGR2RGB) - face_tensor_in = torch.from_numpy(face_rgb.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0).to(device) - - with torch.no_grad(): - normalized_face = normalize(face_tensor_in, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) - parsing_map = parsing_model(normalized_face)[0].argmax(dim=1, keepdim=True) - - parsing_map_np = parsing_map.squeeze().cpu().numpy().astype(np.uint8) - if with_neck: - final_mask_512 = (parsing_map_np != 0).astype(np.uint8) * 255 - else: - parts_to_exclude = [0, 14, 15, 16, 18] - final_mask_512 = np.isin(parsing_map_np, parts_to_exclude, invert=True).astype(np.uint8) * 255 - - if dilation_kernel_size > 0: - kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilation_kernel_size, dilation_kernel_size)) - final_mask_512 = cv2.dilate(final_mask_512, kernel, iterations=1) - - mask_resized_to_crop = cv2.resize(final_mask_512, (face_crop_bgr.shape[1], face_crop_bgr.shape[0]), interpolation=cv2.INTER_NEAREST) - full_mask_np[crop_y1:crop_y2, crop_x1:crop_x2] = mask_resized_to_crop - - target_frame_pil = tensor_to_pil(frame_tensor).copy() - - if confident_boxes.shape[0] > 0 and 'face_crop_bgr' in locals() and face_crop_bgr.size > 0: - - x = crop_x1 - y = crop_y1 - box_w = crop_x2 - crop_x1 - box_h = crop_y2 - crop_y1 + x, y = crop_x1, crop_y1 + box_w, box_h = crop_x2 - crop_x1, crop_y2 - crop_y1 if box_w > 0 and box_h > 0: + # Resize the source face image first face_to_paste_resized = face_to_paste_pil.resize((box_w, box_h), Image.Resampling.LANCZOS) - - target_frame_pil.paste(face_to_paste_resized, (x, y), face_to_paste_resized) + + # --- MODIFIED: Per-frame random flip logic is now INSIDE the loop --- + if random.random() < random_horizontal_flip_chance: + face_to_paste_resized = face_to_paste_resized.transpose(Image.FLIP_LEFT_RIGHT) + # --- END OF MODIFICATION --- + + if not face_only_mode: + target_frame_pil.paste(face_to_paste_resized, (x, y), face_to_paste_resized) + else: + face_crop_bgr = frame_bgr[crop_y1:crop_y2, crop_x1:crop_x2] + + if face_crop_bgr.size > 0: + face_resized = cv2.resize(face_crop_bgr, (512, 512), interpolation=cv2.INTER_AREA) + face_rgb = cv2.cvtColor(face_resized, cv2.COLOR_BGR2RGB) + face_tensor_in = torch.from_numpy(face_rgb.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0).to(device) + + with torch.no_grad(): + normalized_face = normalize(face_tensor_in, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + parsing_map = parsing_model(normalized_face)[0].argmax(dim=1, keepdim=True) + + parsing_map_np = parsing_map.squeeze().cpu().numpy().astype(np.uint8) + + parts_to_exclude = [0, 14, 15, 16, 17, 18] + final_mask_512 = np.isin(parsing_map_np, parts_to_exclude, invert=True).astype(np.uint8) * 255 + + if dilation_kernel_size > 0: + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilation_kernel_size, dilation_kernel_size)) + final_mask_512 = cv2.dilate(final_mask_512, kernel, iterations=1) + + if feather_amount > 0: + if feather_amount % 2 == 0: + feather_amount += 1 + final_mask_512 = cv2.GaussianBlur(final_mask_512, (feather_amount, feather_amount), 0) + + mask_resized_to_crop = cv2.resize(final_mask_512, (box_w, box_h), interpolation=cv2.INTER_LINEAR) + generated_mask_pil = Image.fromarray(mask_resized_to_crop, mode='L') + + target_frame_pil.paste(face_to_paste_resized, (x, y), mask=generated_mask_pil) processed_np = np.array(target_frame_pil.convert("RGB")).astype(np.float32) / 255.0 processed_tensor = torch.from_numpy(processed_np) @@ -121,12 +130,4 @@ class VideoInputPreprocessor: output_image_batch = torch.stack(processed_frames_tensors) - - return (output_image_batch,denoise_strength) - -NODE_CLASS_MAPPINGS = { - "VideoInputPreprocessor": VideoInputPreprocessor, -} -NODE_DISPLAY_NAME_MAPPINGS = { - "VideoInputPreprocessor": "Stand-In VideoInputPreprocessor", -} \ No newline at end of file + return (output_image_batch, denoise_strength) \ No newline at end of file