327 lines
12 KiB
Python
327 lines
12 KiB
Python
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
|
|
class TextOverlay:
|
|
"""
|
|
A node to overlay text on images with various customization options such as
|
|
font size, font type, text color, stroke color, stroke thickness, padding,
|
|
alignment, and position adjustments.
|
|
"""
|
|
|
|
_horizontal_alignments = ["left", "center", "right"]
|
|
_vertical_alignments = ["top", "middle", "bottom"]
|
|
|
|
def __init__(self, device="cpu"):
|
|
"""
|
|
Initializes the TextOverlay node.
|
|
|
|
Parameters:
|
|
- device (str): The computing device to use, default is 'cpu'.
|
|
"""
|
|
|
|
self.device = device
|
|
self._loaded_font = None
|
|
self._full_text = None
|
|
self._x = None
|
|
self._y = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""
|
|
Defines the expected input parameters for the text overlay functionality.
|
|
|
|
Returns:
|
|
- dict: A dictionary specifying required inputs and their attributes.
|
|
"""
|
|
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",), # Input image to overlay text on
|
|
"text": (
|
|
"STRING",
|
|
{"multiline": True, "default": "Hello"},
|
|
), # Text to overlay
|
|
"font_size": (
|
|
"INT",
|
|
{"default": 32, "min": 1, "max": 9999, "step": 1},
|
|
), # Font size
|
|
"font": (
|
|
"STRING",
|
|
{"default": "ariblk.ttf"},
|
|
), # Font name (e.g. arial.ttf)
|
|
"fill_color_hex": (
|
|
"STRING",
|
|
{"default": "#FFFFFF"},
|
|
), # Text fill color in hex
|
|
"stroke_color_hex": (
|
|
"STRING",
|
|
{"default": "#000000"},
|
|
), # Text stroke color in hex
|
|
"stroke_thickness": (
|
|
"FLOAT",
|
|
{"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01},
|
|
), # Stroke thickness
|
|
"padding": (
|
|
"INT",
|
|
{"default": 16, "min": 0, "max": 128, "step": 1},
|
|
), # Padding around text
|
|
"horizontal_alignment": (
|
|
cls._horizontal_alignments,
|
|
{"default": "center"},
|
|
), # Horizontal alignment
|
|
"vertical_alignment": (
|
|
cls._vertical_alignments,
|
|
{"default": "bottom"},
|
|
), # Vertical alignment
|
|
"x_shift": (
|
|
"INT",
|
|
{"default": 0, "min": -128, "max": 128, "step": 1},
|
|
), # Horizontal position adjustment
|
|
"y_shift": (
|
|
"INT",
|
|
{"default": 0, "min": -128, "max": 128, "step": 1},
|
|
), # Vertical position adjustment
|
|
"line_spacing": (
|
|
"FLOAT",
|
|
{"default": 4.0, "min": 0.0, "max": 50.0, "step": 0.5},
|
|
), # Spacing between lines of text
|
|
}
|
|
}
|
|
|
|
# Static attributes defining the return type, function name, and category for the class
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "batch_process"
|
|
CATEGORY = "image/text"
|
|
|
|
def hex_to_rgb(self, hex_color):
|
|
"""
|
|
Converts a hex color string to an RGB tuple.
|
|
|
|
Parameters:
|
|
- hex_color (str): The color in hex format as a string.
|
|
|
|
Returns:
|
|
- tuple: The color converted to an (R, G, B) tuple.
|
|
"""
|
|
|
|
hex_color = hex_color.lstrip("#")
|
|
if len(hex_color) == 3:
|
|
hex_color = hex_color * 2
|
|
return tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4))
|
|
|
|
def draw_text(
|
|
self,
|
|
image,
|
|
text,
|
|
font_size,
|
|
font,
|
|
fill_color_hex,
|
|
stroke_color_hex,
|
|
stroke_thickness,
|
|
padding,
|
|
horizontal_alignment,
|
|
vertical_alignment,
|
|
x_shift,
|
|
y_shift,
|
|
line_spacing,
|
|
use_cache=False,
|
|
):
|
|
"""
|
|
Draws the specified text on the given image with the provided styling and alignment options.
|
|
|
|
Parameters:
|
|
- image (PIL.Image.Image): The image to draw text on.
|
|
- text (str): The text to overlay on the image.
|
|
- font_size (int): The font size of the text.
|
|
- font (str): The font type of the text.
|
|
- fill_color_hex (str): The fill color of the text in hex format.
|
|
- stroke_color_hex (str): The stroke color of the text in hex format.
|
|
- stroke_thickness (float): The thickness of the text stroke.
|
|
- padding (int): The padding around the text.
|
|
- horizontal_alignment (str): The horizontal alignment of the text ('left', 'center', 'right').
|
|
- vertical_alignment (str): The vertical alignment of the text ('top', 'middle', 'bottom').
|
|
- x_shift (int): Horizontal position adjustment for the text.
|
|
- y_shift (int): Vertical position adjustment for the text.
|
|
- line_spacing (float): Spacing between lines of text.
|
|
- use_cache (bool): Flag to use cached font and position calculations to improve performance.
|
|
|
|
Returns:
|
|
- PIL.Image.Image: The image with the text overlay applied.
|
|
"""
|
|
|
|
# Load font from the fonts directory or use default if not found or specified to not use cache
|
|
if self._loaded_font is None or use_cache is False:
|
|
fonts_dir = os.path.join(os.path.dirname(__file__), "fonts")
|
|
font_path = os.path.join(fonts_dir, font)
|
|
|
|
# Check if the font file exists in the fonts directory
|
|
if not os.path.exists(font_path):
|
|
# If not, set path to font name directly - this will cause PIL to search for the font in the system
|
|
font_path = font
|
|
|
|
try:
|
|
self._loaded_font = ImageFont.truetype(font_path, font_size)
|
|
except Exception as e:
|
|
print(f"Error loading font: {e}... Using default font")
|
|
self._loaded_font = ImageFont.load_default(font_size)
|
|
|
|
# Prepare to draw on the image
|
|
draw = ImageDraw.Draw(image)
|
|
|
|
# Process text for multiline support and fit within image dimensions
|
|
words = text.replace("\n", "\n ").split(" ")
|
|
if self._full_text is None or use_cache is False:
|
|
text_lines, line = [], ""
|
|
for word in words:
|
|
extra_line = "\n" in word
|
|
word = word.strip()
|
|
if (
|
|
draw.textlength(line + word, font=self._loaded_font)
|
|
< image.width - 2 * padding
|
|
):
|
|
line += word + " "
|
|
else:
|
|
text_lines.append(line.strip())
|
|
line = word + " "
|
|
if extra_line:
|
|
text_lines.append(line.strip())
|
|
line = ""
|
|
text_lines.append(line.strip())
|
|
self._full_text = "\n".join(text_lines)
|
|
|
|
# Calculate text position based on alignment and position adjustments
|
|
if self._x is None or self._y is None or use_cache is False:
|
|
left, top, right, bottom = draw.multiline_textbbox(
|
|
(0, 0),
|
|
self._full_text,
|
|
font=self._loaded_font,
|
|
stroke_width=int(font_size * stroke_thickness * 0.5),
|
|
align=horizontal_alignment,
|
|
spacing=line_spacing,
|
|
)
|
|
if horizontal_alignment == "left":
|
|
self._x = padding
|
|
elif horizontal_alignment == "center":
|
|
self._x = (image.width - (right - left)) / 2
|
|
elif horizontal_alignment == "right":
|
|
self._x = image.width - (right - left) - padding
|
|
self._x += x_shift
|
|
if vertical_alignment == "middle":
|
|
self._y = (image.height - (bottom - top)) / 2
|
|
elif vertical_alignment == "top":
|
|
self._y = padding
|
|
elif vertical_alignment == "bottom":
|
|
self._y = image.height - (bottom - top) - padding
|
|
self._y += y_shift
|
|
|
|
# Draw the processed text onto the image
|
|
draw.text(
|
|
(self._x, self._y),
|
|
self._full_text,
|
|
fill=self.hex_to_rgb(fill_color_hex),
|
|
stroke_fill=self.hex_to_rgb(stroke_color_hex),
|
|
stroke_width=int(font_size * stroke_thickness * 0.5),
|
|
font=self._loaded_font,
|
|
align=horizontal_alignment,
|
|
spacing=line_spacing,
|
|
)
|
|
return image
|
|
|
|
def batch_process(
|
|
self,
|
|
image,
|
|
text,
|
|
font_size,
|
|
font,
|
|
fill_color_hex,
|
|
stroke_color_hex,
|
|
stroke_thickness,
|
|
padding,
|
|
horizontal_alignment,
|
|
vertical_alignment,
|
|
x_shift,
|
|
y_shift,
|
|
line_spacing,
|
|
):
|
|
"""
|
|
Processes a batch of images or a single image, adding the specified text overlay
|
|
with the given styling and alignment options.
|
|
|
|
Parameters:
|
|
- image (torch.Tensor or numpy.ndarray): The image(s) to process.
|
|
- text (str): The text to overlay on the image(s).
|
|
- font_size (int): The font size of the text.
|
|
- font (str): The font type of the text.
|
|
- fill_color_hex (str): The fill color of the text in hex format.
|
|
- stroke_color_hex (str): The stroke color of the text in hex format.
|
|
- stroke_thickness (float): The thickness of the text stroke.
|
|
- padding (int): The padding around the text.
|
|
- horizontal_alignment (str): The horizontal alignment of the text.
|
|
- vertical_alignment (str): The vertical alignment of the text.
|
|
- x_shift (int): Horizontal position adjustment for the text.
|
|
- y_shift (int): Vertical position adjustment for the text.
|
|
- line_spacing (float): Spacing between lines of text.
|
|
|
|
Returns:
|
|
- tuple: A tuple containing the processed image(s) as a torch.Tensor.
|
|
"""
|
|
|
|
# Handles both single and batch image processing for text overlay
|
|
if len(image.shape) == 3: # Single image
|
|
image_np = image.cpu().numpy()
|
|
image = Image.fromarray((image_np.squeeze(0) * 255).astype(np.uint8))
|
|
image = self.draw_text(
|
|
image,
|
|
text,
|
|
font_size,
|
|
font,
|
|
fill_color_hex,
|
|
stroke_color_hex,
|
|
stroke_thickness,
|
|
padding,
|
|
horizontal_alignment,
|
|
vertical_alignment,
|
|
x_shift,
|
|
y_shift,
|
|
line_spacing,
|
|
)
|
|
image_tensor_out = torch.tensor(np.array(image).astype(np.float32) / 255.0)
|
|
image_tensor_out = torch.unsqueeze(image_tensor_out, 0)
|
|
return (image_tensor_out,)
|
|
else: # Batch of images
|
|
image_np = image.cpu().numpy()
|
|
images = [Image.fromarray((img * 255).astype(np.uint8)) for img in image_np]
|
|
images_out, use_cache = [], False
|
|
for img in images:
|
|
img = self.draw_text(
|
|
img,
|
|
text,
|
|
font_size,
|
|
font,
|
|
fill_color_hex,
|
|
stroke_color_hex,
|
|
stroke_thickness,
|
|
padding,
|
|
horizontal_alignment,
|
|
vertical_alignment,
|
|
x_shift,
|
|
y_shift,
|
|
line_spacing,
|
|
use_cache,
|
|
)
|
|
images_out.append(np.array(img).astype(np.float32) / 255.0)
|
|
use_cache = True
|
|
images_np = np.stack(images_out)
|
|
images_tensor = torch.from_numpy(images_np)
|
|
return (images_tensor,)
|
|
|
|
|
|
# Mapping of node class names to their respective classes
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Text Overlay": TextOverlay,
|
|
}
|