From f9f7db4950421edd0c5c9675cb2116f8a985b866 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20T=C5=82uchowski?= Date: Sun, 23 Feb 2025 10:03:59 +0100 Subject: [PATCH] processor_settings changed to fp_pipe --- nodes/face_fit_and_restore.py | 16 ++++++++-------- nodes/face_wrapper.py | 36 +++++++++++++++++------------------ readme.md | 3 ++- 3 files changed, 28 insertions(+), 27 deletions(-) diff --git a/nodes/face_fit_and_restore.py b/nodes/face_fit_and_restore.py index 0715837..fb2faa5 100644 --- a/nodes/face_fit_and_restore.py +++ b/nodes/face_fit_and_restore.py @@ -34,7 +34,7 @@ class FaceFitAndRestore: }, "optional": { "image": ("IMAGE",), - "processor_settings": ("DICT", { + "fp_pipe": ("DICT", { "default": None }), "frames_data": ("DICT", { @@ -44,15 +44,15 @@ class FaceFitAndRestore: } RETURN_TYPES = ("IMAGE", "DICT", "MASK", "INT") - RETURN_NAMES = ("image", "processor_settings", "mask", "bbox_size") + RETURN_NAMES = ("image", "fp_pipe", "mask", "bbox_size") FUNCTION = "process_image" CATEGORY = "Face Processor" def process_image(self, mode, workflow, padding_percent=0.0, bbox_size="1024", - image=None, processor_settings=None, frames_data=None): + image=None, fp_pipe=None, frames_data=None): - if mode == "Restore" and processor_settings and "workflow" in processor_settings: - workflow = processor_settings["workflow"] + if mode == "Restore" and fp_pipe and "workflow" in fp_pipe: + workflow = fp_pipe["workflow"] print(f"Using workflow from processor settings: {workflow}") if workflow == "single_image": @@ -63,10 +63,10 @@ class FaceFitAndRestore: if mode == "Fit": result = self._fit(image, padding_percent, bbox_size) else: # Restore - if processor_settings is None: + if fp_pipe is None: print("Error: Processor settings are required in Restore mode") return (None, {}, None, int(bbox_size)) - result = self._restore(image, processor_settings) + result = self._restore(image, fp_pipe) # Add workflow info to settings result_settings = result[1] @@ -110,7 +110,7 @@ class FaceFitAndRestore: result = self._fit(frame_tensor, padding_percent, bbox_size) result[1]["original_image_path"] = frame_path elif mode == "Restore": - frame_settings = processor_settings.get("frames", {}).get(f"frame_{frame_idx}") + 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}") continue diff --git a/nodes/face_wrapper.py b/nodes/face_wrapper.py index 85e58a4..95caa69 100644 --- a/nodes/face_wrapper.py +++ b/nodes/face_wrapper.py @@ -34,18 +34,18 @@ class FaceWrapper: "y_transform": ("FLOAT", {"default": 0.0, "min": -0.5, "max": 0.5, "step": 0.01}) }, "optional": { - "processor_settings": ("DICT", {"default": None}), + "fp_pipe": ("DICT", {"default": None}), "mask": ("MASK", {"default": None}) } } RETURN_TYPES = ("IMAGE", "DICT", "MASK") - RETURN_NAMES = ("image", "processor_settings", "mask") + RETURN_NAMES = ("image", "fp_pipe", "mask") FUNCTION = "detect_face" CATEGORY = "Face Processor" def detect_face(self, image, mode, device, show_detection, show_target, refiner, landmark_size, - show_labels, x_scale, y_transform, processor_settings=None, mask=None): + show_labels, x_scale, y_transform, fp_pipe=None, mask=None): # Convert input image to numpy with proper RGB format image_np = self._convert_to_numpy(image) height, width = image_np.shape[:2] @@ -57,7 +57,7 @@ class FaceWrapper: if mode == "Wrap": return self._wrap_mode(image_np, None, width, height, - device, x_scale, y_transform, processor_settings, mask_np) + device, x_scale, y_transform, fp_pipe, mask_np) # Detect facial landmarks landmarks_df = self.face_detector.detect_landmarks(image_np, refiner=(None if refiner == "None" else refiner)) @@ -65,16 +65,16 @@ class FaceWrapper: if landmarks_df is None: print("No face detected") empty_mask = torch.zeros((1, height, width), dtype=torch.float32) if mask is not None else None - return (image, processor_settings or {}, empty_mask) + return (image, fp_pipe or {}, empty_mask) # Handle different modes if mode == "Debug": return self._debug_mode(image_np, landmarks_df, width, height, show_detection, show_target, landmark_size, - show_labels, x_scale, y_transform, processor_settings, mask_np) + show_labels, x_scale, y_transform, fp_pipe, mask_np) elif mode == "Un-Wrap": return self._unwrap_mode(image_np, landmarks_df, width, height, - device, x_scale, y_transform, processor_settings, mask_np) + device, x_scale, y_transform, fp_pipe, mask_np) def _convert_to_numpy(self, image): """Improved image conversion with channel handling""" @@ -123,7 +123,7 @@ class FaceWrapper: def _debug_mode(self, image_np, landmarks_df, width, height, show_detection, show_target, landmark_size, show_labels, x_scale, y_transform, - processor_settings, mask_np): + fp_pipe, mask_np): result_image = image_np.astype(np.float32) / 255.0 overlays = [] @@ -165,10 +165,10 @@ class FaceWrapper: output_mask = self._convert_mask_to_tensor(mask_np) landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks) - return (output_image, self._update_settings(processor_settings, landmarks_data), output_mask) + return (output_image, self._update_settings(fp_pipe, landmarks_data), output_mask) def _unwrap_mode(self, image_np, landmarks_df, width, height, device, - x_scale, y_transform, processor_settings, mask_np): + x_scale, y_transform, fp_pipe, mask_np): base_landmarks = MediapipeBaseLandmarks.get_base_landmarks( (width, height), x_scale=x_scale, y_translation=y_transform ) @@ -189,22 +189,22 @@ class FaceWrapper: output_image = torch.from_numpy(output_image).unsqueeze(0) landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks) - return (output_image, self._update_settings(processor_settings, landmarks_data), warped_mask) + return (output_image, self._update_settings(fp_pipe, landmarks_data), warped_mask) def _wrap_mode(self, image_np, landmarks_df, width, height, device, - x_scale, y_transform, processor_settings, mask_np): - if not processor_settings or 'target_lm' not in processor_settings: + 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") 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) - source_x = processor_settings['target_lm']['x'] - source_y = processor_settings['target_lm']['y'] + source_x = fp_pipe['target_lm']['x'] + source_y = fp_pipe['target_lm']['y'] source_landmarks = np.column_stack((source_x, source_y))[:468] - detected_x = processor_settings['detected_lm']['x'] - detected_y = processor_settings['detected_lm']['y'] + detected_x = fp_pipe['detected_lm']['x'] + detected_y = fp_pipe['detected_lm']['y'] target_landmarks = np.column_stack((detected_x, detected_y))[:468] # Process image @@ -221,7 +221,7 @@ class FaceWrapper: output_image = np.array(warped_image).astype(np.float32) / 255.0 output_image = torch.from_numpy(output_image).unsqueeze(0) - return (output_image, processor_settings, warped_mask) + return (output_image, fp_pipe, warped_mask) def _apply_warping(self, image, source_landmarks, target_landmarks, device): """Apply warping to image using selected device""" diff --git a/readme.md b/readme.md index b432aa1..9744178 100644 --- a/readme.md +++ b/readme.md @@ -138,13 +138,14 @@ normalized_face, settings = face_wrapper.detect_face( 2. Process normalized face with your preferred method 3. Restore face to original position: + ```python # Wrap processed face back final_image, _ = face_wrapper.detect_face( image=processed_face, mode="Wrap", device="CUDA", - processor_settings=settings + fp_pipe=settings ) ```