67 lines
2.3 KiB
Python
67 lines
2.3 KiB
Python
from PIL import Image, ImageDraw, ImageFont
|
|
import torch
|
|
import numpy as np
|
|
|
|
class FatLabels2:
|
|
def __init__(self, device="cpu"):
|
|
self.device = device
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"default": "Hello"}),
|
|
"font_size": ("INT", {"default": 36, "min": 1}), # Font size in pixels
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "create_fat_label_with_cv2"
|
|
CATEGORY = "image/text"
|
|
|
|
def create_fat_label_with_cv2(self, text, font_size):
|
|
# Create a blank grayscale image as canvas with a fixed background color
|
|
bg_color = 0 # Black background (grayscale)
|
|
|
|
# Create a drawing context to calculate text size
|
|
text_width, text_height = self.calculate_text_size(text, font_size)
|
|
|
|
# Calculate canvas dimensions with padding
|
|
canvas_width = text_width + 40 # Add 20px padding on each side
|
|
canvas_height = text_height + 40 # Add 20px padding on each side
|
|
|
|
canvas = Image.new("L", (canvas_width, canvas_height), bg_color)
|
|
|
|
# Font color is always white
|
|
font_color = 255 # White (grayscale)
|
|
|
|
# Create an ImageFont object with the desired font size
|
|
font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", font_size)
|
|
|
|
# Create a drawing context
|
|
draw = ImageDraw.Draw(canvas)
|
|
|
|
# Calculate text position
|
|
x = (canvas_width - text_width) // 2
|
|
y = (canvas_height - text_height) // 2
|
|
|
|
# Draw text on the image with the specified font size
|
|
draw.text((x, y), text, fill=font_color, font=font)
|
|
|
|
# Convert the image to a PyTorch tensor
|
|
image_tensor_out = torch.tensor(np.array(canvas).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
return (image_tensor_out,)
|
|
|
|
def calculate_text_size(self, text, font_size):
|
|
# Create a temporary canvas to calculate text size
|
|
canvas = Image.new("L", (1, 1), 0) # Create a blank 1x1 grayscale image
|
|
draw = ImageDraw.Draw(canvas)
|
|
font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", font_size)
|
|
text_width, text_height = draw.textsize(text, font=font)
|
|
return text_width, text_height
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"FatLabels": FatLabels2,
|
|
}
|