Create fat_labels.py

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,
+}