diff --git a/__init__.py b/__init__.py index 4b1f849..f47c110 100644 --- a/__init__.py +++ b/__init__.py @@ -101,6 +101,7 @@ NODE_CONFIG = { "ImageTransformByNormalizedAmplitude": {"class": ImageTransformByNormalizedAmplitude}, #curve nodes "SplineEditor": {"class": SplineEditor, "name": "Spline Editor"}, + "CreateShapeImageOnPath": {"class": CreateShapeImageOnPath, "name": "Create Shape Image On Path"}, "CreateShapeMaskOnPath": {"class": CreateShapeMaskOnPath, "name": "Create Shape Mask On Path"}, "CreateTextOnPath": {"class": CreateTextOnPath, "name": "Create Text On Path"}, "CreateGradientFromCoords": {"class": CreateGradientFromCoords, "name": "Create Gradient From Coords"}, diff --git a/nodes/curve_nodes.py b/nodes/curve_nodes.py index dea5c03..93112d6 100644 --- a/nodes/curve_nodes.py +++ b/nodes/curve_nodes.py @@ -1,7 +1,7 @@ import torch from torchvision import transforms import json -from PIL import Image, ImageDraw, ImageFont, ImageColor +from PIL import Image, ImageDraw, ImageFont, ImageColor, ImageFilter import numpy as np from ..utility.utility import pil2tensor import folder_paths @@ -302,6 +302,93 @@ Locations are center locations. out.append(mask) outstack = torch.cat(out, dim=0) return (outstack, 1.0 - outstack,) + +class CreateShapeImageOnPath: + + RETURN_TYPES = ("IMAGE", ) + RETURN_NAMES = ("image", ) + FUNCTION = "createshapemask" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Creates an image or batch of images with the specified shape. +Locations are center locations. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "shape": ( + [ 'circle', + 'square', + 'triangle', + ], + { + "default": 'circle' + }), + "coordinates": ("STRING", {"forceInput": True}), + "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "shape_width": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}), + "shape_height": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}), + "shape_color": ("STRING", {"default": 'white'}), + "bg_color": ("STRING", {"default": 'black'}), + "blur_radius": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100, "step": 0.1}), + "intensity": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01}), + }, + "optional": { + "size_multiplier": ("FLOAT", {"default": [1.0], "forceInput": True}), + } + } + + def createshapemask(self, coordinates, frame_width, frame_height, shape_width, shape_height, shape_color, + bg_color, blur_radius, shape, intensity, size_multiplier=[1.0]): + # Define the number of images in the batch + coordinates = coordinates.replace("'", '"') + coordinates = json.loads(coordinates) + + batch_size = len(coordinates) + out = [] + + if len(size_multiplier) != batch_size: + size_multiplier = size_multiplier * (batch_size // len(size_multiplier)) + size_multiplier[:batch_size % len(size_multiplier)] + for i, coord in enumerate(coordinates): + image = Image.new("RGB", (frame_width, frame_height), bg_color) + draw = ImageDraw.Draw(image) + + # Calculate the size for this frame and ensure it's not less than 0 + current_width = max(0, shape_width + i * size_multiplier[i]) + current_height = max(0, shape_height + i * size_multiplier[i]) + + location_x = coord['x'] + location_y = coord['y'] + + if shape == 'circle' or shape == 'square': + # Define the bounding box for the shape + left_up_point = (location_x - current_width // 2, location_y - current_height // 2) + right_down_point = (location_x + current_width // 2, location_y + current_height // 2) + two_points = [left_up_point, right_down_point] + + if shape == 'circle': + draw.ellipse(two_points, fill=shape_color) + elif shape == 'square': + draw.rectangle(two_points, fill=shape_color) + + elif shape == 'triangle': + # Define the points for the triangle + left_up_point = (location_x - current_width // 2, location_y + current_height // 2) # bottom left + right_down_point = (location_x + current_width // 2, location_y + current_height // 2) # bottom right + top_point = (location_x, location_y - current_height // 2) # top point + draw.polygon([top_point, left_up_point, right_down_point], fill=shape_color) + + if blur_radius != 0: + image = image.filter(ImageFilter.GaussianBlur(blur_radius)) + + image = pil2tensor(image) + image = image * intensity + out.append(image) + outstack = torch.cat(out, dim=0) + return (outstack,) class CreateTextOnPath: