|
|
|
@@ -1,196 +1,204 @@
|
|
|
|
|
import torch
|
|
|
|
|
import os
|
|
|
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
|
import torch
|
|
|
|
|
import numpy as np
|
|
|
|
|
import glob
|
|
|
|
|
import platform
|
|
|
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
|
# Font discovery
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
|
|
|
|
|
|
FONT_DIRS = [
|
|
|
|
|
".", # local node dir
|
|
|
|
|
"/usr/share/fonts",
|
|
|
|
|
"/usr/share/fonts/truetype",
|
|
|
|
|
"C:/Windows/Fonts",
|
|
|
|
|
]
|
|
|
|
|
|
|
|
|
|
def find_fonts():
|
|
|
|
|
fonts = []
|
|
|
|
|
for base in FONT_DIRS:
|
|
|
|
|
if not os.path.isdir(base):
|
|
|
|
|
continue
|
|
|
|
|
for root, _, files in os.walk(base):
|
|
|
|
|
for f in files:
|
|
|
|
|
if f.lower().endswith((".ttf", ".otf")):
|
|
|
|
|
fonts.append(os.path.join(root, f))
|
|
|
|
|
return sorted(set(fonts))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
font_list = find_fonts()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
|
# Utility functions
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
|
|
|
|
|
|
def hex_to_rgba(hex_color):
|
|
|
|
|
hex_color = hex_color.lstrip("#")
|
|
|
|
|
if len(hex_color) == 6:
|
|
|
|
|
r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
|
|
|
|
|
return (r, g, b, 255)
|
|
|
|
|
elif len(hex_color) == 8:
|
|
|
|
|
r, g, b, a = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4, 6))
|
|
|
|
|
return (r, g, b, a)
|
|
|
|
|
else:
|
|
|
|
|
raise ValueError(f"Invalid hex color: {hex_color}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def compute_position(position, img_w, img_h, text_w, text_h, padding):
|
|
|
|
|
if position == "top-left":
|
|
|
|
|
return (padding, padding)
|
|
|
|
|
if position == "top-right":
|
|
|
|
|
return (img_w - text_w - padding, padding)
|
|
|
|
|
if position == "bottom-left":
|
|
|
|
|
return (padding, img_h - text_h - padding)
|
|
|
|
|
if position == "bottom-right":
|
|
|
|
|
return (img_w - text_w - padding, img_h - text_h - padding)
|
|
|
|
|
if position == "center":
|
|
|
|
|
return (
|
|
|
|
|
(img_w - text_w) // 2,
|
|
|
|
|
(img_h - text_h) // 2,
|
|
|
|
|
)
|
|
|
|
|
return (padding, padding)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
|
# Node
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
|
|
|
|
|
|
class OverlayFrameNumber:
|
|
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
|
def INPUT_TYPES(cls):
|
|
|
|
|
# Dynamically get available system fonts
|
|
|
|
|
font_list = cls._get_system_fonts()
|
|
|
|
|
if not font_list:
|
|
|
|
|
font_list = ["Arial", "DejaVuSans", "Times New Roman", "Courier New"] # Fallback
|
|
|
|
|
|
|
|
|
|
return {
|
|
|
|
|
"required": {
|
|
|
|
|
"images": ("IMAGE",),
|
|
|
|
|
"font_size": ("INT", {"default": 32, "min": 8, "max": 200}),
|
|
|
|
|
"font_color": (["white", "black", "red", "green", "blue", "yellow", "cyan", "magenta"],),
|
|
|
|
|
"font": (font_list,),
|
|
|
|
|
"h_position": (["left", "center", "right"], {"default": "left"}),
|
|
|
|
|
"v_position": (["top", "center", "bottom"], {"default": "top"}),
|
|
|
|
|
"h_padding": ("INT", {"default": 20, "min": 0, "max": 1000}),
|
|
|
|
|
"v_padding": ("INT", {"default": 20, "min": 0, "max": 1000}),
|
|
|
|
|
"num_padding": ("INT", {"default": 3, "min": 1, "max": 5, "step": 1}),
|
|
|
|
|
"prefix_text": ("STRING", {"default": "Frame", "multiline": False}),
|
|
|
|
|
"outline_enabled": ("BOOLEAN", {"default": False}),
|
|
|
|
|
"outline_color": (["none", "white", "black", "red", "green", "blue", "yellow", "cyan", "magenta"], {"default": "black"}),
|
|
|
|
|
"stroke_width": ("INT", {"default": 1, "min": 0, "max": 5, "step": 1}),
|
|
|
|
|
"frame_number": ("INT", {"default": 1, "min": 1, "max": 99999}),
|
|
|
|
|
"position": (
|
|
|
|
|
["top-left", "top-right", "bottom-left", "bottom-right", "center"],
|
|
|
|
|
{"default": "bottom-right"},
|
|
|
|
|
),
|
|
|
|
|
"font_size": ("INT", {"default": 48, "min": 12, "max": 200}),
|
|
|
|
|
"font_color": ("STRING", {"default": "#FFFFFF"}),
|
|
|
|
|
"background_color": ("STRING", {"default": "#00000000"}),
|
|
|
|
|
"outline_color": ("STRING", {"default": "#000000"}),
|
|
|
|
|
"outline_width": ("INT", {"default": 2, "min": 0, "max": 10}),
|
|
|
|
|
"font_file": (["None"] + sorted(font_list), {"default": "None"}),
|
|
|
|
|
"text_padding": ("INT", {"default": 10, "min": 0, "max": 50}),
|
|
|
|
|
"num_padding": ("INT", {"default": 4, "min": 1, "max": 8}),
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
|
RETURN_NAMES = ("images",)
|
|
|
|
|
FUNCTION = "process"
|
|
|
|
|
CATEGORY = "Image/Processing"
|
|
|
|
|
OUTPUT_IS_LIST = (False,)
|
|
|
|
|
CATEGORY = "image/postprocessing"
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def _get_system_fonts():
|
|
|
|
|
"""Dynamically retrieve available TrueType fonts from the system."""
|
|
|
|
|
font_list = []
|
|
|
|
|
font_paths = []
|
|
|
|
|
def process(
|
|
|
|
|
self,
|
|
|
|
|
images,
|
|
|
|
|
frame_number,
|
|
|
|
|
position,
|
|
|
|
|
font_size,
|
|
|
|
|
font_color,
|
|
|
|
|
background_color,
|
|
|
|
|
outline_color,
|
|
|
|
|
outline_width,
|
|
|
|
|
font_file,
|
|
|
|
|
text_padding,
|
|
|
|
|
num_padding,
|
|
|
|
|
):
|
|
|
|
|
|
|
|
|
|
# Common font directories based on OS
|
|
|
|
|
if platform.system() == "Windows":
|
|
|
|
|
font_dirs = [os.path.join(os.environ.get("WINDIR", "C:\\Windows"), "Fonts")]
|
|
|
|
|
elif platform.system() == "Darwin": # macOS
|
|
|
|
|
font_dirs = ["/System/Library/Fonts", "/Library/Fonts", os.path.expanduser("~/Library/Fonts")]
|
|
|
|
|
else: # Linux/Unix-like
|
|
|
|
|
font_dirs = ["/usr/share/fonts", "/usr/local/share/fonts", os.path.expanduser("~/.fonts")]
|
|
|
|
|
# ----------------------------------------------------
|
|
|
|
|
# Normalize IMAGE input (batch tensor OR list)
|
|
|
|
|
# ----------------------------------------------------
|
|
|
|
|
|
|
|
|
|
# Add fontconfig fonts if available (Linux)
|
|
|
|
|
try:
|
|
|
|
|
import subprocess
|
|
|
|
|
result = subprocess.run(["fc-list", ":file"], capture_output=True, text=True)
|
|
|
|
|
if result.returncode == 0:
|
|
|
|
|
font_paths.extend([line.split(":")[0] for line in result.stdout.splitlines()])
|
|
|
|
|
except (subprocess.SubprocessError, FileNotFoundError):
|
|
|
|
|
pass
|
|
|
|
|
input_was_tensor = False
|
|
|
|
|
image_list = []
|
|
|
|
|
|
|
|
|
|
# Add fonts from common directories
|
|
|
|
|
for font_dir in font_dirs:
|
|
|
|
|
if os.path.exists(font_dir):
|
|
|
|
|
font_paths.extend(glob.glob(os.path.join(font_dir, "*.ttf")))
|
|
|
|
|
if isinstance(images, torch.Tensor) and images.ndim == 4:
|
|
|
|
|
input_was_tensor = True
|
|
|
|
|
for i in range(images.shape[0]):
|
|
|
|
|
image_list.append(images[i])
|
|
|
|
|
elif isinstance(images, (list, tuple)):
|
|
|
|
|
image_list = list(images)
|
|
|
|
|
else:
|
|
|
|
|
raise ValueError("Unsupported IMAGE input type")
|
|
|
|
|
|
|
|
|
|
# Extract font names (without path or extension)
|
|
|
|
|
font_list = sorted(set(os.path.splitext(os.path.basename(f))[0] for f in font_paths))
|
|
|
|
|
# ----------------------------------------------------
|
|
|
|
|
# Font
|
|
|
|
|
# ----------------------------------------------------
|
|
|
|
|
|
|
|
|
|
# Fallback to common fonts if none found
|
|
|
|
|
if not font_list:
|
|
|
|
|
font_list = ["Arial", "DejaVuSans", "Times New Roman", "Courier New"]
|
|
|
|
|
if font_file != "None" and os.path.isfile(font_file):
|
|
|
|
|
font = ImageFont.truetype(font_file, font_size)
|
|
|
|
|
else:
|
|
|
|
|
try:
|
|
|
|
|
font = ImageFont.truetype("arial.ttf", font_size)
|
|
|
|
|
except Exception:
|
|
|
|
|
font = ImageFont.load_default()
|
|
|
|
|
|
|
|
|
|
return font_list
|
|
|
|
|
font_rgba = hex_to_rgba(font_color)
|
|
|
|
|
bg_rgba = hex_to_rgba(background_color)
|
|
|
|
|
outline_rgba = hex_to_rgba(outline_color)
|
|
|
|
|
|
|
|
|
|
def process(self, images, font_size, font_color, font, h_position, v_position, h_padding, v_padding, num_padding, prefix_text, outline_enabled, outline_color, stroke_width):
|
|
|
|
|
# Validate inputs
|
|
|
|
|
if not isinstance(images, torch.Tensor) or len(images.shape) != 4:
|
|
|
|
|
raise ValueError("Input 'images' must be a 4D tensor (batch, height, width, channels)")
|
|
|
|
|
|
|
|
|
|
# Map font_color and outline_color to RGB
|
|
|
|
|
color_map = {
|
|
|
|
|
"white": (255, 255, 255), "black": (0, 0, 0), "red": (255, 0, 0),
|
|
|
|
|
"green": (0, 255, 0), "blue": (0, 0, 255), "yellow": (255, 255, 0),
|
|
|
|
|
"cyan": (0, 255, 255), "magenta": (255, 0, 255), "none": None
|
|
|
|
|
}
|
|
|
|
|
font_color_rgb = color_map[font_color]
|
|
|
|
|
outline_color_rgb = color_map[outline_color] if outline_enabled and outline_color != "none" else None
|
|
|
|
|
|
|
|
|
|
# Map font names to system font files (with common paths or aliases)
|
|
|
|
|
font_map = {
|
|
|
|
|
"Arial": "arial.ttf",
|
|
|
|
|
"DejaVuSans": "DejaVuSans.ttf",
|
|
|
|
|
"Times New Roman": "times.ttf",
|
|
|
|
|
"Courier New": "cour.ttf"
|
|
|
|
|
}
|
|
|
|
|
font_file = font_map.get(font, font + ".ttf") # Try font name directly as fallback
|
|
|
|
|
|
|
|
|
|
# Try to load the font, with fallback to default
|
|
|
|
|
try:
|
|
|
|
|
font_path = ImageFont.truetype(font_file, font_size)
|
|
|
|
|
except (OSError, IOError):
|
|
|
|
|
# Try to find the font in system directories
|
|
|
|
|
font_found = False
|
|
|
|
|
font_dirs = (
|
|
|
|
|
[os.path.join(os.environ.get("WINDIR", "C:\\Windows"), "Fonts")] if platform.system() == "Windows" else
|
|
|
|
|
["/System/Library/Fonts", "/Library/Fonts", os.path.expanduser("~/Library/Fonts")] if platform.system() == "Darwin" else
|
|
|
|
|
["/usr/share/fonts", "/usr/local/share/fonts", os.path.expanduser("~/.fonts")]
|
|
|
|
|
)
|
|
|
|
|
for font_dir in font_dirs:
|
|
|
|
|
font_path_candidate = os.path.join(font_dir, font_file)
|
|
|
|
|
if os.path.exists(font_path_candidate):
|
|
|
|
|
try:
|
|
|
|
|
font_path = ImageFont.truetype(font_path_candidate, font_size)
|
|
|
|
|
font_found = True
|
|
|
|
|
break
|
|
|
|
|
except (OSError, IOError):
|
|
|
|
|
continue
|
|
|
|
|
if not font_found:
|
|
|
|
|
print(f"Font {font_file} not found, using default font")
|
|
|
|
|
font_path = ImageFont.load_default()
|
|
|
|
|
|
|
|
|
|
# Process each frame
|
|
|
|
|
output_images = []
|
|
|
|
|
for i in range(images.shape[0]):
|
|
|
|
|
# Convert tensor to PIL image
|
|
|
|
|
img_tensor = images[i] # Shape: (height, width, channels)
|
|
|
|
|
img_array = (img_tensor.cpu().numpy() * 255).astype(np.uint8) # Convert to uint8
|
|
|
|
|
if img_array.shape[2] == 3: # RGB
|
|
|
|
|
img = Image.fromarray(img_array, mode="RGB")
|
|
|
|
|
else:
|
|
|
|
|
raise ValueError("Input images must have 3 channels (RGB)")
|
|
|
|
|
|
|
|
|
|
# Calculate text position
|
|
|
|
|
draw = ImageDraw.Draw(img)
|
|
|
|
|
# Format text: prefix + zero-padded number, or just number if prefix is empty
|
|
|
|
|
frame_number = f"{i + 1:0{num_padding}d}"
|
|
|
|
|
text = f"{prefix_text} {frame_number}" if prefix_text.strip() else frame_number
|
|
|
|
|
try:
|
|
|
|
|
text_bbox = draw.textbbox((0, 0), text, font=font_path)
|
|
|
|
|
text_width = text_bbox[2] - text_bbox[0]
|
|
|
|
|
text_height = text_bbox[3] - text_bbox[1]
|
|
|
|
|
except AttributeError: # Fallback for older PIL versions
|
|
|
|
|
text_width, text_height = draw.textsize(text, font=font_path)
|
|
|
|
|
# ----------------------------------------------------
|
|
|
|
|
# Process frames
|
|
|
|
|
# ----------------------------------------------------
|
|
|
|
|
|
|
|
|
|
img_width, img_height = img.size
|
|
|
|
|
for i, img_tensor in enumerate(image_list):
|
|
|
|
|
img_np = (img_tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
|
|
|
|
pil_img = Image.fromarray(img_np).convert("RGBA")
|
|
|
|
|
|
|
|
|
|
# Horizontal position
|
|
|
|
|
if h_position == "left":
|
|
|
|
|
x = h_padding
|
|
|
|
|
elif h_position == "center":
|
|
|
|
|
x = (img_width - text_width) // 2
|
|
|
|
|
else: # right
|
|
|
|
|
x = img_width - text_width - h_padding
|
|
|
|
|
draw = ImageDraw.Draw(pil_img)
|
|
|
|
|
|
|
|
|
|
# Vertical position
|
|
|
|
|
if v_position == "top":
|
|
|
|
|
y = v_padding
|
|
|
|
|
elif v_position == "center":
|
|
|
|
|
y = (img_height - text_height) // 2
|
|
|
|
|
else: # bottom
|
|
|
|
|
y = img_height - text_height - v_padding
|
|
|
|
|
label = f"{frame_number + i:0{num_padding}d}"
|
|
|
|
|
|
|
|
|
|
# Draw text with optional outline
|
|
|
|
|
try:
|
|
|
|
|
# Use stroke_width and stroke_fill if outline is enabled and supported (Pillow 9.0+)
|
|
|
|
|
if outline_enabled and outline_color_rgb is not None and stroke_width > 0:
|
|
|
|
|
draw.text((x, y), text, font=font_path, fill=font_color_rgb, stroke_width=stroke_width, stroke_fill=outline_color_rgb)
|
|
|
|
|
else:
|
|
|
|
|
draw.text((x, y), text, font=font_path, fill=font_color_rgb)
|
|
|
|
|
except TypeError: # Fallback for older Pillow versions without stroke_width
|
|
|
|
|
print("Text outline not supported in this Pillow version; rendering without outline")
|
|
|
|
|
draw.text((x, y), text, font=font_path, fill=font_color_rgb)
|
|
|
|
|
text_bbox = draw.textbbox((0, 0), label, font=font, stroke_width=outline_width)
|
|
|
|
|
text_w = text_bbox[2] - text_bbox[0] + text_padding * 2
|
|
|
|
|
text_h = text_bbox[3] - text_bbox[1] + text_padding * 2
|
|
|
|
|
|
|
|
|
|
# Convert back to tensor
|
|
|
|
|
img_array = np.array(img).astype(np.float32) / 255.0
|
|
|
|
|
output_images.append(torch.from_numpy(img_array))
|
|
|
|
|
x, y = compute_position(
|
|
|
|
|
position,
|
|
|
|
|
pil_img.width,
|
|
|
|
|
pil_img.height,
|
|
|
|
|
text_w,
|
|
|
|
|
text_h,
|
|
|
|
|
text_padding,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# Stack images back into a batch tensor
|
|
|
|
|
output_tensor = torch.stack(output_images, dim=0)
|
|
|
|
|
# Background box
|
|
|
|
|
if bg_rgba[3] > 0:
|
|
|
|
|
draw.rectangle(
|
|
|
|
|
[x, y, x + text_w, y + text_h],
|
|
|
|
|
fill=bg_rgba,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
return (output_tensor,)
|
|
|
|
|
# Text
|
|
|
|
|
draw.text(
|
|
|
|
|
(x + text_padding, y + text_padding),
|
|
|
|
|
label,
|
|
|
|
|
fill=font_rgba,
|
|
|
|
|
font=font,
|
|
|
|
|
stroke_width=outline_width,
|
|
|
|
|
stroke_fill=outline_rgba,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
|
def IS_CHANGED(cls, **kwargs):
|
|
|
|
|
# Re-run if inputs change
|
|
|
|
|
return True
|
|
|
|
|
out_np = np.array(pil_img.convert("RGB")).astype(np.float32) / 255.0
|
|
|
|
|
out_tensor = torch.from_numpy(out_np)
|
|
|
|
|
|
|
|
|
|
output_images.append(out_tensor)
|
|
|
|
|
|
|
|
|
|
# ----------------------------------------------------
|
|
|
|
|
# Return
|
|
|
|
|
# ----------------------------------------------------
|
|
|
|
|
|
|
|
|
|
if input_was_tensor:
|
|
|
|
|
return (torch.stack(output_images, dim=0),)
|
|
|
|
|
else:
|
|
|
|
|
return (output_images,)
|
|
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
|
def OUTPUT_UI(cls, outputs):
|
|
|
|
|
# No UI output for image batches (handled by downstream nodes like Preview Image)
|
|
|
|
|
return {}
|
|
|
|
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
|
|
|
"OverlayFrameNumber": OverlayFrameNumber
|
|
|
|
@@ -198,4 +206,4 @@ NODE_CLASS_MAPPINGS = {
|
|
|
|
|
|
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
|
|
|
"OverlayFrameNumber": "Overlay Frame Number"
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|