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

WIP

Browse Source
This commit is contained in:
Rafał Tłuchowski
2025-02-23 11:15:29 +01:00
parent 929679f2d5
commit a31a0d282e
3 changed files with 32 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/image_processor.py
+3 -3
View File
@@ -7,7 +7,7 @@ from typing import Union, Optional, Tuple
class ImageProcessor:
@staticmethod
def _convert_to_numpy(image: Union[torch.Tensor, np.ndarray, Image.Image]) -> np.ndarray:
def convert_to_numpy(image: Union[torch.Tensor, np.ndarray, Image.Image]) -> np.ndarray:
"""Convert different image types to numpy array."""
if torch.is_tensor(image):
image = image.detach().cpu().numpy()
@@ -59,7 +59,7 @@ class ImageProcessor:
if image is None:
return None
image_np = ImageProcessor._convert_to_numpy(image)
image_np = ImageProcessor.convert_to_numpy(image)
h, w = image_np.shape[:2]
# Determine the smaller dimension (height or width)
@@ -98,7 +98,7 @@ class ImageProcessor:
@staticmethod
def rotate_image(image: Union[torch.Tensor, np.ndarray, Image.Image], landmarks_df: pd.DataFrame) -> Tuple[Optional[np.ndarray], Optional[pd.DataFrame]]:
"""Rotate image based on facial landmarks."""
image_np = ImageProcessor._convert_to_numpy(image)
image_np = ImageProcessor.convert_to_numpy(image)
if image_np is None or landmarks_df is None:
return None, None
nodes/face_fit_and_restore.py
+28 -28
View File
@@ -53,19 +53,19 @@ class FaceFitAndRestore:
if mode == "Restore" and fp_pipe and "workflow" in fp_pipe:
workflow = fp_pipe["workflow"]
print(f"Using workflow from processor settings: {workflow}")
print(f"Using workflow from face processor pipe: {workflow}")
if workflow == "single_image":
if image is None:
print("Error: Image is required for single_image workflow")
return (None, {}, None, int(bbox_size))
return None, {}, None, int(bbox_size)
if mode == "Fit":
result = self._fit(image, padding_percent, bbox_size)
else: # Restore
if fp_pipe is None:
print("Error: Processor settings are required in Restore mode")
return (None, {}, None, int(bbox_size))
return None, {}, None, int(bbox_size)
result = self._restore(image, fp_pipe)
# Add workflow info to settings
@@ -79,7 +79,7 @@ class FaceFitAndRestore:
if frames_data is None or not frames_data.get("frames"):
print("Error: Valid frames_data is required for image_sequence workflow")
return (None, {}, None, int(bbox_size))
return None, {}, None, int(bbox_size)
sequence_settings = {
@@ -112,7 +112,7 @@ class FaceFitAndRestore:
elif mode == "Restore":
frame_settings = fp_pipe.get("frames", {}).get(f"frame_{frame_idx}")
if frame_settings is None:
print(f"Warning: No processor settings for frame {frame_idx}")
print(f"Warning: No face processor pipe for frame {frame_idx}")
continue
print(f"Processing restore for frame {frame_idx} with settings: {frame_settings}")
@@ -131,48 +131,48 @@ class FaceFitAndRestore:
continue
if first_result is None:
return (None, sequence_settings, None, int(bbox_size))
return None, sequence_settings, None, int(bbox_size)
return (first_result[0], sequence_settings, first_result[2], int(bbox_size))
return first_result[0], sequence_settings, first_result[2], int(bbox_size)
def _fit(self, image, padding_percent, bbox_size):
"""Fit mode: Crop and process the face."""
image_np = self.image_processor._convert_to_numpy(image)
image_np = self.image_processor.convert_to_numpy(image)
if image_np is None:
return (image, {}, self._create_empty_mask(image), int(bbox_size))
return image, {}, self._create_empty_mask(image), int(bbox_size)
# Detect facial landmarks using mediapipe
landmarks_df = self.face_detector.detect_landmarks_mp(image_np)
if landmarks_df is None:
print("No face detected, returning original image")
return (image, {}, self._create_empty_mask(image), int(bbox_size))
return image, {}, self._create_empty_mask(image), int(bbox_size)
# Calculate rotation angle and rotate the image
rotation_angle = self.image_processor.calculate_rotation_angle(landmarks_df)
rotated_image, updated_landmarks = self.image_processor.rotate_image(image_np, landmarks_df)
if rotated_image is None:
print("Failed to rotate image, returning original")
return (image, {}, self._create_empty_mask(image), int(bbox_size))
return image, {}, self._create_empty_mask(image), int(bbox_size)
# Crop face region to a square (1:1)
cropped_face, crop_bbox = self.image_processor.crop_face_to_square(rotated_image, updated_landmarks,
padding_percent)
if cropped_face is None:
print("Failed to crop face, returning original image")
return (image, {}, self._create_empty_mask(image), int(bbox_size))
return image, {}, self._create_empty_mask(image), int(bbox_size)
# Resize to target size
target_size = int(bbox_size)
final_image = self.image_processor.resize_image(cropped_face, target_size)
if final_image is None:
print("Failed to resize image, returning original")
return (image, {}, self._create_empty_mask(image), int(bbox_size))
return image, {}, self._create_empty_mask(image), int(bbox_size)
# Convert back to float32 format expected by ComfyUI
final_image = final_image.astype(np.float32) / 255.0
final_image = torch.from_numpy(final_image).unsqueeze(0)
# Save processor settings for restoration
# Save face processor pipe for restoration
processor_settings = {
"original_image_shape": image_np.shape,
"rotation_angle": rotation_angle,
@@ -198,25 +198,25 @@ class FaceFitAndRestore:
mask = resized_mask
mask = torch.from_numpy(mask).unsqueeze(0)
return (final_image, processor_settings, mask, int(bbox_size))
return final_image, processor_settings, mask, int(bbox_size)
def _restore(self, image, processor_settings):
def _restore(self, image, fp_pipe):
"""Restore mode: Restore the face to the original image."""
if processor_settings is None:
if fp_pipe is None:
print("Error in restore: processor_settings is None")
return (image, {}, self._create_empty_mask(image))
return image, {}, self._create_empty_mask(image)
processed_face_np = self.image_processor._convert_to_numpy(image)
processed_face_np = self.image_processor.convert_to_numpy(image)
if processed_face_np is None:
print("Error in restore: Failed to convert image to numpy array")
return (image, {}, self._create_empty_mask(image))
return image, {}, self._create_empty_mask(image)
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")
original_image_shape = fp_pipe.get("original_image_shape")
rotation_angle = fp_pipe.get("rotation_angle")
crop_bbox = fp_pipe.get("crop_bbox")
padding_percent = fp_pipe.get("padding_percent")
bbox_size = fp_pipe.get("bbox_size")
# Debug info
print(f"Restore settings: shape={original_image_shape}, rotation={rotation_angle}, bbox={crop_bbox}")
@@ -228,7 +228,7 @@ class FaceFitAndRestore:
if not crop_bbox:
missing.append("crop_bbox")
print(f"Error in restore: Missing required settings: {', '.join(missing)}")
return (image, {}, self._create_empty_mask(image))
return image, {}, self._create_empty_mask(image)
try:
restored_image = np.zeros(original_image_shape, dtype=np.uint8)
@@ -250,11 +250,11 @@ class FaceFitAndRestore:
mask = self._create_mask(original_image_shape, crop_bbox, rotation_angle)
return (restored_image, processor_settings, mask)
return restored_image, fp_pipe, mask
except Exception as e:
print(f"Error in restore: {str(e)}")
return (image, processor_settings, self._create_empty_mask(image))
return image, fp_pipe, self._create_empty_mask(image)
def _create_mask(self, image_shape, crop_bbox, rotation_angle):
"""Create a mask for the face area."""
nodes/face_wrapper.py
+1 -1
View File
@@ -194,7 +194,7 @@ class FaceWrapper:
def _wrap_mode(self, image_np, landmarks_df, width, height, device,
x_scale, y_transform, fp_pipe, mask_np):
if not fp_pipe or 'target_lm' not in fp_pipe:
print("No landmarks found in processor settings")
print("No landmarks found in face processor pipe")
output_image = torch.from_numpy(image_np.astype(np.float32) / 255.0).unsqueeze(0)
output_mask = self._convert_mask_to_tensor(mask_np)
return output_image, {}, output_mask