TravelContentCreator/core/simple_collage.py

# -*- coding: utf-8 -*-
import os
import random
import traceback
import math
from pathlib import Path
from PIL import Image, ImageDraw, ImageEnhance, ImageFilter, ImageOps
import logging # Import logging module

class ImageCollageCreator:
    def __init__(self, ):
        """初始化拼图创建器"""

        # 定义可用拼接样式
        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):
        """创建指定样式的拼接画布"""
        logging.info(f"--- Starting Collage Creation for Directory: {input_dir} ---") # Start Log
        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)
            logging.info(f"Using collage style: {style} with target size: {target_size}")
            
            # 检查目录是否存在
            if not os.path.exists(input_dir):
                logging.error(f"Input directory does not exist: {input_dir}")
                return None
            
            # 支持的图片格式
            image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
            
            # 获取目录中的所有文件
            try:
                all_files = os.listdir(input_dir)
                logging.info(f"Files found in directory: {all_files}")
            except Exception as e:
                logging.exception(f"Error listing directory {input_dir}: {e}")
                return None
                
            # 过滤图片文件
            all_images_names = [f for f in all_files 
                               if f.lower().endswith(image_extensions) and os.path.isfile(os.path.join(input_dir, f))]
            logging.info(f"Filtered image files: {all_images_names}")
            
            if not all_images_names:
                 logging.warning(f"No valid image files found in directory: {input_dir}")
                 return None # Return None if no images found
            
            # 根据不同样式，确定需要的图片数量
            # ... (logic for num_images based on style) ...
            num_images = 4
            if style == "mosaic":
                num_images = 9
            elif style == "filmstrip":
                num_images = 5
            elif style == "fullscreen":
                num_images = 6
            elif style == "vertical_stack":
                num_images = 2
            logging.info(f"Style '{style}' requires {num_images} images.")
            
            # 确保有足够的图像 (或重复使用)
            selected_images_names = []
            if len(all_images_names) < num_images:
                logging.warning(f"Need {num_images} images for style '{style}', but only found {len(all_images_names)}. Will repeat images.")
                if len(all_images_names) > 0:
                    # Repeat available images to meet the count
                    selected_images_names = (all_images_names * (num_images // len(all_images_names) + 1))[:num_images]
                else:
                     logging.error("Cannot select images, none were found.") # Should not happen due to earlier check
                     return None
            else:
                # 随机选择指定数量的图片
                selected_images_names = random.sample(all_images_names, num_images)
            
            logging.info(f"Selected image files for collage: {selected_images_names}")
            
            # 加载图片
            images = []
            loaded_image_paths = set()
            for img_name in selected_images_names:
                img_path = os.path.join(input_dir, img_name)
                try:
                    img = Image.open(img_path).convert('RGBA')
                    images.append(img)
                    loaded_image_paths.add(img_path)
                    logging.info(f"Successfully loaded image: {img_path}")
                except Exception as e:
                    logging.error(f"Failed to load image {img_path}: {e}", exc_info=True) # Log exception info
                    # Optionally: try to replace failed image (or just log and continue)
                    # For simplicity now, just log and continue; the check below handles insufficient images.
            
            # 再次检查实际加载成功的图片数量
            if len(images) < num_images:
                 logging.error(f"Needed {num_images} images, but only successfully loaded {len(images)}. Cannot create collage.")
                 # Log which images failed if possible (from error logs above)
                 return None
                 
            logging.info(f"Successfully loaded {len(images)} images for collage.")

            # 创建空白画布 (moved after image loading success check)
            # collage_image = Image.new('RGBA', target_size, (0, 0, 0, 0)) # This line seems unused as styles create their own canvas
            
            # 应用所选样式
            logging.info(f"Applying style '{style}'...")
            result_collage = None
            if style == "grid_2x2":
                result_collage = self._create_grid_2x2_collage(images, target_size)
            # ... (elif for all other styles) ...
            elif style == "asymmetric":
                result_collage = self._create_asymmetric_collage(images, target_size)
            elif style == "filmstrip":
                result_collage = self._create_filmstrip_collage(images, target_size)
            # elif style == "circles":
            #     result_collage = self._create_circles_collage(images, target_size)
            elif style == "polaroid":
                result_collage = self._create_polaroid_collage(images, target_size)
            elif style == "overlap":
                result_collage = self._create_overlap_collage(images, target_size)
            elif style == "mosaic":
                result_collage = self._create_mosaic_collage(images, target_size)
            elif style == "fullscreen":
                result_collage = self._create_fullscreen_collage(images, target_size)
            elif style == "vertical_stack":
                result_collage = self._create_vertical_stack_collage(images, target_size)
            else:
                logging.warning(f"Unknown style '{style}', defaulting to grid_2x2.")
                result_collage = self._create_grid_2x2_collage(images, target_size)
            
            if result_collage is None:
                logging.error(f"Collage creation failed during style application ('{style}').")
                return None
            else:
                logging.info(f"--- Collage Creation Successful for Directory: {input_dir} ---")
                return result_collage # Return the created collage image
            
        except Exception as e:
            logging.exception(f"An unexpected error occurred during collage creation for {input_dir}: {e}") # Log full traceback
            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_collage_style(self, collage_style):
        """设置拼贴画样式"""
        self.collage_style = collage_style
        return self.collage_style

def process_directory(directory_path, style=None, target_size=(900, 1200), output_count=1):
    """
    处理指定目录中的图片，创建指定数量的拼贴图。
    
    参数:
        directory_path: 包含图片的目录路径
        target_size: 拼贴图目标尺寸，默认为 (900, 1200)
        output_count: 需要生成的拼贴图数量，默认为 1
        
    返回:
        list: 生成的拼贴图列表（PIL.Image 对象）；如果生成失败，返回空列表
    """
    logging.info(f"处理目录中的图片并创建 {output_count} 个拼贴图: {directory_path}")
    
    # 创建 ImageCollageCreator 实例
    collage_creator = ImageCollageCreator()
    collage_images = []
    
    # 检查目录是否存在
    if not os.path.exists(directory_path):
        logging.error(f"目录不存在: {directory_path}")
        return []
        
    # 尝试创建请求数量的拼贴图
    for i in range(output_count):
        try:
            # 随机选择一个样式（由 create_collage_with_style 内部实现）
            # 传入 None 作为 style 参数，让函数内部随机选择
            collage = collage_creator.create_collage_with_style(
                directory_path, 
                style=style,  # 让方法内部随机选择样式
                target_size=target_size
            )
            
            if collage:
                collage_images.append(collage)
                logging.info(f"成功创建拼贴图 {i+1}/{output_count}")
            else:
                logging.error(f"无法创建拼贴图 {i+1}/{output_count}")
        except Exception as e:
            logging.exception(f"创建拼贴图 {i+1}/{output_count} 时发生异常: {e}")
    
    logging.info(f"已处理目录 {directory_path}，成功创建 {len(collage_images)}/{output_count} 个拼贴图")
    return collage_images

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():
    """展示如何使用 ImageCollageCreator 和 process_directory 函数的示例。"""
    logging.basicConfig(level=logging.INFO, 
                      format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s')
    
    # 示例目录路径 - 根据实际情况修改
    test_directory = "/root/autodl-tmp/sanming_img/modify/古田会议旧址"  # 修改为你实际的图片目录
    
    logging.info(f"测试目录: {test_directory}")
    
    # 方法 1: 使用 process_directory 函数 (推荐用于外部调用)
    logging.info("方法 1: 使用 process_directory 函数生成拼贴图...")
    collages_1 = process_directory(
        directory_path=test_directory,
        target_size=(900, 1200),  # 默认 3:4 比例
        output_count=2  # 创建 2 张不同的拼贴图
    )
    
    if collages_1:
        logging.info(f"成功创建了 {len(collages_1)} 张拼贴图 (使用 process_directory)")
        # 可选: 保存图片到文件
        for i, collage in enumerate(collages_1):
            output_path = f"/tmp/collage_method1_{i}.png"
            collage.save(output_path)
            logging.info(f"拼贴图已保存到: {output_path}")
    else:
        logging.error("使用 process_directory 创建拼贴图失败")
    
    # 方法 2: 直接使用 ImageCollageCreator 类 (用于更精细的控制)
    logging.info("方法 2: 直接使用 ImageCollageCreator 类...")
    creator = ImageCollageCreator()
    
    # 指定样式创建拼贴图 (可选样式: grid_2x2, asymmetric, filmstrip, overlap, mosaic, fullscreen, vertical_stack)
    styles_to_try = ["grid_2x2", "overlap", "mosaic"]
    collages_2 = []
    
    for style in styles_to_try:
        logging.info(f"尝试使用样式: {style}")
        collage = creator.create_collage_with_style(
            input_dir=test_directory,
            style=style,
            target_size=(800, 1000)  # 自定义尺寸
        )
        
        if collage:
            collages_2.append(collage)
            # 可选: 保存图片到文件
            output_path = f"/tmp/collage_method2_{style}.png"
            collage.save(output_path)
            logging.info(f"使用样式 '{style}' 的拼贴图已保存到: {output_path}")
        else:
            logging.error(f"使用样式 '{style}' 创建拼贴图失败")
    
    logging.info(f"总共成功创建了 {len(collages_2)} 张拼贴图 (使用 ImageCollageCreator)")
    
    # 比较两种方法
    logging.info("===== 拼贴图创建测试完成 =====")
    logging.info(f"方法 1 (process_directory): {len(collages_1)} 张拼贴图")
    logging.info(f"方法 2 (直接使用 ImageCollageCreator): {len(collages_2)} 张拼贴图")
    
if __name__ == "__main__":
    main()