From 0a048da5dd88281582f727afe19cb7f23b51aa2e Mon Sep 17 00:00:00 2001 From: jinye_huang Date: Tue, 22 Apr 2025 14:30:25 +0800 Subject: [PATCH] =?UTF-8?q?=E8=A7=A3=E8=80=A6=E4=BA=86=E6=B5=B7=E6=8A=A5?= =?UTF-8?q?=E7=94=9F=E6=88=90=E4=B8=AD=E7=9A=84=E9=83=A8=E5=88=86=E5=8A=9F?= =?UTF-8?q?=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- example_config.json | 3 +- main.py | 245 +++++++++------------------------------ utils/tweet_generator.py | 209 +++++++++++++++++++++++++++++++++ 3 files changed, 267 insertions(+), 190 deletions(-) diff --git a/example_config.json b/example_config.json index e0c6e94..e05a748 100644 --- a/example_config.json +++ b/example_config.json @@ -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 } \ No newline at end of file diff --git a/main.py b/main.py index ae9e28d..e2b22e1 100644 --- a/main.py +++ b/main.py @@ -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"]) + # 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 validity once + image_base_dir = config.get("image_base_dir", None) + if not image_base_dir or not os.path.isdir(image_base_dir): + 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 AI Agent: {e}. Cannot proceed with content generation.") + print(f"Error initializing resources for content/poster generation: {e}") traceback.print_exc() - return # Cannot continue without AI agent - - # Check image base directory - 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 - 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 --- + if ai_agent: ai_agent.close() + return + + # --- 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 = [] - - # --- 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}") + print(f"\nProcessing Topic {topic_index}/{total_topics}: {topic.get('title', 'N/A')}") + + # 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 + ) + + # 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" 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 - ) - - 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}") - 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 diff --git a/utils/tweet_generator.py b/utils/tweet_generator.py index 210e7eb..489cdc0 100644 --- a/utils/tweet_generator.py +++ b/utils/tweet_generator.py @@ -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 = {