nodes/WeChatCV-Stand-In_Preprocessor_ComfyUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add only face mode

Browse Source
This commit is contained in:
KBRASK
2025-08-17 14:43:50 +08:00
parent 131ed382a1
commit c8feda62cc
3 changed files with 127 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+17 -10
View File
@@ -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,
videocombine.py
+58 -23
View File
@@ -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",
}
return (output_image_batch,)
videoinput.py
+52 -51
View File
@@ -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",
}
return (output_image_batch, denoise_strength)