TravelContentCreator/core/simple_collage.py

754 lines
31 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# -*- coding: utf-8 -*-
import os
import random
import traceback
import math
from pathlib import Path
from PIL import Image, ImageDraw, ImageEnhance, ImageFilter, ImageOps
class ImageCollageCreator:
def __init__(self):
"""初始化拼图创建器"""
self.output_dir = "/root/autodl-tmp/poster_baseboard_0403/output_collage"
# 创建输出目录
os.makedirs(self.output_dir, exist_ok=True)
# 定义可用拼接样式
self.collage_styles = [
"grid_2x2", # 标准2x2网格
"asymmetric", # 非对称布局
"filmstrip", # 胶片条布局
# "circles", # 圆形布局
"overlap", # 重叠风格
"mosaic", # 马赛克风格 3x3
"fullscreen", # 全覆盖拼图样式
"vertical_stack" # 新增:上下拼图样式
# "polaroid", # 宝丽来风格
]
def resize_and_crop(self, img, target_size):
"""调整图片大小并居中裁剪为指定尺寸"""
width, height = img.size
target_width, target_height = target_size
# 计算宽高比
img_ratio = width / height
target_ratio = target_width / target_height
if img_ratio > target_ratio:
# 图片较宽,以高度为基准调整
new_height = target_height
new_width = int(width * target_height / height)
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
# 居中裁剪
left = (new_width - target_width) // 2
img = img.crop((left, 0, left + target_width, target_height))
else:
# 图片较高,以宽度为基准调整
new_width = target_width
new_height = int(height * target_width / width)
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
# 居中裁剪
top = (new_height - target_height) // 2
img = img.crop((0, top, target_width, top + target_height))
return img
def add_border(self, img, color=(255, 255, 255, 200), width=2, no_border=True):
"""给图像添加边框,可选择不添加边框"""
if no_border:
return img # 如果设置为无边框,直接返回原图
try:
w, h = img.size
new_img = Image.new('RGBA', (w, h), (0, 0, 0, 0))
draw = ImageDraw.Draw(new_img)
# 绘制边框(在四条边上)
for i in range(width):
# 上边框
draw.line([(i, i), (w-i-1, i)], fill=color, width=1)
# 右边框
draw.line([(w-i-1, i), (w-i-1, h-i-1)], fill=color, width=1)
# 下边框
draw.line([(i, h-i-1), (w-i-1, h-i-1)], fill=color, width=1)
# 左边框
draw.line([(i, i), (i, h-i-1)], fill=color, width=1)
# 合并原图和边框
result = img.copy()
result.alpha_composite(new_img)
return result
except Exception as e:
print(f"添加边框时出错: {str(e)}")
return img
def add_polaroid_frame(self, img, margin=20, bottom_margin=60, background_color=(255, 255, 255, 255)):
"""添加宝丽来风格的相框"""
try:
w, h = img.size
frame_width = w + 2 * margin
frame_height = h + margin + bottom_margin
# 创建白色背景
frame = Image.new('RGBA', (frame_width, frame_height), background_color)
# 将图像粘贴到框架中
frame.paste(img, (margin, margin))
# 添加稍微的阴影效果
shadow = Image.new('RGBA', frame.size, (0, 0, 0, 0))
shadow_draw = ImageDraw.Draw(shadow)
shadow_draw.rectangle([2, 2, frame_width-2, frame_height-2], fill=(0, 0, 0, 40))
# 模糊阴影
shadow = shadow.filter(ImageFilter.GaussianBlur(3))
# 创建最终图像
final = Image.new('RGBA', (frame_width+6, frame_height+6), (0, 0, 0, 0))
final.paste(shadow, (6, 6))
final.paste(frame, (0, 0), frame)
return final
except Exception as e:
print(f"添加宝丽来相框时出错: {str(e)}")
return img
def apply_image_effect(self, img, effect="none"):
"""应用各种图像效果 - 所有图片适度增强对比度和亮度"""
try:
# 适度增强对比度
contrast = ImageEnhance.Contrast(img)
enhanced = contrast.enhance(1.1) # 降低对比度系数从1.6降至1.3
# 轻微增强亮度
brightness = ImageEnhance.Brightness(enhanced)
enhanced = brightness.enhance(1.1) # 保持轻微增加亮度
# 轻微增强色彩饱和度
color = ImageEnhance.Color(enhanced)
enhanced = color.enhance(1.15) # 轻微降低饱和度从1.2降至1.15
return enhanced
except Exception as e:
print(f"增强图片效果时出错: {str(e)}")
return img
def create_collage_with_style(self, input_dir, style=None, target_size=None):
"""创建指定样式的拼接画布"""
try:
# 设置默认尺寸为3:4比例
if target_size is None:
target_size = (900, 1200) # 3:4比例
# 如果没有指定样式,随机选择一种
if style is None or style not in self.collage_styles:
style = random.choice(self.collage_styles)
print(f"使用拼接样式: {style}")
# 检查目录是否存在
if not os.path.exists(input_dir):
print(f"目录不存在: {input_dir}")
return None
# 支持的图片格式
image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
# 获取目录中的所有图片文件
all_images = [f for f in os.listdir(input_dir)
if f.lower().endswith(image_extensions) and os.path.isfile(os.path.join(input_dir, f))]
if len(all_images) < 4:
print(f"目录中图片不足四张: {input_dir}")
return None
# 根据不同样式,可能需要的图片数量不同
num_images = 4
if style == "mosaic":
num_images = 9
elif style == "filmstrip":
num_images = 5
elif style == "fullscreen":
num_images = 6 # 全覆盖样式使用6张图片
elif style == "vertical_stack":
num_images = 2 # 上下拼图样式只需要2张图片
# 确保有足够的图像
if len(all_images) < num_images:
print(f"样式'{style}'需要至少{num_images}张图片,但目录只有{len(all_images)}")
# 多次使用相同图片
if len(all_images) > 0:
all_images = all_images * (num_images // len(all_images) + 1)
# 随机选择指定数量的图片
selected_images = random.sample(all_images, num_images)
print(f"随机选择的图片: {selected_images}")
# 创建空白画布
collage_image = Image.new('RGBA', target_size, (0, 0, 0, 0))
# 加载图片
images = []
for img_name in selected_images:
img_path = os.path.join(input_dir, img_name)
try:
img = Image.open(img_path).convert('RGBA')
images.append(img)
print(f"已加载图片: {img_name}")
except Exception as e:
print(f"加载图片 {img_name} 时出错: {e}")
# 如果某张图片加载失败,随机选择另一张图片代替
remaining_images = [f for f in all_images if f not in selected_images]
if remaining_images:
replacement = random.choice(remaining_images)
selected_images.append(replacement)
try:
replacement_path = os.path.join(input_dir, replacement)
replacement_img = Image.open(replacement_path).convert('RGBA')
images.append(replacement_img)
print(f"使用替代图片: {replacement}")
except:
print(f"替代图片 {replacement} 也加载失败")
# 确保图片数量足够
while len(images) < num_images:
if images:
images.append(random.choice(images).copy())
else:
print("没有可用的图片来创建拼贴画")
return None
# 应用所选样式
if style == "grid_2x2":
return self._create_grid_2x2_collage(images, target_size)
elif style == "asymmetric":
return self._create_asymmetric_collage(images, target_size)
elif style == "filmstrip":
return self._create_filmstrip_collage(images, target_size)
elif style == "circles":
return self._create_circles_collage(images, target_size)
elif style == "polaroid":
return self._create_polaroid_collage(images, target_size)
elif style == "overlap":
return self._create_overlap_collage(images, target_size)
elif style == "mosaic":
return self._create_mosaic_collage(images, target_size)
elif style == "fullscreen":
return self._create_fullscreen_collage(images, target_size)
elif style == "vertical_stack":
return self._create_vertical_stack_collage(images, target_size)
else:
# 默认使用2x2网格
return self._create_grid_2x2_collage(images, target_size)
except Exception as e:
print(f"创建拼贴画时出错: {str(e)}")
traceback.print_exc()
return None
def _create_grid_2x2_collage(self, images, target_size):
"""创建2x2网格拼贴画"""
collage = Image.new('RGBA', target_size, (255, 255, 255, 255)) # 使用白色背景
# 计算每个块的大小
block_width = target_size[0] // 2
block_height = target_size[1] // 2
# 定义四个区域位置
positions = [
(0, 0), # 左上
(block_width, 0), # 右上
(0, block_height), # 左下
(block_width, block_height) # 右下
]
# 将图像粘贴到拼贴画位置
for i, position in enumerate(positions):
if i < len(images):
img = images[i].copy()
# 调整大小
img = self.resize_and_crop(img, (block_width, block_height))
# 不应用边框,实现无缝拼接
# 粘贴到拼贴画
collage.paste(img, position, img)
print(f"添加拼贴画块 {i+1} 到位置: {position}")
print(f"无缝2x2网格拼贴画创建成功尺寸: {target_size}")
return collage
def _create_asymmetric_collage(self, images, target_size):
"""创建非对称布局拼贴画"""
collage = Image.new('RGBA', target_size, (255, 255, 255, 255))
width, height = target_size
# 定义非对称区域位置
positions = [
(0, 0, width*2//3, height//2), # 左上 (大)
(width*2//3, 0, width, height//3), # 右上
(width*2//3, height//3, width, height//2), # 右中
(0, height//2, width, height) # 底部 (全宽)
]
# 定义不同的效果
effects = ["none", "grayscale", "vintage", "color_boost"]
random.shuffle(effects)
# 将图像粘贴到拼贴画位置
for i, (x1, y1, x2, y2) in enumerate(positions):
if i < len(images):
img = images[i].copy()
# 调整大小
img = self.resize_and_crop(img, (x2-x1, y2-y1))
# 应用效果
img = self.apply_image_effect(img, effects[i % len(effects)])
# 不添加边框
# 粘贴到拼贴画
collage.paste(img, (x1, y1), img)
print(f"添加非对称拼贴画块 {i+1} 到位置: ({x1},{y1},{x2},{y2})")
print(f"无缝非对称拼贴画创建成功,尺寸: {target_size}")
return collage
def _create_filmstrip_collage(self, images, target_size):
"""创建胶片条布局拼贴画"""
collage = Image.new('RGBA', target_size, (0, 0, 0, 0))
width, height = target_size
# 胶片条中每个图像的高度
strip_height = height // 5
# 添加黑条边框
film_border_width = 15
# 将图像粘贴为胶片条
for i in range(5):
if i < len(images):
img = images[i].copy()
# 调整大小,考虑边框
img = self.resize_and_crop(img, (width - 2*film_border_width, strip_height - 2*film_border_width))
# 不应用效果,保持原始颜色
# 创建黑色胶片边框
film_frame = Image.new('RGBA', (width, strip_height), (0, 0, 0, 255))
# 在边框中间贴上图片
film_frame.paste(img, (film_border_width, film_border_width), img)
# 添加胶片冲孔
draw = ImageDraw.Draw(film_frame)
hole_spacing = 30
hole_radius = 5
num_holes = width // hole_spacing
for h in range(num_holes):
hole_center_x = h * hole_spacing + hole_spacing // 2
# 顶部和底部的冲孔
draw.ellipse((hole_center_x - hole_radius, 3, hole_center_x + hole_radius, 13), fill=(50, 50, 50, 255))
draw.ellipse((hole_center_x - hole_radius, strip_height - 13, hole_center_x + hole_radius, strip_height - 3), fill=(50, 50, 50, 255))
# 粘贴到拼贴画
y_position = i * strip_height
collage.paste(film_frame, (0, y_position), film_frame)
print(f"添加胶片条拼贴画块 {i+1} 到位置 y={y_position}")
print(f"胶片条拼贴画创建成功,尺寸: {target_size}")
return collage
def _create_circles_collage(self, images, target_size):
"""创建圆形布局拼贴画"""
collage = Image.new('RGBA', target_size, (0, 0, 0, 0))
width, height = target_size
# 定义圆形的位置和大小
circle_positions = [
(width//4, height//4, width//2.5), # 左上
(width*3//4, height//4, width//3), # 右上
(width//4, height*3//4, width//3), # 左下
(width*3//4, height*3//4, width//2.5) # 右下
]
# 为每个圆形创建蒙版
for i, (center_x, center_y, radius) in enumerate(circle_positions):
if i < len(images):
img = images[i].copy()
# 应用效果
img = self.apply_image_effect(img)
# 调整图像大小为圆的直径 - 确保是整数
diam = int(radius*2)
img = self.resize_and_crop(img, (diam, diam))
# 创建圆形蒙版
mask = Image.new('L', img.size, 0)
draw = ImageDraw.Draw(mask)
draw.ellipse((0, 0, mask.width, mask.height), fill=255)
# 模糊边缘
mask = mask.filter(ImageFilter.GaussianBlur(radius=5))
# 应用蒙版
img.putalpha(mask)
# 计算粘贴位置,使圆心在定义的位置
paste_x = int(center_x - radius)
paste_y = int(center_y - radius)
# 粘贴到拼贴画
collage.paste(img, (paste_x, paste_y), img)
print(f"添加圆形拼贴画块 {i+1} 到位置: ({paste_x},{paste_y})")
# 添加轻微的渐变背景
background = Image.new('RGBA', target_size, (245, 245, 245, 100))
collage = Image.alpha_composite(background, collage)
print(f"圆形拼贴画创建成功,尺寸: {target_size}")
return collage
def _create_polaroid_collage(self, images, target_size):
"""创建宝丽来风格拼贴画 - 最小化图片重叠"""
collage = Image.new('RGBA', target_size, (240, 240, 240, 255))
width, height = target_size
# 宝丽来照片的大小 - 适当调整尺寸,减少重叠
polaroid_sizes = [
(int(width//2.2), int(height//2.8)), # 大号
(int(width//2.5), int(height//3)), # 中大号
(int(width//2.8), int(height//3.5)), # 中号
(int(width//3), int(height//4)) # 中小号
]
# 随机打乱尺寸
random.shuffle(polaroid_sizes)
# 创建网格布局,降低重叠概率
grid_cells = [
(0, 0, width//2, height//2), # 左上
(width//2, 0, width, height//2), # 右上
(0, height//2, width//2, height), # 左下
(width//2, height//2, width, height) # 右下
]
# 随机打乱网格单元
random.shuffle(grid_cells)
# 用于记录已放置的区域
placed_areas = []
for i, img_size in enumerate(polaroid_sizes):
if i < len(images) and i < len(grid_cells):
img = images[i].copy()
# 调整大小
img = self.resize_and_crop(img, img_size)
# 应用效果
img = self.apply_image_effect(img)
# 添加宝丽来相框
img = self.add_polaroid_frame(img)
# 轻微旋转(-3到3度之间进一步减小旋转角度
rotation = random.uniform(-3, 3)
img = img.rotate(rotation, expand=True, resample=Image.Resampling.BICUBIC)
# 从当前网格单元获取可用区域
cell = grid_cells[i]
cell_x1, cell_y1, cell_x2, cell_y2 = cell
# 确保照片至少有80%在当前网格单元内
cell_width = cell_x2 - cell_x1
cell_height = cell_y2 - cell_y1
# 计算可用的粘贴位置范围
min_x = max(10, cell_x1 - img.width * 0.2) # 允许20%超出左边
max_x = min(width - img.width - 10, cell_x2 - img.width * 0.8) # 确保至少80%在单元内
min_y = max(10, cell_y1 - img.height * 0.2) # 允许20%超出上边
max_y = min(height - img.height - 10, cell_y2 - img.height * 0.8) # 确保至少80%在单元内
# 确保坐标范围有效,如果无效则使用单元中心
if min_x >= max_x:
center_x = (cell_x1 + cell_x2) // 2
min_x = max(10, center_x - img.width // 2)
max_x = min_x + 1
if min_y >= max_y:
center_y = (cell_y1 + cell_y2) // 2
min_y = max(10, center_y - img.height // 2)
max_y = min_y + 1
# 在可用范围内随机选择位置
paste_x = random.randint(int(min_x), int(max_x))
paste_y = random.randint(int(min_y), int(max_y))
# 记录这个位置
placed_areas.append((paste_x, paste_y, paste_x + img.width, paste_y + img.height))
# 粘贴到拼贴画
collage.paste(img, (paste_x, paste_y), img)
print(f"添加宝丽来风格块 {i+1} 到位置: ({paste_x},{paste_y}),尺寸: {img.size},单元: {cell}")
print(f"宝丽来风格拼贴画创建成功,尺寸: {target_size}")
return collage
def _create_overlap_collage(self, images, target_size):
"""创建无缝重叠风格拼贴画"""
collage = Image.new('RGBA', target_size, (255, 255, 255, 255))
width, height = target_size
# 计算每个图像确切填充一个区域的大小
img_width = width // 2
img_height = height // 2
# 网格位置 - 完全填充画布
positions = [
(0, 0), # 左上
(img_width, 0), # 右上
(0, img_height), # 左下
(img_width, img_height) # 右下
]
# 添加图片到位置
for i, position in enumerate(positions):
if i < len(images):
img = images[i].copy()
# 调整大小
img = self.resize_and_crop(img, (img_width, img_height))
# 应用效果
img = self.apply_image_effect(img)
# 不添加边框
# 粘贴到拼贴画
collage.paste(img, position)
print(f"添加无缝拼贴画块 {i+1} 到位置: {position}")
print(f"无缝拼贴画创建成功,尺寸: {target_size}")
return collage
def _create_mosaic_collage(self, images, target_size):
"""创建马赛克风格拼贴画需要9张图片"""
collage = Image.new('RGBA', target_size, (255, 255, 255, 255))
width, height = target_size
# 创建3x3网格
grid_width = width // 3
grid_height = height // 3
# 生成网格位置
positions = []
for row in range(3):
for col in range(3):
positions.append((col * grid_width, row * grid_height))
# 定义不同的效果
effects = ["none", "grayscale", "sepia", "vintage", "high_contrast",
"color_boost", "vibrant", "warm", "none"]
# 保证所有效果都能使用
if len(effects) > len(images):
effects = effects[:len(images)]
# 随机打乱效果
random.shuffle(effects)
# 将图像粘贴到马赛克位置
for i, position in enumerate(positions):
if i < len(images):
img = images[i].copy()
# 调整大小
img = self.resize_and_crop(img, (grid_width, grid_height))
# 应用效果
img = self.apply_image_effect(img, effects[i % len(effects)])
# 不添加边框
# 粘贴到拼贴画
collage.paste(img, position)
print(f"添加马赛克拼贴画块 {i+1} 到位置: {position}")
print(f"无缝马赛克拼贴画创建成功,尺寸: {target_size}")
return collage
def _create_fullscreen_collage(self, images, target_size):
"""创建全覆盖拼图样式(完全填满画布,无空白)"""
width, height = target_size
collage = Image.new('RGBA', target_size, (255, 255, 255, 255)) # 白色背景
# 确保至少有6张图片
while len(images) < 6:
if images:
images.append(random.choice(images).copy())
else:
return None
# 定义区域划分 - 按照完全填满画布设计
regions = [
# 左列 - 上中下三块
(0, 0, width//2, height//3), # 左上
(0, height//3, width//2, height*2//3), # 左中
(0, height*2//3, width//2, height), # 左下
# 右列 - 上中下三块
(width//2, 0, width, height//3), # 右上
(width//2, height//3, width, height*2//3), # 右中
(width//2, height*2//3, width, height) # 右下
]
# 添加图片到各个区域,确保完全覆盖
for i, (x1, y1, x2, y2) in enumerate(regions):
if i < len(images):
img = images[i].copy()
# 调整大小以完全填充区域
region_width = x2 - x1
region_height = y2 - y1
img = self.resize_and_crop(img, (region_width, region_height))
# 应用轻微的图像效果,确保视觉统一性
img = self.apply_image_effect(img)
# 不添加边框,实现无缝拼接
# 粘贴到画布上
collage.paste(img, (x1, y1))
print(f"添加全屏拼图块 {i+1} 到位置: ({x1}, {y1}, {x2}, {y2})")
# 不添加分隔线,保持无缝
print(f"无缝全覆盖拼图创建成功,尺寸: {target_size}")
return collage
def _create_vertical_stack_collage(self, images, target_size):
"""创建上下拼图样式(两张图片上下排列)"""
collage = Image.new('RGBA', target_size, (255, 255, 255, 255)) # 白色背景
width, height = target_size
# 确保至少有2张图片
while len(images) < 2:
if images:
images.append(images[0].copy())
else:
print("没有可用的图片来创建上下拼图")
return None
# 设置间隙(可选)
gap = 0 # 无间隙拼接设置为0
# 计算每张图片的高度
img_height = (height - gap) // 2
# 定义图片位置
positions = [
(0, 0), # 上方图片
(0, img_height + gap) # 下方图片
]
# 添加图片
for i, position in enumerate(positions):
if i < len(images) and i < 2: # 只使用前两张图片
img = images[i].copy()
# 调整大小以适应宽度
img = self.resize_and_crop(img, (width, img_height))
# 应用轻微的图像效果
img = self.apply_image_effect(img)
# 粘贴到画布上
collage.paste(img, position, img)
print(f"添加上下拼图块 {i+1} 到位置: {position}")
# 可选:添加分隔线
if gap > 0:
draw = ImageDraw.Draw(collage)
line_y = img_height + gap // 2
draw.line([(0, line_y), (width, line_y)], fill=(200, 200, 200, 255), width=gap)
print(f"上下拼图创建成功,尺寸: {target_size}")
return collage
def save_collage(self, collage, output_path):
"""保存拼贴画"""
if collage:
# 确保有背景 - 创建白色背景并将拼贴画合并上去
background = Image.new('RGB', collage.size, (255, 255, 255))
# 如果拼贴画有透明通道,将其合并到白色背景上
if collage.mode == 'RGBA':
background.paste(collage, (0, 0), collage)
final_image = background
else:
final_image = collage.convert('RGB')
final_image.save(output_path)
print(f"无缝拼贴画已保存: {output_path}")
return output_path
return None
def set_output_dir(self, output_dir):
"""设置输出目录"""
self.output_dir = output_dir
os.makedirs(self.output_dir, exist_ok=True)
return self.output_dir
def set_collage_style(self, collage_style):
"""设置拼贴画样式"""
self.collage_style = collage_style
return self.collage_style
def process_directory(directory_path, target_size=(900, 1200), output_count=1):
"""
Processes images in a directory: finds main subject, adjusts contrast/saturation,
performs smart cropping/resizing, creates a collage, and returns PIL Image objects.
Args:
directory_path: Path to the directory containing images.
target_size: Tuple (width, height) for the final collage.
output_count: Number of collages to generate.
Returns:
A list containing the generated PIL Image objects for the collages,
or an empty list if processing fails.
"""
image_files = [os.path.join(directory_path, f) for f in os.listdir(directory_path)
if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
if not image_files:
print(f"No images found in {directory_path}")
return []
# Create collage
collage_images = []
for i in range(output_count):
collage = create_collage(image_files, target_size)
if collage:
# collage_filename = f"collage_{i}.png"
# save_path = os.path.join(output_dir, collage_filename)
# collage.save(save_path)
# print(f"Collage saved to {save_path}")
# collage_images.append({'path': save_path, 'image': collage})
collage_images.append(collage) # Return the PIL Image object directly
else:
print(f"Failed to create collage {i}")
return collage_images
def create_collage(image_paths, target_size=(900, 1200)):
# ... (internal logic, including find_main_subject, adjust_image, smart_crop_and_resize) ...
pass
def find_main_subject(image):
# ... (keep the existing implementation) ...
pass
def adjust_image(image, contrast=1.0, saturation=1.0):
# ... (keep the existing implementation) ...
pass
def smart_crop_and_resize(image, target_aspect_ratio):
# ... (keep the existing implementation) ...
pass
def main():
# 设置基础路径
base_path = "/root/autodl-tmp"
# 默认图片目录
input_dir = os.path.join(base_path, "陈家祠")
## 考虑一下,是否需要直接传递图片结果
# 处理目录中的图片生成10个随机风格拼贴画
process_directory(input_dir, output_count=10)
if __name__ == "__main__":
main()