解耦了海报生成中的部分功能

This commit is contained in:
jinye_huang 2025-04-22 14:30:25 +08:00
parent f415abaec3
commit 0a048da5dd
3 changed files with 267 additions and 190 deletions

View File

@ -31,5 +31,6 @@
"variants": 3,
"topic_temperature": 0.2,
"content_temperature": 0.3,
"poster_target_size": [900, 1200]
"poster_target_size": [900, 1200],
"text_possibility": 0.3
}

231
main.py
View File

@ -12,12 +12,14 @@ import core.contentGen as contentGen
import core.posterGen as posterGen
import core.simple_collage as simple_collage
from utils.resource_loader import ResourceLoader
from utils.tweet_generator import generate_single_content, run_topic_generation_pipeline # Import the new pipeline function
from utils.tweet_generator import ( # Import the moved functions
run_topic_generation_pipeline,
generate_content_for_topic,
generate_posters_for_topic
)
from utils.prompt_manager import PromptManager # Import PromptManager
import random
TEXT_POSBILITY = 0.3 # Consider moving this to config if it varies
def load_config(config_path="poster_gen_config.json"):
"""Loads configuration from a JSON file."""
if not os.path.exists(config_path):
@ -46,209 +48,74 @@ def load_config(config_path="poster_gen_config.json"):
# Removed generate_topics_step function definition from here
# Its logic is now in utils.tweet_generator.run_topic_generation_pipeline
# --- Main Orchestration Step (Remains in main.py) ---
def generate_content_and_posters_step(config, run_id, tweet_topic_record):
"""Generates content and posters based on generated topics."""
if not run_id or not tweet_topic_record:
print("Missing run_id or topics data. Skipping content and poster generation.")
"""Generates content and posters by calling decoupled functions for each topic."""
if not run_id or not tweet_topic_record or not tweet_topic_record.topics_list:
print("Missing run_id or valid topics data. Skipping content and poster generation.")
return
print("\nStep 2: Generating Content and Posters...")
print("\nStep 2: Processing Topics for Content and Posters...")
base_output_dir = config["output_dir"]
output_dir = os.path.join(base_output_dir, run_id) # Directory for this specific run
# --- Pre-load resources and initialize shared objects ---
# Initialize PromptManager once for this phase
# --- Initialize shared resources once ---
prompt_manager = None
ai_agent = None
# We might not need to initialize these here if the moved functions handle it
# content_gen = None
# poster_gen_instance = None
initialized_ok = False
try:
prompt_manager = PromptManager(config)
except Exception as e:
print(f"Error initializing PromptManager: {e}")
traceback.print_exc()
return
# Load content generation system prompt once
content_system_prompt = ResourceLoader.load_system_prompt(config["content_system_prompt"])
if not content_system_prompt:
print("Warning: Content generation system prompt is empty. Using default logic if available or might fail.")
# Initialize AI Agent once for the entire content generation phase
ai_agent = None
try:
print(f"Initializing AI Agent ({config['model']})...")
ai_agent = AI_Agent(config["api_url"], config["model"], config["api_key"])
except Exception as e:
print(f"Error initializing AI Agent: {e}. Cannot proceed with content generation.")
traceback.print_exc()
return # Cannot continue without AI agent
# content_gen = core_contentGen.ContentGenerator() # Instantiation moved to generate_posters_for_topic
# poster_gen_instance = core_posterGen.PosterGenerator() # Instantiation moved to generate_posters_for_topic
# Check image base directory
# Check image base directory validity once
image_base_dir = config.get("image_base_dir", None)
if not image_base_dir or not os.path.isdir(image_base_dir):
print(f"Error: 'image_base_dir' ({image_base_dir}) not specified or not a valid directory in config. Cannot locate images.")
if ai_agent: ai_agent.close() # Close agent if initialized
raise ValueError(f"'image_base_dir' ({image_base_dir}) not specified or not a valid directory in config.")
initialized_ok = True
except Exception as e:
print(f"Error initializing resources for content/poster generation: {e}")
traceback.print_exc()
if ai_agent: ai_agent.close()
return
camera_image_subdir = config.get("camera_image_subdir", "相机") # Default '相机'
modify_image_subdir = config.get("modify_image_subdir", "modify") # Default 'modify'
# Initialize ContentGenerator and PosterGenerator once if they are stateless
# Assuming they are stateless for now
content_gen = contentGen.ContentGenerator()
poster_gen_instance = posterGen.PosterGenerator()
# --- Process each topic ---
# --- Process each topic using decoupled functions ---
total_topics = len(tweet_topic_record.topics_list)
for i, topic in enumerate(tweet_topic_record.topics_list):
topic_index = i + 1
print(f"\nProcessing Topic {topic_index}/{len(tweet_topic_record.topics_list)}: {topic.get('title', 'N/A')}")
tweet_content_list = []
print(f"\nProcessing Topic {topic_index}/{total_topics}: {topic.get('title', 'N/A')}")
# --- Content Generation Loop (using the single AI Agent) ---
for j in range(config["variants"]):
variant_index = j + 1
print(f" Generating Variant {variant_index}/{config['variants']}...")
# Get prompts for this specific topic item using PromptManager
content_system_prompt, content_user_prompt = prompt_manager.get_content_prompts(topic)
if not content_system_prompt or not content_user_prompt:
print(f" Skipping Variant {variant_index} due to missing content prompts.")
continue # Skip this variant if prompts failed
time.sleep(random.random() * 0.5) # Slightly reduced delay
try:
# Use the pre-initialized AI Agent and pass the specific prompts
tweet_content, gen_result = generate_single_content(
ai_agent, content_system_prompt, content_user_prompt, topic,
output_dir, run_id, topic_index, variant_index, config.get("content_temperature", 0.3)
)
if tweet_content:
# Assuming get_json_file() returns a dictionary or similar structure
tweet_content_data = tweet_content.get_json_file()
if tweet_content_data:
tweet_content_list.append(tweet_content_data)
else:
print(f" Warning: generate_single_content for Topic {topic_index}, Variant {variant_index} returned empty data.")
else:
print(f" Failed to generate content for Topic {topic_index}, Variant {variant_index}. Skipping.")
except Exception as e:
print(f" Error during content generation for Topic {topic_index}, Variant {variant_index}: {e}")
# Decide if traceback is needed here, might be too verbose for loop errors
# traceback.print_exc()
# Do NOT close the agent here
if not tweet_content_list:
print(f" No valid content generated for Topic {topic_index}. Skipping poster generation.")
continue
# --- Poster Generation Setup ---
object_name = topic.get("object", "")
if not object_name:
print(" Warning: Topic object name is missing. Cannot determine image path.")
continue
# Clean object name (consider making this a utility function)
try:
# More robust cleaning might be needed depending on actual object name formats
object_name_cleaned = object_name.split(".")[0].replace("景点信息-", "").strip()
if not object_name_cleaned:
print(f" Warning: Object name '{object_name}' resulted in empty string after cleaning.")
continue
object_name = object_name_cleaned
except Exception as e:
print(f" Warning: Could not fully clean object name '{object_name}': {e}")
# Continue with potentially unclean name? Or skip?
# Let's continue for now, path checks below might catch issues.
# Construct and check image paths using config base dir
input_img_dir_path = os.path.join(image_base_dir, modify_image_subdir, object_name)
camera_img_dir_path = os.path.join(image_base_dir, camera_image_subdir, object_name)
description_file_path = os.path.join(camera_img_dir_path, "description.txt")
if not os.path.exists(input_img_dir_path) or not os.path.isdir(input_img_dir_path):
print(f" Image directory not found or not a directory: '{input_img_dir_path}'. Skipping poster generation for this topic.")
continue
info_directory = []
if os.path.exists(description_file_path):
info_directory = [description_file_path]
print(f" Using description file: {description_file_path}")
else:
print(f" Description file not found: '{description_file_path}'. Using generated content for poster text.")
# --- Generate Text Configurations for Posters ---
try:
# Pass the list of content data directly
poster_text_configs_raw = content_gen.run(info_directory, config["variants"], tweet_content_list)
# print(f" Raw poster text configs: {poster_text_configs_raw}") # For debugging
if not poster_text_configs_raw:
print(" Warning: ContentGenerator returned empty configuration data.")
continue # Skip if no text configs generated
poster_config_summary = posterGen.PosterConfig(poster_text_configs_raw)
except Exception as e:
print(f" Error running ContentGenerator or parsing poster configs: {e}")
traceback.print_exc()
continue # Skip poster generation for this topic
# --- Poster Generation Loop ---
poster_num = config["variants"] # Same as content variants
target_size = tuple(config.get("poster_target_size", [900, 1200]))
for j_index in range(poster_num):
variant_index = j_index + 1
print(f" Generating Poster {variant_index}/{poster_num}...")
try:
poster_config = poster_config_summary.get_config_by_index(j_index)
if not poster_config:
print(f" Warning: Could not get poster config for index {j_index}. Skipping.")
continue
variant_output_dir = os.path.join(output_dir, f"{topic_index}_{variant_index}")
collage_output_dir = os.path.join(variant_output_dir, "collage_img")
poster_output_dir = os.path.join(variant_output_dir, "poster")
os.makedirs(collage_output_dir, exist_ok=True)
os.makedirs(poster_output_dir, exist_ok=True)
# --- Image Collage ---
img_list = simple_collage.process_directory(
input_img_dir_path,
target_size=target_size,
output_count=1, # Assuming 1 collage image per poster variant
output_dir=collage_output_dir
# 1. Generate Content for this topic (Call function from utils.tweet_generator)
tweet_content_list = generate_content_for_topic(
ai_agent, prompt_manager, config, topic,
base_output_dir, run_id, topic_index
)
if not img_list or len(img_list) == 0 or not img_list[0].get('path'):
print(f" Failed to generate collage image for Variant {variant_index}. Skipping poster.")
continue
collage_img_path = img_list[0]['path']
print(f" Using collage image: {collage_img_path}")
# --- Create Poster (using the single poster_gen_instance) ---
text_data = {
"title": poster_config.get('main_title', 'Default Title'),
"subtitle": "",
"additional_texts": []
}
texts = poster_config.get('texts', [])
if texts:
text_data["additional_texts"].append({"text": texts[0], "position": "bottom", "size_factor": 0.5})
if len(texts) > 1 and random.random() < TEXT_POSBILITY:
text_data["additional_texts"].append({"text": texts[1], "position": "bottom", "size_factor": 0.5})
final_poster_path = os.path.join(poster_output_dir, "poster.jpg")
result_path = poster_gen_instance.create_poster(collage_img_path, text_data, final_poster_path)
if result_path:
print(f" Successfully generated poster: {result_path}")
# 2. Generate Posters for this topic (Call function from utils.tweet_generator)
if tweet_content_list:
# Pass only necessary parts of config? Or the whole config?
# Pass config for now, the function extracts what it needs.
generate_posters_for_topic(
config, # Pass config
topic, tweet_content_list,
base_output_dir, run_id, topic_index
# Instances are now created inside generate_posters_for_topic
# poster_gen_instance, content_gen,
)
else:
print(f" Poster generation function did not return a valid path.")
except Exception as e:
print(f" Error during poster generation for Variant {variant_index}: {e}")
traceback.print_exc()
continue # Continue to next variant
print(f" Skipping poster generation for Topic {topic_index} due to content generation failure.")
# --- Cleanup ---
# Close the AI Agent after processing all topics
if ai_agent:
print("\nClosing AI Agent...")
print("\nClosing shared AI Agent...")
ai_agent.close()
print("\nFinished processing all topics.")
def main():
# No argparse for now, directly load default config

View File

@ -16,6 +16,9 @@ from TravelContentCreator.core.topic_parser import TopicParser
# ResourceLoader is now used implicitly via PromptManager
# from TravelContentCreator.utils.resource_loader import ResourceLoader
from TravelContentCreator.utils.prompt_manager import PromptManager # Import PromptManager
from ..core import contentGen as core_contentGen # Use relative imports for core
from ..core import posterGen as core_posterGen
from ..core import simple_collage as core_simple_collage
class tweetTopic:
def __init__(self, index, date, logic, object, product, product_logic, style, style_logic, target_audience, target_audience_logic):
@ -308,6 +311,212 @@ def run_topic_generation_pipeline(config):
print(f"Topics generated successfully. Run ID: {run_id}")
return run_id, tweet_topic_record
# --- Decoupled Functional Units (Moved from main.py) ---
def generate_content_for_topic(ai_agent, prompt_manager, config, topic_item, output_dir, run_id, topic_index):
"""Generates all content variants for a single topic item.
Args:
ai_agent: An initialized AI_Agent instance.
prompt_manager: An initialized PromptManager instance.
config: The global configuration dictionary.
topic_item: The dictionary representing a single topic.
output_dir: The base output directory for the entire run (e.g., ./result).
run_id: The ID for the current run.
topic_index: The 1-based index of the current topic.
Returns:
A list of tweet content data (dictionaries) generated for the topic,
or None if generation failed.
"""
print(f" Generating content for Topic {topic_index}...")
tweet_content_list = []
variants = config.get("variants", 1)
for j in range(variants):
variant_index = j + 1
print(f" Generating Variant {variant_index}/{variants}...")
# Get prompts for this specific topic item
# Assuming prompt_manager is correctly initialized and passed
content_system_prompt, content_user_prompt = prompt_manager.get_content_prompts(topic_item)
if not content_system_prompt or not content_user_prompt:
print(f" Skipping Variant {variant_index} due to missing content prompts.")
continue # Skip this variant
time.sleep(random.random() * 0.5)
try:
# Call the core generation function (generate_single_content is in this file)
tweet_content, gen_result = generate_single_content(
ai_agent, content_system_prompt, content_user_prompt, topic_item,
# Pass the base output_dir, run_id etc. generate_single_content creates subdirs
output_dir, run_id, topic_index, variant_index,
config.get("content_temperature", 0.3)
)
if tweet_content:
tweet_content_data = tweet_content.get_json_file()
if tweet_content_data:
tweet_content_list.append(tweet_content_data)
else:
print(f" Warning: generate_single_content for Topic {topic_index}, Variant {variant_index} returned empty data.")
else:
print(f" Failed to generate content for Topic {topic_index}, Variant {variant_index}. Skipping.")
except Exception as e:
print(f" Error during content generation for Topic {topic_index}, Variant {variant_index}: {e}")
# traceback.print_exc()
if not tweet_content_list:
print(f" No valid content generated for Topic {topic_index}.")
return None
else:
print(f" Successfully generated {len(tweet_content_list)} content variants for Topic {topic_index}.")
return tweet_content_list
def generate_posters_for_topic(config, topic_item, tweet_content_list, output_dir, run_id, topic_index):
"""Generates all posters for a single topic item based on its generated content.
Args:
config: The global configuration dictionary.
topic_item: The dictionary representing a single topic.
tweet_content_list: List of content data generated by generate_content_for_topic.
output_dir: The base output directory for the entire run (e.g., ./result).
run_id: The ID for the current run.
topic_index: The 1-based index of the current topic.
Returns:
True if poster generation was attempted (regardless of individual variant success),
False if setup failed before attempting variants.
"""
print(f" Generating posters for Topic {topic_index}...")
# Initialize necessary generators here, assuming they are stateless or cheap to create
# Alternatively, pass initialized instances if they hold state or are expensive
try:
content_gen_instance = core_contentGen.ContentGenerator()
poster_gen_instance = core_posterGen.PosterGenerator()
except Exception as e:
print(f" Error initializing generators for poster creation: {e}")
return False
# --- Setup: Paths and Object Name ---
image_base_dir = config.get("image_base_dir")
if not image_base_dir:
print(" Error: image_base_dir missing in config for poster generation.")
return False
modify_image_subdir = config.get("modify_image_subdir", "modify")
camera_image_subdir = config.get("camera_image_subdir", "相机")
object_name = topic_item.get("object", "")
if not object_name:
print(" Warning: Topic object name is missing. Cannot generate posters.")
return False
# Clean object name
try:
object_name_cleaned = object_name.split(".")[0].replace("景点信息-", "").strip()
if not object_name_cleaned:
print(f" Warning: Object name '{object_name}' resulted in empty string after cleaning. Skipping posters.")
return False
object_name = object_name_cleaned
except Exception as e:
print(f" Warning: Could not fully clean object name '{object_name}': {e}. Skipping posters.")
return False
# Construct and check image paths
input_img_dir_path = os.path.join(image_base_dir, modify_image_subdir, object_name)
camera_img_dir_path = os.path.join(image_base_dir, camera_image_subdir, object_name)
description_file_path = os.path.join(camera_img_dir_path, "description.txt")
if not os.path.exists(input_img_dir_path) or not os.path.isdir(input_img_dir_path):
print(f" Image directory not found or not a directory: '{input_img_dir_path}'. Skipping posters for this topic.")
return False
info_directory = []
if os.path.exists(description_file_path):
info_directory = [description_file_path]
print(f" Using description file: {description_file_path}")
else:
print(f" Description file not found: '{description_file_path}'. Using generated content for poster text.")
# --- Generate Text Configurations for All Variants ---
try:
poster_text_configs_raw = content_gen_instance.run(info_directory, config["variants"], tweet_content_list)
if not poster_text_configs_raw:
print(" Warning: ContentGenerator returned empty configuration data. Skipping posters.")
return False
poster_config_summary = core_posterGen.PosterConfig(poster_text_configs_raw)
except Exception as e:
print(f" Error running ContentGenerator or parsing poster configs: {e}")
traceback.print_exc()
return False # Cannot proceed if text config fails
# --- Poster Generation Loop for each variant ---
poster_num = config.get("variants", 1)
target_size = tuple(config.get("poster_target_size", [900, 1200]))
any_poster_attempted = False
text_possibility = config.get("text_possibility", 0.3) # Get from config
for j_index in range(poster_num):
variant_index = j_index + 1
print(f" Generating Poster {variant_index}/{poster_num}...")
any_poster_attempted = True
try:
poster_config = poster_config_summary.get_config_by_index(j_index)
if not poster_config:
print(f" Warning: Could not get poster config for index {j_index}. Skipping.")
continue
# Define output directories for this specific variant
run_output_dir = os.path.join(output_dir, run_id) # Base dir for the run
variant_output_dir = os.path.join(run_output_dir, f"{topic_index}_{variant_index}")
collage_output_dir = os.path.join(variant_output_dir, "collage_img")
poster_output_dir = os.path.join(variant_output_dir, "poster")
os.makedirs(collage_output_dir, exist_ok=True)
os.makedirs(poster_output_dir, exist_ok=True)
# --- Image Collage ---
img_list = core_simple_collage.process_directory(
input_img_dir_path,
target_size=target_size,
output_count=1,
output_dir=collage_output_dir
)
if not img_list or len(img_list) == 0 or not img_list[0].get('path'):
print(f" Failed to generate collage image for Variant {variant_index}. Skipping poster.")
continue
collage_img_path = img_list[0]['path']
print(f" Using collage image: {collage_img_path}")
# --- Create Poster ---
text_data = {
"title": poster_config.get('main_title', 'Default Title'),
"subtitle": "",
"additional_texts": []
}
texts = poster_config.get('texts', [])
if texts:
# Ensure TEXT_POSBILITY is accessible, maybe pass via config?
# text_possibility = config.get("text_possibility", 0.3)
text_data["additional_texts"].append({"text": texts[0], "position": "bottom", "size_factor": 0.5})
if len(texts) > 1 and random.random() < text_possibility: # Use variable from config
text_data["additional_texts"].append({"text": texts[1], "position": "bottom", "size_factor": 0.5})
final_poster_path = os.path.join(poster_output_dir, "poster.jpg")
result_path = poster_gen_instance.create_poster(collage_img_path, text_data, final_poster_path)
if result_path:
print(f" Successfully generated poster: {result_path}")
else:
print(f" Poster generation function did not return a valid path.")
except Exception as e:
print(f" Error during poster generation for Variant {variant_index}: {e}")
traceback.print_exc()
continue # Continue to next variant
return any_poster_attempted
def main():
"""主函数入口"""
config_file = {