xinyun/ComfyUI-KJNodes
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git synced 2025-12-08 20:34:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add CreateTextMask -node

Browse Source 
CreateTextMask allows creating text masks
This commit is contained in:
kijai 2023-10-07 19:29:20 +03:00
parent 57e8efb7f7
commit d526683f25
2 changed files with 58 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

BIN
fonts/TTNorms-Black.otf Normal file
View File

Binary file not shown.

60
nodes.py
View File

@ -3,10 +3,12 @@ import torch
import torch.nn.functional as F import torch.nn.functional as F
import scipy.ndimage import scipy.ndimage
import numpy as np import numpy as np
from PIL import ImageColor, Image from PIL import ImageColor, Image, ImageDraw, ImageFont
import os
from nodes import MAX_RESOLUTION from nodes import MAX_RESOLUTION
script_dir = os.path.dirname(os.path.abspath(__file__))
class INTConstant: class INTConstant:
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
@ -50,7 +52,6 @@ class CreateGradientMask:
def createmask(self, frames, width, height, invert): def createmask(self, frames, width, height, invert):
# Define the number of images in the batch # Define the number of images in the batch
batch_size = frames batch_size = frames
out = [] out = []
# Create an empty array to store the image batch # Create an empty array to store the image batch
image_batch = np.zeros((batch_size, height, width), dtype=np.float32) image_batch = np.zeros((batch_size, height, width), dtype=np.float32)
@ -66,6 +67,59 @@ class CreateGradientMask:
if invert: if invert:
return (1.0 - torch.stack(out, dim=0),) return (1.0 - torch.stack(out, dim=0),)
return (torch.stack(out, dim=0),) return (torch.stack(out, dim=0),)
class CreateTextMask:
RETURN_TYPES = ("MASK",)
FUNCTION = "createtextmask"
CATEGORY = "KJNodes"
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"invert": ("BOOLEAN", {"default": False}),
"frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}),
"text_x": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}),
"text_y": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}),
"font_size": ("INT", {"default": 32,"min": 8, "max": 4096, "step": 1}),
"text": ("STRING", {"default": "HELLO!"}),
"font_path": ("STRING", {"default": "fonts\\TTNorms-Black.otf"}),
"width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}),
"height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}),
"start_rotation": ("INT", {"default": 0,"min": 0, "max": 359, "step": 1}),
"end_rotation": ("INT", {"default": 359,"min": 0, "max": 359, "step": 1}),
},
}
def createtextmask(self, frames, width, height, invert, text_x, text_y, text, font_size, font_path, start_rotation, end_rotation):
# Define the number of images in the batch
batch_size = frames
out = []
rotation = start_rotation
rotation_increment = (end_rotation - start_rotation) / (batch_size - 1)
# Create an empty array to store the image batch
image_batch = np.zeros((batch_size, height, width), dtype=np.float32)
if font_path == "fonts\\TTNorms-Black.otf": #I don't know why relative path won't work otherwise...
font_path = os.path.join(script_dir, font_path)
# Generate the text
for i in range(batch_size):
image = Image.new("RGB", (width, height), "black")
draw = ImageDraw.Draw(image)
font = ImageFont.truetype(font_path, font_size)
text_width, text_height = draw.textsize(text, font=font)
text_center_x = text_x + text_width / 2
text_center_y = text_y + text_height / 2
draw.text((text_x, text_y), text, font=font, fill="white")
image = image.rotate(rotation, center=(text_center_x, text_center_y))
rotation += rotation_increment
image_batch[i] = np.array(image.convert("L"))
output = torch.from_numpy(image_batch)
rotation += 10
out.append(output)
if invert:
return (1.0 - torch.stack(out, dim=0),)
return (torch.stack(out, dim=0),)
class GrowMaskWithBlur: class GrowMaskWithBlur:
@classmethod @classmethod
@ -293,6 +347,7 @@ NODE_CLASS_MAPPINGS = {
"GrowMaskWithBlur": GrowMaskWithBlur, "GrowMaskWithBlur": GrowMaskWithBlur,
"ColorToMask": ColorToMask, "ColorToMask": ColorToMask,
"CreateGradientMask": CreateGradientMask, "CreateGradientMask": CreateGradientMask,
"CreateTextMask": CreateTextMask,
} }
NODE_DISPLAY_NAME_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = {
"INTConstant": "INT Constant", "INTConstant": "INT Constant",
@ -301,4 +356,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"GrowMaskWithBlur": "GrowMaskWithBlur", "GrowMaskWithBlur": "GrowMaskWithBlur",
"ColorToMask": "ColorToMask", "ColorToMask": "ColorToMask",
"CreateGradientMask": "CreateGradientMask", "CreateGradientMask": "CreateGradientMask",
"CreateTextMask" : "CreateTextMask",
} }