mirror of
https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git
synced 2025-12-08 20:34:35 +08:00
Continue restructuring
This commit is contained in:
kijai
parent
6510f7eda9
commit
98add2716b
11
README.md
11
README.md
@ -17,7 +17,8 @@ This is still work in progress, like everything else.
|
||||
## Javascript
|
||||
|
||||
### browserstatus.js
|
||||
Sets the favicon to green circle when not processing anything, sets it to red when processing and shows progress percentage and the lenghth of your queue. Might clash with other scripts that affect the page title, delete this file to disable (until I figure out how to add options).
|
||||
Sets the favicon to green circle when not processing anything, sets it to red when processing and shows progress percentage and the lenghth of your queue.
|
||||
Default off, needs to be enabled from options, overrides Custom-Scripts favicon when enabled.
|
||||
|
||||
## Nodes:
|
||||
|
||||
@ -47,14 +48,6 @@ Mask and combine two sets of conditions, saves space.
|
||||
|
||||
Grows or shrinks (with negative values) mask, option to invert input, returns mask and inverted mask. Additionally Blurs the mask, this is a slow operation especially with big batches.
|
||||
|
||||
### CreateFadeMask
|
||||
|
||||
This node creates batch of single color images by interpolating between white/black levels. Useful to control mask strengths or QR code controlnet input weight when combined with MaskComposite node.
|
||||
|
||||
### CreateAudioMask
|
||||
|
||||
Work in progress, currently creates a sphere that's size is synced with audio input.
|
||||
|
||||
### RoundMask
|
||||
|
||||
|
||||
|
||||
117
__init__.py
117
__init__.py
@ -1,15 +1,26 @@
|
||||
from .nodes import *
|
||||
from .curve_nodes import *
|
||||
from .nodes.nodes import *
|
||||
from .nodes.curve_nodes import *
|
||||
from .nodes.batchcrop_nodes import *
|
||||
from .nodes.audioscheduler_nodes import *
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
#constants
|
||||
"INTConstant": INTConstant,
|
||||
"FloatConstant": FloatConstant,
|
||||
"ImageBatchMulti": ImageBatchMulti,
|
||||
"MaskBatchMulti": MaskBatchMulti,
|
||||
"StringConstant": StringConstant,
|
||||
"StringConstantMultiline": StringConstantMultiline,
|
||||
#conditioning
|
||||
"ConditioningMultiCombine": ConditioningMultiCombine,
|
||||
"ConditioningSetMaskAndCombine": ConditioningSetMaskAndCombine,
|
||||
"ConditioningSetMaskAndCombine3": ConditioningSetMaskAndCombine3,
|
||||
"ConditioningSetMaskAndCombine4": ConditioningSetMaskAndCombine4,
|
||||
"ConditioningSetMaskAndCombine5": ConditioningSetMaskAndCombine5,
|
||||
"CondPassThrough": CondPassThrough,
|
||||
#masking
|
||||
"BatchCLIPSeg": BatchCLIPSeg,
|
||||
"RoundMask": RoundMask,
|
||||
"ResizeMask": ResizeMask,
|
||||
"OffsetMask": OffsetMask,
|
||||
"MaskBatchMulti": MaskBatchMulti,
|
||||
"GrowMaskWithBlur": GrowMaskWithBlur,
|
||||
"ColorToMask": ColorToMask,
|
||||
"CreateGradientMask": CreateGradientMask,
|
||||
@ -18,11 +29,14 @@ NODE_CLASS_MAPPINGS = {
|
||||
"CreateFadeMask": CreateFadeMask,
|
||||
"CreateFadeMaskAdvanced": CreateFadeMaskAdvanced,
|
||||
"CreateFluidMask" :CreateFluidMask,
|
||||
"VRAM_Debug" : VRAM_Debug,
|
||||
"SomethingToString" : SomethingToString,
|
||||
"CrossFadeImages": CrossFadeImages,
|
||||
"EmptyLatentImagePresets": EmptyLatentImagePresets,
|
||||
"CreateShapeMask": CreateShapeMask,
|
||||
"CreateVoronoiMask": CreateVoronoiMask,
|
||||
"CreateMagicMask": CreateMagicMask,
|
||||
"RemapMaskRange": RemapMaskRange,
|
||||
#images
|
||||
"ImageBatchMulti": ImageBatchMulti,
|
||||
"ColorMatch": ColorMatch,
|
||||
"CrossFadeImages": CrossFadeImages,
|
||||
"GetImageRangeFromBatch": GetImageRangeFromBatch,
|
||||
"SaveImageWithAlpha": SaveImageWithAlpha,
|
||||
"ReverseImageBatch": ReverseImageBatch,
|
||||
@ -31,65 +45,63 @@ NODE_CLASS_MAPPINGS = {
|
||||
"ImageConcanate": ImageConcanate,
|
||||
"ImageBatchTestPattern": ImageBatchTestPattern,
|
||||
"ReplaceImagesInBatch": ReplaceImagesInBatch,
|
||||
"ImageGrabPIL": ImageGrabPIL,
|
||||
"AddLabel": AddLabel,
|
||||
"ImageUpscaleWithModelBatched": ImageUpscaleWithModelBatched,
|
||||
"GetImagesFromBatchIndexed": GetImagesFromBatchIndexed,
|
||||
"InsertImagesToBatchIndexed": InsertImagesToBatchIndexed,
|
||||
"ImageBatchRepeatInterleaving": ImageBatchRepeatInterleaving,
|
||||
"ImageNormalize_Neg1_To_1": ImageNormalize_Neg1_To_1,
|
||||
"RemapImageRange": RemapImageRange,
|
||||
"ImagePass": ImagePass,
|
||||
"ImagePadForOutpaintMasked": ImagePadForOutpaintMasked,
|
||||
"ImageAndMaskPreview": ImageAndMaskPreview,
|
||||
#batch cropping
|
||||
"BatchCropFromMask": BatchCropFromMask,
|
||||
"BatchCropFromMaskAdvanced": BatchCropFromMaskAdvanced,
|
||||
"FilterZeroMasksAndCorrespondingImages": FilterZeroMasksAndCorrespondingImages,
|
||||
"InsertImageBatchByIndexes": InsertImageBatchByIndexes,
|
||||
"BatchUncrop": BatchUncrop,
|
||||
"BatchUncropAdvanced": BatchUncropAdvanced,
|
||||
"BatchCLIPSeg": BatchCLIPSeg,
|
||||
"RoundMask": RoundMask,
|
||||
"ResizeMask": ResizeMask,
|
||||
"OffsetMask": OffsetMask,
|
||||
"WidgetToString": WidgetToString,
|
||||
"CreateShapeMask": CreateShapeMask,
|
||||
"CreateVoronoiMask": CreateVoronoiMask,
|
||||
"CreateMagicMask": CreateMagicMask,
|
||||
"BboxToInt": BboxToInt,
|
||||
"SplitBboxes": SplitBboxes,
|
||||
"ImageGrabPIL": ImageGrabPIL,
|
||||
"DummyLatentOut": DummyLatentOut,
|
||||
"BboxToInt": BboxToInt,
|
||||
"BboxVisualize": BboxVisualize,
|
||||
#noise
|
||||
"GenerateNoise": GenerateNoise,
|
||||
"FlipSigmasAdjusted": FlipSigmasAdjusted,
|
||||
"InjectNoiseToLatent": InjectNoiseToLatent,
|
||||
"AddLabel": AddLabel,
|
||||
"SoundReactive": SoundReactive,
|
||||
"GenerateNoise": GenerateNoise,
|
||||
"StableZero123_BatchSchedule": StableZero123_BatchSchedule,
|
||||
"SV3D_BatchSchedule": SV3D_BatchSchedule,
|
||||
"GetImagesFromBatchIndexed": GetImagesFromBatchIndexed,
|
||||
"InsertImagesToBatchIndexed": InsertImagesToBatchIndexed,
|
||||
"ImageBatchRepeatInterleaving": ImageBatchRepeatInterleaving,
|
||||
"CustomSigmas": CustomSigmas,
|
||||
#utility
|
||||
"WidgetToString": WidgetToString,
|
||||
"DummyLatentOut": DummyLatentOut,
|
||||
"GetLatentsFromBatchIndexed": GetLatentsFromBatchIndexed,
|
||||
"ScaleBatchPromptSchedule": ScaleBatchPromptSchedule,
|
||||
"CameraPoseVisualizer": CameraPoseVisualizer,
|
||||
"JoinStrings": JoinStrings,
|
||||
"Sleep": Sleep,
|
||||
"VRAM_Debug" : VRAM_Debug,
|
||||
"SomethingToString" : SomethingToString,
|
||||
"EmptyLatentImagePresets": EmptyLatentImagePresets,
|
||||
#audioscheduler stuff
|
||||
"NormalizedAmplitudeToMask": NormalizedAmplitudeToMask,
|
||||
"OffsetMaskByNormalizedAmplitude": OffsetMaskByNormalizedAmplitude,
|
||||
"ImageTransformByNormalizedAmplitude": ImageTransformByNormalizedAmplitude,
|
||||
"GetLatentsFromBatchIndexed": GetLatentsFromBatchIndexed,
|
||||
"StringConstant": StringConstant,
|
||||
"GLIGENTextBoxApplyBatch": GLIGENTextBoxApplyBatch,
|
||||
"CondPassThrough": CondPassThrough,
|
||||
"ImageUpscaleWithModelBatched": ImageUpscaleWithModelBatched,
|
||||
"ScaleBatchPromptSchedule": ScaleBatchPromptSchedule,
|
||||
"ImageNormalize_Neg1_To_1": ImageNormalize_Neg1_To_1,
|
||||
"Intrinsic_lora_sampling": Intrinsic_lora_sampling,
|
||||
"RemapMaskRange": RemapMaskRange,
|
||||
"LoadResAdapterNormalization": LoadResAdapterNormalization,
|
||||
"Superprompt": Superprompt,
|
||||
"RemapImageRange": RemapImageRange,
|
||||
"CameraPoseVisualizer": CameraPoseVisualizer,
|
||||
"BboxVisualize": BboxVisualize,
|
||||
"StringConstantMultiline": StringConstantMultiline,
|
||||
"JoinStrings": JoinStrings,
|
||||
"Sleep": Sleep,
|
||||
"ImagePadForOutpaintMasked": ImagePadForOutpaintMasked,
|
||||
"ImageAndMaskPreview": ImageAndMaskPreview,
|
||||
"StabilityAPI_SD3": StabilityAPI_SD3,
|
||||
#curve nodes
|
||||
"SplineEditor": SplineEditor,
|
||||
"CreateShapeMaskOnPath": CreateShapeMaskOnPath,
|
||||
"WeightScheduleExtend": WeightScheduleExtend,
|
||||
"MaskOrImageToWeight": MaskOrImageToWeight,
|
||||
"WeightScheduleConvert": WeightScheduleConvert,
|
||||
"FloatToMask": FloatToMask,
|
||||
"CustomSigmas": CustomSigmas,
|
||||
"ImagePass": ImagePass,
|
||||
"SplineEditor": SplineEditor,
|
||||
"CreateShapeMaskOnPath": CreateShapeMaskOnPath,
|
||||
"WeightScheduleExtend": WeightScheduleExtend
|
||||
#experimental
|
||||
"StabilityAPI_SD3": StabilityAPI_SD3,
|
||||
"SoundReactive": SoundReactive,
|
||||
"StableZero123_BatchSchedule": StableZero123_BatchSchedule,
|
||||
"SV3D_BatchSchedule": SV3D_BatchSchedule,
|
||||
"LoadResAdapterNormalization": LoadResAdapterNormalization,
|
||||
"Superprompt": Superprompt,
|
||||
"GLIGENTextBoxApplyBatch": GLIGENTextBoxApplyBatch,
|
||||
"Intrinsic_lora_sampling": Intrinsic_lora_sampling,
|
||||
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@ -179,6 +191,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FloatToMask": "Float To Mask",
|
||||
"CustomSigmas": "Custom Sigmas",
|
||||
"ImagePass": "ImagePass",
|
||||
#curve nodes
|
||||
"SplineEditor": "Spline Editor",
|
||||
"CreateShapeMaskOnPath": "Create Shape Mask On Path",
|
||||
"WeightScheduleExtend": "Weight Schedule Extend"
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 1.0 KiB |
BIN
favicon.ico
BIN
favicon.ico
Binary file not shown.
|
Before Width: | Height: | Size: 1006 B |
2
kjweb_async/d3.v6.min.js
vendored
2
kjweb_async/d3.v6.min.js
vendored
File diff suppressed because one or more lines are too long
230
nodes/audioscheduler_nodes.py
Normal file
230
nodes/audioscheduler_nodes.py
Normal file
@ -0,0 +1,230 @@
|
||||
# to be used with https://github.com/a1lazydog/ComfyUI-AudioScheduler
|
||||
import torch
|
||||
from torchvision.transforms import functional as TF
|
||||
from PIL import Image, ImageDraw
|
||||
import numpy as np
|
||||
from ..utility.utility import pil2tensor
|
||||
from nodes import MAX_RESOLUTION
|
||||
|
||||
class NormalizedAmplitudeToMask:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"normalized_amp": ("NORMALIZED_AMPLITUDE",),
|
||||
"width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}),
|
||||
"height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}),
|
||||
"frame_offset": ("INT", {"default": 0,"min": -255, "max": 255, "step": 1}),
|
||||
"location_x": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}),
|
||||
"location_y": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}),
|
||||
"size": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}),
|
||||
"shape": (
|
||||
[
|
||||
'none',
|
||||
'circle',
|
||||
'square',
|
||||
'triangle',
|
||||
],
|
||||
{
|
||||
"default": 'none'
|
||||
}),
|
||||
"color": (
|
||||
[
|
||||
'white',
|
||||
'amplitude',
|
||||
],
|
||||
{
|
||||
"default": 'amplitude'
|
||||
}),
|
||||
},}
|
||||
|
||||
CATEGORY = "KJNodes/audio"
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "convert"
|
||||
DESCRIPTION = """
|
||||
Works as a bridge to the AudioScheduler -nodes:
|
||||
https://github.com/a1lazydog/ComfyUI-AudioScheduler
|
||||
Creates masks based on the normalized amplitude.
|
||||
"""
|
||||
|
||||
def convert(self, normalized_amp, width, height, frame_offset, shape, location_x, location_y, size, color):
|
||||
# Ensure normalized_amp is an array and within the range [0, 1]
|
||||
normalized_amp = np.clip(normalized_amp, 0.0, 1.0)
|
||||
|
||||
# Offset the amplitude values by rolling the array
|
||||
normalized_amp = np.roll(normalized_amp, frame_offset)
|
||||
|
||||
# Initialize an empty list to hold the image tensors
|
||||
out = []
|
||||
# Iterate over each amplitude value to create an image
|
||||
for amp in normalized_amp:
|
||||
# Scale the amplitude value to cover the full range of grayscale values
|
||||
if color == 'amplitude':
|
||||
grayscale_value = int(amp * 255)
|
||||
elif color == 'white':
|
||||
grayscale_value = 255
|
||||
# Convert the grayscale value to an RGB format
|
||||
gray_color = (grayscale_value, grayscale_value, grayscale_value)
|
||||
finalsize = size * amp
|
||||
|
||||
if shape == 'none':
|
||||
shapeimage = Image.new("RGB", (width, height), gray_color)
|
||||
else:
|
||||
shapeimage = Image.new("RGB", (width, height), "black")
|
||||
|
||||
draw = ImageDraw.Draw(shapeimage)
|
||||
if shape == 'circle' or shape == 'square':
|
||||
# Define the bounding box for the shape
|
||||
left_up_point = (location_x - finalsize, location_y - finalsize)
|
||||
right_down_point = (location_x + finalsize,location_y + finalsize)
|
||||
two_points = [left_up_point, right_down_point]
|
||||
|
||||
if shape == 'circle':
|
||||
draw.ellipse(two_points, fill=gray_color)
|
||||
elif shape == 'square':
|
||||
draw.rectangle(two_points, fill=gray_color)
|
||||
|
||||
elif shape == 'triangle':
|
||||
# Define the points for the triangle
|
||||
left_up_point = (location_x - finalsize, location_y + finalsize) # bottom left
|
||||
right_down_point = (location_x + finalsize, location_y + finalsize) # bottom right
|
||||
top_point = (location_x, location_y) # top point
|
||||
draw.polygon([top_point, left_up_point, right_down_point], fill=gray_color)
|
||||
|
||||
shapeimage = pil2tensor(shapeimage)
|
||||
mask = shapeimage[:, :, :, 0]
|
||||
out.append(mask)
|
||||
|
||||
return (torch.cat(out, dim=0),)
|
||||
|
||||
class OffsetMaskByNormalizedAmplitude:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"normalized_amp": ("NORMALIZED_AMPLITUDE",),
|
||||
"mask": ("MASK",),
|
||||
"x": ("INT", { "default": 0, "min": -4096, "max": MAX_RESOLUTION, "step": 1, "display": "number" }),
|
||||
"y": ("INT", { "default": 0, "min": -4096, "max": MAX_RESOLUTION, "step": 1, "display": "number" }),
|
||||
"rotate": ("BOOLEAN", { "default": False }),
|
||||
"angle_multiplier": ("FLOAT", { "default": 0.0, "min": -1.0, "max": 1.0, "step": 0.001, "display": "number" }),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
RETURN_NAMES = ("mask",)
|
||||
FUNCTION = "offset"
|
||||
CATEGORY = "KJNodes/audio"
|
||||
DESCRIPTION = """
|
||||
Works as a bridge to the AudioScheduler -nodes:
|
||||
https://github.com/a1lazydog/ComfyUI-AudioScheduler
|
||||
Offsets masks based on the normalized amplitude.
|
||||
"""
|
||||
|
||||
def offset(self, mask, x, y, angle_multiplier, rotate, normalized_amp):
|
||||
|
||||
# Ensure normalized_amp is an array and within the range [0, 1]
|
||||
offsetmask = mask.clone()
|
||||
normalized_amp = np.clip(normalized_amp, 0.0, 1.0)
|
||||
|
||||
batch_size, height, width = mask.shape
|
||||
|
||||
if rotate:
|
||||
for i in range(batch_size):
|
||||
rotation_amp = int(normalized_amp[i] * (360 * angle_multiplier))
|
||||
rotation_angle = rotation_amp
|
||||
offsetmask[i] = TF.rotate(offsetmask[i].unsqueeze(0), rotation_angle).squeeze(0)
|
||||
if x != 0 or y != 0:
|
||||
for i in range(batch_size):
|
||||
offset_amp = normalized_amp[i] * 10
|
||||
shift_x = min(x*offset_amp, width-1)
|
||||
shift_y = min(y*offset_amp, height-1)
|
||||
if shift_x != 0:
|
||||
offsetmask[i] = torch.roll(offsetmask[i], shifts=int(shift_x), dims=1)
|
||||
if shift_y != 0:
|
||||
offsetmask[i] = torch.roll(offsetmask[i], shifts=int(shift_y), dims=0)
|
||||
|
||||
return offsetmask,
|
||||
|
||||
class ImageTransformByNormalizedAmplitude:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"normalized_amp": ("NORMALIZED_AMPLITUDE",),
|
||||
"zoom_scale": ("FLOAT", { "default": 0.0, "min": -1.0, "max": 1.0, "step": 0.001, "display": "number" }),
|
||||
"x_offset": ("INT", { "default": 0, "min": (1 -MAX_RESOLUTION), "max": MAX_RESOLUTION, "step": 1, "display": "number" }),
|
||||
"y_offset": ("INT", { "default": 0, "min": (1 -MAX_RESOLUTION), "max": MAX_RESOLUTION, "step": 1, "display": "number" }),
|
||||
"cumulative": ("BOOLEAN", { "default": False }),
|
||||
"image": ("IMAGE",),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "amptransform"
|
||||
CATEGORY = "KJNodes/audio"
|
||||
DESCRIPTION = """
|
||||
Works as a bridge to the AudioScheduler -nodes:
|
||||
https://github.com/a1lazydog/ComfyUI-AudioScheduler
|
||||
Transforms image based on the normalized amplitude.
|
||||
"""
|
||||
|
||||
def amptransform(self, image, normalized_amp, zoom_scale, cumulative, x_offset, y_offset):
|
||||
# Ensure normalized_amp is an array and within the range [0, 1]
|
||||
normalized_amp = np.clip(normalized_amp, 0.0, 1.0)
|
||||
transformed_images = []
|
||||
|
||||
# Initialize the cumulative zoom factor
|
||||
prev_amp = 0.0
|
||||
|
||||
for i in range(image.shape[0]):
|
||||
img = image[i] # Get the i-th image in the batch
|
||||
amp = normalized_amp[i] # Get the corresponding amplitude value
|
||||
|
||||
# Incrementally increase the cumulative zoom factor
|
||||
if cumulative:
|
||||
prev_amp += amp
|
||||
amp += prev_amp
|
||||
|
||||
# Convert the image tensor from BxHxWxC to CxHxW format expected by torchvision
|
||||
img = img.permute(2, 0, 1)
|
||||
|
||||
# Convert PyTorch tensor to PIL Image for processing
|
||||
pil_img = TF.to_pil_image(img)
|
||||
|
||||
# Calculate the crop size based on the amplitude
|
||||
width, height = pil_img.size
|
||||
crop_size = int(min(width, height) * (1 - amp * zoom_scale))
|
||||
crop_size = max(crop_size, 1)
|
||||
|
||||
# Calculate the crop box coordinates (centered crop)
|
||||
left = (width - crop_size) // 2
|
||||
top = (height - crop_size) // 2
|
||||
right = (width + crop_size) // 2
|
||||
bottom = (height + crop_size) // 2
|
||||
|
||||
# Crop and resize back to original size
|
||||
cropped_img = TF.crop(pil_img, top, left, crop_size, crop_size)
|
||||
resized_img = TF.resize(cropped_img, (height, width))
|
||||
|
||||
# Convert back to tensor in CxHxW format
|
||||
tensor_img = TF.to_tensor(resized_img)
|
||||
|
||||
# Convert the tensor back to BxHxWxC format
|
||||
tensor_img = tensor_img.permute(1, 2, 0)
|
||||
|
||||
# Offset the image based on the amplitude
|
||||
offset_amp = amp * 10 # Calculate the offset magnitude based on the amplitude
|
||||
shift_x = min(x_offset * offset_amp, img.shape[1] - 1) # Calculate the shift in x direction
|
||||
shift_y = min(y_offset * offset_amp, img.shape[0] - 1) # Calculate the shift in y direction
|
||||
|
||||
# Apply the offset to the image tensor
|
||||
if shift_x != 0:
|
||||
tensor_img = torch.roll(tensor_img, shifts=int(shift_x), dims=1)
|
||||
if shift_y != 0:
|
||||
tensor_img = torch.roll(tensor_img, shifts=int(shift_y), dims=0)
|
||||
|
||||
# Add to the list
|
||||
transformed_images.append(tensor_img)
|
||||
|
||||
# Stack all transformed images into a batch
|
||||
transformed_batch = torch.stack(transformed_images)
|
||||
|
||||
return (transformed_batch,)
|
||||
677
nodes/batchcrop_nodes.py
Normal file
677
nodes/batchcrop_nodes.py
Normal file
@ -0,0 +1,677 @@
|
||||
from ..utility.utility import tensor2pil, pil2tensor
|
||||
from PIL import Image, ImageDraw, ImageFilter
|
||||
import numpy as np
|
||||
import torch
|
||||
from torchvision.transforms import Resize, CenterCrop, InterpolationMode
|
||||
import math
|
||||
|
||||
#based on nodes from mtb https://github.com/melMass/comfy_mtb
|
||||
class BatchCropFromMask:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"original_images": ("IMAGE",),
|
||||
"masks": ("MASK",),
|
||||
"crop_size_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}),
|
||||
"bbox_smooth_alpha": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"IMAGE",
|
||||
"IMAGE",
|
||||
"BBOX",
|
||||
"INT",
|
||||
"INT",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"original_images",
|
||||
"cropped_images",
|
||||
"bboxes",
|
||||
"width",
|
||||
"height",
|
||||
)
|
||||
FUNCTION = "crop"
|
||||
CATEGORY = "KJNodes/masking"
|
||||
|
||||
def smooth_bbox_size(self, prev_bbox_size, curr_bbox_size, alpha):
|
||||
if alpha == 0:
|
||||
return prev_bbox_size
|
||||
return round(alpha * curr_bbox_size + (1 - alpha) * prev_bbox_size)
|
||||
|
||||
def smooth_center(self, prev_center, curr_center, alpha=0.5):
|
||||
if alpha == 0:
|
||||
return prev_center
|
||||
return (
|
||||
round(alpha * curr_center[0] + (1 - alpha) * prev_center[0]),
|
||||
round(alpha * curr_center[1] + (1 - alpha) * prev_center[1])
|
||||
)
|
||||
|
||||
def crop(self, masks, original_images, crop_size_mult, bbox_smooth_alpha):
|
||||
|
||||
bounding_boxes = []
|
||||
cropped_images = []
|
||||
|
||||
self.max_bbox_width = 0
|
||||
self.max_bbox_height = 0
|
||||
|
||||
# First, calculate the maximum bounding box size across all masks
|
||||
curr_max_bbox_width = 0
|
||||
curr_max_bbox_height = 0
|
||||
for mask in masks:
|
||||
_mask = tensor2pil(mask)[0]
|
||||
non_zero_indices = np.nonzero(np.array(_mask))
|
||||
min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1])
|
||||
min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0])
|
||||
width = max_x - min_x
|
||||
height = max_y - min_y
|
||||
curr_max_bbox_width = max(curr_max_bbox_width, width)
|
||||
curr_max_bbox_height = max(curr_max_bbox_height, height)
|
||||
|
||||
# Smooth the changes in the bounding box size
|
||||
self.max_bbox_width = self.smooth_bbox_size(self.max_bbox_width, curr_max_bbox_width, bbox_smooth_alpha)
|
||||
self.max_bbox_height = self.smooth_bbox_size(self.max_bbox_height, curr_max_bbox_height, bbox_smooth_alpha)
|
||||
|
||||
# Apply the crop size multiplier
|
||||
self.max_bbox_width = round(self.max_bbox_width * crop_size_mult)
|
||||
self.max_bbox_height = round(self.max_bbox_height * crop_size_mult)
|
||||
bbox_aspect_ratio = self.max_bbox_width / self.max_bbox_height
|
||||
|
||||
# Then, for each mask and corresponding image...
|
||||
for i, (mask, img) in enumerate(zip(masks, original_images)):
|
||||
_mask = tensor2pil(mask)[0]
|
||||
non_zero_indices = np.nonzero(np.array(_mask))
|
||||
min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1])
|
||||
min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0])
|
||||
|
||||
# Calculate center of bounding box
|
||||
center_x = np.mean(non_zero_indices[1])
|
||||
center_y = np.mean(non_zero_indices[0])
|
||||
curr_center = (round(center_x), round(center_y))
|
||||
|
||||
# If this is the first frame, initialize prev_center with curr_center
|
||||
if not hasattr(self, 'prev_center'):
|
||||
self.prev_center = curr_center
|
||||
|
||||
# Smooth the changes in the center coordinates from the second frame onwards
|
||||
if i > 0:
|
||||
center = self.smooth_center(self.prev_center, curr_center, bbox_smooth_alpha)
|
||||
else:
|
||||
center = curr_center
|
||||
|
||||
# Update prev_center for the next frame
|
||||
self.prev_center = center
|
||||
|
||||
# Create bounding box using max_bbox_width and max_bbox_height
|
||||
half_box_width = round(self.max_bbox_width / 2)
|
||||
half_box_height = round(self.max_bbox_height / 2)
|
||||
min_x = max(0, center[0] - half_box_width)
|
||||
max_x = min(img.shape[1], center[0] + half_box_width)
|
||||
min_y = max(0, center[1] - half_box_height)
|
||||
max_y = min(img.shape[0], center[1] + half_box_height)
|
||||
|
||||
# Append bounding box coordinates
|
||||
bounding_boxes.append((min_x, min_y, max_x - min_x, max_y - min_y))
|
||||
|
||||
# Crop the image from the bounding box
|
||||
cropped_img = img[min_y:max_y, min_x:max_x, :]
|
||||
|
||||
# Calculate the new dimensions while maintaining the aspect ratio
|
||||
new_height = min(cropped_img.shape[0], self.max_bbox_height)
|
||||
new_width = round(new_height * bbox_aspect_ratio)
|
||||
|
||||
# Resize the image
|
||||
resize_transform = Resize((new_height, new_width))
|
||||
resized_img = resize_transform(cropped_img.permute(2, 0, 1))
|
||||
|
||||
# Perform the center crop to the desired size
|
||||
crop_transform = CenterCrop((self.max_bbox_height, self.max_bbox_width)) # swap the order here if necessary
|
||||
cropped_resized_img = crop_transform(resized_img)
|
||||
|
||||
cropped_images.append(cropped_resized_img.permute(1, 2, 0))
|
||||
|
||||
cropped_out = torch.stack(cropped_images, dim=0)
|
||||
|
||||
return (original_images, cropped_out, bounding_boxes, self.max_bbox_width, self.max_bbox_height, )
|
||||
|
||||
|
||||
def bbox_to_region(bbox, target_size=None):
|
||||
bbox = bbox_check(bbox, target_size)
|
||||
return (bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3])
|
||||
|
||||
def bbox_check(bbox, target_size=None):
|
||||
if not target_size:
|
||||
return bbox
|
||||
|
||||
new_bbox = (
|
||||
bbox[0],
|
||||
bbox[1],
|
||||
min(target_size[0] - bbox[0], bbox[2]),
|
||||
min(target_size[1] - bbox[1], bbox[3]),
|
||||
)
|
||||
return new_bbox
|
||||
|
||||
class BatchUncrop:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"original_images": ("IMAGE",),
|
||||
"cropped_images": ("IMAGE",),
|
||||
"bboxes": ("BBOX",),
|
||||
"border_blending": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}, ),
|
||||
"crop_rescale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
||||
"border_top": ("BOOLEAN", {"default": True}),
|
||||
"border_bottom": ("BOOLEAN", {"default": True}),
|
||||
"border_left": ("BOOLEAN", {"default": True}),
|
||||
"border_right": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "uncrop"
|
||||
|
||||
CATEGORY = "KJNodes/masking"
|
||||
|
||||
def uncrop(self, original_images, cropped_images, bboxes, border_blending, crop_rescale, border_top, border_bottom, border_left, border_right):
|
||||
def inset_border(image, border_width, border_color, border_top, border_bottom, border_left, border_right):
|
||||
draw = ImageDraw.Draw(image)
|
||||
width, height = image.size
|
||||
if border_top:
|
||||
draw.rectangle((0, 0, width, border_width), fill=border_color)
|
||||
if border_bottom:
|
||||
draw.rectangle((0, height - border_width, width, height), fill=border_color)
|
||||
if border_left:
|
||||
draw.rectangle((0, 0, border_width, height), fill=border_color)
|
||||
if border_right:
|
||||
draw.rectangle((width - border_width, 0, width, height), fill=border_color)
|
||||
return image
|
||||
|
||||
if len(original_images) != len(cropped_images):
|
||||
raise ValueError(f"The number of original_images ({len(original_images)}) and cropped_images ({len(cropped_images)}) should be the same")
|
||||
|
||||
# Ensure there are enough bboxes, but drop the excess if there are more bboxes than images
|
||||
if len(bboxes) > len(original_images):
|
||||
print(f"Warning: Dropping excess bounding boxes. Expected {len(original_images)}, but got {len(bboxes)}")
|
||||
bboxes = bboxes[:len(original_images)]
|
||||
elif len(bboxes) < len(original_images):
|
||||
raise ValueError("There should be at least as many bboxes as there are original and cropped images")
|
||||
|
||||
input_images = tensor2pil(original_images)
|
||||
crop_imgs = tensor2pil(cropped_images)
|
||||
|
||||
out_images = []
|
||||
for i in range(len(input_images)):
|
||||
img = input_images[i]
|
||||
crop = crop_imgs[i]
|
||||
bbox = bboxes[i]
|
||||
|
||||
# uncrop the image based on the bounding box
|
||||
bb_x, bb_y, bb_width, bb_height = bbox
|
||||
|
||||
paste_region = bbox_to_region((bb_x, bb_y, bb_width, bb_height), img.size)
|
||||
|
||||
# scale factors
|
||||
scale_x = crop_rescale
|
||||
scale_y = crop_rescale
|
||||
|
||||
# scaled paste_region
|
||||
paste_region = (round(paste_region[0]*scale_x), round(paste_region[1]*scale_y), round(paste_region[2]*scale_x), round(paste_region[3]*scale_y))
|
||||
|
||||
# rescale the crop image to fit the paste_region
|
||||
crop = crop.resize((round(paste_region[2]-paste_region[0]), round(paste_region[3]-paste_region[1])))
|
||||
crop_img = crop.convert("RGB")
|
||||
|
||||
if border_blending > 1.0:
|
||||
border_blending = 1.0
|
||||
elif border_blending < 0.0:
|
||||
border_blending = 0.0
|
||||
|
||||
blend_ratio = (max(crop_img.size) / 2) * float(border_blending)
|
||||
|
||||
blend = img.convert("RGBA")
|
||||
mask = Image.new("L", img.size, 0)
|
||||
|
||||
mask_block = Image.new("L", (paste_region[2]-paste_region[0], paste_region[3]-paste_region[1]), 255)
|
||||
mask_block = inset_border(mask_block, round(blend_ratio / 2), (0), border_top, border_bottom, border_left, border_right)
|
||||
|
||||
mask.paste(mask_block, paste_region)
|
||||
blend.paste(crop_img, paste_region)
|
||||
|
||||
mask = mask.filter(ImageFilter.BoxBlur(radius=blend_ratio / 4))
|
||||
mask = mask.filter(ImageFilter.GaussianBlur(radius=blend_ratio / 4))
|
||||
|
||||
blend.putalpha(mask)
|
||||
img = Image.alpha_composite(img.convert("RGBA"), blend)
|
||||
out_images.append(img.convert("RGB"))
|
||||
|
||||
return (pil2tensor(out_images),)
|
||||
|
||||
class BatchCropFromMaskAdvanced:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"original_images": ("IMAGE",),
|
||||
"masks": ("MASK",),
|
||||
"crop_size_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
||||
"bbox_smooth_alpha": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"IMAGE",
|
||||
"IMAGE",
|
||||
"MASK",
|
||||
"IMAGE",
|
||||
"MASK",
|
||||
"BBOX",
|
||||
"BBOX",
|
||||
"INT",
|
||||
"INT",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"original_images",
|
||||
"cropped_images",
|
||||
"cropped_masks",
|
||||
"combined_crop_image",
|
||||
"combined_crop_masks",
|
||||
"bboxes",
|
||||
"combined_bounding_box",
|
||||
"bbox_width",
|
||||
"bbox_height",
|
||||
)
|
||||
FUNCTION = "crop"
|
||||
CATEGORY = "KJNodes/masking"
|
||||
|
||||
def smooth_bbox_size(self, prev_bbox_size, curr_bbox_size, alpha):
|
||||
return round(alpha * curr_bbox_size + (1 - alpha) * prev_bbox_size)
|
||||
|
||||
def smooth_center(self, prev_center, curr_center, alpha=0.5):
|
||||
return (round(alpha * curr_center[0] + (1 - alpha) * prev_center[0]),
|
||||
round(alpha * curr_center[1] + (1 - alpha) * prev_center[1]))
|
||||
|
||||
def crop(self, masks, original_images, crop_size_mult, bbox_smooth_alpha):
|
||||
bounding_boxes = []
|
||||
combined_bounding_box = []
|
||||
cropped_images = []
|
||||
cropped_masks = []
|
||||
cropped_masks_out = []
|
||||
combined_crop_out = []
|
||||
combined_cropped_images = []
|
||||
combined_cropped_masks = []
|
||||
|
||||
def calculate_bbox(mask):
|
||||
non_zero_indices = np.nonzero(np.array(mask))
|
||||
|
||||
# handle empty masks
|
||||
min_x, max_x, min_y, max_y = 0, 0, 0, 0
|
||||
if len(non_zero_indices[1]) > 0 and len(non_zero_indices[0]) > 0:
|
||||
min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1])
|
||||
min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0])
|
||||
|
||||
width = max_x - min_x
|
||||
height = max_y - min_y
|
||||
bbox_size = max(width, height)
|
||||
return min_x, max_x, min_y, max_y, bbox_size
|
||||
|
||||
combined_mask = torch.max(masks, dim=0)[0]
|
||||
_mask = tensor2pil(combined_mask)[0]
|
||||
new_min_x, new_max_x, new_min_y, new_max_y, combined_bbox_size = calculate_bbox(_mask)
|
||||
center_x = (new_min_x + new_max_x) / 2
|
||||
center_y = (new_min_y + new_max_y) / 2
|
||||
half_box_size = round(combined_bbox_size // 2)
|
||||
new_min_x = max(0, round(center_x - half_box_size))
|
||||
new_max_x = min(original_images[0].shape[1], round(center_x + half_box_size))
|
||||
new_min_y = max(0, round(center_y - half_box_size))
|
||||
new_max_y = min(original_images[0].shape[0], round(center_y + half_box_size))
|
||||
|
||||
combined_bounding_box.append((new_min_x, new_min_y, new_max_x - new_min_x, new_max_y - new_min_y))
|
||||
|
||||
self.max_bbox_size = 0
|
||||
|
||||
# First, calculate the maximum bounding box size across all masks
|
||||
curr_max_bbox_size = max(calculate_bbox(tensor2pil(mask)[0])[-1] for mask in masks)
|
||||
# Smooth the changes in the bounding box size
|
||||
self.max_bbox_size = self.smooth_bbox_size(self.max_bbox_size, curr_max_bbox_size, bbox_smooth_alpha)
|
||||
# Apply the crop size multiplier
|
||||
self.max_bbox_size = round(self.max_bbox_size * crop_size_mult)
|
||||
# Make sure max_bbox_size is divisible by 16, if not, round it upwards so it is
|
||||
self.max_bbox_size = math.ceil(self.max_bbox_size / 16) * 16
|
||||
|
||||
if self.max_bbox_size > original_images[0].shape[0] or self.max_bbox_size > original_images[0].shape[1]:
|
||||
# max_bbox_size can only be as big as our input's width or height, and it has to be even
|
||||
self.max_bbox_size = math.floor(min(original_images[0].shape[0], original_images[0].shape[1]) / 2) * 2
|
||||
|
||||
# Then, for each mask and corresponding image...
|
||||
for i, (mask, img) in enumerate(zip(masks, original_images)):
|
||||
_mask = tensor2pil(mask)[0]
|
||||
non_zero_indices = np.nonzero(np.array(_mask))
|
||||
|
||||
# check for empty masks
|
||||
if len(non_zero_indices[0]) > 0 and len(non_zero_indices[1]) > 0:
|
||||
min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1])
|
||||
min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0])
|
||||
|
||||
# Calculate center of bounding box
|
||||
center_x = np.mean(non_zero_indices[1])
|
||||
center_y = np.mean(non_zero_indices[0])
|
||||
curr_center = (round(center_x), round(center_y))
|
||||
|
||||
# If this is the first frame, initialize prev_center with curr_center
|
||||
if not hasattr(self, 'prev_center'):
|
||||
self.prev_center = curr_center
|
||||
|
||||
# Smooth the changes in the center coordinates from the second frame onwards
|
||||
if i > 0:
|
||||
center = self.smooth_center(self.prev_center, curr_center, bbox_smooth_alpha)
|
||||
else:
|
||||
center = curr_center
|
||||
|
||||
# Update prev_center for the next frame
|
||||
self.prev_center = center
|
||||
|
||||
# Create bounding box using max_bbox_size
|
||||
half_box_size = self.max_bbox_size // 2
|
||||
min_x = max(0, center[0] - half_box_size)
|
||||
max_x = min(img.shape[1], center[0] + half_box_size)
|
||||
min_y = max(0, center[1] - half_box_size)
|
||||
max_y = min(img.shape[0], center[1] + half_box_size)
|
||||
|
||||
# Append bounding box coordinates
|
||||
bounding_boxes.append((min_x, min_y, max_x - min_x, max_y - min_y))
|
||||
|
||||
# Crop the image from the bounding box
|
||||
cropped_img = img[min_y:max_y, min_x:max_x, :]
|
||||
cropped_mask = mask[min_y:max_y, min_x:max_x]
|
||||
|
||||
# Resize the cropped image to a fixed size
|
||||
new_size = max(cropped_img.shape[0], cropped_img.shape[1])
|
||||
resize_transform = Resize(new_size, interpolation=InterpolationMode.NEAREST, max_size=max(img.shape[0], img.shape[1]))
|
||||
resized_mask = resize_transform(cropped_mask.unsqueeze(0).unsqueeze(0)).squeeze(0).squeeze(0)
|
||||
resized_img = resize_transform(cropped_img.permute(2, 0, 1))
|
||||
# Perform the center crop to the desired size
|
||||
# Constrain the crop to the smaller of our bbox or our image so we don't expand past the image dimensions.
|
||||
crop_transform = CenterCrop((min(self.max_bbox_size, resized_img.shape[1]), min(self.max_bbox_size, resized_img.shape[2])))
|
||||
|
||||
cropped_resized_img = crop_transform(resized_img)
|
||||
cropped_images.append(cropped_resized_img.permute(1, 2, 0))
|
||||
|
||||
cropped_resized_mask = crop_transform(resized_mask)
|
||||
cropped_masks.append(cropped_resized_mask)
|
||||
|
||||
combined_cropped_img = original_images[i][new_min_y:new_max_y, new_min_x:new_max_x, :]
|
||||
combined_cropped_images.append(combined_cropped_img)
|
||||
|
||||
combined_cropped_mask = masks[i][new_min_y:new_max_y, new_min_x:new_max_x]
|
||||
combined_cropped_masks.append(combined_cropped_mask)
|
||||
else:
|
||||
bounding_boxes.append((0, 0, img.shape[1], img.shape[0]))
|
||||
cropped_images.append(img)
|
||||
cropped_masks.append(mask)
|
||||
combined_cropped_images.append(img)
|
||||
combined_cropped_masks.append(mask)
|
||||
|
||||
cropped_out = torch.stack(cropped_images, dim=0)
|
||||
combined_crop_out = torch.stack(combined_cropped_images, dim=0)
|
||||
cropped_masks_out = torch.stack(cropped_masks, dim=0)
|
||||
combined_crop_mask_out = torch.stack(combined_cropped_masks, dim=0)
|
||||
|
||||
return (original_images, cropped_out, cropped_masks_out, combined_crop_out, combined_crop_mask_out, bounding_boxes, combined_bounding_box, self.max_bbox_size, self.max_bbox_size)
|
||||
|
||||
class FilterZeroMasksAndCorrespondingImages:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"masks": ("MASK",),
|
||||
},
|
||||
"optional": {
|
||||
"original_images": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK", "IMAGE", "IMAGE", "INDEXES",)
|
||||
RETURN_NAMES = ("non_zero_masks_out", "non_zero_mask_images_out", "zero_mask_images_out", "zero_mask_images_out_indexes",)
|
||||
FUNCTION = "filter"
|
||||
CATEGORY = "KJNodes/masking"
|
||||
DESCRIPTION = """
|
||||
Filter out all the empty (i.e. all zero) mask in masks
|
||||
Also filter out all the corresponding images in original_images by indexes if provide
|
||||
|
||||
original_images (optional): If provided, need have same length as masks.
|
||||
"""
|
||||
|
||||
def filter(self, masks, original_images=None):
|
||||
non_zero_masks = []
|
||||
non_zero_mask_images = []
|
||||
zero_mask_images = []
|
||||
zero_mask_images_indexes = []
|
||||
|
||||
masks_num = len(masks)
|
||||
also_process_images = False
|
||||
if original_images is not None:
|
||||
imgs_num = len(original_images)
|
||||
if len(original_images) == masks_num:
|
||||
also_process_images = True
|
||||
else:
|
||||
print(f"[WARNING] ignore input: original_images, due to number of original_images ({imgs_num}) is not equal to number of masks ({masks_num})")
|
||||
|
||||
for i in range(masks_num):
|
||||
non_zero_num = np.count_nonzero(np.array(masks[i]))
|
||||
if non_zero_num > 0:
|
||||
non_zero_masks.append(masks[i])
|
||||
if also_process_images:
|
||||
non_zero_mask_images.append(original_images[i])
|
||||
else:
|
||||
zero_mask_images.append(original_images[i])
|
||||
zero_mask_images_indexes.append(i)
|
||||
|
||||
non_zero_masks_out = torch.stack(non_zero_masks, dim=0)
|
||||
non_zero_mask_images_out = zero_mask_images_out = zero_mask_images_out_indexes = None
|
||||
|
||||
if also_process_images:
|
||||
non_zero_mask_images_out = torch.stack(non_zero_mask_images, dim=0)
|
||||
if len(zero_mask_images) > 0:
|
||||
zero_mask_images_out = torch.stack(zero_mask_images, dim=0)
|
||||
zero_mask_images_out_indexes = zero_mask_images_indexes
|
||||
|
||||
return (non_zero_masks_out, non_zero_mask_images_out, zero_mask_images_out, zero_mask_images_out_indexes)
|
||||
|
||||
class InsertImageBatchByIndexes:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"images_to_insert": ("IMAGE",),
|
||||
"insert_indexes": ("INDEXES",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", )
|
||||
RETURN_NAMES = ("images_after_insert", )
|
||||
FUNCTION = "insert"
|
||||
CATEGORY = "KJNodes/image"
|
||||
DESCRIPTION = """
|
||||
This node is designed to be use with node FilterZeroMasksAndCorrespondingImages
|
||||
It inserts the images_to_insert into images according to insert_indexes
|
||||
|
||||
Returns:
|
||||
images_after_insert: updated original images with origonal sequence order
|
||||
"""
|
||||
|
||||
def insert(self, images, images_to_insert, insert_indexes):
|
||||
images_after_insert = images
|
||||
|
||||
if images_to_insert is not None and insert_indexes is not None:
|
||||
images_to_insert_num = len(images_to_insert)
|
||||
insert_indexes_num = len(insert_indexes)
|
||||
if images_to_insert_num == insert_indexes_num:
|
||||
images_after_insert = []
|
||||
|
||||
i_images = 0
|
||||
for i in range(len(images) + images_to_insert_num):
|
||||
if i in insert_indexes:
|
||||
images_after_insert.append(images_to_insert[insert_indexes.index(i)])
|
||||
else:
|
||||
images_after_insert.append(images[i_images])
|
||||
i_images += 1
|
||||
|
||||
images_after_insert = torch.stack(images_after_insert, dim=0)
|
||||
|
||||
else:
|
||||
print(f"[WARNING] skip this node, due to number of images_to_insert ({images_to_insert_num}) is not equal to number of insert_indexes ({insert_indexes_num})")
|
||||
|
||||
|
||||
return (images_after_insert, )
|
||||
|
||||
def bbox_to_region(bbox, target_size=None):
|
||||
bbox = bbox_check(bbox, target_size)
|
||||
return (bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3])
|
||||
|
||||
def bbox_check(bbox, target_size=None):
|
||||
if not target_size:
|
||||
return bbox
|
||||
|
||||
new_bbox = (
|
||||
bbox[0],
|
||||
bbox[1],
|
||||
min(target_size[0] - bbox[0], bbox[2]),
|
||||
min(target_size[1] - bbox[1], bbox[3]),
|
||||
)
|
||||
return new_bbox
|
||||
|
||||
class BatchUncropAdvanced:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"original_images": ("IMAGE",),
|
||||
"cropped_images": ("IMAGE",),
|
||||
"cropped_masks": ("MASK",),
|
||||
"combined_crop_mask": ("MASK",),
|
||||
"bboxes": ("BBOX",),
|
||||
"border_blending": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}, ),
|
||||
"crop_rescale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
||||
"use_combined_mask": ("BOOLEAN", {"default": False}),
|
||||
"use_square_mask": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
"optional": {
|
||||
"combined_bounding_box": ("BBOX", {"default": None}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "uncrop"
|
||||
CATEGORY = "KJNodes/masking"
|
||||
|
||||
|
||||
def uncrop(self, original_images, cropped_images, cropped_masks, combined_crop_mask, bboxes, border_blending, crop_rescale, use_combined_mask, use_square_mask, combined_bounding_box = None):
|
||||
|
||||
def inset_border(image, border_width=20, border_color=(0)):
|
||||
width, height = image.size
|
||||
bordered_image = Image.new(image.mode, (width, height), border_color)
|
||||
bordered_image.paste(image, (0, 0))
|
||||
draw = ImageDraw.Draw(bordered_image)
|
||||
draw.rectangle((0, 0, width - 1, height - 1), outline=border_color, width=border_width)
|
||||
return bordered_image
|
||||
|
||||
if len(original_images) != len(cropped_images):
|
||||
raise ValueError(f"The number of original_images ({len(original_images)}) and cropped_images ({len(cropped_images)}) should be the same")
|
||||
|
||||
# Ensure there are enough bboxes, but drop the excess if there are more bboxes than images
|
||||
if len(bboxes) > len(original_images):
|
||||
print(f"Warning: Dropping excess bounding boxes. Expected {len(original_images)}, but got {len(bboxes)}")
|
||||
bboxes = bboxes[:len(original_images)]
|
||||
elif len(bboxes) < len(original_images):
|
||||
raise ValueError("There should be at least as many bboxes as there are original and cropped images")
|
||||
|
||||
crop_imgs = tensor2pil(cropped_images)
|
||||
input_images = tensor2pil(original_images)
|
||||
out_images = []
|
||||
|
||||
for i in range(len(input_images)):
|
||||
img = input_images[i]
|
||||
crop = crop_imgs[i]
|
||||
bbox = bboxes[i]
|
||||
|
||||
if use_combined_mask:
|
||||
bb_x, bb_y, bb_width, bb_height = combined_bounding_box[0]
|
||||
paste_region = bbox_to_region((bb_x, bb_y, bb_width, bb_height), img.size)
|
||||
mask = combined_crop_mask[i]
|
||||
else:
|
||||
bb_x, bb_y, bb_width, bb_height = bbox
|
||||
paste_region = bbox_to_region((bb_x, bb_y, bb_width, bb_height), img.size)
|
||||
mask = cropped_masks[i]
|
||||
|
||||
# scale paste_region
|
||||
scale_x = scale_y = crop_rescale
|
||||
paste_region = (round(paste_region[0]*scale_x), round(paste_region[1]*scale_y), round(paste_region[2]*scale_x), round(paste_region[3]*scale_y))
|
||||
|
||||
# rescale the crop image to fit the paste_region
|
||||
crop = crop.resize((round(paste_region[2]-paste_region[0]), round(paste_region[3]-paste_region[1])))
|
||||
crop_img = crop.convert("RGB")
|
||||
|
||||
#border blending
|
||||
if border_blending > 1.0:
|
||||
border_blending = 1.0
|
||||
elif border_blending < 0.0:
|
||||
border_blending = 0.0
|
||||
|
||||
blend_ratio = (max(crop_img.size) / 2) * float(border_blending)
|
||||
blend = img.convert("RGBA")
|
||||
|
||||
if use_square_mask:
|
||||
mask = Image.new("L", img.size, 0)
|
||||
mask_block = Image.new("L", (paste_region[2]-paste_region[0], paste_region[3]-paste_region[1]), 255)
|
||||
mask_block = inset_border(mask_block, round(blend_ratio / 2), (0))
|
||||
mask.paste(mask_block, paste_region)
|
||||
else:
|
||||
original_mask = tensor2pil(mask)[0]
|
||||
original_mask = original_mask.resize((paste_region[2]-paste_region[0], paste_region[3]-paste_region[1]))
|
||||
mask = Image.new("L", img.size, 0)
|
||||
mask.paste(original_mask, paste_region)
|
||||
|
||||
mask = mask.filter(ImageFilter.BoxBlur(radius=blend_ratio / 4))
|
||||
mask = mask.filter(ImageFilter.GaussianBlur(radius=blend_ratio / 4))
|
||||
|
||||
blend.paste(crop_img, paste_region)
|
||||
blend.putalpha(mask)
|
||||
|
||||
img = Image.alpha_composite(img.convert("RGBA"), blend)
|
||||
out_images.append(img.convert("RGB"))
|
||||
|
||||
return (pil2tensor(out_images),)
|
||||
|
||||
class SplitBboxes:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"bboxes": ("BBOX",),
|
||||
"index": ("INT", {"default": 0,"min": 0, "max": 99999999, "step": 1}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BBOX","BBOX",)
|
||||
RETURN_NAMES = ("bboxes_a","bboxes_b",)
|
||||
FUNCTION = "splitbbox"
|
||||
CATEGORY = "KJNodes/masking"
|
||||
DESCRIPTION = """
|
||||
Splits the specified bbox list at the given index into two lists.
|
||||
"""
|
||||
|
||||
def splitbbox(self, bboxes, index):
|
||||
bboxes_a = bboxes[:index] # Sub-list from the start of bboxes up to (but not including) the index
|
||||
bboxes_b = bboxes[index:] # Sub-list from the index to the end of bboxes
|
||||
|
||||
return (bboxes_a, bboxes_b,)
|
||||
@ -2,7 +2,7 @@ import torch
|
||||
import json
|
||||
from PIL import Image, ImageDraw
|
||||
import numpy as np
|
||||
from .utility import pil2tensor
|
||||
from ..utility.utility import pil2tensor
|
||||
|
||||
class SplineEditor:
|
||||
|
||||
@ -307,7 +307,6 @@ Converts different value lists/series to another type.
|
||||
import pandas as pd
|
||||
input_type = self.detect_input_type(input_values)
|
||||
|
||||
# Convert input_values to a list of floats
|
||||
if input_type == 'list of lists':
|
||||
float_values = [item for sublist in input_values for item in sublist]
|
||||
elif input_type == 'pandas series':
|
||||
@ -1,6 +1,5 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torchvision.transforms import Resize, CenterCrop, InterpolationMode
|
||||
from torchvision.transforms import functional as TF
|
||||
|
||||
import scipy.ndimage
|
||||
@ -17,9 +16,11 @@ import random
|
||||
import math
|
||||
|
||||
import model_management
|
||||
from nodes import MAX_RESOLUTION, SaveImage
|
||||
from nodes import MAX_RESOLUTION, SaveImage, CLIPTextEncode
|
||||
from comfy_extras.nodes_mask import ImageCompositeMasked
|
||||
import comfy.sample
|
||||
import folder_paths
|
||||
from ..utility.utility import tensor2pil, pil2tensor
|
||||
|
||||
script_directory = os.path.dirname(os.path.abspath(__file__))
|
||||
folder_paths.add_model_folder_path("kjnodes_fonts", os.path.join(script_directory, "fonts"))
|
||||
@ -141,7 +142,7 @@ class CreateFluidMask:
|
||||
}
|
||||
#using code from https://github.com/GregTJ/stable-fluids
|
||||
def createfluidmask(self, frames, width, height, invert, inflow_count, inflow_velocity, inflow_radius, inflow_padding, inflow_duration):
|
||||
from .fluid import Fluid
|
||||
from ..utility.fluid import Fluid
|
||||
from scipy.spatial import erf
|
||||
out = []
|
||||
masks = []
|
||||
@ -1845,655 +1846,6 @@ class ImageBatchTestPattern:
|
||||
|
||||
return (out_tensor,)
|
||||
|
||||
#based on nodes from mtb https://github.com/melMass/comfy_mtb
|
||||
|
||||
from .utility import tensor2pil, pil2tensor
|
||||
|
||||
class BatchCropFromMask:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"original_images": ("IMAGE",),
|
||||
"masks": ("MASK",),
|
||||
"crop_size_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}),
|
||||
"bbox_smooth_alpha": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"IMAGE",
|
||||
"IMAGE",
|
||||
"BBOX",
|
||||
"INT",
|
||||
"INT",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"original_images",
|
||||
"cropped_images",
|
||||
"bboxes",
|
||||
"width",
|
||||
"height",
|
||||
)
|
||||
FUNCTION = "crop"
|
||||
CATEGORY = "KJNodes/masking"
|
||||
|
||||
def smooth_bbox_size(self, prev_bbox_size, curr_bbox_size, alpha):
|
||||
if alpha == 0:
|
||||
return prev_bbox_size
|
||||
return round(alpha * curr_bbox_size + (1 - alpha) * prev_bbox_size)
|
||||
|
||||
def smooth_center(self, prev_center, curr_center, alpha=0.5):
|
||||
if alpha == 0:
|
||||
return prev_center
|
||||
return (
|
||||
round(alpha * curr_center[0] + (1 - alpha) * prev_center[0]),
|
||||
round(alpha * curr_center[1] + (1 - alpha) * prev_center[1])
|
||||
)
|
||||
|
||||
def crop(self, masks, original_images, crop_size_mult, bbox_smooth_alpha):
|
||||
|
||||
bounding_boxes = []
|
||||
cropped_images = []
|
||||
|
||||
self.max_bbox_width = 0
|
||||
self.max_bbox_height = 0
|
||||
|
||||
# First, calculate the maximum bounding box size across all masks
|
||||
curr_max_bbox_width = 0
|
||||
curr_max_bbox_height = 0
|
||||
for mask in masks:
|
||||
_mask = tensor2pil(mask)[0]
|
||||
non_zero_indices = np.nonzero(np.array(_mask))
|
||||
min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1])
|
||||
min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0])
|
||||
width = max_x - min_x
|
||||
height = max_y - min_y
|
||||
curr_max_bbox_width = max(curr_max_bbox_width, width)
|
||||
curr_max_bbox_height = max(curr_max_bbox_height, height)
|
||||
|
||||
# Smooth the changes in the bounding box size
|
||||
self.max_bbox_width = self.smooth_bbox_size(self.max_bbox_width, curr_max_bbox_width, bbox_smooth_alpha)
|
||||
self.max_bbox_height = self.smooth_bbox_size(self.max_bbox_height, curr_max_bbox_height, bbox_smooth_alpha)
|
||||
|
||||
# Apply the crop size multiplier
|
||||
self.max_bbox_width = round(self.max_bbox_width * crop_size_mult)
|
||||
self.max_bbox_height = round(self.max_bbox_height * crop_size_mult)
|
||||
bbox_aspect_ratio = self.max_bbox_width / self.max_bbox_height
|
||||
|
||||
# Then, for each mask and corresponding image...
|
||||
for i, (mask, img) in enumerate(zip(masks, original_images)):
|
||||
_mask = tensor2pil(mask)[0]
|
||||
non_zero_indices = np.nonzero(np.array(_mask))
|
||||
min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1])
|
||||
min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0])
|
||||
|
||||
# Calculate center of bounding box
|
||||
center_x = np.mean(non_zero_indices[1])
|
||||
center_y = np.mean(non_zero_indices[0])
|
||||
curr_center = (round(center_x), round(center_y))
|
||||
|
||||
# If this is the first frame, initialize prev_center with curr_center
|
||||
if not hasattr(self, 'prev_center'):
|
||||
self.prev_center = curr_center
|
||||
|
||||
# Smooth the changes in the center coordinates from the second frame onwards
|
||||
if i > 0:
|
||||
center = self.smooth_center(self.prev_center, curr_center, bbox_smooth_alpha)
|
||||
else:
|
||||
center = curr_center
|
||||
|
||||
# Update prev_center for the next frame
|
||||
self.prev_center = center
|
||||
|
||||
# Create bounding box using max_bbox_width and max_bbox_height
|
||||
half_box_width = round(self.max_bbox_width / 2)
|
||||
half_box_height = round(self.max_bbox_height / 2)
|
||||
min_x = max(0, center[0] - half_box_width)
|
||||
max_x = min(img.shape[1], center[0] + half_box_width)
|
||||
min_y = max(0, center[1] - half_box_height)
|
||||
max_y = min(img.shape[0], center[1] + half_box_height)
|
||||
|
||||
# Append bounding box coordinates
|
||||
bounding_boxes.append((min_x, min_y, max_x - min_x, max_y - min_y))
|
||||
|
||||
# Crop the image from the bounding box
|
||||
cropped_img = img[min_y:max_y, min_x:max_x, :]
|
||||
|
||||
# Calculate the new dimensions while maintaining the aspect ratio
|
||||
new_height = min(cropped_img.shape[0], self.max_bbox_height)
|
||||
new_width = round(new_height * bbox_aspect_ratio)
|
||||
|
||||
# Resize the image
|
||||
resize_transform = Resize((new_height, new_width))
|
||||
resized_img = resize_transform(cropped_img.permute(2, 0, 1))
|
||||
|
||||
# Perform the center crop to the desired size
|
||||
crop_transform = CenterCrop((self.max_bbox_height, self.max_bbox_width)) # swap the order here if necessary
|
||||
cropped_resized_img = crop_transform(resized_img)
|
||||
|
||||
cropped_images.append(cropped_resized_img.permute(1, 2, 0))
|
||||
|
||||
cropped_out = torch.stack(cropped_images, dim=0)
|
||||
|
||||
return (original_images, cropped_out, bounding_boxes, self.max_bbox_width, self.max_bbox_height, )
|
||||
|
||||
|
||||
def bbox_to_region(bbox, target_size=None):
|
||||
bbox = bbox_check(bbox, target_size)
|
||||
return (bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3])
|
||||
|
||||
def bbox_check(bbox, target_size=None):
|
||||
if not target_size:
|
||||
return bbox
|
||||
|
||||
new_bbox = (
|
||||
bbox[0],
|
||||
bbox[1],
|
||||
min(target_size[0] - bbox[0], bbox[2]),
|
||||
min(target_size[1] - bbox[1], bbox[3]),
|
||||
)
|
||||
return new_bbox
|
||||
|
||||
class BatchUncrop:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"original_images": ("IMAGE",),
|
||||
"cropped_images": ("IMAGE",),
|
||||
"bboxes": ("BBOX",),
|
||||
"border_blending": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}, ),
|
||||
"crop_rescale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
||||
"border_top": ("BOOLEAN", {"default": True}),
|
||||
"border_bottom": ("BOOLEAN", {"default": True}),
|
||||
"border_left": ("BOOLEAN", {"default": True}),
|
||||
"border_right": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "uncrop"
|
||||
|
||||
CATEGORY = "KJNodes/masking"
|
||||
|
||||
def uncrop(self, original_images, cropped_images, bboxes, border_blending, crop_rescale, border_top, border_bottom, border_left, border_right):
|
||||
def inset_border(image, border_width, border_color, border_top, border_bottom, border_left, border_right):
|
||||
draw = ImageDraw.Draw(image)
|
||||
width, height = image.size
|
||||
if border_top:
|
||||
draw.rectangle((0, 0, width, border_width), fill=border_color)
|
||||
if border_bottom:
|
||||
draw.rectangle((0, height - border_width, width, height), fill=border_color)
|
||||
if border_left:
|
||||
draw.rectangle((0, 0, border_width, height), fill=border_color)
|
||||
if border_right:
|
||||
draw.rectangle((width - border_width, 0, width, height), fill=border_color)
|
||||
return image
|
||||
|
||||
if len(original_images) != len(cropped_images):
|
||||
raise ValueError(f"The number of original_images ({len(original_images)}) and cropped_images ({len(cropped_images)}) should be the same")
|
||||
|
||||
# Ensure there are enough bboxes, but drop the excess if there are more bboxes than images
|
||||
if len(bboxes) > len(original_images):
|
||||
print(f"Warning: Dropping excess bounding boxes. Expected {len(original_images)}, but got {len(bboxes)}")
|
||||
bboxes = bboxes[:len(original_images)]
|
||||
elif len(bboxes) < len(original_images):
|
||||
raise ValueError("There should be at least as many bboxes as there are original and cropped images")
|
||||
|
||||
input_images = tensor2pil(original_images)
|
||||
crop_imgs = tensor2pil(cropped_images)
|
||||
|
||||
out_images = []
|
||||
for i in range(len(input_images)):
|
||||
img = input_images[i]
|
||||
crop = crop_imgs[i]
|
||||
bbox = bboxes[i]
|
||||
|
||||
# uncrop the image based on the bounding box
|
||||
bb_x, bb_y, bb_width, bb_height = bbox
|
||||
|
||||
paste_region = bbox_to_region((bb_x, bb_y, bb_width, bb_height), img.size)
|
||||
|
||||
# scale factors
|
||||
scale_x = crop_rescale
|
||||
scale_y = crop_rescale
|
||||
|
||||
# scaled paste_region
|
||||
paste_region = (round(paste_region[0]*scale_x), round(paste_region[1]*scale_y), round(paste_region[2]*scale_x), round(paste_region[3]*scale_y))
|
||||
|
||||
# rescale the crop image to fit the paste_region
|
||||
crop = crop.resize((round(paste_region[2]-paste_region[0]), round(paste_region[3]-paste_region[1])))
|
||||
crop_img = crop.convert("RGB")
|
||||
|
||||
if border_blending > 1.0:
|
||||
border_blending = 1.0
|
||||
elif border_blending < 0.0:
|
||||
border_blending = 0.0
|
||||
|
||||
blend_ratio = (max(crop_img.size) / 2) * float(border_blending)
|
||||
|
||||
blend = img.convert("RGBA")
|
||||
mask = Image.new("L", img.size, 0)
|
||||
|
||||
mask_block = Image.new("L", (paste_region[2]-paste_region[0], paste_region[3]-paste_region[1]), 255)
|
||||
mask_block = inset_border(mask_block, round(blend_ratio / 2), (0), border_top, border_bottom, border_left, border_right)
|
||||
|
||||
mask.paste(mask_block, paste_region)
|
||||
blend.paste(crop_img, paste_region)
|
||||
|
||||
mask = mask.filter(ImageFilter.BoxBlur(radius=blend_ratio / 4))
|
||||
mask = mask.filter(ImageFilter.GaussianBlur(radius=blend_ratio / 4))
|
||||
|
||||
blend.putalpha(mask)
|
||||
img = Image.alpha_composite(img.convert("RGBA"), blend)
|
||||
out_images.append(img.convert("RGB"))
|
||||
|
||||
return (pil2tensor(out_images),)
|
||||
|
||||
class BatchCropFromMaskAdvanced:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"original_images": ("IMAGE",),
|
||||
"masks": ("MASK",),
|
||||
"crop_size_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
||||
"bbox_smooth_alpha": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"IMAGE",
|
||||
"IMAGE",
|
||||
"MASK",
|
||||
"IMAGE",
|
||||
"MASK",
|
||||
"BBOX",
|
||||
"BBOX",
|
||||
"INT",
|
||||
"INT",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"original_images",
|
||||
"cropped_images",
|
||||
"cropped_masks",
|
||||
"combined_crop_image",
|
||||
"combined_crop_masks",
|
||||
"bboxes",
|
||||
"combined_bounding_box",
|
||||
"bbox_width",
|
||||
"bbox_height",
|
||||
)
|
||||
FUNCTION = "crop"
|
||||
CATEGORY = "KJNodes/masking"
|
||||
|
||||
def smooth_bbox_size(self, prev_bbox_size, curr_bbox_size, alpha):
|
||||
return round(alpha * curr_bbox_size + (1 - alpha) * prev_bbox_size)
|
||||
|
||||
def smooth_center(self, prev_center, curr_center, alpha=0.5):
|
||||
return (round(alpha * curr_center[0] + (1 - alpha) * prev_center[0]),
|
||||
round(alpha * curr_center[1] + (1 - alpha) * prev_center[1]))
|
||||
|
||||
def crop(self, masks, original_images, crop_size_mult, bbox_smooth_alpha):
|
||||
bounding_boxes = []
|
||||
combined_bounding_box = []
|
||||
cropped_images = []
|
||||
cropped_masks = []
|
||||
cropped_masks_out = []
|
||||
combined_crop_out = []
|
||||
combined_cropped_images = []
|
||||
combined_cropped_masks = []
|
||||
|
||||
def calculate_bbox(mask):
|
||||
non_zero_indices = np.nonzero(np.array(mask))
|
||||
|
||||
# handle empty masks
|
||||
min_x, max_x, min_y, max_y = 0, 0, 0, 0
|
||||
if len(non_zero_indices[1]) > 0 and len(non_zero_indices[0]) > 0:
|
||||
min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1])
|
||||
min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0])
|
||||
|
||||
width = max_x - min_x
|
||||
height = max_y - min_y
|
||||
bbox_size = max(width, height)
|
||||
return min_x, max_x, min_y, max_y, bbox_size
|
||||
|
||||
combined_mask = torch.max(masks, dim=0)[0]
|
||||
_mask = tensor2pil(combined_mask)[0]
|
||||
new_min_x, new_max_x, new_min_y, new_max_y, combined_bbox_size = calculate_bbox(_mask)
|
||||
center_x = (new_min_x + new_max_x) / 2
|
||||
center_y = (new_min_y + new_max_y) / 2
|
||||
half_box_size = round(combined_bbox_size // 2)
|
||||
new_min_x = max(0, round(center_x - half_box_size))
|
||||
new_max_x = min(original_images[0].shape[1], round(center_x + half_box_size))
|
||||
new_min_y = max(0, round(center_y - half_box_size))
|
||||
new_max_y = min(original_images[0].shape[0], round(center_y + half_box_size))
|
||||
|
||||
combined_bounding_box.append((new_min_x, new_min_y, new_max_x - new_min_x, new_max_y - new_min_y))
|
||||
|
||||
self.max_bbox_size = 0
|
||||
|
||||
# First, calculate the maximum bounding box size across all masks
|
||||
curr_max_bbox_size = max(calculate_bbox(tensor2pil(mask)[0])[-1] for mask in masks)
|
||||
# Smooth the changes in the bounding box size
|
||||
self.max_bbox_size = self.smooth_bbox_size(self.max_bbox_size, curr_max_bbox_size, bbox_smooth_alpha)
|
||||
# Apply the crop size multiplier
|
||||
self.max_bbox_size = round(self.max_bbox_size * crop_size_mult)
|
||||
# Make sure max_bbox_size is divisible by 16, if not, round it upwards so it is
|
||||
self.max_bbox_size = math.ceil(self.max_bbox_size / 16) * 16
|
||||
|
||||
if self.max_bbox_size > original_images[0].shape[0] or self.max_bbox_size > original_images[0].shape[1]:
|
||||
# max_bbox_size can only be as big as our input's width or height, and it has to be even
|
||||
self.max_bbox_size = math.floor(min(original_images[0].shape[0], original_images[0].shape[1]) / 2) * 2
|
||||
|
||||
# Then, for each mask and corresponding image...
|
||||
for i, (mask, img) in enumerate(zip(masks, original_images)):
|
||||
_mask = tensor2pil(mask)[0]
|
||||
non_zero_indices = np.nonzero(np.array(_mask))
|
||||
|
||||
# check for empty masks
|
||||
if len(non_zero_indices[0]) > 0 and len(non_zero_indices[1]) > 0:
|
||||
min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1])
|
||||
min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0])
|
||||
|
||||
# Calculate center of bounding box
|
||||
center_x = np.mean(non_zero_indices[1])
|
||||
center_y = np.mean(non_zero_indices[0])
|
||||
curr_center = (round(center_x), round(center_y))
|
||||
|
||||
# If this is the first frame, initialize prev_center with curr_center
|
||||
if not hasattr(self, 'prev_center'):
|
||||
self.prev_center = curr_center
|
||||
|
||||
# Smooth the changes in the center coordinates from the second frame onwards
|
||||
if i > 0:
|
||||
center = self.smooth_center(self.prev_center, curr_center, bbox_smooth_alpha)
|
||||
else:
|
||||
center = curr_center
|
||||
|
||||
# Update prev_center for the next frame
|
||||
self.prev_center = center
|
||||
|
||||
# Create bounding box using max_bbox_size
|
||||
half_box_size = self.max_bbox_size // 2
|
||||
min_x = max(0, center[0] - half_box_size)
|
||||
max_x = min(img.shape[1], center[0] + half_box_size)
|
||||
min_y = max(0, center[1] - half_box_size)
|
||||
max_y = min(img.shape[0], center[1] + half_box_size)
|
||||
|
||||
# Append bounding box coordinates
|
||||
bounding_boxes.append((min_x, min_y, max_x - min_x, max_y - min_y))
|
||||
|
||||
# Crop the image from the bounding box
|
||||
cropped_img = img[min_y:max_y, min_x:max_x, :]
|
||||
cropped_mask = mask[min_y:max_y, min_x:max_x]
|
||||
|
||||
# Resize the cropped image to a fixed size
|
||||
new_size = max(cropped_img.shape[0], cropped_img.shape[1])
|
||||
resize_transform = Resize(new_size, interpolation=InterpolationMode.NEAREST, max_size=max(img.shape[0], img.shape[1]))
|
||||
resized_mask = resize_transform(cropped_mask.unsqueeze(0).unsqueeze(0)).squeeze(0).squeeze(0)
|
||||
resized_img = resize_transform(cropped_img.permute(2, 0, 1))
|
||||
# Perform the center crop to the desired size
|
||||
# Constrain the crop to the smaller of our bbox or our image so we don't expand past the image dimensions.
|
||||
crop_transform = CenterCrop((min(self.max_bbox_size, resized_img.shape[1]), min(self.max_bbox_size, resized_img.shape[2])))
|
||||
|
||||
cropped_resized_img = crop_transform(resized_img)
|
||||
cropped_images.append(cropped_resized_img.permute(1, 2, 0))
|
||||
|
||||
cropped_resized_mask = crop_transform(resized_mask)
|
||||
cropped_masks.append(cropped_resized_mask)
|
||||
|
||||
combined_cropped_img = original_images[i][new_min_y:new_max_y, new_min_x:new_max_x, :]
|
||||
combined_cropped_images.append(combined_cropped_img)
|
||||
|
||||
combined_cropped_mask = masks[i][new_min_y:new_max_y, new_min_x:new_max_x]
|
||||
combined_cropped_masks.append(combined_cropped_mask)
|
||||
else:
|
||||
bounding_boxes.append((0, 0, img.shape[1], img.shape[0]))
|
||||
cropped_images.append(img)
|
||||
cropped_masks.append(mask)
|
||||
combined_cropped_images.append(img)
|
||||
combined_cropped_masks.append(mask)
|
||||
|
||||
cropped_out = torch.stack(cropped_images, dim=0)
|
||||
combined_crop_out = torch.stack(combined_cropped_images, dim=0)
|
||||
cropped_masks_out = torch.stack(cropped_masks, dim=0)
|
||||
combined_crop_mask_out = torch.stack(combined_cropped_masks, dim=0)
|
||||
|
||||
return (original_images, cropped_out, cropped_masks_out, combined_crop_out, combined_crop_mask_out, bounding_boxes, combined_bounding_box, self.max_bbox_size, self.max_bbox_size)
|
||||
|
||||
class FilterZeroMasksAndCorrespondingImages:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"masks": ("MASK",),
|
||||
},
|
||||
"optional": {
|
||||
"original_images": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK", "IMAGE", "IMAGE", "INDEXES",)
|
||||
RETURN_NAMES = ("non_zero_masks_out", "non_zero_mask_images_out", "zero_mask_images_out", "zero_mask_images_out_indexes",)
|
||||
FUNCTION = "filter"
|
||||
CATEGORY = "KJNodes/masking"
|
||||
DESCRIPTION = """
|
||||
Filter out all the empty (i.e. all zero) mask in masks
|
||||
Also filter out all the corresponding images in original_images by indexes if provide
|
||||
|
||||
original_images (optional): If provided, need have same length as masks.
|
||||
"""
|
||||
|
||||
def filter(self, masks, original_images=None):
|
||||
non_zero_masks = []
|
||||
non_zero_mask_images = []
|
||||
zero_mask_images = []
|
||||
zero_mask_images_indexes = []
|
||||
|
||||
masks_num = len(masks)
|
||||
also_process_images = False
|
||||
if original_images is not None:
|
||||
imgs_num = len(original_images)
|
||||
if len(original_images) == masks_num:
|
||||
also_process_images = True
|
||||
else:
|
||||
print(f"[WARNING] ignore input: original_images, due to number of original_images ({imgs_num}) is not equal to number of masks ({masks_num})")
|
||||
|
||||
for i in range(masks_num):
|
||||
non_zero_num = np.count_nonzero(np.array(masks[i]))
|
||||
if non_zero_num > 0:
|
||||
non_zero_masks.append(masks[i])
|
||||
if also_process_images:
|
||||
non_zero_mask_images.append(original_images[i])
|
||||
else:
|
||||
zero_mask_images.append(original_images[i])
|
||||
zero_mask_images_indexes.append(i)
|
||||
|
||||
non_zero_masks_out = torch.stack(non_zero_masks, dim=0)
|
||||
non_zero_mask_images_out = zero_mask_images_out = zero_mask_images_out_indexes = None
|
||||
|
||||
if also_process_images:
|
||||
non_zero_mask_images_out = torch.stack(non_zero_mask_images, dim=0)
|
||||
if len(zero_mask_images) > 0:
|
||||
zero_mask_images_out = torch.stack(zero_mask_images, dim=0)
|
||||
zero_mask_images_out_indexes = zero_mask_images_indexes
|
||||
|
||||
return (non_zero_masks_out, non_zero_mask_images_out, zero_mask_images_out, zero_mask_images_out_indexes)
|
||||
|
||||
class InsertImageBatchByIndexes:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"images_to_insert": ("IMAGE",),
|
||||
"insert_indexes": ("INDEXES",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", )
|
||||
RETURN_NAMES = ("images_after_insert", )
|
||||
FUNCTION = "insert"
|
||||
CATEGORY = "KJNodes/image"
|
||||
DESCRIPTION = """
|
||||
This node is designed to be use with node FilterZeroMasksAndCorrespondingImages
|
||||
It inserts the images_to_insert into images according to insert_indexes
|
||||
|
||||
Returns:
|
||||
images_after_insert: updated original images with origonal sequence order
|
||||
"""
|
||||
|
||||
def insert(self, images, images_to_insert, insert_indexes):
|
||||
images_after_insert = images
|
||||
|
||||
if images_to_insert is not None and insert_indexes is not None:
|
||||
images_to_insert_num = len(images_to_insert)
|
||||
insert_indexes_num = len(insert_indexes)
|
||||
if images_to_insert_num == insert_indexes_num:
|
||||
images_after_insert = []
|
||||
|
||||
i_images = 0
|
||||
for i in range(len(images) + images_to_insert_num):
|
||||
if i in insert_indexes:
|
||||
images_after_insert.append(images_to_insert[insert_indexes.index(i)])
|
||||
else:
|
||||
images_after_insert.append(images[i_images])
|
||||
i_images += 1
|
||||
|
||||
images_after_insert = torch.stack(images_after_insert, dim=0)
|
||||
|
||||
else:
|
||||
print(f"[WARNING] skip this node, due to number of images_to_insert ({images_to_insert_num}) is not equal to number of insert_indexes ({insert_indexes_num})")
|
||||
|
||||
|
||||
return (images_after_insert, )
|
||||
|
||||
def bbox_to_region(bbox, target_size=None):
|
||||
bbox = bbox_check(bbox, target_size)
|
||||
return (bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3])
|
||||
|
||||
def bbox_check(bbox, target_size=None):
|
||||
if not target_size:
|
||||
return bbox
|
||||
|
||||
new_bbox = (
|
||||
bbox[0],
|
||||
bbox[1],
|
||||
min(target_size[0] - bbox[0], bbox[2]),
|
||||
min(target_size[1] - bbox[1], bbox[3]),
|
||||
)
|
||||
return new_bbox
|
||||
|
||||
class BatchUncropAdvanced:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"original_images": ("IMAGE",),
|
||||
"cropped_images": ("IMAGE",),
|
||||
"cropped_masks": ("MASK",),
|
||||
"combined_crop_mask": ("MASK",),
|
||||
"bboxes": ("BBOX",),
|
||||
"border_blending": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}, ),
|
||||
"crop_rescale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
||||
"use_combined_mask": ("BOOLEAN", {"default": False}),
|
||||
"use_square_mask": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
"optional": {
|
||||
"combined_bounding_box": ("BBOX", {"default": None}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "uncrop"
|
||||
CATEGORY = "KJNodes/masking"
|
||||
|
||||
|
||||
def uncrop(self, original_images, cropped_images, cropped_masks, combined_crop_mask, bboxes, border_blending, crop_rescale, use_combined_mask, use_square_mask, combined_bounding_box = None):
|
||||
|
||||
def inset_border(image, border_width=20, border_color=(0)):
|
||||
width, height = image.size
|
||||
bordered_image = Image.new(image.mode, (width, height), border_color)
|
||||
bordered_image.paste(image, (0, 0))
|
||||
draw = ImageDraw.Draw(bordered_image)
|
||||
draw.rectangle((0, 0, width - 1, height - 1), outline=border_color, width=border_width)
|
||||
return bordered_image
|
||||
|
||||
if len(original_images) != len(cropped_images):
|
||||
raise ValueError(f"The number of original_images ({len(original_images)}) and cropped_images ({len(cropped_images)}) should be the same")
|
||||
|
||||
# Ensure there are enough bboxes, but drop the excess if there are more bboxes than images
|
||||
if len(bboxes) > len(original_images):
|
||||
print(f"Warning: Dropping excess bounding boxes. Expected {len(original_images)}, but got {len(bboxes)}")
|
||||
bboxes = bboxes[:len(original_images)]
|
||||
elif len(bboxes) < len(original_images):
|
||||
raise ValueError("There should be at least as many bboxes as there are original and cropped images")
|
||||
|
||||
crop_imgs = tensor2pil(cropped_images)
|
||||
input_images = tensor2pil(original_images)
|
||||
out_images = []
|
||||
|
||||
for i in range(len(input_images)):
|
||||
img = input_images[i]
|
||||
crop = crop_imgs[i]
|
||||
bbox = bboxes[i]
|
||||
|
||||
if use_combined_mask:
|
||||
bb_x, bb_y, bb_width, bb_height = combined_bounding_box[0]
|
||||
paste_region = bbox_to_region((bb_x, bb_y, bb_width, bb_height), img.size)
|
||||
mask = combined_crop_mask[i]
|
||||
else:
|
||||
bb_x, bb_y, bb_width, bb_height = bbox
|
||||
paste_region = bbox_to_region((bb_x, bb_y, bb_width, bb_height), img.size)
|
||||
mask = cropped_masks[i]
|
||||
|
||||
# scale paste_region
|
||||
scale_x = scale_y = crop_rescale
|
||||
paste_region = (round(paste_region[0]*scale_x), round(paste_region[1]*scale_y), round(paste_region[2]*scale_x), round(paste_region[3]*scale_y))
|
||||
|
||||
# rescale the crop image to fit the paste_region
|
||||
crop = crop.resize((round(paste_region[2]-paste_region[0]), round(paste_region[3]-paste_region[1])))
|
||||
crop_img = crop.convert("RGB")
|
||||
|
||||
#border blending
|
||||
if border_blending > 1.0:
|
||||
border_blending = 1.0
|
||||
elif border_blending < 0.0:
|
||||
border_blending = 0.0
|
||||
|
||||
blend_ratio = (max(crop_img.size) / 2) * float(border_blending)
|
||||
blend = img.convert("RGBA")
|
||||
|
||||
if use_square_mask:
|
||||
mask = Image.new("L", img.size, 0)
|
||||
mask_block = Image.new("L", (paste_region[2]-paste_region[0], paste_region[3]-paste_region[1]), 255)
|
||||
mask_block = inset_border(mask_block, round(blend_ratio / 2), (0))
|
||||
mask.paste(mask_block, paste_region)
|
||||
else:
|
||||
original_mask = tensor2pil(mask)[0]
|
||||
original_mask = original_mask.resize((paste_region[2]-paste_region[0], paste_region[3]-paste_region[1]))
|
||||
mask = Image.new("L", img.size, 0)
|
||||
mask.paste(original_mask, paste_region)
|
||||
|
||||
mask = mask.filter(ImageFilter.BoxBlur(radius=blend_ratio / 4))
|
||||
mask = mask.filter(ImageFilter.GaussianBlur(radius=blend_ratio / 4))
|
||||
|
||||
blend.paste(crop_img, paste_region)
|
||||
blend.putalpha(mask)
|
||||
|
||||
img = Image.alpha_composite(img.convert("RGBA"), blend)
|
||||
out_images.append(img.convert("RGB"))
|
||||
|
||||
return (pil2tensor(out_images),)
|
||||
|
||||
class BatchCLIPSeg:
|
||||
|
||||
def __init__(self):
|
||||
@ -2950,7 +2302,7 @@ class CreateMagicMask:
|
||||
}
|
||||
|
||||
def createmagicmask(self, frames, transitions, depth, distortion, seed, frame_width, frame_height):
|
||||
from .magictex import coordinate_grid, random_transform, magic
|
||||
from ..utility.magictex import coordinate_grid, random_transform, magic
|
||||
rng = np.random.default_rng(seed)
|
||||
out = []
|
||||
coords = coordinate_grid((frame_width, frame_height))
|
||||
@ -3080,31 +2432,6 @@ Visualizes the specified bbox on the image.
|
||||
|
||||
return (torch.cat(image_list, dim=0),)
|
||||
|
||||
class SplitBboxes:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"bboxes": ("BBOX",),
|
||||
"index": ("INT", {"default": 0,"min": 0, "max": 99999999, "step": 1}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BBOX","BBOX",)
|
||||
RETURN_NAMES = ("bboxes_a","bboxes_b",)
|
||||
FUNCTION = "splitbbox"
|
||||
CATEGORY = "KJNodes/masking"
|
||||
DESCRIPTION = """
|
||||
Splits the specified bbox list at the given index into two lists.
|
||||
"""
|
||||
|
||||
def splitbbox(self, bboxes, index):
|
||||
bboxes_a = bboxes[:index] # Sub-list from the start of bboxes up to (but not including) the index
|
||||
bboxes_b = bboxes[index:] # Sub-list from the index to the end of bboxes
|
||||
|
||||
return (bboxes_a, bboxes_b,)
|
||||
|
||||
from PIL import ImageGrab
|
||||
import time
|
||||
class ImageGrabPIL:
|
||||
@ -3775,229 +3102,6 @@ with repeats 2 becomes batch of 10 images: 0, 0, 1, 1, 2, 2, 3, 3, 4, 4
|
||||
repeated_images = torch.repeat_interleave(images, repeats=repeats, dim=0)
|
||||
return (repeated_images, )
|
||||
|
||||
class NormalizedAmplitudeToMask:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"normalized_amp": ("NORMALIZED_AMPLITUDE",),
|
||||
"width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}),
|
||||
"height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}),
|
||||
"frame_offset": ("INT", {"default": 0,"min": -255, "max": 255, "step": 1}),
|
||||
"location_x": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}),
|
||||
"location_y": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}),
|
||||
"size": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}),
|
||||
"shape": (
|
||||
[
|
||||
'none',
|
||||
'circle',
|
||||
'square',
|
||||
'triangle',
|
||||
],
|
||||
{
|
||||
"default": 'none'
|
||||
}),
|
||||
"color": (
|
||||
[
|
||||
'white',
|
||||
'amplitude',
|
||||
],
|
||||
{
|
||||
"default": 'amplitude'
|
||||
}),
|
||||
},}
|
||||
|
||||
CATEGORY = "KJNodes/audio"
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "convert"
|
||||
DESCRIPTION = """
|
||||
Works as a bridge to the AudioScheduler -nodes:
|
||||
https://github.com/a1lazydog/ComfyUI-AudioScheduler
|
||||
Creates masks based on the normalized amplitude.
|
||||
"""
|
||||
|
||||
def convert(self, normalized_amp, width, height, frame_offset, shape, location_x, location_y, size, color):
|
||||
# Ensure normalized_amp is an array and within the range [0, 1]
|
||||
normalized_amp = np.clip(normalized_amp, 0.0, 1.0)
|
||||
|
||||
# Offset the amplitude values by rolling the array
|
||||
normalized_amp = np.roll(normalized_amp, frame_offset)
|
||||
|
||||
# Initialize an empty list to hold the image tensors
|
||||
out = []
|
||||
# Iterate over each amplitude value to create an image
|
||||
for amp in normalized_amp:
|
||||
# Scale the amplitude value to cover the full range of grayscale values
|
||||
if color == 'amplitude':
|
||||
grayscale_value = int(amp * 255)
|
||||
elif color == 'white':
|
||||
grayscale_value = 255
|
||||
# Convert the grayscale value to an RGB format
|
||||
gray_color = (grayscale_value, grayscale_value, grayscale_value)
|
||||
finalsize = size * amp
|
||||
|
||||
if shape == 'none':
|
||||
shapeimage = Image.new("RGB", (width, height), gray_color)
|
||||
else:
|
||||
shapeimage = Image.new("RGB", (width, height), "black")
|
||||
|
||||
draw = ImageDraw.Draw(shapeimage)
|
||||
if shape == 'circle' or shape == 'square':
|
||||
# Define the bounding box for the shape
|
||||
left_up_point = (location_x - finalsize, location_y - finalsize)
|
||||
right_down_point = (location_x + finalsize,location_y + finalsize)
|
||||
two_points = [left_up_point, right_down_point]
|
||||
|
||||
if shape == 'circle':
|
||||
draw.ellipse(two_points, fill=gray_color)
|
||||
elif shape == 'square':
|
||||
draw.rectangle(two_points, fill=gray_color)
|
||||
|
||||
elif shape == 'triangle':
|
||||
# Define the points for the triangle
|
||||
left_up_point = (location_x - finalsize, location_y + finalsize) # bottom left
|
||||
right_down_point = (location_x + finalsize, location_y + finalsize) # bottom right
|
||||
top_point = (location_x, location_y) # top point
|
||||
draw.polygon([top_point, left_up_point, right_down_point], fill=gray_color)
|
||||
|
||||
shapeimage = pil2tensor(shapeimage)
|
||||
mask = shapeimage[:, :, :, 0]
|
||||
out.append(mask)
|
||||
|
||||
return (torch.cat(out, dim=0),)
|
||||
|
||||
class OffsetMaskByNormalizedAmplitude:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"normalized_amp": ("NORMALIZED_AMPLITUDE",),
|
||||
"mask": ("MASK",),
|
||||
"x": ("INT", { "default": 0, "min": -4096, "max": MAX_RESOLUTION, "step": 1, "display": "number" }),
|
||||
"y": ("INT", { "default": 0, "min": -4096, "max": MAX_RESOLUTION, "step": 1, "display": "number" }),
|
||||
"rotate": ("BOOLEAN", { "default": False }),
|
||||
"angle_multiplier": ("FLOAT", { "default": 0.0, "min": -1.0, "max": 1.0, "step": 0.001, "display": "number" }),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
RETURN_NAMES = ("mask",)
|
||||
FUNCTION = "offset"
|
||||
CATEGORY = "KJNodes/audio"
|
||||
DESCRIPTION = """
|
||||
Works as a bridge to the AudioScheduler -nodes:
|
||||
https://github.com/a1lazydog/ComfyUI-AudioScheduler
|
||||
Offsets masks based on the normalized amplitude.
|
||||
"""
|
||||
|
||||
def offset(self, mask, x, y, angle_multiplier, rotate, normalized_amp):
|
||||
|
||||
# Ensure normalized_amp is an array and within the range [0, 1]
|
||||
offsetmask = mask.clone()
|
||||
normalized_amp = np.clip(normalized_amp, 0.0, 1.0)
|
||||
|
||||
batch_size, height, width = mask.shape
|
||||
|
||||
if rotate:
|
||||
for i in range(batch_size):
|
||||
rotation_amp = int(normalized_amp[i] * (360 * angle_multiplier))
|
||||
rotation_angle = rotation_amp
|
||||
offsetmask[i] = TF.rotate(offsetmask[i].unsqueeze(0), rotation_angle).squeeze(0)
|
||||
if x != 0 or y != 0:
|
||||
for i in range(batch_size):
|
||||
offset_amp = normalized_amp[i] * 10
|
||||
shift_x = min(x*offset_amp, width-1)
|
||||
shift_y = min(y*offset_amp, height-1)
|
||||
if shift_x != 0:
|
||||
offsetmask[i] = torch.roll(offsetmask[i], shifts=int(shift_x), dims=1)
|
||||
if shift_y != 0:
|
||||
offsetmask[i] = torch.roll(offsetmask[i], shifts=int(shift_y), dims=0)
|
||||
|
||||
return offsetmask,
|
||||
|
||||
class ImageTransformByNormalizedAmplitude:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"normalized_amp": ("NORMALIZED_AMPLITUDE",),
|
||||
"zoom_scale": ("FLOAT", { "default": 0.0, "min": -1.0, "max": 1.0, "step": 0.001, "display": "number" }),
|
||||
"x_offset": ("INT", { "default": 0, "min": (1 -MAX_RESOLUTION), "max": MAX_RESOLUTION, "step": 1, "display": "number" }),
|
||||
"y_offset": ("INT", { "default": 0, "min": (1 -MAX_RESOLUTION), "max": MAX_RESOLUTION, "step": 1, "display": "number" }),
|
||||
"cumulative": ("BOOLEAN", { "default": False }),
|
||||
"image": ("IMAGE",),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "amptransform"
|
||||
CATEGORY = "KJNodes/audio"
|
||||
DESCRIPTION = """
|
||||
Works as a bridge to the AudioScheduler -nodes:
|
||||
https://github.com/a1lazydog/ComfyUI-AudioScheduler
|
||||
Transforms image based on the normalized amplitude.
|
||||
"""
|
||||
|
||||
def amptransform(self, image, normalized_amp, zoom_scale, cumulative, x_offset, y_offset):
|
||||
# Ensure normalized_amp is an array and within the range [0, 1]
|
||||
normalized_amp = np.clip(normalized_amp, 0.0, 1.0)
|
||||
transformed_images = []
|
||||
|
||||
# Initialize the cumulative zoom factor
|
||||
prev_amp = 0.0
|
||||
|
||||
for i in range(image.shape[0]):
|
||||
img = image[i] # Get the i-th image in the batch
|
||||
amp = normalized_amp[i] # Get the corresponding amplitude value
|
||||
|
||||
# Incrementally increase the cumulative zoom factor
|
||||
if cumulative:
|
||||
prev_amp += amp
|
||||
amp += prev_amp
|
||||
|
||||
# Convert the image tensor from BxHxWxC to CxHxW format expected by torchvision
|
||||
img = img.permute(2, 0, 1)
|
||||
|
||||
# Convert PyTorch tensor to PIL Image for processing
|
||||
pil_img = TF.to_pil_image(img)
|
||||
|
||||
# Calculate the crop size based on the amplitude
|
||||
width, height = pil_img.size
|
||||
crop_size = int(min(width, height) * (1 - amp * zoom_scale))
|
||||
crop_size = max(crop_size, 1)
|
||||
|
||||
# Calculate the crop box coordinates (centered crop)
|
||||
left = (width - crop_size) // 2
|
||||
top = (height - crop_size) // 2
|
||||
right = (width + crop_size) // 2
|
||||
bottom = (height + crop_size) // 2
|
||||
|
||||
# Crop and resize back to original size
|
||||
cropped_img = TF.crop(pil_img, top, left, crop_size, crop_size)
|
||||
resized_img = TF.resize(cropped_img, (height, width))
|
||||
|
||||
# Convert back to tensor in CxHxW format
|
||||
tensor_img = TF.to_tensor(resized_img)
|
||||
|
||||
# Convert the tensor back to BxHxWxC format
|
||||
tensor_img = tensor_img.permute(1, 2, 0)
|
||||
|
||||
# Offset the image based on the amplitude
|
||||
offset_amp = amp * 10 # Calculate the offset magnitude based on the amplitude
|
||||
shift_x = min(x_offset * offset_amp, img.shape[1] - 1) # Calculate the shift in x direction
|
||||
shift_y = min(y_offset * offset_amp, img.shape[0] - 1) # Calculate the shift in y direction
|
||||
|
||||
# Apply the offset to the image tensor
|
||||
if shift_x != 0:
|
||||
tensor_img = torch.roll(tensor_img, shifts=int(shift_x), dims=1)
|
||||
if shift_y != 0:
|
||||
tensor_img = torch.roll(tensor_img, shifts=int(shift_y), dims=0)
|
||||
|
||||
# Add to the list
|
||||
transformed_images.append(tensor_img)
|
||||
|
||||
# Stack all transformed images into a batch
|
||||
transformed_batch = torch.stack(transformed_images)
|
||||
|
||||
return (transformed_batch,)
|
||||
|
||||
def parse_coordinates(coordinates_str):
|
||||
coordinates = {}
|
||||
pattern = r'(\d+):\((\d+),(\d+)\)'
|
||||
@ -4203,8 +3307,6 @@ Normalize the images to be in the range [-1, 1]
|
||||
images = images * 2.0 - 1.0
|
||||
return (images,)
|
||||
|
||||
import comfy.sample
|
||||
from nodes import CLIPTextEncode
|
||||
folder_paths.add_model_folder_path("intristic_loras", os.path.join(script_directory, "intristic_loras"))
|
||||
|
||||
class Intrinsic_lora_sampling:
|
||||
@ -211,6 +211,10 @@ const create_documentation_stylesheet = () => {
|
||||
this.show_doc = !this.show_doc
|
||||
docElement.parentNode.removeChild(docElement)
|
||||
docElement = null
|
||||
if (contentWrapper) {
|
||||
contentWrapper.remove()
|
||||
contentWrapper = null
|
||||
}
|
||||
},
|
||||
{ signal: this.docCtrl.signal },
|
||||
);
|
||||
@ -247,7 +251,7 @@ const create_documentation_stylesheet = () => {
|
||||
const transform = new DOMMatrix()
|
||||
.scaleSelf(scaleX, scaleY)
|
||||
.multiplySelf(ctx.getTransform())
|
||||
.translateSelf(this.size[0] * scaleX * window.devicePixelRatio, 0)
|
||||
.translateSelf(this.size[0] * scaleX * Math.max(1.0,window.devicePixelRatio) , 0)
|
||||
.translateSelf(10, -32)
|
||||
|
||||
const scale = new DOMMatrix()
|
||||
@ -301,4 +305,20 @@ const create_documentation_stylesheet = () => {
|
||||
}
|
||||
return r;
|
||||
}
|
||||
const onRem = nodeType.prototype.onRemoved
|
||||
|
||||
nodeType.prototype.onRemoved = function () {
|
||||
const r = onRem ? onRem.apply(this, []) : undefined
|
||||
|
||||
if (docElement) {
|
||||
docElement.remove()
|
||||
docElement = null
|
||||
}
|
||||
|
||||
if (contentWrapper) {
|
||||
contentWrapper.remove()
|
||||
contentWrapper = null
|
||||
}
|
||||
return r
|
||||
}
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user