351 lines
14 KiB
Python
351 lines
14 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFilter
|
|
import math
|
|
|
|
class SmartBorderFrame:
|
|
"""
|
|
A ComfyUI node for adding intelligent borders and frames to images.
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"input_image": ("IMAGE",),
|
|
"frame_style": (["none", "simple", "classic", "modern", "vintage", "polaroid", "film"], {"default": "simple"}),
|
|
"border_width": ("INT", {"default": 20, "min": 0, "max": 200, "step": 2}),
|
|
"border_color": (["white", "black", "gray", "custom"], {"default": "white"}),
|
|
"corner_style": (["square", "rounded", "beveled"], {"default": "square"}),
|
|
"shadow_enabled": ("BOOLEAN", {"default": False}),
|
|
"shadow_blur": ("INT", {"default": 10, "min": 0, "max": 50, "step": 1}),
|
|
"shadow_offset": ("INT", {"default": 5, "min": 0, "max": 30, "step": 1}),
|
|
"shadow_opacity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1}),
|
|
},
|
|
"optional": {
|
|
"custom_color_r": ("INT", {"default": 128, "min": 0, "max": 255}),
|
|
"custom_color_g": ("INT", {"default": 128, "min": 0, "max": 255}),
|
|
"custom_color_b": ("INT", {"default": 128, "min": 0, "max": 255}),
|
|
"corner_radius": ("INT", {"default": 10, "min": 0, "max": 50, "step": 1}),
|
|
"inner_border_width": ("INT", {"default": 5, "min": 0, "max": 50, "step": 1}),
|
|
"inner_border_color": (["white", "black", "gold", "silver"], {"default": "gold"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("framed_image",)
|
|
FUNCTION = "add_frame"
|
|
CATEGORY = "Niutonian/Image Processing"
|
|
|
|
def tensor_to_pil(self, tensor):
|
|
"""Convert tensor to PIL Image"""
|
|
if tensor is None:
|
|
return None
|
|
|
|
if len(tensor.shape) == 4:
|
|
tensor = tensor[0]
|
|
|
|
np_image = tensor.cpu().numpy()
|
|
|
|
if np_image.dtype == np.float32 or np_image.dtype == np.float64:
|
|
np_image = (np_image * 255).astype(np.uint8)
|
|
|
|
if len(np_image.shape) == 3:
|
|
return Image.fromarray(np_image, 'RGB')
|
|
else:
|
|
return Image.fromarray(np_image, 'L')
|
|
|
|
def pil_to_tensor(self, pil_image):
|
|
"""Convert PIL Image to tensor"""
|
|
if pil_image.mode != 'RGB':
|
|
pil_image = pil_image.convert('RGB')
|
|
|
|
np_image = np.array(pil_image).astype(np.float32) / 255.0
|
|
tensor = torch.from_numpy(np_image).unsqueeze(0)
|
|
return tensor
|
|
|
|
def get_border_color(self, color_name, custom_r=128, custom_g=128, custom_b=128):
|
|
"""Get border color from name or custom values"""
|
|
colors = {
|
|
"white": (255, 255, 255),
|
|
"black": (0, 0, 0),
|
|
"gray": (128, 128, 128),
|
|
"gold": (255, 215, 0),
|
|
"silver": (192, 192, 192),
|
|
}
|
|
|
|
if color_name == "custom":
|
|
return (custom_r, custom_g, custom_b)
|
|
|
|
return colors.get(color_name, (255, 255, 255))
|
|
|
|
def create_shadow(self, image, blur, offset, opacity):
|
|
"""Create drop shadow for image"""
|
|
# Create shadow image
|
|
shadow = Image.new('RGBA',
|
|
(image.width + offset * 2, image.height + offset * 2),
|
|
(0, 0, 0, 0))
|
|
|
|
# Create shadow shape
|
|
shadow_color = (0, 0, 0, int(255 * opacity))
|
|
shadow_img = Image.new('RGBA', image.size, shadow_color)
|
|
|
|
# Paste shadow at offset position
|
|
shadow.paste(shadow_img, (offset, offset))
|
|
|
|
# Apply blur
|
|
if blur > 0:
|
|
shadow = shadow.filter(ImageFilter.GaussianBlur(radius=blur))
|
|
|
|
return shadow
|
|
|
|
def add_rounded_corners(self, image, radius):
|
|
"""Add rounded corners to image"""
|
|
# Create mask for rounded corners
|
|
mask = Image.new('L', image.size, 0)
|
|
draw = ImageDraw.Draw(mask)
|
|
|
|
# Draw rounded rectangle
|
|
draw.rounded_rectangle([0, 0, image.width, image.height], radius, fill=255)
|
|
|
|
# Apply mask
|
|
if image.mode != 'RGBA':
|
|
image = image.convert('RGBA')
|
|
|
|
# Create output with transparency
|
|
output = Image.new('RGBA', image.size, (0, 0, 0, 0))
|
|
output.paste(image, mask=mask)
|
|
|
|
return output
|
|
|
|
def add_beveled_corners(self, image, bevel_size=10):
|
|
"""Add beveled corners to image"""
|
|
if image.mode != 'RGBA':
|
|
image = image.convert('RGBA')
|
|
|
|
# Create mask for beveled corners
|
|
mask = Image.new('L', image.size, 255)
|
|
draw = ImageDraw.Draw(mask)
|
|
|
|
# Draw beveled corners (cut off corners)
|
|
w, h = image.size
|
|
|
|
# Top-left corner
|
|
draw.polygon([(0, 0), (bevel_size, 0), (0, bevel_size)], fill=0)
|
|
# Top-right corner
|
|
draw.polygon([(w, 0), (w - bevel_size, 0), (w, bevel_size)], fill=0)
|
|
# Bottom-left corner
|
|
draw.polygon([(0, h), (bevel_size, h), (0, h - bevel_size)], fill=0)
|
|
# Bottom-right corner
|
|
draw.polygon([(w, h), (w - bevel_size, h), (w, h - bevel_size)], fill=0)
|
|
|
|
# Apply mask
|
|
output = Image.new('RGBA', image.size, (0, 0, 0, 0))
|
|
output.paste(image, mask=mask)
|
|
|
|
return output
|
|
|
|
def create_simple_frame(self, image, border_width, border_color):
|
|
"""Create simple border frame"""
|
|
if border_width == 0:
|
|
return image
|
|
|
|
# Create new image with border
|
|
new_width = image.width + 2 * border_width
|
|
new_height = image.height + 2 * border_width
|
|
|
|
framed = Image.new('RGB', (new_width, new_height), border_color)
|
|
framed.paste(image, (border_width, border_width))
|
|
|
|
return framed
|
|
|
|
def create_classic_frame(self, image, border_width, border_color, inner_width, inner_color):
|
|
"""Create classic double-border frame"""
|
|
if border_width == 0:
|
|
return image
|
|
|
|
# Create outer border
|
|
outer_width = image.width + 2 * border_width
|
|
outer_height = image.height + 2 * border_width
|
|
|
|
framed = Image.new('RGB', (outer_width, outer_height), border_color)
|
|
|
|
# Add inner border if specified
|
|
if inner_width > 0 and inner_width < border_width:
|
|
inner_x = border_width - inner_width
|
|
inner_y = border_width - inner_width
|
|
inner_w = image.width + 2 * inner_width
|
|
inner_h = image.height + 2 * inner_width
|
|
|
|
# Draw inner border
|
|
draw = ImageDraw.Draw(framed)
|
|
draw.rectangle([inner_x, inner_y, inner_x + inner_w, inner_y + inner_h],
|
|
fill=inner_color)
|
|
|
|
# Paste original image
|
|
framed.paste(image, (border_width, border_width))
|
|
|
|
return framed
|
|
|
|
def create_modern_frame(self, image, border_width, border_color):
|
|
"""Create modern minimalist frame"""
|
|
if border_width == 0:
|
|
return image
|
|
|
|
# Create asymmetric modern border (thicker on bottom)
|
|
top_border = border_width // 2
|
|
side_border = border_width // 2
|
|
bottom_border = border_width * 2
|
|
|
|
new_width = image.width + 2 * side_border
|
|
new_height = image.height + top_border + bottom_border
|
|
|
|
framed = Image.new('RGB', (new_width, new_height), border_color)
|
|
framed.paste(image, (side_border, top_border))
|
|
|
|
return framed
|
|
|
|
def create_vintage_frame(self, image, border_width, border_color):
|
|
"""Create vintage-style frame with texture"""
|
|
if border_width == 0:
|
|
return image
|
|
|
|
# Create base frame
|
|
framed = self.create_classic_frame(image, border_width, border_color,
|
|
border_width // 4, (139, 69, 19)) # Brown inner
|
|
|
|
# Add vintage texture effect
|
|
draw = ImageDraw.Draw(framed)
|
|
|
|
# Add some decorative elements in corners
|
|
corner_size = border_width // 2
|
|
|
|
# Simple corner decorations
|
|
for x, y in [(corner_size//2, corner_size//2),
|
|
(framed.width - corner_size//2, corner_size//2),
|
|
(corner_size//2, framed.height - corner_size//2),
|
|
(framed.width - corner_size//2, framed.height - corner_size//2)]:
|
|
draw.ellipse([x-3, y-3, x+3, y+3], fill=(255, 215, 0)) # Gold dots
|
|
|
|
return framed
|
|
|
|
def create_polaroid_frame(self, image, border_width):
|
|
"""Create Polaroid-style frame"""
|
|
# Polaroid has white border with extra space at bottom
|
|
top_border = border_width
|
|
side_border = border_width
|
|
bottom_border = border_width * 3 # Extra space for "photo" effect
|
|
|
|
new_width = image.width + 2 * side_border
|
|
new_height = image.height + top_border + bottom_border
|
|
|
|
framed = Image.new('RGB', (new_width, new_height), (255, 255, 255))
|
|
framed.paste(image, (side_border, top_border))
|
|
|
|
return framed
|
|
|
|
def create_film_frame(self, image, border_width):
|
|
"""Create film strip style frame"""
|
|
# Film frame with perforations
|
|
new_width = image.width + 2 * border_width
|
|
new_height = image.height + 2 * border_width
|
|
|
|
framed = Image.new('RGB', (new_width, new_height), (0, 0, 0))
|
|
framed.paste(image, (border_width, border_width))
|
|
|
|
# Add film perforations
|
|
draw = ImageDraw.Draw(framed)
|
|
perf_size = border_width // 4
|
|
|
|
# Top and bottom perforations
|
|
for x in range(perf_size, new_width - perf_size, perf_size * 2):
|
|
# Top perforations
|
|
draw.rectangle([x, perf_size//2, x + perf_size, perf_size + perf_size//2],
|
|
fill=(64, 64, 64))
|
|
# Bottom perforations
|
|
draw.rectangle([x, new_height - perf_size - perf_size//2,
|
|
x + perf_size, new_height - perf_size//2],
|
|
fill=(64, 64, 64))
|
|
|
|
return framed
|
|
|
|
def add_frame(self, input_image, frame_style, border_width, border_color, corner_style,
|
|
shadow_enabled, shadow_blur, shadow_offset, shadow_opacity,
|
|
custom_color_r=128, custom_color_g=128, custom_color_b=128,
|
|
corner_radius=10, inner_border_width=5, inner_border_color="gold"):
|
|
"""Main function to add frame to image"""
|
|
|
|
# Convert input to PIL
|
|
pil_image = self.tensor_to_pil(input_image)
|
|
|
|
if pil_image is None:
|
|
default_img = Image.new('RGB', (512, 512), (128, 128, 128))
|
|
return (self.pil_to_tensor(default_img),)
|
|
|
|
# Get colors
|
|
main_color = self.get_border_color(border_color, custom_color_r, custom_color_g, custom_color_b)
|
|
inner_color = self.get_border_color(inner_border_color)
|
|
|
|
# Apply frame style
|
|
if frame_style == "none":
|
|
framed = pil_image
|
|
elif frame_style == "simple":
|
|
framed = self.create_simple_frame(pil_image, border_width, main_color)
|
|
elif frame_style == "classic":
|
|
framed = self.create_classic_frame(pil_image, border_width, main_color,
|
|
inner_border_width, inner_color)
|
|
elif frame_style == "modern":
|
|
framed = self.create_modern_frame(pil_image, border_width, main_color)
|
|
elif frame_style == "vintage":
|
|
framed = self.create_vintage_frame(pil_image, border_width, main_color)
|
|
elif frame_style == "polaroid":
|
|
framed = self.create_polaroid_frame(pil_image, border_width)
|
|
elif frame_style == "film":
|
|
framed = self.create_film_frame(pil_image, border_width)
|
|
else:
|
|
framed = self.create_simple_frame(pil_image, border_width, main_color)
|
|
|
|
# Apply corner style
|
|
if corner_style == "rounded":
|
|
framed = self.add_rounded_corners(framed, corner_radius)
|
|
elif corner_style == "beveled":
|
|
framed = self.add_beveled_corners(framed, corner_radius)
|
|
|
|
# Add shadow if enabled
|
|
if shadow_enabled:
|
|
shadow = self.create_shadow(framed, shadow_blur, shadow_offset, shadow_opacity)
|
|
|
|
# Create final image with shadow
|
|
final_width = shadow.width
|
|
final_height = shadow.height
|
|
final_image = Image.new('RGBA', (final_width, final_height), (255, 255, 255, 0))
|
|
|
|
# Paste shadow first
|
|
final_image.paste(shadow, (0, 0), shadow)
|
|
|
|
# Paste framed image on top
|
|
if framed.mode != 'RGBA':
|
|
framed = framed.convert('RGBA')
|
|
final_image.paste(framed, (0, 0), framed)
|
|
|
|
# Convert back to RGB
|
|
if final_image.mode == 'RGBA':
|
|
rgb_final = Image.new('RGB', final_image.size, (255, 255, 255))
|
|
rgb_final.paste(final_image, mask=final_image.split()[-1])
|
|
framed = rgb_final
|
|
else:
|
|
framed = final_image
|
|
|
|
return (self.pil_to_tensor(framed),)
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SmartBorderFrame": SmartBorderFrame
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SmartBorderFrame": "Smart Border & Frame"
|
|
} |