Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update mask_nodes.py #86

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Update mask_nodes.py #86

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 68 additions & 19 deletions nodes/mask_nodes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -771,8 +771,8 @@ class CreateShapeMask:
FUNCTION = "createshapemask"
CATEGORY = "KJNodes/masking/generate"
DESCRIPTION = """
Creates a mask or batch of masks with the specified shape.
Locations are center locations.
Creates a mask or batch of masks with the specified shape and movement trajectory.
Locations are the starting center locations.
Grow value is the amount to grow the shape on each frame, creating animated masks.
"""

Expand All @@ -781,25 +781,71 @@ def INPUT_TYPES(s):
return {
"required": {
"shape": (
[ 'circle',
'square',
'triangle',
],
{
"default": 'circle'
}),
"frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}),
"location_x": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}),
"location_y": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}),
['circle', 'square', 'triangle'],
{"default": 'circle'}
),
"frames": ("INT", {"default": 1, "min": 1, "max": 4096, "step": 1}),
"start_location_x": ("INT", {"default": 256, "min": 0, "max": 4096, "step": 1}),
"start_location_y": ("INT", {"default": 256, "min": 0, "max": 4096, "step": 1}),
"grow": ("INT", {"default": 0, "min": -512, "max": 512, "step": 1}),
"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}),
},
}
"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}),
"movement_type": (
['none', 'linear', 'circular', 'zigzag', 'bounce', 'sinusoidal', 'spiral', 'random'],
{"default": 'none'}
),
"end_location_x": ("INT", {"default": 256, "min": 0, "max": 4096, "step": 1}),
"end_location_y": ("INT", {"default": 256, "min": 0, "max": 4096, "step": 1}),
"radius": ("INT", {"default": 128, "min": 0, "max": 4096, "step": 1}),
"angle_step": ("FLOAT", {"default": 10.0, "min": 0.1, "max": 360.0, "step": 0.1}),
"zigzag_amplitude": ("INT", {"default": 20, "min": 1, "max": 512, "step": 1}),
"bounce_height": ("INT", {"default": 50, "min": 1, "max": 512, "step": 1}),
"sin_amplitude": ("INT", {"default": 50, "min": 1, "max": 512, "step": 1}),
"spiral_tightness": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 1.0, "step": 0.01}),
},
}

def createshapemask(self, frames, frame_width, frame_height, location_x, location_y, shape_width, shape_height, grow, shape):
def createshapemask(self, frames, frame_width, frame_height, start_location_x, start_location_y, shape_width, shape_height, grow, shape, movement_type, end_location_x, end_location_y, radius, angle_step, zigzag_amplitude, bounce_height, sin_amplitude, spiral_tightness):
from math import sin, cos, radians, pi
import random

def get_position(i, movement_type, start_x, start_y, end_x, end_y, radius, angle_step, zigzag_amplitude, bounce_height, sin_amplitude, spiral_tightness):
if movement_type == 'linear':
fraction = i / (frames - 1)
return start_x + fraction * (end_x - start_x), start_y + fraction * (end_y - start_y)
elif movement_type == 'circular':
angle = angle_step * i
return start_x + radius * cos(radians(angle)), start_y + radius * sin(radians(angle))
elif movement_type == 'zigzag':
fraction = i / (frames - 1)
zigzag_offset = zigzag_amplitude * sin(2 * pi * fraction * 10)
return start_x + fraction * (end_x - start_x), start_y + zigzag_offset
elif movement_type == 'bounce':
fraction = i / (frames - 1)
bounce_offset = bounce_height * abs(sin(pi * fraction * 2))
return start_x + fraction * (end_x - start_x), start_y + bounce_offset
elif movement_type == 'sinusoidal':
fraction = i / (frames - 1)
sin_offset = sin_amplitude * sin(2 * pi * fraction)
return start_x + fraction * (end_x - start_x), start_y + sin_offset
elif movement_type == 'spiral':
angle = i * spiral_tightness
r = radius * (i / frames)
return start_x + r * cos(angle), start_y + r * sin(angle)
elif movement_type == 'random':
if i == 0: # initialize previous_positions if first frame
self.previous_positions = (start_x, start_y)
else:
prev_x, prev_y = self.previous_positions
new_x = max(0, min(frame_width, prev_x + random.randint(-10, 10)))
new_y = max(0, min(frame_height, prev_y + random.randint(-10, 10)))
self.previous_positions = (new_x, new_y)
return self.previous_positions
else:
return start_x, start_y

# Define the number of images in the batch
batch_size = frames
out = []
Expand All @@ -811,6 +857,9 @@ def createshapemask(self, frames, frame_width, frame_height, location_x, locatio
# Calculate the size for this frame and ensure it's not less than 0
current_width = max(0, shape_width + i*grow)
current_height = max(0, shape_height + i*grow)

# Get the current position based on the movement type
location_x, location_y = get_position(i, movement_type, start_location_x, start_location_y, end_location_x, end_location_y, radius, angle_step, zigzag_amplitude, bounce_height, sin_amplitude, spiral_tightness)

if shape == 'circle' or shape == 'square':
# Define the bounding box for the shape
Expand Down