nodes/SykkoAtHome-ComfyUI_FaceProcessor
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

processor_settings - pass through with optional extension

Browse Source
This commit is contained in:
Rafał Tłuchowski
2025-01-26 11:35:32 +01:00
parent 773d308cba
commit 501c176c09
3 changed files with 33 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
contex_generator.py
+1 -1
View File
@@ -57,7 +57,7 @@ def generate_context_file(
# Example usage
if __name__ == "__main__":
output_file = "D:/code/code_base/face_processor_05.txt"
output_file = "D:/code/code_base/face_processor_06.txt"
generate_context_file(
root_dir="E:/ComfyUI-DEV/custom_nodes/ComfyUI_FaceProcessor",
output_file=output_file,
nodes/face_fit_and_restore.py
+18 -38
View File
@@ -32,25 +32,25 @@ class FaceFitAndRestore:
}),
},
"optional": {
"face_settings": ("DICT", {
"processor_settings": ("DICT", {
"default": None
}),
}
}
RETURN_TYPES = ("IMAGE", "DICT", "MASK")
RETURN_NAMES = ("image", "face_settings", "mask")
RETURN_NAMES = ("image", "processor_settings", "mask")
FUNCTION = "process_image"
CATEGORY = "Face Processor"
def process_image(self, mode, image, padding_percent=0.0, bbox_size="1024", face_settings=None):
def process_image(self, mode, image, padding_percent=0.0, bbox_size="1024", processor_settings=None):
if mode == "Fit":
return self._fit(image, padding_percent, bbox_size)
elif mode == "Restore":
if face_settings is None:
print("Face settings are required in Restore mode, returning original image")
if processor_settings is None:
print("Processor settings are required in Restore mode, returning original image")
return (image, {}, self._create_empty_mask(image))
return self._restore(image, face_settings)
return self._restore(image, processor_settings)
else:
print(f"Invalid mode: {mode}, returning original image")
return (image, {}, self._create_empty_mask(image))
@@ -92,8 +92,8 @@ class FaceFitAndRestore:
final_image = final_image.astype(np.float32) / 255.0
final_image = torch.from_numpy(final_image).unsqueeze(0)
# Save face settings for restoration
face_settings = {
# Save processor settings for restoration
processor_settings = {
"original_image_shape": image_np.shape,
"rotation_angle": rotation_angle,
"crop_bbox": crop_bbox, # (x, y, w, h)
@@ -101,76 +101,56 @@ class FaceFitAndRestore:
"bbox_size": target_size,
}
# Create a mask for the face area (1:1)
# Initialize mask with ones (all pixels are part of the face)
# Create a mask for the face area
mask = np.ones((target_size, target_size), dtype=np.float32)
# If there was a rotation, exclude areas that are outside the original image
if rotation_angle != 0:
# Create a blank image with the original size
original_mask = np.ones(image_np.shape[:2], dtype=np.float32)
# Rotate the original mask to match the rotated image
height, width = image_np.shape[:2]
center = (width // 2, height // 2)
rotation_matrix = cv2.getRotationMatrix2D(center, rotation_angle, 1.0)
rotated_mask = cv2.warpAffine(original_mask, rotation_matrix, (width, height), flags=cv2.INTER_LINEAR)
# Crop the rotated mask to the face region
x1, y1, w, h = crop_bbox
cropped_mask = rotated_mask[y1:y1 + h, x1:x1 + w]
# Resize the cropped mask to the target size
resized_mask = cv2.resize(cropped_mask, (target_size, target_size), interpolation=cv2.INTER_LINEAR)
# Update the final mask to exclude areas outside the original image
mask = resized_mask
# Convert mask to torch tensor
mask = torch.from_numpy(mask).unsqueeze(0)
return (final_image, processor_settings, mask)
return (final_image, face_settings, mask)
def _restore(self, image, face_settings):
def _restore(self, image, processor_settings):
"""Restore mode: Restore the face to the original image."""
# Convert tensor to numpy
processed_face_np = self._convert_to_numpy(image)
if processed_face_np is None:
return (image, {}, self._create_empty_mask(image))
# Extract face settings
original_image_shape = face_settings.get("original_image_shape")
rotation_angle = face_settings.get("rotation_angle")
crop_bbox = face_settings.get("crop_bbox")
padding_percent = face_settings.get("padding_percent")
bbox_size = face_settings.get("bbox_size")
original_image_shape = processor_settings.get("original_image_shape")
rotation_angle = processor_settings.get("rotation_angle")
crop_bbox = processor_settings.get("crop_bbox")
padding_percent = processor_settings.get("padding_percent")
bbox_size = processor_settings.get("bbox_size")
if not all([original_image_shape, crop_bbox]):
print("Invalid face settings, returning processed face")
print("Invalid processor settings, returning processed face")
return (image, {}, self._create_empty_mask(image))
# Create a blank image with the original size
restored_image = np.zeros(original_image_shape, dtype=np.uint8)
# Resize the processed face back to the original crop size
x1, y1, w, h = crop_bbox
resized_face = cv2.resize(processed_face_np, (w, h), interpolation=cv2.INTER_LANCZOS4)
# Place the resized face back into the rotated image
restored_image[y1:y1 + h, x1:x1 + w] = resized_face
# Rotate the image back to the original orientation (reverse the rotation)
if rotation_angle != 0:
height, width = original_image_shape[:2]
center = (width // 2, height // 2)
rotation_matrix = cv2.getRotationMatrix2D(center, -rotation_angle, 1.0) # Reverse the rotation
rotation_matrix = cv2.getRotationMatrix2D(center, -rotation_angle, 1.0)
restored_image = cv2.warpAffine(restored_image, rotation_matrix, (width, height), flags=cv2.INTER_LANCZOS4)
# Convert back to float32 format expected by ComfyUI
restored_image = restored_image.astype(np.float32) / 255.0
restored_image = torch.from_numpy(restored_image).unsqueeze(0)
# Create a mask for the face area (original image size)
mask = self._create_mask(original_image_shape, crop_bbox, rotation_angle)
return (restored_image, {}, mask)
nodes/face_wrapper.py
+14 -14
View File
@@ -45,19 +45,23 @@ class FaceWrapper:
"max": 0.5,
"step": 0.05
})
},
"optional": {
"processor_settings": ("DICT", {
"default": None
})
}
}
RETURN_TYPES = ("IMAGE", "DICT")
RETURN_NAMES = ("image", "landmarks")
RETURN_NAMES = ("image", "processor_settings")
FUNCTION = "detect_face"
CATEGORY = "Face Processor"
def detect_face(self, image, show_detection, show_target, landmark_size=2,
show_labels=False, x_scale=1.0, y_transform=0.0):
show_labels=False, x_scale=1.0, y_transform=0.0, processor_settings=None):
"""Detect facial landmarks and optionally visualize them."""
# Convert tensor to numpy for detection
if torch.is_tensor(image):
image_np = image.detach().cpu().numpy()
if len(image_np.shape) == 4:
@@ -66,18 +70,15 @@ class FaceWrapper:
height, width = image_np.shape[:2]
# Detect landmarks
landmarks_df = self.face_detector.detect_landmarks(image_np)
if landmarks_df is None:
print("No face detected in the image")
return (image, {})
# Start with original image
result_image = image_np.astype(np.float32) / 255.0
overlays = []
# Create detected landmarks overlay if requested
if show_detection:
landmark_overlay = ImageProcessor.draw_landmarks(
(width, height),
@@ -90,14 +91,12 @@ class FaceWrapper:
if landmark_overlay is not None:
overlays.append(landmark_overlay)
# Get base landmarks with transformations
base_landmarks = MediapipeBaseLandmarks.get_base_landmarks(
(width, height),
x_scale=x_scale,
y_translation=y_transform
)
# Create target landmarks overlay if requested
if show_target:
base_df = pd.DataFrame({
'x': base_landmarks[:, 0],
@@ -117,18 +116,15 @@ class FaceWrapper:
if target_overlay is not None:
overlays.append(target_overlay)
# Blend all overlays with the image
for overlay in overlays:
overlay = overlay.astype(np.float32) / 255.0
alpha = overlay[:, :, 3:]
rgb = overlay[:, :, :3]
result_image = result_image * (1 - alpha) + rgb * alpha
# Convert final image back to torch tensor
image = torch.from_numpy(result_image).unsqueeze(0)
# Create landmarks dictionary
landmarks = {
landmarks_data = {
'detected_lm': {
'x': landmarks_df['x'].tolist(),
'y': landmarks_df['y'].tolist(),
@@ -141,5 +137,9 @@ class FaceWrapper:
}
}
return (image, landmarks)
if processor_settings is not None:
updated_settings = processor_settings.copy()
updated_settings.update(landmarks_data)
return (image, updated_settings)
else:
return (image, landmarks_data)