Initial commit: Complete Niutonian Smart Image Suite with 7 nodes
This commit is contained in:
Niutonian
@@ -0,0 +1,351 @@
|
||||
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"
|
||||
}
|
||||
Reference in New Issue
Block a user