diff --git a/__init__.py b/__init__.py index f6b6485..6c63b7f 100644 --- a/__init__.py +++ b/__init__.py @@ -2,7 +2,8 @@ import os import sys -from .nodes.tools import HighPassFilter +from .nodes.image_feeder import ImageFeeder +from .nodes.image_filters import HighPassFilter from .nodes.face_wrapper import FaceWrapper from .nodes.face_fit_and_restore import FaceFitAndRestore @@ -15,13 +16,15 @@ NODE_PATH = os.path.dirname(os.path.realpath(__file__)) NODE_CLASS_MAPPINGS = { "FaceFitAndRestore": FaceFitAndRestore, "FaceWrapper": FaceWrapper, - "HighPassFilter": HighPassFilter + "HighPassFilter": HighPassFilter, + "ImageFeeder": ImageFeeder } NODE_DISPLAY_NAME_MAPPINGS = { "FaceFitAndRestore": "Face Fit or Restore", "FaceWrapper": "Face Wrapper", - "HighPassFilter": "High Pass Filter (HPF)" + "HighPassFilter": "High Pass Filter (HPF)", + "ImageFeeder": "Image Feeder" } -__version__ = "1.0.1" +__version__ = "1.1.0" diff --git a/nodes/image_feeder.py b/nodes/image_feeder.py new file mode 100644 index 0000000..f0e185f --- /dev/null +++ b/nodes/image_feeder.py @@ -0,0 +1,157 @@ +import os +import torch +import numpy as np +from PIL import Image +from typing import List, Dict, Union, Tuple + + +class ImageFeeder: + """ComfyUI node for feeding images from a directory.""" + + def __init__(self): + self.current_dir = None + self.image_files = [] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "directory": ("STRING", { + "default": "", + "multiline": False + }), + "frame_number": ("INT", { + "default": 0, + "min": 0, + "step": 1 + }), + } + } + + RETURN_TYPES = ("IMAGE", "DICT") + RETURN_NAMES = ("image", "frames_data") + FUNCTION = "feed_images" + CATEGORY = "Face Processor/Image" + + def _load_image(self, image_path: str) -> torch.Tensor: + """ + Load an image from path and convert it to ComfyUI format + + Args: + image_path: Path to the image file + + Returns: + torch.Tensor: Image in ComfyUI format (B,H,W,C) normalized to 0-1 range + """ + try: + image = Image.open(image_path) + + # Convert to RGB if necessary + 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) + + except Exception as e: + print(f"Error loading image {image_path}: {str(e)}") + return None + + def _scan_directory(self, directory: str) -> List[str]: + """ + Scan directory for image files + + Args: + directory: Path to directory to scan + + Returns: + List of image file paths + """ + # List of supported image extensions + valid_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.webp'} + + image_files = [] + + try: + # Get absolute path + abs_dir = os.path.abspath(directory) + + if not os.path.exists(abs_dir): + print(f"Directory not found: {abs_dir}") + return image_files + + # Scan directory for image files + print(f"Scanning directory: {abs_dir}") + + for file in os.listdir(abs_dir): + # Check file extension + if os.path.splitext(file)[1].lower() in valid_extensions: + full_path = os.path.join(abs_dir, file) + image_files.append(full_path) + + # Sort files for consistent ordering + image_files.sort() + + print(f"Found {len(image_files)} image files") + + except Exception as e: + print(f"Error scanning directory {directory}: {str(e)}") + + return image_files + + def feed_images(self, directory: str, frame_number: int) -> Tuple[torch.Tensor, Dict]: + """ + Main processing function + + Args: + directory: Path to images directory + frame_number: Index of frame to return + + Returns: + Tuple containing: + - Selected frame as tensor + - Dictionary with frames data + """ + # Check if directory changed + if directory != self.current_dir: + print(f"New directory detected, scanning: {directory}") + self.image_files = self._scan_directory(directory) + self.current_dir = directory + + # Prepare return data with frame to file mapping + frames_data = { + "frames": { + idx: file_path + for idx, file_path in enumerate(self.image_files) + }, + "current_frame": frame_number + } + + # Handle empty directory case + if not self.image_files: + print("No images found in directory") + # Return empty image and data + empty_image = torch.zeros((1, 64, 64, 3)) + return (empty_image, frames_data) + + # Validate frame number + max_frame = len(self.image_files) - 1 + if frame_number > max_frame: + print(f"Warning: Requested frame {frame_number} is out of range. Using last available frame ({max_frame})") + frame_number = max_frame + elif frame_number < 0: + print(f"Warning: Requested frame {frame_number} is negative. Using first frame (0)") + frame_number = 0 + + # Load selected frame + selected_frame = self._load_image(self.image_files[frame_number]) + + 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 (selected_frame, frames_data) diff --git a/nodes/image_filters.py b/nodes/image_filters.py new file mode 100644 index 0000000..2174fe0 --- /dev/null +++ b/nodes/image_filters.py @@ -0,0 +1,114 @@ +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""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "blur_radius": ("INT", { + "default": 2, + "min": 1, + "max": 20, + "step": 1 + }), + "blur_iterations": ("INT", { + "default": 6, + "min": 1, + "max": 20, + "step": 1 + }), + "blend_opacity": ("FLOAT", { + "default": 0.5, + "min": 0.45, + "max": 0.55, + "step": 0.01 + }), + "input_black": ("INT", { + "default": 117, + "min": 0, + "max": 255, + "step": 1 + }), + "input_white": ("INT", { + "default": 137, + "min": 0, + "max": 255, + "step": 1 + }), + "gamma": ("FLOAT", { + "default": 0, + "min": 0.5, + "max": 2, + "step": 0.01 + }), + "show_histogram": ("BOOLEAN", { + "default": False + }), + } + } + + RETURN_TYPES = ("IMAGE", "DICT") + RETURN_NAMES = ("image", "HPF Settings") + FUNCTION = "apply_hpf" + CATEGORY = "Face Processor/Tools" + + 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) + + # 1. Apply box blur with multiple iterations + kernel_size = 2 * blur_radius + 1 + kernel = np.ones((kernel_size, kernel_size), np.float32) / (kernel_size ** 2) + blurred = np_img.copy() + for _ in range(blur_iterations): + blurred = cv2.filter2D(blurred, -1, kernel, borderType=cv2.BORDER_REPLICATE) + + # 2. Invert blurred image + inverted = cv2.bitwise_not(blurred) + + # 3. Blend original with inverted using specified opacity + composite = cv2.addWeighted(np_img, 1 - blend_opacity, inverted, blend_opacity, 0) + + # 4. Apply levels adjustment + in_black = np.clip(input_black, 0, 255) + in_white = np.clip(input_white, in_black + 1, 255) + levels = np.clip((composite.astype(np.float32) - in_black) / (in_white - in_black), 0, 1) + levels = np.power(levels, 1.0 / gamma) if gamma > 0 else levels + result = (levels * 255).astype(np.uint8) + + # Add histogram visualization if enabled + if show_histogram: + hist_img = ImageProcessor.draw_dynamic_histogram(composite, in_black, in_white, gamma) + target_width = result.shape[1] + hist_img = cv2.resize(hist_img, (target_width, 200)) + result = np.vstack([result, hist_img]) + + # Convert back to tensor format + result = torch.from_numpy(result.astype(np.float32) / 255.0).unsqueeze(0) + + # Prepare settings dictionary + settings = { + "high_pass_filter": { + "blur_radius": blur_radius, + "blur_iterations": blur_iterations, + "blend_opacity": blend_opacity, + "input_black": input_black, + "input_white": input_white, + "gamma": gamma, + "show_histogram": False + } + } + + return result, settings diff --git a/nodes/tools.py b/nodes/tools.py deleted file mode 100644 index 3c0c7f8..0000000 --- a/nodes/tools.py +++ /dev/null @@ -1,71 +0,0 @@ -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""" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - "blur_radius": ("INT", {"default": 2, "min": 1, "max": 20, "step": 1}), - "blur_iterations": ("INT", {"default": 6, "min": 1, "max": 20, "step": 1}), - "blend_opacity": ("FLOAT", {"default": 0.5, "min": 0.45, "max": 0.55, "step": 0.01}), - "input_black": ("INT", {"default": 117, "min": 0, "max": 255, "step": 1}), - "input_white": ("INT", {"default": 137, "min": 0, "max": 255, "step": 1}), - "gamma": ("FLOAT", {"default": 0, "min": 0.5, "max": 2, "step": 0.01}), - "show_histogram": ("BOOLEAN", {"default": False}), - } - } - - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("image",) - FUNCTION = "apply_effect" - CATEGORY = "Face Processor/Tools" - - def apply_effect(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]: - """Main processing function implementing the AE-style effect workflow""" - # Convert tensor to numpy array - np_img = image[0].numpy() * 255 - np_img = np.clip(np_img, 0, 255).astype(np.uint8) - - # 1. Apply box blur with multiple iterations - kernel_size = 2 * blur_radius + 1 - kernel = np.ones((kernel_size, kernel_size), np.float32) / (kernel_size ** 2) - blurred = np_img.copy() - for _ in range(blur_iterations): - blurred = cv2.filter2D(blurred, -1, kernel, borderType=cv2.BORDER_REPLICATE) - - # 2. Invert blurred image - inverted = cv2.bitwise_not(blurred) - - # 3. Blend original with inverted using specified opacity - composite = cv2.addWeighted(np_img, 1 - blend_opacity, inverted, blend_opacity, 0) - - # 4. Apply levels adjustment - in_black = np.clip(input_black, 0, 255) - in_white = np.clip(input_white, in_black + 1, 255) - levels = np.clip((composite.astype(np.float32) - in_black) / (in_white - in_black), 0, 1) - levels = np.power(levels, 1.0 / gamma) - result = (levels * 255).astype(np.uint8) - - # Add histogram visualization if enabled - if show_histogram: - hist_img = ImageProcessor.draw_dynamic_histogram(composite, in_black, in_white, gamma) - target_width = result.shape[1] - hist_img = cv2.resize(hist_img, (target_width, 200)) - result = np.vstack([result, hist_img]) - - # Convert back to tensor format - result = torch.from_numpy(result.astype(np.float32) / 255.0).unsqueeze(0) - - return (result,) \ No newline at end of file