diff --git a/fat_labels.py b/fat_labels.py new file mode 100644 index 0000000..e621c30 --- /dev/null +++ b/fat_labels.py @@ -0,0 +1,66 @@ +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, +}