From 929679f2d58a61491a0ed72c027920591e48c9e0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20T=C5=82uchowski?= Date: Sun, 23 Feb 2025 10:49:39 +0100 Subject: [PATCH 1/3] WIP --- core/image_processor.py | 14 +------------- nodes/face_wrapper.py | 10 +++++----- nodes/image_feeder.py | 18 +++++++++++------- nodes/image_filters.py | 1 + 4 files changed, 18 insertions(+), 25 deletions(-) diff --git a/core/image_processor.py b/core/image_processor.py index 50fc1df..c244b56 100644 --- a/core/image_processor.py +++ b/core/image_processor.py @@ -51,19 +51,7 @@ class ImageProcessor: x2 = max_x + pad_x y2 = max_y + pad_y - return (int(x1), int(y1), int(x2 - x1), int(y2 - y1)) - - @staticmethod - def crop_face(image: Union[torch.Tensor, np.ndarray, Image.Image], landmarks_df: pd.DataFrame, padding_percent: float = 0.0) -> Optional[np.ndarray]: - """Crop face region from image based on landmarks.""" - image_np = ImageProcessor._convert_to_numpy(image) - bbox = ImageProcessor.calculate_face_bbox(landmarks_df, padding_percent) - - if bbox is None: - return None - - x, y, w, h = bbox - return image_np[y:y + h, x:x + w] + return int(x1), int(y1), int(x2 - x1), int(y2 - y1) @staticmethod def resize_image(image: Union[torch.Tensor, np.ndarray, Image.Image], target_size: int) -> Optional[np.ndarray]: diff --git a/nodes/face_wrapper.py b/nodes/face_wrapper.py index 95caa69..4288807 100644 --- a/nodes/face_wrapper.py +++ b/nodes/face_wrapper.py @@ -65,7 +65,7 @@ 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, fp_pipe or {}, empty_mask) + return image, fp_pipe or {}, empty_mask # Handle different modes if mode == "Debug": @@ -165,7 +165,7 @@ 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(fp_pipe, 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, fp_pipe, mask_np): @@ -189,7 +189,7 @@ 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(fp_pipe, 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, fp_pipe, mask_np): @@ -197,7 +197,7 @@ class FaceWrapper: 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) + return output_image, {}, output_mask source_x = fp_pipe['target_lm']['x'] source_y = fp_pipe['target_lm']['y'] @@ -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, fp_pipe, 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/nodes/image_feeder.py b/nodes/image_feeder.py index f0e185f..2b1194e 100644 --- a/nodes/image_feeder.py +++ b/nodes/image_feeder.py @@ -1,8 +1,10 @@ import os -import torch +from typing import List, Dict, Tuple + import numpy as np +import torch from PIL import Image -from typing import List, Dict, Union, Tuple +from torch import Tensor class ImageFeeder: @@ -33,7 +35,8 @@ class ImageFeeder: FUNCTION = "feed_images" CATEGORY = "Face Processor/Image" - def _load_image(self, image_path: str) -> torch.Tensor: + @staticmethod + def _load_image(image_path: str) -> Tensor | None: """ Load an image from path and convert it to ComfyUI format @@ -60,7 +63,8 @@ class ImageFeeder: print(f"Error loading image {image_path}: {str(e)}") return None - def _scan_directory(self, directory: str) -> List[str]: + @staticmethod + def _scan_directory(directory: str) -> List[str]: """ Scan directory for image files @@ -135,7 +139,7 @@ class ImageFeeder: print("No images found in directory") # Return empty image and data empty_image = torch.zeros((1, 64, 64, 3)) - return (empty_image, frames_data) + return empty_image, frames_data # Validate frame number max_frame = len(self.image_files) - 1 @@ -152,6 +156,6 @@ class ImageFeeder: if selected_frame is None: print(f"Failed to load frame {frame_number}") empty_image = torch.zeros((1, 64, 64, 3)) - return (empty_image, frames_data) + return empty_image, frames_data - return (selected_frame, frames_data) + return selected_frame, frames_data diff --git a/nodes/image_filters.py b/nodes/image_filters.py index 22721b8..3e0aabf 100644 --- a/nodes/image_filters.py +++ b/nodes/image_filters.py @@ -60,6 +60,7 @@ class HighPassFilter: FUNCTION = "apply_hpf" CATEGORY = "Face Processor/Tools" + @staticmethod def apply_hpf(self, image: torch.Tensor, blur_radius: int, blur_iterations: int, blend_opacity: float, input_black: int, input_white: int, gamma: float, show_histogram: bool) -> tuple[Tensor, dict]: 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 2/3] 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 From 085dbc278ec7104bab5ed4ba71921a1e6fff5e76 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20T=C5=82uchowski?= Date: Sun, 23 Feb 2025 11:57:02 +0100 Subject: [PATCH 3/3] refactor: moved all image converters to ImageProcessor class --- core/face_detector.py | 30 +---- core/image_processor.py | 216 ++++++++++++++++++++++++++++++---- nodes/face_fit_and_restore.py | 49 +++----- nodes/face_wrapper.py | 94 ++++----------- nodes/image_feeder.py | 20 ++-- nodes/image_filters.py | 13 +- 6 files changed, 253 insertions(+), 169 deletions(-) diff --git a/core/face_detector.py b/core/face_detector.py index 23d2f90..02c0312 100644 --- a/core/face_detector.py +++ b/core/face_detector.py @@ -9,6 +9,7 @@ from pandas import DataFrame from ..core.lm_mapping import LandmarkMappings from ..core.resources.model_loader import ModelDlib, ModelMediaPipe +from ..core.image_processor import ImageProcessor class FaceDetector: @@ -17,7 +18,7 @@ class FaceDetector: # Models will be initialized on demand self.mediapipe_model = None self.dlib_model = None - self.insight_model = None + self.image_processor = ImageProcessor() def detect_landmarks(self, image: Union[torch.Tensor, np.ndarray, Image.Image], @@ -68,7 +69,7 @@ class FaceDetector: if self.mediapipe_model is None: self.mediapipe_model = ModelMediaPipe() - image_np = self._convert_to_numpy(image) + image_np = self.image_processor.convert_to_numpy(image) if image_np is None: return None @@ -111,7 +112,7 @@ class FaceDetector: self.dlib_model = ModelDlib() # Convert image to numpy format suitable for dlib - image_np = self._convert_to_numpy(image) + image_np = self.image_processor.convert_to_numpy(image) if image_np is None: print("Failed to convert image for dlib processing") return None @@ -138,7 +139,7 @@ class FaceDetector: 'index': [] } - # Extract landmarks (adding 1 to index to match dlib's 1-based indexing) + # Extract landmarks for i in range(68): point = shape.part(i) landmarks_data['x'].append(float(point.x)) @@ -223,27 +224,6 @@ class FaceDetector: print(f"Error during landmark interpolation: {str(e)}") return mp_landmarks - def _convert_to_numpy(self, image: Union[torch.Tensor, np.ndarray, Image.Image]) -> Optional[np.ndarray]: - """Convert different image types to numpy array.""" - if torch.is_tensor(image): - image = image.detach().cpu().numpy() - if len(image.shape) == 4: - image = image[0] - image = (image * 255).astype(np.uint8) - elif isinstance(image, Image.Image): - image = np.array(image) - elif isinstance(image, np.ndarray): - if image.dtype == np.float32 and image.max() <= 1.0: - image = (image * 255).astype(np.uint8) - if len(image.shape) == 4: - image = image[0] - - if len(image.shape) == 2: - image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) - elif image.shape[2] == 4: - image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB) - - return image def __del__(self): """Clean up resources.""" diff --git a/core/image_processor.py b/core/image_processor.py index b1c7e5c..ea8d638 100644 --- a/core/image_processor.py +++ b/core/image_processor.py @@ -6,28 +6,6 @@ from PIL import Image from typing import Union, Optional, Tuple class ImageProcessor: - @staticmethod - 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() - if len(image.shape) == 4: - image = image[0] - image = (image * 255).astype(np.uint8) - elif isinstance(image, Image.Image): - image = np.array(image) - elif isinstance(image, np.ndarray): - if image.dtype == np.float32 and image.max() <= 1.0: - image = (image * 255).astype(np.uint8) - if len(image.shape) == 4: - image = image[0] - - if len(image.shape) == 2: - image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) - elif image.shape[2] == 4: - image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB) - - return image @staticmethod def calculate_face_bbox(landmarks_df: pd.DataFrame, padding_percent: float = 0.0) -> Optional[Tuple[int, int, int, int]]: @@ -288,3 +266,197 @@ class ImageProcessor: cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) return hist_img + + # Image Converters + @staticmethod + def tensor_to_numpy(tensor: torch.Tensor) -> np.ndarray: + """Convert PyTorch tensor to numpy array. + + Args: + tensor: Input tensor in format (B,H,W,C) or (B,C,H,W), values 0-1 + + Returns: + np.ndarray: RGB image (H,W,3), uint8 values 0-255 + """ + if tensor is None: + return None + + # Handle different tensor formats + img = tensor.detach().cpu().numpy() + if len(img.shape) == 4: + img = img[0] # Remove batch dimension + + # Handle channel dimension if in (C,H,W) format + if img.shape[0] in [1, 3, 4]: + img = img.transpose(1, 2, 0) + + # Handle single channel + if len(img.shape) == 2 or img.shape[2] == 1: + img = np.stack((img,) * 3, axis=-1) + + # Handle RGBA + elif img.shape[2] == 4: + img = img[..., :3] + + # Convert to uint8 if float + if img.dtype == np.float32 or img.dtype == np.float64: + img = (img * 255).clip(0, 255).astype(np.uint8) + + return img + + @staticmethod + def numpy_to_tensor(array: np.ndarray) -> torch.Tensor: + """Convert numpy array to PyTorch tensor. + + Args: + array: RGB image (H,W,3), uint8 or float32 + + Returns: + torch.Tensor: (1,H,W,3) tensor, values 0-1 + """ + if array is None: + return None + + # Convert to float32 if uint8 + if array.dtype == np.uint8: + array = array.astype(np.float32) / 255.0 + + # Add batch dimension + tensor = torch.from_numpy(array).unsqueeze(0) + return tensor + + @staticmethod + def pil_to_numpy(image: Image.Image) -> np.ndarray: + """Convert PIL image to numpy array. + + Args: + image: PIL Image + + Returns: + np.ndarray: RGB image (H,W,3), uint8 values 0-255 + """ + if image is None: + return None + + # Convert to RGB if needed + if image.mode != 'RGB': + image = image.convert('RGB') + + return np.array(image) + + @staticmethod + def numpy_to_pil(array: np.ndarray) -> Image.Image: + """Convert numpy array to PIL image. + + Args: + array: RGB image (H,W,3), uint8 or float32 + + Returns: + PIL.Image: RGB PIL image + """ + if array is None: + return None + + # Convert to uint8 if float + if array.dtype == np.float32 or array.dtype == np.float64: + array = (array * 255).clip(0, 255).astype(np.uint8) + + return Image.fromarray(array) + + @staticmethod + def tensor_to_pil(tensor: torch.Tensor) -> Image.Image: + """Convert PyTorch tensor to PIL image. + + Args: + tensor: Input tensor in format (B,H,W,C), values 0-1 + + Returns: + PIL.Image: RGB PIL image + """ + return ImageProcessor.numpy_to_pil(ImageProcessor.tensor_to_numpy(tensor)) + + @staticmethod + def pil_to_tensor(image: Image.Image) -> torch.Tensor: + """Convert PIL image to PyTorch tensor. + + Args: + image: PIL Image + + Returns: + torch.Tensor: (1,H,W,3) tensor, values 0-1 + """ + return ImageProcessor.numpy_to_tensor(ImageProcessor.pil_to_numpy(image)) + + @staticmethod + def convert_to_numpy(image: Union[torch.Tensor, np.ndarray, Image.Image]) -> np.ndarray: + """Universal converter to numpy array. + + Args: + image: Input in any supported format + + Returns: + np.ndarray: RGB image (H,W,3), uint8 values 0-255 + """ + if image is None: + return None + + if torch.is_tensor(image): + return ImageProcessor.tensor_to_numpy(image) + elif isinstance(image, Image.Image): + return ImageProcessor.pil_to_numpy(image) + elif isinstance(image, np.ndarray): + # Handle grayscale + if len(image.shape) == 2: + image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) + # Handle RGBA + elif image.shape[2] == 4: + image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB) + # Convert to uint8 if float + if image.dtype == np.float32 or image.dtype == np.float64: + image = (image * 255).clip(0, 255).astype(np.uint8) + return image + else: + raise ValueError(f"Unsupported image type: {type(image)}") + + # Mask conversion methods + @staticmethod + def convert_mask_to_numpy(mask: Union[torch.Tensor, np.ndarray]) -> np.ndarray: + """Convert mask to numpy array. + + Args: + mask: Input mask as tensor (1,H,W) or array + + Returns: + np.ndarray: Mask as (H,W) uint8 array, values 0-255 + """ + if mask is None: + return None + + if torch.is_tensor(mask): + mask = mask.detach().cpu().numpy() + if len(mask.shape) == 3: + mask = mask[0] # Remove batch dimension + + if mask.dtype == np.float32 or mask.dtype == np.float64: + mask = (mask * 255).clip(0, 255).astype(np.uint8) + + return mask + + @staticmethod + def convert_mask_to_tensor(mask: np.ndarray) -> torch.Tensor: + """Convert numpy mask to tensor. + + Args: + mask: Input mask as (H,W) array + + Returns: + torch.Tensor: Mask as (1,H,W) tensor, values 0-1 + """ + if mask is None: + return None + + if mask.dtype == np.uint8: + mask = mask.astype(np.float32) / 255.0 + + return torch.from_numpy(mask).unsqueeze(0) + diff --git a/nodes/face_fit_and_restore.py b/nodes/face_fit_and_restore.py index be0225e..11d0f1a 100644 --- a/nodes/face_fit_and_restore.py +++ b/nodes/face_fit_and_restore.py @@ -53,7 +53,7 @@ class FaceFitAndRestore: if mode == "Restore" and fp_pipe and "workflow" in fp_pipe: workflow = fp_pipe["workflow"] - print(f"Using workflow from face processor pipe: {workflow}") + print(f"Using workflow from pipe: {workflow}") if workflow == "single_image": if image is None: @@ -64,7 +64,7 @@ class FaceFitAndRestore: result = self._fit(image, padding_percent, bbox_size) else: # Restore if fp_pipe is None: - print("Error: Processor settings are required in Restore mode") + print("Error: fp_pipe is required in Restore mode") return None, {}, None, int(bbox_size) result = self._restore(image, fp_pipe) @@ -73,34 +73,24 @@ class FaceFitAndRestore: result_settings["workflow"] = "single_image" return result - elif workflow == "image_sequence": - 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) sequence_settings = { - "workflow": "image_sequence", - "frames": {} - } first_result = None - total_frames = len(frames_data["frames"]) - print(f"Processing sequence of {total_frames} frames...") for frame_idx, frame_path in frames_data["frames"].items(): - print(f"Processing frame {frame_idx + 1}/{total_frames}") try: - # Load image from path frame_tensor = self._load_image_from_path(frame_path) if frame_tensor is None: @@ -112,7 +102,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 face processor pipe for frame {frame_idx}") + print(f"Warning: No processor settings for frame {frame_idx}") continue print(f"Processing restore for frame {frame_idx} with settings: {frame_settings}") @@ -123,11 +113,8 @@ class FaceFitAndRestore: sequence_settings["frames"][f"frame_{frame_idx}"] = result[1] - except Exception as e: - print(f"Error processing frame {frame_idx}: {str(e)}") - continue if first_result is None: @@ -168,12 +155,11 @@ class FaceFitAndRestore: print("Failed to resize image, returning original") 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) + # Convert to tensor format + final_image = self.image_processor.numpy_to_tensor(final_image) - # Save face processor pipe for restoration - processor_settings = { + # Save processor settings for restoration + fp_pipe = { "original_image_shape": image_np.shape, "rotation_angle": rotation_angle, "crop_bbox": crop_bbox, # (x, y, w, h) @@ -197,14 +183,14 @@ class FaceFitAndRestore: resized_mask = cv2.resize(cropped_mask, (target_size, target_size), interpolation=cv2.INTER_LINEAR) mask = resized_mask - mask = torch.from_numpy(mask).unsqueeze(0) - return final_image, processor_settings, mask, int(bbox_size) + mask = self.image_processor.convert_mask_to_tensor(mask) + return final_image, fp_pipe, mask, int(bbox_size) def _restore(self, image, fp_pipe): """Restore mode: Restore the face to the original image.""" if fp_pipe is None: - print("Error in restore: processor_settings is None") + print("Error in restore: fp_pipe is None") return image, {}, self._create_empty_mask(image) processed_face_np = self.image_processor.convert_to_numpy(image) @@ -245,9 +231,8 @@ class FaceFitAndRestore: restored_image = cv2.warpAffine(restored_image, rotation_matrix, (width, height), flags=cv2.INTER_LANCZOS4) - restored_image = restored_image.astype(np.float32) / 255.0 - restored_image = torch.from_numpy(restored_image).unsqueeze(0) - + # Convert to tensor format + restored_image = self.image_processor.numpy_to_tensor(restored_image) mask = self._create_mask(original_image_shape, crop_bbox, rotation_angle) return restored_image, fp_pipe, mask @@ -271,8 +256,7 @@ class FaceFitAndRestore: rotation_matrix = cv2.getRotationMatrix2D(center, -rotation_angle, 1.0) mask = cv2.warpAffine(mask, rotation_matrix, (width, height), flags=cv2.INTER_LINEAR) - mask = torch.from_numpy(mask).unsqueeze(0) - return mask + return self.image_processor.convert_mask_to_tensor(mask) def _create_empty_mask(self, image): """Create an empty mask with the same shape as the input image.""" @@ -286,9 +270,10 @@ class FaceFitAndRestore: def _load_image_from_path(self, image_path): """Load and convert image from file path to tensor format.""" try: - frame_image = Image.open(image_path) - frame_tensor = torch.from_numpy(np.array(frame_image).astype(np.float32) / 255.0).unsqueeze(0) - return frame_tensor + image = Image.open(image_path) + if image.mode != 'RGB': + image = image.convert('RGB') + return self.image_processor.pil_to_tensor(image) except Exception as e: print(f"Error loading image from {image_path}: {str(e)}") return None \ No newline at end of file diff --git a/nodes/face_wrapper.py b/nodes/face_wrapper.py index 43471c0..38eaf00 100644 --- a/nodes/face_wrapper.py +++ b/nodes/face_wrapper.py @@ -1,14 +1,12 @@ -import torch import numpy as np import pandas as pd -from PIL import Image -import cv2 +import torch -from ..core.face_detector import FaceDetector -from ..core.image_processor import ImageProcessor from ..core.base_mesh import MediapipeBaseLandmarks from ..core.cpu_deformer import CPUDeformer +from ..core.face_detector import FaceDetector from ..core.gpu_deformer import GPUDeformer +from ..core.image_processor import ImageProcessor class FaceWrapper: @@ -47,13 +45,13 @@ class FaceWrapper: def detect_face(self, image, mode, device, show_detection, show_target, refiner, landmark_size, 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) + image_np = self.image_processor.convert_to_numpy(image) height, width = image_np.shape[:2] # Convert mask if provided mask_np = None if mask is not None: - mask_np = self._convert_mask_to_numpy(mask) + mask_np = self.image_processor.convert_mask_to_numpy(mask) if mode == "Wrap": return self._wrap_mode(image_np, None, width, height, @@ -76,50 +74,6 @@ class FaceWrapper: return self._unwrap_mode(image_np, landmarks_df, width, height, device, x_scale, y_transform, fp_pipe, mask_np) - def _convert_to_numpy(self, image): - """Improved image conversion with channel handling""" - if torch.is_tensor(image): - image = image.detach().cpu().numpy() - if len(image.shape) == 4: # Batch dimension - image = image[0] - - # Handle different channel formats - if image.shape[0] == 1: # Grayscale - image = np.stack((image[0],) * 3, axis=-1) - elif image.shape[0] == 3: # RGB - image = image.transpose(1, 2, 0) - elif image.shape[0] == 4: # RGBA - image = image[:3].transpose(1, 2, 0) - - image = (image * 255).astype(np.uint8) - - # Ensure 3 channels for numpy arrays - if isinstance(image, np.ndarray): - if len(image.shape) == 2: # Grayscale - image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) - elif image.shape[2] == 4: # RGBA - image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB) - elif image.shape[2] == 1: # Single channel - image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) - - return image - - def _convert_mask_to_numpy(self, mask): - """Convert mask tensor to numpy array""" - if torch.is_tensor(mask): - mask = mask.detach().cpu().numpy() - if len(mask.shape) == 3: - mask = mask[0] # Remove batch dimension - # Scale to 0-255 range - mask = (mask * 255).astype(np.uint8) - return mask - - def _convert_mask_to_tensor(self, mask_np): - """Convert numpy mask back to tensor format""" - if mask_np is not None: - mask = mask_np.astype(np.float32) / 255.0 - return torch.from_numpy(mask).unsqueeze(0) - return None def _debug_mode(self, image_np, landmarks_df, width, height, show_detection, show_target, landmark_size, show_labels, x_scale, y_transform, @@ -128,7 +82,7 @@ class FaceWrapper: overlays = [] if show_detection: - det_overlay = ImageProcessor.draw_landmarks( + det_overlay = self.image_processor.draw_landmarks( (width, height), landmarks_df, transparency=0.4, color=(0, 255, 0), radius=landmark_size, label=show_labels @@ -147,7 +101,7 @@ class FaceWrapper: 'z': np.zeros(len(base_landmarks)), 'index': range(len(base_landmarks)) }) - target_overlay = ImageProcessor.draw_landmarks( + target_overlay = self.image_processor.draw_landmarks( (width, height), base_df, transparency=0.4, color=(255, 0, 0), radius=landmark_size, label=show_labels @@ -161,11 +115,11 @@ class FaceWrapper: rgb = overlay[:, :, :3] result_image = result_image * (1 - alpha) + rgb * alpha - output_image = torch.from_numpy(result_image).unsqueeze(0) - output_mask = self._convert_mask_to_tensor(mask_np) + output_image = self.image_processor.numpy_to_tensor(result_image) + output_mask = self.image_processor.convert_mask_to_tensor(mask_np) landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks) - return output_image, self._update_settings(fp_pipe, landmarks_data), output_mask + return output_image, self._update_pipe(fp_pipe, landmarks_data), output_mask def _unwrap_mode(self, image_np, landmarks_df, width, height, device, x_scale, y_transform, fp_pipe, mask_np): @@ -175,28 +129,27 @@ class FaceWrapper: source_landmarks = landmarks_df.iloc[:468][['x', 'y']].values.astype(np.float32) # Process image - pil_image = Image.fromarray(image_np) + pil_image = self.image_processor.numpy_to_pil(image_np) warped_image = self._apply_warping(pil_image, source_landmarks, base_landmarks, device) # Process mask if provided warped_mask = None if mask_np is not None: - pil_mask = Image.fromarray(mask_np) + pil_mask = self.image_processor.numpy_to_pil(mask_np) warped_mask = self._apply_warping(pil_mask, source_landmarks, base_landmarks, device) - warped_mask = self._convert_mask_to_tensor(np.array(warped_mask)) + warped_mask = self.image_processor.convert_mask_to_tensor(np.array(warped_mask)) - output_image = np.array(warped_image).astype(np.float32) / 255.0 - output_image = torch.from_numpy(output_image).unsqueeze(0) + output_image = self.image_processor.pil_to_tensor(warped_image) landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks) - return output_image, self._update_settings(fp_pipe, landmarks_data), warped_mask + return output_image, self._update_pipe(fp_pipe, landmarks_data), warped_mask 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 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) + output_image = self.image_processor.numpy_to_tensor(image_np) + output_mask = self.image_processor.convert_mask_to_tensor(mask_np) return output_image, {}, output_mask source_x = fp_pipe['target_lm']['x'] @@ -208,18 +161,17 @@ class FaceWrapper: target_landmarks = np.column_stack((detected_x, detected_y))[:468] # Process image - pil_image = Image.fromarray(image_np) + pil_image = self.image_processor.numpy_to_pil(image_np) warped_image = self._apply_warping(pil_image, source_landmarks, target_landmarks, device) # Process mask if provided warped_mask = None if mask_np is not None: - pil_mask = Image.fromarray(mask_np) + pil_mask = self.image_processor.numpy_to_pil(mask_np) warped_mask = self._apply_warping(pil_mask, source_landmarks, target_landmarks, device) - warped_mask = self._convert_mask_to_tensor(np.array(warped_mask)) + warped_mask = self.image_processor.convert_mask_to_tensor(np.array(warped_mask)) - output_image = np.array(warped_image).astype(np.float32) / 255.0 - output_image = torch.from_numpy(output_image).unsqueeze(0) + output_image = self.image_processor.pil_to_tensor(warped_image) return output_image, fp_pipe, warped_mask @@ -244,5 +196,5 @@ class FaceWrapper: } } - def _update_settings(self, settings, new_data): - return {**(settings or {}), **new_data} + def _update_pipe(self, pipe, new_data): + return {**(pipe or {}), **new_data} diff --git a/nodes/image_feeder.py b/nodes/image_feeder.py index 2b1194e..999a428 100644 --- a/nodes/image_feeder.py +++ b/nodes/image_feeder.py @@ -1,10 +1,10 @@ import os from typing import List, Dict, Tuple -import numpy as np import torch from PIL import Image -from torch import Tensor + +from ..core.image_processor import ImageProcessor class ImageFeeder: @@ -13,6 +13,7 @@ class ImageFeeder: def __init__(self): self.current_dir = None self.image_files = [] + self.image_processor = ImageProcessor() @classmethod def INPUT_TYPES(cls): @@ -35,8 +36,7 @@ class ImageFeeder: FUNCTION = "feed_images" CATEGORY = "Face Processor/Image" - @staticmethod - def _load_image(image_path: str) -> Tensor | None: + def _load_image(self, image_path: str) -> torch.Tensor: """ Load an image from path and convert it to ComfyUI format @@ -53,18 +53,14 @@ class ImageFeeder: if image.mode != 'RGB': image = image.convert('RGB') - # Convert to numpy array and normalize - image_np = np.array(image).astype(np.float32) / 255.0 - - # Convert to torch tensor and add batch dimension - return torch.from_numpy(image_np).unsqueeze(0) + # Use ImageProcessor to convert to tensor + return self.image_processor.pil_to_tensor(image) except Exception as e: print(f"Error loading image {image_path}: {str(e)}") return None - @staticmethod - def _scan_directory(directory: str) -> List[str]: + def _scan_directory(self, directory: str) -> List[str]: """ Scan directory for image files @@ -158,4 +154,4 @@ class ImageFeeder: empty_image = torch.zeros((1, 64, 64, 3)) return empty_image, frames_data - return selected_frame, frames_data + return selected_frame, frames_data \ No newline at end of file diff --git a/nodes/image_filters.py b/nodes/image_filters.py index 3e0aabf..5791151 100644 --- a/nodes/image_filters.py +++ b/nodes/image_filters.py @@ -1,10 +1,10 @@ import torch import numpy as np import cv2 - from torch import Tensor from ..core.image_processor import ImageProcessor + class HighPassFilter: """ComfyUI node implementing AE-style high-pass filter with dynamic histogram visualization""" @@ -60,14 +60,13 @@ class HighPassFilter: FUNCTION = "apply_hpf" CATEGORY = "Face Processor/Tools" - @staticmethod def apply_hpf(self, image: torch.Tensor, blur_radius: int, blur_iterations: int, blend_opacity: float, input_black: int, input_white: int, gamma: float, show_histogram: bool) -> tuple[Tensor, dict]: """Main processing function implementing the High Pass Filter""" - # Convert tensor to numpy array - np_img = image[0].numpy() * 255 - np_img = np.clip(np_img, 0, 255).astype(np.uint8) + + # Convert input tensor to numpy array + np_img = ImageProcessor.tensor_to_numpy(image) # 1. Apply box blur with multiple iterations kernel_size = 2 * blur_radius + 1 @@ -97,7 +96,7 @@ class HighPassFilter: result = np.vstack([result, hist_img]) # Convert back to tensor format - result = torch.from_numpy(result.astype(np.float32) / 255.0).unsqueeze(0) + result = ImageProcessor.numpy_to_tensor(result) # Prepare settings dictionary settings = { @@ -112,4 +111,4 @@ class HighPassFilter: } } - return result, settings + return result, settings \ No newline at end of file