commit 6eba632c96a714bb9671e0fd46d751b048e05d8d Author: Serge Katzmann Date: Sat Nov 25 22:39:39 2023 +0100 Initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ed8ebf5 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__ \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..dac0528 --- /dev/null +++ b/__init__.py @@ -0,0 +1,37 @@ +import os +import subprocess +import importlib.util +import sys + +import __main__ + +python = sys.executable + +def is_installed(package, package_overwrite=None): + try: + spec = importlib.util.find_spec(package) + except ModuleNotFoundError: + pass + + package = package_overwrite or package + + if spec is None: + print(f"Installing {package}...") + command = f'"{python}" -m pip install {package}' + + result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, env=os.environ) + + if result.returncode != 0: + print(f"Couldn't install\nCommand: {command}\nError code: {result.returncode}") + +from .image_fitting_node import ImageSquareFittingNode + +NODE_CLASS_MAPPINGS = { + "ImageSquareAdapterNode": ImageSquareAdapterNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageSquareAdapterNode": "Image Square Adapter Node" +} + +print('\033[34mNimbus Nodes: \033[92mLoaded\033[0m') diff --git a/image_fitting_node.py b/image_fitting_node.py new file mode 100644 index 0000000..27dc14a --- /dev/null +++ b/image_fitting_node.py @@ -0,0 +1,99 @@ +from PIL import Image +import torch +import numpy as np + +from .utils import pil2tensor, tensor2pil + + +class ImageSquareFittingNode: + """ + A custom node for ComfyUI to fit an image into a square frame, + resizing and padding it as necessary, with options for resampling, supersampling, + and various fitting modes. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "target_size": ("INT", {"default": 224, "min": 1, "max": 10000, "step": 1}), + "fill_color": ("STRING", {"default": "255,255,255"}), + "resampling": (["lanczos", "nearest", "bilinear", "bicubic"], {"default": "lanczos"}), + "supersample": (["true", "false"], {"default": "false"}), + "fitting_mode": (["none", "top", "bottom", "center"], {"default": "none"}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "image_fit_in_square" + + CATEGORY = "Nimbus-Pack/Image" + + def image_fit_in_square(self, image, target_size=224, fill_color='255,255,255', resampling='lanczos', + supersample='false', fitting_mode='none'): + scaled_images = [] + for img in image: + scaled_images.append( + self.apply_fit_image(tensor2pil(img), target_size, fill_color, resampling, supersample, fitting_mode)) + + scaled_images = torch.cat(scaled_images, dim=0) + + return (scaled_images,) + + def apply_fit_image(self, image: Image.Image, target_size: int, fill_color: str, resample: str, supersample: str, + fitting_mode: str): + # Convert fill_color string to tuple + fill_color = tuple(map(int, fill_color.split(','))) + + # Define a dictionary of resampling filters + resample_filters = { + 'nearest': Image.NEAREST, + 'bilinear': Image.BILINEAR, + 'bicubic': Image.BICUBIC, + 'lanczos': Image.LANCZOS + } + + # Calculate scaling factor and new size + scaling_factor = target_size / float(max(image.size)) + new_size = tuple([int(x * scaling_factor) for x in image.size]) + + # Apply supersample if needed + if supersample == 'true': + image = image.resize((new_size[0] * 8, new_size[1] * 8), resample=resample_filters[resample]) + + # Resize the image + image = image.resize(new_size, resample=resample_filters[resample]) + + # Adjust image fitting based on the mode + if fitting_mode == 'none': + # Current behavior - centering the image + new_img = Image.new("RGB", (target_size, target_size), fill_color) + position = ((target_size - new_size[0]) // 2, (target_size - new_size[1]) // 2) + new_img.paste(image, position) + else: + # Resize width to target size, adjust height placement based on the fitting_mode + width, height = image.size + new_height = int(height * (target_size / float(width))) + image = image.resize((target_size, new_height), resample=resample_filters[resample]) + new_img = Image.new("RGB", (target_size, target_size), fill_color) + + if fitting_mode == 'top': + position = (0, 0) + elif fitting_mode == 'bottom': + position = (0, target_size - new_height) + elif fitting_mode == 'center': + position = (0, (target_size - new_height) // 2) + + new_img.paste(image, position) + + return pil2tensor(new_img) + + +NODE_CLASS_MAPPINGS = { + "ImageSquareAdapterNode": ImageSquareAdapterNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageSquareAdapterNode": "Image Square Adapter Node" +} diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..9fd3234 --- /dev/null +++ b/utils.py @@ -0,0 +1,11 @@ +from PIL import Image +import torch +import numpy as np + +# Tensor to PIL +def tensor2pil(img): + return Image.fromarray(np.clip(255. * img.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + +# PIL to Tensor +def pil2tensor(img): + return torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0) \ No newline at end of file