241 lines
6.4 KiB
Python
241 lines
6.4 KiB
Python
from PIL import Image, ImageDraw, ImageFont
|
|
import os
|
|
import io
|
|
from datetime import datetime
|
|
import textwrap
|
|
import string
|
|
from Common.Utils import (
|
|
get_system_font_files,
|
|
images_data_to_tensor,
|
|
)
|
|
|
|
|
|
def get_default_font():
|
|
try:
|
|
return ImageFont.load_default()
|
|
except IOError:
|
|
possible_fonts = [
|
|
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", # Linux
|
|
"/Library/Fonts/Arial.ttf", # macOS
|
|
"C:\\Windows\\Fonts\\arial.ttf", # Windows
|
|
]
|
|
for font_path in possible_fonts:
|
|
if os.path.exists(font_path):
|
|
return font_path
|
|
raise IOError("No usable default font found.")
|
|
|
|
|
|
def create_text_image(
|
|
text,
|
|
width,
|
|
height,
|
|
font,
|
|
font_size=None,
|
|
text_color="white",
|
|
bg_color="black",
|
|
offset_x=0,
|
|
offset_y=0,
|
|
):
|
|
img = Image.new("RGB", (width, height), color=bg_color)
|
|
draw = ImageDraw.Draw(img)
|
|
|
|
try:
|
|
if isinstance(font, str):
|
|
font = ImageFont.truetype(font, font_size or 20)
|
|
elif font_size:
|
|
font = font.font_variant(size=font_size)
|
|
|
|
# Calculate the average character width
|
|
avg_char_width = (
|
|
sum(font.getbbox(char)[2] for char in string.ascii_lowercase) / 26
|
|
)
|
|
|
|
# Calculate the maximum characters per line
|
|
max_char_count = int(width / avg_char_width)
|
|
|
|
# Wrap the text
|
|
lines = textwrap.wrap(text, width=max_char_count)
|
|
|
|
# Calculate total text height
|
|
line_height = font.getbbox("hg")[3] - font.getbbox("hg")[1]
|
|
text_height = len(lines) * line_height
|
|
|
|
# Calculate starting Y position to center the text block
|
|
y = offset_y + (height - text_height) / 2
|
|
|
|
for line in lines:
|
|
# Get line width
|
|
line_width = font.getbbox(line)[2]
|
|
|
|
# Calculate starting X position to center this line
|
|
x = offset_x + (width - line_width) / 2
|
|
|
|
# Draw the line
|
|
draw.text((x, y), line, font=font, fill=text_color)
|
|
|
|
# Move to next line
|
|
y += line_height
|
|
|
|
except Exception as e:
|
|
print(f"Error creating image: {str(e)}")
|
|
return None
|
|
|
|
return img
|
|
|
|
|
|
def create_text_image_pil(
|
|
text,
|
|
width,
|
|
height,
|
|
font,
|
|
font_size=None,
|
|
text_color="white",
|
|
bg_color="black",
|
|
offset_x=0,
|
|
offset_y=0,
|
|
):
|
|
img = create_text_image(
|
|
text, width, height, font, font_size, text_color, bg_color, offset_x, offset_y
|
|
)
|
|
img_byte_arr = io.BytesIO()
|
|
img.save(img_byte_arr, format="PNG")
|
|
img_byte_arr.seek(0)
|
|
return Image.open(img_byte_arr)
|
|
|
|
|
|
class MLTaskUtilsTextImageGenerator:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"text": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"default": "text here",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"width": (
|
|
"INT",
|
|
{
|
|
"default": 512,
|
|
},
|
|
),
|
|
"height": (
|
|
"INT",
|
|
{
|
|
"default": 512,
|
|
},
|
|
),
|
|
"font_size": (
|
|
"INT",
|
|
{
|
|
"default": 100,
|
|
},
|
|
),
|
|
# TODO move this to a widget
|
|
"font_name": (sorted(get_system_font_files()),),
|
|
"offset_x": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
},
|
|
),
|
|
"offset_y": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
},
|
|
),
|
|
# TODO color widget
|
|
# "text_color": (
|
|
# "INT",
|
|
# {
|
|
# "default": 0,
|
|
# "min": 0,
|
|
# "max": 0xFFFFFF,
|
|
# "step": 1,
|
|
# "display": "color",
|
|
# },
|
|
# ),
|
|
# "bg_color": (
|
|
# "INT",
|
|
# {
|
|
# "default": 0,
|
|
# "min": 0,
|
|
# "max": 0x000000,
|
|
# "step": 1,
|
|
# "display": "color",
|
|
# },
|
|
# ),
|
|
},
|
|
}
|
|
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "generate_text_image"
|
|
CATEGORY = "MLTask/SocialMan/Utils"
|
|
RETURN_TYPES = (
|
|
"IMAGE",
|
|
"MASK",
|
|
"IMAGE",
|
|
"MASK",
|
|
)
|
|
RETURN_NAMES = (
|
|
"text_image",
|
|
"text_image_mask",
|
|
"text_image_inverted",
|
|
"text_image_mask_inverted",
|
|
)
|
|
|
|
def generate_text_image(
|
|
self, text, width, height, font_size, font_name, offset_x, offset_y
|
|
):
|
|
|
|
# font = get_default_font() if args.font_path is None else args.font_path
|
|
font = font_name # "Arial Rounded Bold.ttf"
|
|
text_color = "black"
|
|
bg_color = "white"
|
|
img = create_text_image_pil(
|
|
text,
|
|
width,
|
|
height,
|
|
font,
|
|
font_size,
|
|
text_color,
|
|
bg_color,
|
|
offset_x,
|
|
offset_y,
|
|
)
|
|
# INVERTED
|
|
text_color = "white"
|
|
bg_color = "black"
|
|
img_inverted = create_text_image_pil(
|
|
text,
|
|
width,
|
|
height,
|
|
font,
|
|
font_size,
|
|
text_color,
|
|
bg_color,
|
|
offset_x,
|
|
offset_y,
|
|
)
|
|
if img:
|
|
# script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
# timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
# file_name = f"text_image_{timestamp}.png"
|
|
# # file_path = os.path.join(script_dir, file_name)
|
|
# file_path = f"{folder_paths.get_output_directory()}/{file_name}"
|
|
# img.save(file_path)
|
|
# print(f"Image saved as {file_path}")
|
|
# return images_file_to_tensor(file_path)
|
|
|
|
return images_data_to_tensor(img) + images_data_to_tensor(img_inverted)
|
|
else:
|
|
print("Failed to create image.")
|