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 50fc1df..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]]: @@ -51,19 +29,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]: @@ -71,7 +37,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) @@ -110,7 +76,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 @@ -300,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 fb2faa5..11d0f1a 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 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)) + print("Error: fp_pipe is required in Restore mode") + return None, {}, None, int(bbox_size) result = self._restore(image, fp_pipe) # Add workflow info to settings @@ -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)) + 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: @@ -123,57 +113,53 @@ 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: - 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) + # Convert to tensor format + final_image = self.image_processor.numpy_to_tensor(final_image) # Save processor settings for restoration - processor_settings = { + fp_pipe = { "original_image_shape": image_np.shape, "rotation_angle": rotation_angle, "crop_bbox": crop_bbox, # (x, y, w, h) @@ -197,26 +183,26 @@ 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, processor_settings): + def _restore(self, image, fp_pipe): """Restore mode: Restore the face to the original image.""" - if processor_settings is None: - print("Error in restore: processor_settings is None") - return (image, {}, self._create_empty_mask(image)) + if fp_pipe 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) + 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 +214,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) @@ -245,16 +231,15 @@ 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, 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.""" @@ -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 95caa69..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, @@ -65,7 +63,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": @@ -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,29 +129,28 @@ 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 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) + print("No landmarks found in face processor pipe") + 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'] source_y = fp_pipe['target_lm']['y'] @@ -208,20 +161,19 @@ 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) + return output_image, fp_pipe, warped_mask def _apply_warping(self, image, source_landmarks, target_landmarks, device): """Apply warping to image using selected device""" @@ -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 f0e185f..999a428 100644 --- a/nodes/image_feeder.py +++ b/nodes/image_feeder.py @@ -1,8 +1,10 @@ import os +from typing import List, Dict, Tuple + import torch -import numpy as np from PIL import Image -from typing import List, Dict, Union, Tuple + +from ..core.image_processor import ImageProcessor class ImageFeeder: @@ -11,6 +13,7 @@ class ImageFeeder: def __init__(self): self.current_dir = None self.image_files = [] + self.image_processor = ImageProcessor() @classmethod def INPUT_TYPES(cls): @@ -50,11 +53,8 @@ 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)}") @@ -135,7 +135,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 +152,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 \ No newline at end of file diff --git a/nodes/image_filters.py b/nodes/image_filters.py index 22721b8..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""" @@ -64,9 +64,9 @@ class HighPassFilter: 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 @@ -96,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 = { @@ -111,4 +111,4 @@ class HighPassFilter: } } - return result, settings + return result, settings \ No newline at end of file