Merge pull request #7 from SykkoAtHome/image_processor

Image processor

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."""

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)
+

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

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}

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

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