200 lines
9.5 KiB
Python
200 lines
9.5 KiB
Python
import os
|
|
import time
|
|
from datetime import datetime
|
|
import argparse
|
|
import sys
|
|
import traceback
|
|
|
|
from core.ai_agent import AI_Agent
|
|
from core.topic_parser import TopicParser
|
|
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 prepare_topic_generation, generate_topics, generate_single_content
|
|
import random
|
|
|
|
TEXT_POSBILITY = 0.3
|
|
|
|
def main():
|
|
config_file = {
|
|
"date": "4月24日",
|
|
"num": 10,
|
|
"model": "qwenQWQ",
|
|
"api_url": "vllm",
|
|
"api_key": "EMPTY",
|
|
"topic_system_prompt": "/root/autodl-tmp/TravelContentCreator/SelectPrompt/systemPrompt.txt",
|
|
"topic_user_prompt": "/root/autodl-tmp/TravelContentCreator/SelectPrompt/userPrompt.txt",
|
|
"content_system_prompt": "/root/autodl-tmp/TravelContentCreator/genPrompts/systemPrompt.txt",
|
|
"resource_dir": [{
|
|
"type": "Object",
|
|
"num": 4,
|
|
"file_path": ["/root/autodl-tmp/TravelContentCreator/resource/Object/景点信息-尚书第.txt",
|
|
"/root/autodl-tmp/TravelContentCreator/resource/Object/景点信息-明清园.txt",
|
|
"/root/autodl-tmp/TravelContentCreator/resource/Object/景点信息-泰宁古城.txt",
|
|
"/root/autodl-tmp/TravelContentCreator/resource/Object/景点信息-甘露寺.txt"
|
|
]},
|
|
{
|
|
"type": "Product",
|
|
"num": 0,
|
|
"file_path": []
|
|
}
|
|
],
|
|
"prompts_dir": "/root/autodl-tmp/TravelContentCreator/genPrompts",
|
|
"output_dir": "/root/autodl-tmp/TravelContentCreator/result",
|
|
"variants": 10,
|
|
"topic_temperature": 0.2,
|
|
"content_temperature": 0.3
|
|
}
|
|
|
|
if True:
|
|
# 1. 首先生成选题
|
|
ai_agent, system_prompt, user_prompt, output_dir = prepare_topic_generation(
|
|
config_file["date"], config_file["num"], config_file["topic_system_prompt"], config_file["topic_user_prompt"],
|
|
config_file["api_url"], config_file["model"], config_file["api_key"], config_file["prompts_dir"], config_file["resource_dir"], config_file["output_dir"]
|
|
)
|
|
|
|
run_id, tweet_topic_record = generate_topics(
|
|
ai_agent, system_prompt, user_prompt, config_file["output_dir"],
|
|
config_file["topic_temperature"], 0.5, 1.5
|
|
)
|
|
|
|
output_dir = os.path.join(config_file["output_dir"], run_id)
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
tweet_topic_record.save_topics(os.path.join(output_dir, "tweet_topic.json"))
|
|
tweet_topic_record.save_prompt(os.path.join(output_dir, "tweet_prompt.txt"))
|
|
ai_agent.close()
|
|
# raise Exception("选题生成失败，退出程序")
|
|
if not run_id or not tweet_topic_record:
|
|
print("选题生成失败，退出程序")
|
|
return
|
|
|
|
# 2. 然后生成内容
|
|
print("\n开始根据选题生成内容...")
|
|
|
|
# 加载内容生成的系统提示词
|
|
content_system_prompt = ResourceLoader.load_system_prompt(config_file["content_system_prompt"])
|
|
|
|
|
|
if not content_system_prompt:
|
|
print("内容生成系统提示词为空，使用选题生成的系统提示词")
|
|
content_system_prompt = system_prompt
|
|
|
|
# 直接使用同一个AI Agent实例
|
|
for i in range(len(tweet_topic_record.topics_list)):
|
|
tweet_content_list = []
|
|
for j in range(config_file["variants"]):
|
|
time.sleep(random.random())
|
|
ai_agent = AI_Agent(config_file["api_url"], config_file["model"], config_file["api_key"])
|
|
tweet_content, gen_result = generate_single_content(
|
|
ai_agent, content_system_prompt, tweet_topic_record.topics_list[i],
|
|
config_file["prompts_dir"], config_file["resource_dir"],
|
|
output_dir, run_id, i+1, j+1, config_file["content_temperature"]
|
|
)
|
|
tweet_content_list.append(tweet_content.get_json_file())
|
|
ai_agent.close()
|
|
if not tweet_content:
|
|
print(f"生成第{i+1}篇文章的第{j+1}个变体失败，跳过")
|
|
continue
|
|
object_name = tweet_topic_record.topics_list[i]["object"]
|
|
try:
|
|
object_name = object_name.split(".")[0]
|
|
except:
|
|
pass
|
|
try:
|
|
object_name = object_name.split("景点信息-")[1]
|
|
except:
|
|
pass
|
|
# 处理对象名称中可能包含的"+"等特殊字符
|
|
# if "+" in object_name:
|
|
# # 取第一个景点名称作为主要对象
|
|
# object_name = object_name.split("+")[0].strip()
|
|
# print(f"对象名称包含多个景点，使用第一个景点：{object_name}")
|
|
|
|
# 检查图片路径是否存在
|
|
img_dir_path = f"/root/autodl-tmp/sanming_img/modify/{object_name}"
|
|
if not os.path.exists(img_dir_path):
|
|
print(f"图片目录不存在：{img_dir_path}，跳过该对象")
|
|
continue
|
|
info_directory = [
|
|
f"/root/autodl-tmp/sanming_img/相机/{object_name}/description.txt"
|
|
]
|
|
# 检查描述文件是否存在
|
|
if not os.path.exists(info_directory[0]):
|
|
print(f"描述文件不存在：{info_directory[0]}，使用生成的内容替代")
|
|
info_directory = []
|
|
|
|
poster_num = config_file["variants"]
|
|
|
|
input_dir = img_dir_path # 使用前面检查过的目录路径
|
|
target_size = (900, 1200)
|
|
result_path = []
|
|
|
|
content_gen = contentGen.ContentGenerator()
|
|
response = content_gen.run(info_directory, poster_num, tweet_content_list)
|
|
print(response)
|
|
try:
|
|
poster_config_summary = posterGen.PosterConfig(response)
|
|
for j_index in range(config_file["variants"]):
|
|
poster_config = poster_config_summary.get_config_by_index(j_index)
|
|
img_dir = os.path.join(output_dir, f"{i+1}_{j_index+1}")
|
|
try:
|
|
# 创建输出目录
|
|
collage_output_dir = os.path.join(img_dir, "collage_img")
|
|
os.makedirs(collage_output_dir, exist_ok=True)
|
|
poster_output_dir = os.path.join(img_dir, "poster")
|
|
os.makedirs(poster_output_dir, exist_ok=True)
|
|
|
|
# 处理图片目录
|
|
img_list = simple_collage.process_directory(
|
|
input_dir,
|
|
target_size=target_size,
|
|
output_count=1,
|
|
output_dir=collage_output_dir
|
|
)
|
|
print(img_list)
|
|
|
|
if not img_list or len(img_list) == 0:
|
|
print(f"未能生成拼贴图片，跳过海报生成")
|
|
continue
|
|
|
|
# 生成海报
|
|
poster_gen = posterGen.PosterGenerator()
|
|
|
|
if random.random() < TEXT_POSBILITY:
|
|
text_data = {
|
|
"title": f"{poster_config['main_title']}",
|
|
"subtitle": "",
|
|
"additional_texts": [
|
|
{"text": f"{poster_config['texts'][0]}", "position": "bottom", "size_factor": 0.5},
|
|
{"text": f"{poster_config['texts'][1]}", "position": "bottom", "size_factor": 0.5}
|
|
]
|
|
}
|
|
else:
|
|
text_data = {
|
|
"title": f"{poster_config['main_title']}",
|
|
"subtitle": "",
|
|
"additional_texts": [
|
|
{"text": f"{poster_config['texts'][0]}", "position": "bottom", "size_factor": 0.5},
|
|
# {"text": f"{poster_config['texts'][1]}", "position": "bottom", "size_factor": 0.5}
|
|
]
|
|
}
|
|
print(text_data)
|
|
img_path = img_list[0]['path']
|
|
print(f"使用图片路径: {img_path}")
|
|
output_path = os.path.join(poster_output_dir, f"poster.jpg")
|
|
result_path.append(poster_gen.create_poster(img_path, text_data, output_path))
|
|
except Exception as e:
|
|
print(f"海报生成过程中出错: {e}")
|
|
traceback.print_exc() # 打印完整的堆栈跟踪信息
|
|
continue
|
|
except Exception as e:
|
|
print(f"配置解析失败，跳过当前项目: {e}")
|
|
traceback.print_exc() # 打印完整的堆栈跟踪信息
|
|
continue
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
|
|
|
|