From a31a0d282eca3f23a4a161fb2090df64b21f50de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20T=C5=82uchowski?= Date: Sun, 23 Feb 2025 11:15:29 +0100 Subject: [PATCH] WIP --- core/image_processor.py | 6 ++-- nodes/face_fit_and_restore.py | 56 +++++++++++++++++------------------ nodes/face_wrapper.py | 2 +- 3 files changed, 32 insertions(+), 32 deletions(-) diff --git a/core/image_processor.py b/core/image_processor.py index c244b56..b1c7e5c 100644 --- a/core/image_processor.py +++ b/core/image_processor.py @@ -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 diff --git a/nodes/face_fit_and_restore.py b/nodes/face_fit_and_restore.py index fb2faa5..be0225e 100644 --- a/nodes/face_fit_and_restore.py +++ b/nodes/face_fit_and_restore.py @@ -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.""" diff --git a/nodes/face_wrapper.py b/nodes/face_wrapper.py index 4288807..43471c0 100644 --- a/nodes/face_wrapper.py +++ b/nodes/face_wrapper.py @@ -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