399 lines
15 KiB
Python
399 lines
15 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw
|
|
import random
|
|
import math
|
|
|
|
class SmartCollageMaker:
|
|
"""
|
|
A ComfyUI node for creating artistic collages with varied image sizes and layouts.
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"layout_style": (["magazine", "photo_wall", "mosaic", "scattered", "spiral", "custom"], {"default": "magazine"}),
|
|
"canvas_width": ("INT", {"default": 1024, "min": 512, "max": 4096, "step": 64}),
|
|
"canvas_height": ("INT", {"default": 1024, "min": 512, "max": 4096, "step": 64}),
|
|
"size_variation": (["none", "small", "medium", "large", "random"], {"default": "medium"}),
|
|
"rotation_enabled": ("BOOLEAN", {"default": False}),
|
|
"max_rotation": ("INT", {"default": 15, "min": 0, "max": 45, "step": 5}),
|
|
"overlap_allowed": ("BOOLEAN", {"default": False}),
|
|
"background_color": (["white", "black", "transparent", "random"], {"default": "white"}),
|
|
"border_style": (["none", "thin", "thick", "shadow"], {"default": "none"}),
|
|
"spacing": ("INT", {"default": 10, "min": 0, "max": 50, "step": 2}),
|
|
},
|
|
"optional": {
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 999999}),
|
|
"image_1": ("IMAGE",),
|
|
"image_2": ("IMAGE",),
|
|
"image_3": ("IMAGE",),
|
|
"image_4": ("IMAGE",),
|
|
"image_5": ("IMAGE",),
|
|
"image_6": ("IMAGE",),
|
|
"image_7": ("IMAGE",),
|
|
"image_8": ("IMAGE",),
|
|
"image_9": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("collage_image",)
|
|
FUNCTION = "create_collage"
|
|
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_valid_images(self, **kwargs):
|
|
"""Extract valid images from inputs"""
|
|
valid_images = []
|
|
|
|
for i in range(1, 10): # image_1 through image_9
|
|
image_key = f"image_{i}"
|
|
if image_key in kwargs and kwargs[image_key] is not None:
|
|
pil_img = self.tensor_to_pil(kwargs[image_key])
|
|
if pil_img is not None:
|
|
valid_images.append(pil_img)
|
|
|
|
return valid_images
|
|
|
|
def calculate_size_variations(self, images, canvas_size, variation_level):
|
|
"""Calculate varied sizes for images based on variation level"""
|
|
canvas_width, canvas_height = canvas_size
|
|
base_area = (canvas_width * canvas_height) / len(images)
|
|
base_size = int(math.sqrt(base_area))
|
|
|
|
sizes = []
|
|
|
|
if variation_level == "none":
|
|
# All same size
|
|
for _ in images:
|
|
sizes.append((base_size, base_size))
|
|
|
|
elif variation_level == "small":
|
|
# 20% variation
|
|
for i, img in enumerate(images):
|
|
factor = 0.9 + (i % 3) * 0.1 # 0.9, 1.0, 1.1
|
|
new_size = int(base_size * factor)
|
|
sizes.append((new_size, new_size))
|
|
|
|
elif variation_level == "medium":
|
|
# 50% variation
|
|
for i, img in enumerate(images):
|
|
factor = 0.7 + (i % 4) * 0.2 # 0.7, 0.9, 1.1, 1.3
|
|
new_size = int(base_size * factor)
|
|
sizes.append((new_size, new_size))
|
|
|
|
elif variation_level == "large":
|
|
# 100% variation
|
|
for i, img in enumerate(images):
|
|
factor = 0.5 + (i % 5) * 0.3 # 0.5, 0.8, 1.1, 1.4, 1.7
|
|
new_size = int(base_size * factor)
|
|
sizes.append((new_size, new_size))
|
|
|
|
elif variation_level == "random":
|
|
# Random sizes
|
|
for _ in images:
|
|
factor = random.uniform(0.4, 1.8)
|
|
new_size = int(base_size * factor)
|
|
sizes.append((new_size, new_size))
|
|
|
|
return sizes
|
|
|
|
def resize_image_proportional(self, img, target_size):
|
|
"""Resize image maintaining aspect ratio"""
|
|
target_width, target_height = target_size
|
|
|
|
# Calculate scale to fit within target size
|
|
scale_w = target_width / img.width
|
|
scale_h = target_height / img.height
|
|
scale = min(scale_w, scale_h)
|
|
|
|
new_width = int(img.width * scale)
|
|
new_height = int(img.height * scale)
|
|
|
|
return img.resize((new_width, new_height), Image.LANCZOS)
|
|
|
|
def rotate_image(self, img, angle):
|
|
"""Rotate image with transparent background"""
|
|
if angle == 0:
|
|
return img
|
|
|
|
# Convert to RGBA for rotation
|
|
if img.mode != 'RGBA':
|
|
img = img.convert('RGBA')
|
|
|
|
# Rotate with transparent background
|
|
rotated = img.rotate(angle, expand=True, fillcolor=(255, 255, 255, 0))
|
|
return rotated
|
|
|
|
def add_border(self, img, border_style, border_color=(0, 0, 0)):
|
|
"""Add border to image"""
|
|
if border_style == "none":
|
|
return img
|
|
|
|
if border_style == "thin":
|
|
border_width = 2
|
|
elif border_style == "thick":
|
|
border_width = 5
|
|
elif border_style == "shadow":
|
|
# Add drop shadow effect
|
|
shadow_offset = 3
|
|
shadow_color = (128, 128, 128, 180)
|
|
|
|
# Create shadow
|
|
shadow = Image.new('RGBA',
|
|
(img.width + shadow_offset, img.height + shadow_offset),
|
|
shadow_color)
|
|
|
|
# Create result image
|
|
result = Image.new('RGBA', shadow.size, (255, 255, 255, 0))
|
|
result.paste(shadow, (shadow_offset, shadow_offset))
|
|
result.paste(img, (0, 0), img if img.mode == 'RGBA' else None)
|
|
|
|
return result
|
|
else:
|
|
border_width = 2
|
|
|
|
# Add simple border
|
|
bordered = Image.new('RGB',
|
|
(img.width + 2 * border_width, img.height + 2 * border_width),
|
|
border_color)
|
|
bordered.paste(img, (border_width, border_width))
|
|
|
|
return bordered
|
|
|
|
def check_overlap(self, pos1, size1, pos2, size2):
|
|
"""Check if two rectangles overlap"""
|
|
x1, y1 = pos1
|
|
w1, h1 = size1
|
|
x2, y2 = pos2
|
|
w2, h2 = size2
|
|
|
|
return not (x1 + w1 <= x2 or x2 + w2 <= x1 or y1 + h1 <= y2 or y2 + h2 <= y1)
|
|
|
|
def find_non_overlapping_position(self, canvas_size, img_size, placed_items, max_attempts=100):
|
|
"""Find a position that doesn't overlap with existing items"""
|
|
canvas_width, canvas_height = canvas_size
|
|
img_width, img_height = img_size
|
|
|
|
for _ in range(max_attempts):
|
|
x = random.randint(0, max(0, canvas_width - img_width))
|
|
y = random.randint(0, max(0, canvas_height - img_height))
|
|
|
|
# Check for overlaps
|
|
overlap_found = False
|
|
for placed_pos, placed_size in placed_items:
|
|
if self.check_overlap((x, y), img_size, placed_pos, placed_size):
|
|
overlap_found = True
|
|
break
|
|
|
|
if not overlap_found:
|
|
return (x, y)
|
|
|
|
# If no non-overlapping position found, return random position
|
|
x = random.randint(0, max(0, canvas_width - img_width))
|
|
y = random.randint(0, max(0, canvas_height - img_height))
|
|
return (x, y)
|
|
|
|
def create_magazine_layout(self, images, sizes, canvas_size, spacing):
|
|
"""Create magazine-style layout"""
|
|
positions = []
|
|
canvas_width, canvas_height = canvas_size
|
|
|
|
# Simple grid-based magazine layout
|
|
cols = int(math.sqrt(len(images))) + 1
|
|
rows = math.ceil(len(images) / cols)
|
|
|
|
cell_width = canvas_width // cols
|
|
cell_height = canvas_height // rows
|
|
|
|
for i, (img, size) in enumerate(zip(images, sizes)):
|
|
row = i // cols
|
|
col = i % cols
|
|
|
|
# Center in cell
|
|
cell_x = col * cell_width
|
|
cell_y = row * cell_height
|
|
|
|
img_width, img_height = size
|
|
x = cell_x + (cell_width - img_width) // 2
|
|
y = cell_y + (cell_height - img_height) // 2
|
|
|
|
# Ensure within bounds
|
|
x = max(0, min(x, canvas_width - img_width))
|
|
y = max(0, min(y, canvas_height - img_height))
|
|
|
|
positions.append((x, y))
|
|
|
|
return positions
|
|
|
|
def create_scattered_layout(self, images, sizes, canvas_size, overlap_allowed):
|
|
"""Create scattered layout"""
|
|
positions = []
|
|
placed_items = []
|
|
|
|
for img, size in zip(images, sizes):
|
|
if overlap_allowed:
|
|
# Random position
|
|
x = random.randint(0, max(0, canvas_size[0] - size[0]))
|
|
y = random.randint(0, max(0, canvas_size[1] - size[1]))
|
|
pos = (x, y)
|
|
else:
|
|
# Find non-overlapping position
|
|
pos = self.find_non_overlapping_position(canvas_size, size, placed_items)
|
|
|
|
positions.append(pos)
|
|
placed_items.append((pos, size))
|
|
|
|
return positions
|
|
|
|
def create_spiral_layout(self, images, sizes, canvas_size):
|
|
"""Create spiral layout"""
|
|
positions = []
|
|
center_x, center_y = canvas_size[0] // 2, canvas_size[1] // 2
|
|
|
|
angle = 0
|
|
radius = 50
|
|
|
|
for i, (img, size) in enumerate(zip(images, sizes)):
|
|
# Calculate position on spiral
|
|
x = center_x + int(radius * math.cos(angle)) - size[0] // 2
|
|
y = center_y + int(radius * math.sin(angle)) - size[1] // 2
|
|
|
|
# Ensure within bounds
|
|
x = max(0, min(x, canvas_size[0] - size[0]))
|
|
y = max(0, min(y, canvas_size[1] - size[1]))
|
|
|
|
positions.append((x, y))
|
|
|
|
# Update for next position
|
|
angle += math.pi / 3 # 60 degrees
|
|
radius += 30
|
|
|
|
return positions
|
|
|
|
def get_background_color(self, bg_type):
|
|
"""Get background color"""
|
|
if bg_type == "white":
|
|
return (255, 255, 255)
|
|
elif bg_type == "black":
|
|
return (0, 0, 0)
|
|
elif bg_type == "random":
|
|
return (random.randint(200, 255), random.randint(200, 255), random.randint(200, 255))
|
|
else: # transparent
|
|
return (255, 255, 255, 0)
|
|
|
|
def create_collage(self, layout_style, canvas_width, canvas_height, size_variation,
|
|
rotation_enabled, max_rotation, overlap_allowed, background_color,
|
|
border_style, spacing, seed=42, **kwargs):
|
|
"""Main function to create collage"""
|
|
|
|
# Set random seed for reproducible results
|
|
random.seed(seed)
|
|
|
|
# Get valid images
|
|
valid_images = self.get_valid_images(**kwargs)
|
|
|
|
if not valid_images:
|
|
default_img = Image.new('RGB', (canvas_width, canvas_height), (128, 128, 128))
|
|
return (self.pil_to_tensor(default_img),)
|
|
|
|
canvas_size = (canvas_width, canvas_height)
|
|
|
|
# Calculate size variations
|
|
target_sizes = self.calculate_size_variations(valid_images, canvas_size, size_variation)
|
|
|
|
# Resize images
|
|
resized_images = []
|
|
for img, target_size in zip(valid_images, target_sizes):
|
|
resized_img = self.resize_image_proportional(img, target_size)
|
|
|
|
# Add rotation if enabled
|
|
if rotation_enabled:
|
|
angle = random.randint(-max_rotation, max_rotation)
|
|
resized_img = self.rotate_image(resized_img, angle)
|
|
|
|
# Add border
|
|
resized_img = self.add_border(resized_img, border_style)
|
|
|
|
resized_images.append(resized_img)
|
|
|
|
# Get actual sizes after processing
|
|
actual_sizes = [(img.width, img.height) for img in resized_images]
|
|
|
|
# Calculate positions based on layout style
|
|
if layout_style == "magazine":
|
|
positions = self.create_magazine_layout(resized_images, actual_sizes, canvas_size, spacing)
|
|
elif layout_style == "scattered":
|
|
positions = self.create_scattered_layout(resized_images, actual_sizes, canvas_size, overlap_allowed)
|
|
elif layout_style == "spiral":
|
|
positions = self.create_spiral_layout(resized_images, actual_sizes, canvas_size)
|
|
else: # photo_wall, mosaic, custom - use scattered for now
|
|
positions = self.create_scattered_layout(resized_images, actual_sizes, canvas_size, overlap_allowed)
|
|
|
|
# Create canvas
|
|
bg_color = self.get_background_color(background_color)
|
|
if background_color == "transparent":
|
|
canvas = Image.new('RGBA', canvas_size, bg_color)
|
|
else:
|
|
canvas = Image.new('RGB', canvas_size, bg_color)
|
|
|
|
# Place images on canvas
|
|
for img, pos in zip(resized_images, positions):
|
|
x, y = pos
|
|
|
|
# Convert modes if necessary
|
|
if canvas.mode == 'RGBA' and img.mode != 'RGBA':
|
|
img = img.convert('RGBA')
|
|
elif canvas.mode == 'RGB' and img.mode == 'RGBA':
|
|
# Create white background for RGBA image
|
|
white_bg = Image.new('RGB', img.size, (255, 255, 255))
|
|
white_bg.paste(img, mask=img.split()[-1] if img.mode == 'RGBA' else None)
|
|
img = white_bg
|
|
|
|
# Paste image
|
|
if img.mode == 'RGBA' and canvas.mode == 'RGBA':
|
|
canvas.paste(img, (x, y), img)
|
|
else:
|
|
canvas.paste(img, (x, y))
|
|
|
|
return (self.pil_to_tensor(canvas),)
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SmartCollageMaker": SmartCollageMaker
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SmartCollageMaker": "Smart Collage Maker"
|
|
} |