183 lines
6.2 KiB
Python
183 lines
6.2 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
测试图片拼贴功能
|
||
|
||
该脚本演示如何独立使用simple_collage模块创建图片拼贴。
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import argparse
|
||
import time
|
||
import random
|
||
from pathlib import Path
|
||
from PIL import Image
|
||
|
||
# 将项目根目录添加到PATH
|
||
project_root = str(Path(__file__).parent.parent.absolute())
|
||
if project_root not in sys.path:
|
||
sys.path.append(project_root)
|
||
|
||
from core import simple_collage
|
||
|
||
def parse_arguments():
|
||
"""解析命令行参数"""
|
||
parser = argparse.ArgumentParser(description='测试图片拼贴功能')
|
||
parser.add_argument('--input_dir', '-i', type=str,
|
||
help='输入图片目录路径')
|
||
parser.add_argument('--output_dir', '-o', type=str, default='collage_output',
|
||
help='输出拼贴图保存目录，默认为"collage_output"')
|
||
parser.add_argument('--width', '-w', type=int, default=900,
|
||
help='拼贴图宽度，默认为900')
|
||
parser.add_argument('--height', '-h', type=int, default=1200,
|
||
help='拼贴图高度，默认为1200')
|
||
parser.add_argument('--count', '-c', type=int, default=3,
|
||
help='生成拼贴图数量，默认为3')
|
||
parser.add_argument('--seed', '-s', type=int, default=None,
|
||
help='随机种子，用于重现结果')
|
||
return parser.parse_args()
|
||
|
||
def get_image_files(directory):
|
||
"""获取目录中的所有图片文件"""
|
||
if not os.path.exists(directory):
|
||
raise FileNotFoundError(f"目录不存在: {directory}")
|
||
|
||
image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.webp', '.tiff']
|
||
image_files = []
|
||
|
||
for root, _, files in os.walk(directory):
|
||
for file in files:
|
||
if any(file.lower().endswith(ext) for ext in image_extensions):
|
||
image_files.append(os.path.join(root, file))
|
||
|
||
if not image_files:
|
||
raise ValueError(f"目录中没有找到有效的图片文件: {directory}")
|
||
|
||
return image_files
|
||
|
||
def create_test_images(directory, count=5, size=(640, 480)):
|
||
"""如果没有提供输入目录，创建测试图片"""
|
||
if not os.path.exists(directory):
|
||
os.makedirs(directory)
|
||
|
||
colors = [
|
||
(255, 0, 0), # 红色
|
||
(0, 255, 0), # 绿色
|
||
(0, 0, 255), # 蓝色
|
||
(255, 255, 0), # 黄色
|
||
(255, 0, 255), # 紫色
|
||
(0, 255, 255), # 青色
|
||
(255, 165, 0), # 橙色
|
||
(128, 0, 128), # 紫色
|
||
(210, 105, 30), # 巧克力色
|
||
(0, 128, 128) # 墨绿色
|
||
]
|
||
|
||
image_files = []
|
||
for i in range(count):
|
||
# 随机选择颜色
|
||
color = colors[i % len(colors)]
|
||
|
||
# 创建纯色图片
|
||
img = Image.new('RGB', size, color)
|
||
|
||
# 保存图片
|
||
file_path = os.path.join(directory, f"test_image_{i+1}.jpg")
|
||
img.save(file_path)
|
||
image_files.append(file_path)
|
||
print(f"创建测试图片: {file_path}")
|
||
|
||
return image_files
|
||
|
||
def main():
|
||
"""主函数"""
|
||
args = parse_arguments()
|
||
|
||
# 设置随机种子
|
||
if args.seed is not None:
|
||
random.seed(args.seed)
|
||
|
||
# 确保输出目录存在
|
||
if not os.path.exists(args.output_dir):
|
||
os.makedirs(args.output_dir)
|
||
|
||
# 获取输入图片
|
||
if not args.input_dir:
|
||
print("未提供输入目录，创建测试图片...")
|
||
test_img_dir = os.path.join(args.output_dir, "test_images")
|
||
image_files = create_test_images(test_img_dir, count=7)
|
||
input_dir = test_img_dir
|
||
else:
|
||
input_dir = args.input_dir
|
||
try:
|
||
# 尝试获取图片文件列表，验证目录有效
|
||
get_image_files(input_dir)
|
||
except Exception as e:
|
||
print(f"错误: {e}")
|
||
return
|
||
|
||
# 设置目标尺寸
|
||
target_size = (args.width, args.height)
|
||
print(f"使用输入目录: {input_dir}")
|
||
print(f"拼贴图尺寸: {target_size}")
|
||
|
||
# 生成拼贴图
|
||
print(f"开始生成{args.count}张拼贴图...")
|
||
start_time = time.time()
|
||
|
||
try:
|
||
# 创建与测试不同样式的拼贴图
|
||
collage_creator = simple_collage.ImageCollageCreator()
|
||
|
||
# 1. 处理整个目录，获取多个拼贴图
|
||
print("方法1: 使用process_directory批量生成")
|
||
collages, used_image_names = simple_collage.process_directory(
|
||
input_dir,
|
||
target_size=target_size,
|
||
output_count=args.count
|
||
)
|
||
|
||
if not collages:
|
||
print("拼贴图生成失败！")
|
||
return
|
||
|
||
# 保存拼贴图
|
||
for i, collage in enumerate(collages):
|
||
output_path = os.path.join(args.output_dir, f"collage_auto_{i+1}.png")
|
||
collage.save(output_path)
|
||
# 输出使用的图片名称
|
||
if i < len(used_image_names):
|
||
print(f"拼贴图已保存: {output_path}，使用图片: {used_image_names[i]}")
|
||
else:
|
||
print(f"拼贴图已保存: {output_path}")
|
||
|
||
# 2. 使用不同风格创建拼贴图
|
||
print("\n方法2: 测试不同风格")
|
||
styles = ["grid", "asymmetrical", "filmstrip", "overlap", "mosaic"]
|
||
|
||
for style in styles:
|
||
print(f"创建 {style} 风格拼贴图...")
|
||
try:
|
||
collage, selected_images = collage_creator.create_collage_with_style(input_dir, style, target_size)
|
||
if collage:
|
||
output_path = os.path.join(args.output_dir, f"collage_style_{style}.png")
|
||
collage.save(output_path)
|
||
print(f"风格拼贴图已保存: {output_path}，使用图片: {selected_images}")
|
||
else:
|
||
print(f"创建 {style} 风格失败: 未返回有效拼贴图")
|
||
except Exception as e:
|
||
print(f"创建 {style} 风格失败: {e}")
|
||
|
||
elapsed = time.time() - start_time
|
||
print(f"\n拼贴图生成完成，总耗时: {elapsed:.2f}秒")
|
||
print(f"输出目录: {os.path.abspath(args.output_dir)}")
|
||
|
||
except Exception as e:
|
||
print(f"拼贴图生成出错: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
if __name__ == "__main__":
|
||
main() |