210 lines
6.6 KiB
Python
210 lines
6.6 KiB
Python
import os
|
|
import torch
|
|
import numpy as np
|
|
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):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"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",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "image/postprocessing"
|
|
|
|
def process(
|
|
self,
|
|
images,
|
|
frame_number,
|
|
position,
|
|
font_size,
|
|
font_color,
|
|
background_color,
|
|
outline_color,
|
|
outline_width,
|
|
font_file,
|
|
text_padding,
|
|
num_padding,
|
|
):
|
|
|
|
# ----------------------------------------------------
|
|
# Normalize IMAGE input (batch tensor OR list)
|
|
# ----------------------------------------------------
|
|
|
|
input_was_tensor = False
|
|
image_list = []
|
|
|
|
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")
|
|
|
|
# ----------------------------------------------------
|
|
# Font
|
|
# ----------------------------------------------------
|
|
|
|
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()
|
|
|
|
font_rgba = hex_to_rgba(font_color)
|
|
bg_rgba = hex_to_rgba(background_color)
|
|
outline_rgba = hex_to_rgba(outline_color)
|
|
|
|
output_images = []
|
|
|
|
# ----------------------------------------------------
|
|
# Process frames
|
|
# ----------------------------------------------------
|
|
|
|
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")
|
|
|
|
draw = ImageDraw.Draw(pil_img)
|
|
|
|
label = f"{frame_number + i:0{num_padding}d}"
|
|
|
|
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
|
|
|
|
x, y = compute_position(
|
|
position,
|
|
pil_img.width,
|
|
pil_img.height,
|
|
text_w,
|
|
text_h,
|
|
text_padding,
|
|
)
|
|
|
|
# Background box
|
|
if bg_rgba[3] > 0:
|
|
draw.rectangle(
|
|
[x, y, x + text_w, y + text_h],
|
|
fill=bg_rgba,
|
|
)
|
|
|
|
# Text
|
|
draw.text(
|
|
(x + text_padding, y + text_padding),
|
|
label,
|
|
fill=font_rgba,
|
|
font=font,
|
|
stroke_width=outline_width,
|
|
stroke_fill=outline_rgba,
|
|
)
|
|
|
|
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,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"OverlayFrameNumber": OverlayFrameNumber
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"OverlayFrameNumber": "Overlay Frame Number"
|
|
}
|