From 64dd52cffd33c9259d9bc87fceb217bb19866f20 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20T=C5=82uchowski?= Date: Mon, 27 Jan 2025 21:05:09 +0100 Subject: [PATCH] new tool - High Pass Filter (HPF) --- __init__.py | 10 +++-- nodes/tools.py | 107 +++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 4 deletions(-) create mode 100644 nodes/tools.py diff --git a/__init__.py b/__init__.py index f13a62a..f6b6485 100644 --- a/__init__.py +++ b/__init__.py @@ -2,6 +2,7 @@ import os import sys +from .nodes.tools import HighPassFilter from .nodes.face_wrapper import FaceWrapper from .nodes.face_fit_and_restore import FaceFitAndRestore @@ -11,15 +12,16 @@ sys.path.append(os.path.dirname(os.path.realpath(__file__))) # Get the path to the current directory NODE_PATH = os.path.dirname(os.path.realpath(__file__)) - NODE_CLASS_MAPPINGS = { "FaceFitAndRestore": FaceFitAndRestore, - "FaceWrapper": FaceWrapper + "FaceWrapper": FaceWrapper, + "HighPassFilter": HighPassFilter } NODE_DISPLAY_NAME_MAPPINGS = { - "FaceFitAndRestore": "Face Fit and Restore", - "FaceWrapper": "Face Wrapper" + "FaceFitAndRestore": "Face Fit or Restore", + "FaceWrapper": "Face Wrapper", + "HighPassFilter": "High Pass Filter (HPF)" } __version__ = "1.0.1" diff --git a/nodes/tools.py b/nodes/tools.py new file mode 100644 index 0000000..a198eec --- /dev/null +++ b/nodes/tools.py @@ -0,0 +1,107 @@ +import torch +import numpy as np +import cv2 +from typing import Tuple + +from torch import Tensor + + +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 draw_dynamic_histogram(self, composite: np.ndarray, + input_black: int, input_white: int, + gamma: float) -> np.ndarray: + """Creates histogram visualization with dynamic control lines""" + # Convert to grayscale for histogram calculation + gray = cv2.cvtColor(composite, cv2.COLOR_RGB2GRAY) if len(composite.shape) == 3 else composite + + # Calculate histogram + hist = cv2.calcHist([gray], [0], None, [256], [0, 256]) + hist_img = np.zeros((200, 256, 3), dtype=np.uint8) + cv2.normalize(hist, hist, 0, hist_img.shape[0], cv2.NORM_MINMAX) + + # Draw histogram baseline + for i in range(1, 256): + cv2.line(hist_img, + (i - 1, 200 - int(hist[i - 1])), + (i, 200 - int(hist[i])), + (128, 128, 128), 1) + + # Calculate gamma position + gamma_pos = int(255 * (0.5 ** (1 / gamma))) + + # Draw control lines + cv2.line(hist_img, (input_black, 0), (input_black, 200), (255, 0, 0), 2) # Blue - black level + cv2.line(hist_img, (input_white, 0), (input_white, 200), (0, 255, 0), 2) # Green - white level + cv2.line(hist_img, (gamma_pos, 0), (gamma_pos, 200), (255, 255, 255), 2) # White - gamma + + # Add legend + cv2.putText(hist_img, f"Black: {input_black}", (10, 20), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) + cv2.putText(hist_img, f"White: {input_white}", (10, 40), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) + cv2.putText(hist_img, f"Gamma: {gamma:.2f}", (10, 60), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) + + return hist_img + + 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 = self.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