201 lines
8.9 KiB
Python
201 lines
8.9 KiB
Python
import torch
|
|
import os
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
import numpy as np
|
|
import glob
|
|
import platform
|
|
|
|
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}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Image/Processing"
|
|
OUTPUT_IS_LIST = (False,)
|
|
|
|
@staticmethod
|
|
def _get_system_fonts():
|
|
"""Dynamically retrieve available TrueType fonts from the system."""
|
|
font_list = []
|
|
font_paths = []
|
|
|
|
# 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")]
|
|
|
|
# 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
|
|
|
|
# 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")))
|
|
|
|
# Extract font names (without path or extension)
|
|
font_list = sorted(set(os.path.splitext(os.path.basename(f))[0] for f in font_paths))
|
|
|
|
# Fallback to common fonts if none found
|
|
if not font_list:
|
|
font_list = ["Arial", "DejaVuSans", "Times New Roman", "Courier New"]
|
|
|
|
return font_list
|
|
|
|
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)
|
|
|
|
img_width, img_height = img.size
|
|
|
|
# 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
|
|
|
|
# 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
|
|
|
|
# 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)
|
|
|
|
# Convert back to tensor
|
|
img_array = np.array(img).astype(np.float32) / 255.0
|
|
output_images.append(torch.from_numpy(img_array))
|
|
|
|
# Stack images back into a batch tensor
|
|
output_tensor = torch.stack(output_images, dim=0)
|
|
|
|
return (output_tensor,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
# Re-run if inputs change
|
|
return True
|
|
|
|
@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
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"OverlayFrameNumber": "Overlay Frame Number"
|
|
} |