372 lines
13 KiB
Python
372 lines
13 KiB
Python
|
import os
|
|||
|
|
import random
|
|||
|
|
import logging
|
|||
|
|
import json
|
|||
|
|
from PIL import Image
|
|||
|
|
import traceback
|
|||
|
|
from typing import List, Tuple, Dict, Any, Optional
|
|||
|
|
|
|||
|
|
from .output_handler import OutputHandler
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
class PosterNotesCreator:
|
|||
|
|
"""
|
|||
|
|
处理原始海报作为主图，并随机选择额外的图片作为笔记图片。
|
|||
|
|
确保选择的笔记图片与海报中使用的图片不重复。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, output_handler: OutputHandler):
|
|||
|
|
"""
|
|||
|
|
初始化 PosterNotesCreator
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
output_handler: 可选的 OutputHandler 实例，用于处理输出
|
|||
|
|
"""
|
|||
|
|
self.output_handler = output_handler
|
|||
|
|
logging.info("PosterNotesCreator 初始化完成")
|
|||
|
|
|
|||
|
|
def create_notes_images(
|
|||
|
|
self,
|
|||
|
|
run_id: str,
|
|||
|
|
topic_index: int,
|
|||
|
|
variant_index: int,
|
|||
|
|
poster_image_path: str,
|
|||
|
|
poster_metadata_path: str,
|
|||
|
|
source_image_dir: str,
|
|||
|
|
num_additional_images: int,
|
|||
|
|
output_filename_template: str = "note_{index}.jpg"
|
|||
|
|
) -> List[str]:
|
|||
|
|
"""
|
|||
|
|
创建笔记图像
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
run_id: 运行ID
|
|||
|
|
topic_index: 主题索引
|
|||
|
|
variant_index: 变体索引
|
|||
|
|
poster_image_path: 海报图像路径
|
|||
|
|
poster_metadata_path: 海报元数据路径
|
|||
|
|
source_image_dir: 源图像目录
|
|||
|
|
num_additional_images: 要使用的额外图像数量
|
|||
|
|
output_filename_template: 输出文件名模板
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
List[str]: 保存的笔记图像路径列表
|
|||
|
|
"""
|
|||
|
|
# 检查输入路径是否存在
|
|||
|
|
if not os.path.exists(poster_image_path):
|
|||
|
|
logger.error(f"海报图像不存在: {poster_image_path}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
if not os.path.exists(poster_metadata_path):
|
|||
|
|
logger.error(f"海报元数据不存在: {poster_metadata_path}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
if not os.path.exists(source_image_dir) or not os.path.isdir(source_image_dir):
|
|||
|
|
logger.error(f"源图像目录不存在: {source_image_dir}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 从元数据文件中读取已使用的图像信息
|
|||
|
|
try:
|
|||
|
|
with open(poster_metadata_path, 'r', encoding='utf-8') as f:
|
|||
|
|
poster_metadata = json.load(f)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"无法读取海报元数据: {e}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 获取已经在海报中使用的图像
|
|||
|
|
used_images = []
|
|||
|
|
if 'collage_images' in poster_metadata:
|
|||
|
|
used_images = poster_metadata['collage_images']
|
|||
|
|
logger.info(f"海报中已使用 {len(used_images)} 张图像: {', '.join(used_images)}")
|
|||
|
|
|
|||
|
|
# 列出源目录中的所有图像文件
|
|||
|
|
image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
|
|||
|
|
available_images = [
|
|||
|
|
f for f in os.listdir(source_image_dir)
|
|||
|
|
if os.path.isfile(os.path.join(source_image_dir, f)) and
|
|||
|
|
f.lower().endswith(image_extensions)
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
if not available_images:
|
|||
|
|
logger.error(f"源目录中没有找到图像: {source_image_dir}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
logger.info(f"源目录中找到 {len(available_images)} 张图像")
|
|||
|
|
|
|||
|
|
# 过滤掉已经在海报中使用的图像
|
|||
|
|
available_images = [img for img in available_images if img not in used_images]
|
|||
|
|
|
|||
|
|
if not available_images:
|
|||
|
|
logger.warning("所有图像都已在海报中使用，无法创建额外笔记")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
logger.info(f"过滤后可用图像数量: {len(available_images)}")
|
|||
|
|
|
|||
|
|
# 如果可用图像少于请求数量，进行警告但继续处理
|
|||
|
|
if len(available_images) < num_additional_images:
|
|||
|
|
logger.warning(
|
|||
|
|
f"可用图像数量 ({len(available_images)}) 少于请求的笔记数量 ({num_additional_images})，"
|
|||
|
|
f"将使用所有可用图像"
|
|||
|
|
)
|
|||
|
|
selected_images = available_images
|
|||
|
|
else:
|
|||
|
|
# 随机选择额外图像
|
|||
|
|
selected_images = random.sample(available_images, num_additional_images)
|
|||
|
|
|
|||
|
|
logger.info(f"已选择 {len(selected_images)} 张图像作为笔记")
|
|||
|
|
|
|||
|
|
# 保存选择的笔记图像
|
|||
|
|
saved_paths = []
|
|||
|
|
for i, image_filename in enumerate(selected_images):
|
|||
|
|
try:
|
|||
|
|
# 加载图像
|
|||
|
|
image_path = os.path.join(source_image_dir, image_filename)
|
|||
|
|
image = Image.open(image_path)
|
|||
|
|
|
|||
|
|
# 生成输出文件名
|
|||
|
|
output_filename = output_filename_template.format(index=i+1)
|
|||
|
|
|
|||
|
|
# 创建元数据
|
|||
|
|
note_metadata = {
|
|||
|
|
"original_image": image_filename,
|
|||
|
|
"note_index": i + 1,
|
|||
|
|
"source_dir": source_image_dir,
|
|||
|
|
"associated_poster": os.path.basename(poster_image_path)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 使用输出处理器保存图像
|
|||
|
|
saved_path = self.output_handler.handle_generated_image(
|
|||
|
|
run_id,
|
|||
|
|
topic_index,
|
|||
|
|
variant_index,
|
|||
|
|
'note', # 图像类型为note
|
|||
|
|
image,
|
|||
|
|
output_filename,
|
|||
|
|
note_metadata
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
saved_paths.append(saved_path)
|
|||
|
|
logger.info(f"已保存笔记图像 {i+1}/{len(selected_images)}: {saved_path}")
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"处理图像时出错 '{image_filename}': {e}")
|
|||
|
|
|
|||
|
|
return saved_paths
|
|||
|
|
|
|||
|
|
def create_additional_images(
|
|||
|
|
self,
|
|||
|
|
run_id: str,
|
|||
|
|
topic_index: int,
|
|||
|
|
variant_index: int,
|
|||
|
|
poster_metadata_path: str,
|
|||
|
|
source_image_dir: str,
|
|||
|
|
num_additional_images: int = 3,
|
|||
|
|
output_filename_template: str = "additional_{index}.jpg"
|
|||
|
|
) -> List[str]:
|
|||
|
|
"""
|
|||
|
|
选择未被海报使用的图像作为额外配图
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
run_id: 运行ID
|
|||
|
|
topic_index: 主题索引
|
|||
|
|
variant_index: 变体索引
|
|||
|
|
poster_metadata_path: 海报元数据路径
|
|||
|
|
source_image_dir: 源图像目录
|
|||
|
|
num_additional_images: 要选择的额外图像数量
|
|||
|
|
output_filename_template: 输出文件名模板
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
List[str]: 保存的额外配图路径列表
|
|||
|
|
"""
|
|||
|
|
logger.info(f"开始为主题 {topic_index} 变体 {variant_index} 选择额外配图")
|
|||
|
|
|
|||
|
|
# 检查输入路径是否存在
|
|||
|
|
if not os.path.exists(poster_metadata_path):
|
|||
|
|
logger.error(f"海报元数据不存在: {poster_metadata_path}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
if not os.path.exists(source_image_dir) or not os.path.isdir(source_image_dir):
|
|||
|
|
logger.error(f"源图像目录不存在: {source_image_dir}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 从元数据文件中读取已使用的图像信息
|
|||
|
|
try:
|
|||
|
|
with open(poster_metadata_path, 'r', encoding='utf-8') as f:
|
|||
|
|
poster_metadata = json.load(f)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"无法读取海报元数据: {e}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 获取已经在海报中使用的图像
|
|||
|
|
used_images = []
|
|||
|
|
if 'collage_images' in poster_metadata:
|
|||
|
|
used_images = poster_metadata['collage_images']
|
|||
|
|
logger.info(f"海报中已使用 {len(used_images)} 张图像: {', '.join(used_images)}")
|
|||
|
|
|
|||
|
|
# 列出源目录中的所有图像文件
|
|||
|
|
image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
|
|||
|
|
available_images = [
|
|||
|
|
f for f in os.listdir(source_image_dir)
|
|||
|
|
if os.path.isfile(os.path.join(source_image_dir, f)) and
|
|||
|
|
f.lower().endswith(image_extensions)
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
if not available_images:
|
|||
|
|
logger.error(f"源目录中没有找到图像: {source_image_dir}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
logger.info(f"源目录中找到 {len(available_images)} 张图像")
|
|||
|
|
|
|||
|
|
# 过滤掉已经在海报中使用的图像
|
|||
|
|
available_images = [img for img in available_images if img not in used_images]
|
|||
|
|
|
|||
|
|
if not available_images:
|
|||
|
|
logger.warning("所有图像都已在海报中使用，无法创建额外配图")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
logger.info(f"过滤后可用图像数量: {len(available_images)}")
|
|||
|
|
|
|||
|
|
# 如果可用图像少于请求数量，进行警告但继续处理
|
|||
|
|
if len(available_images) < num_additional_images:
|
|||
|
|
logger.warning(
|
|||
|
|
f"可用图像数量 ({len(available_images)}) 少于请求的配图数量 ({num_additional_images})，"
|
|||
|
|
f"将使用所有可用图像"
|
|||
|
|
)
|
|||
|
|
selected_images = available_images
|
|||
|
|
else:
|
|||
|
|
# 随机选择额外图像
|
|||
|
|
selected_images = random.sample(available_images, num_additional_images)
|
|||
|
|
|
|||
|
|
logger.info(f"已选择 {len(selected_images)} 张图像作为额外配图")
|
|||
|
|
|
|||
|
|
# 保存选择的额外配图
|
|||
|
|
saved_paths = []
|
|||
|
|
for i, image_filename in enumerate(selected_images):
|
|||
|
|
try:
|
|||
|
|
# 加载图像
|
|||
|
|
image_path = os.path.join(source_image_dir, image_filename)
|
|||
|
|
image = Image.open(image_path)
|
|||
|
|
|
|||
|
|
# 生成输出文件名
|
|||
|
|
output_filename = output_filename_template.format(index=i+1)
|
|||
|
|
|
|||
|
|
# 创建元数据
|
|||
|
|
additional_metadata = {
|
|||
|
|
"original_image": image_filename,
|
|||
|
|
"additional_index": i + 1,
|
|||
|
|
"source_dir": source_image_dir,
|
|||
|
|
"is_additional_image": True
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 使用输出处理器保存图像
|
|||
|
|
saved_path = self.output_handler.handle_generated_image(
|
|||
|
|
run_id,
|
|||
|
|
topic_index,
|
|||
|
|
variant_index,
|
|||
|
|
'additional', # 图像类型为additional
|
|||
|
|
image,
|
|||
|
|
output_filename,
|
|||
|
|
additional_metadata
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
saved_paths.append(saved_path)
|
|||
|
|
logger.info(f"已保存额外配图 {i+1}/{len(selected_images)}: {saved_path}")
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"处理图像时出错 '{image_filename}': {e}")
|
|||
|
|
|
|||
|
|
return saved_paths
|
|||
|
|
|
|||
|
|
def process_poster_for_notes(
|
|||
|
|
run_id: str,
|
|||
|
|
topic_index: int,
|
|||
|
|
variant_index: int,
|
|||
|
|
poster_image_path: str,
|
|||
|
|
poster_metadata_path: str,
|
|||
|
|
source_image_dir: str,
|
|||
|
|
num_additional_images: int,
|
|||
|
|
output_handler: OutputHandler,
|
|||
|
|
output_filename_template: str = "note_{index}.jpg"
|
|||
|
|
) -> List[str]:
|
|||
|
|
"""
|
|||
|
|
处理海报并创建笔记图像
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
run_id: 运行ID
|
|||
|
|
topic_index: 主题索引
|
|||
|
|
variant_index: 变体索引
|
|||
|
|
poster_image_path: 海报图像路径
|
|||
|
|
poster_metadata_path: 海报元数据路径
|
|||
|
|
source_image_dir: 源图像目录
|
|||
|
|
num_additional_images: 要使用的额外图像数量
|
|||
|
|
output_handler: 输出处理器
|
|||
|
|
output_filename_template: 输出文件名模板
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
List[str]: 保存的笔记图像路径列表
|
|||
|
|
"""
|
|||
|
|
logger.info(f"开始为海报创建笔记图像: {poster_image_path}")
|
|||
|
|
|
|||
|
|
# 验证输入
|
|||
|
|
if not os.path.exists(poster_image_path):
|
|||
|
|
logger.error(f"海报图像不存在: {poster_image_path}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 创建处理器实例并处理
|
|||
|
|
creator = PosterNotesCreator(output_handler)
|
|||
|
|
return creator.create_notes_images(
|
|||
|
|
run_id,
|
|||
|
|
topic_index,
|
|||
|
|
variant_index,
|
|||
|
|
poster_image_path,
|
|||
|
|
poster_metadata_path,
|
|||
|
|
source_image_dir,
|
|||
|
|
num_additional_images,
|
|||
|
|
output_filename_template
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def select_additional_images(
|
|||
|
|
run_id: str,
|
|||
|
|
topic_index: int,
|
|||
|
|
variant_index: int,
|
|||
|
|
poster_metadata_path: str,
|
|||
|
|
source_image_dir: str,
|
|||
|
|
num_additional_images: int,
|
|||
|
|
output_handler: OutputHandler,
|
|||
|
|
output_filename_template: str = "additional_{index}.jpg"
|
|||
|
|
) -> List[str]:
|
|||
|
|
"""
|
|||
|
|
选择未被海报使用的图像作为额外配图
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
run_id: 运行ID
|
|||
|
|
topic_index: 主题索引
|
|||
|
|
variant_index: 变体索引
|
|||
|
|
poster_metadata_path: 海报元数据路径
|
|||
|
|
source_image_dir: 源图像目录
|
|||
|
|
num_additional_images: 要选择的额外图像数量
|
|||
|
|
output_handler: 输出处理器
|
|||
|
|
output_filename_template: 输出文件名模板
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
List[str]: 保存的额外配图路径列表
|
|||
|
|
"""
|
|||
|
|
logger.info(f"开始为主题 {topic_index} 变体 {variant_index} 选择额外配图")
|
|||
|
|
|
|||
|
|
# 验证输入
|
|||
|
|
if not os.path.exists(poster_metadata_path):
|
|||
|
|
logger.error(f"海报元数据不存在: {poster_metadata_path}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 创建处理器实例并处理
|
|||
|
|
creator = PosterNotesCreator(output_handler)
|
|||
|
|
return creator.create_additional_images(
|
|||
|
|
run_id,
|
|||
|
|
topic_index,
|
|||
|
|
variant_index,
|
|||
|
|
poster_metadata_path,
|
|||
|
|
source_image_dir,
|
|||
|
|
num_additional_images,
|
|||
|
|
output_filename_template
|
|||
|
|
)
|