377 lines
15 KiB
Python
377 lines
15 KiB
Python
import os
|
||
from openai import OpenAI
|
||
import pandas as pd
|
||
from datetime import datetime
|
||
import cv2
|
||
import time
|
||
import random
|
||
import json
|
||
|
||
class ContentGenerator:
|
||
def __init__(self,
|
||
model_name="qwenQWQ",
|
||
api_base_url="http://localhost:8000/v1",
|
||
api_key="EMPTY",
|
||
output_dir="/root/autodl-tmp/poster_generate_result",
|
||
):
|
||
"""
|
||
初始化海报生成器
|
||
|
||
参数:
|
||
csv_path: CSV文件路径
|
||
img_base_dir: 图片基础目录
|
||
output_dir: 输出结果保存目录
|
||
model_name: 使用的模型名称
|
||
api_base_url: API基础URL
|
||
api_key: API密钥
|
||
"""
|
||
self.output_dir = output_dir
|
||
self.model_name = model_name
|
||
self.api_base_url = api_base_url
|
||
self.api_key = api_key
|
||
|
||
# 不在初始化时创建OpenAI客户端，而是在需要时临时创建
|
||
self.client = None
|
||
|
||
# 初始化数据
|
||
self.df = None
|
||
self.all_images_info = []
|
||
self.structured_prompt = ""
|
||
self.current_img_info = None
|
||
self.add_description = ""
|
||
|
||
self.temperature = 0.7
|
||
self.top_p = 0.8
|
||
self.presence_penalty = 1.2
|
||
|
||
|
||
def load_infomation(self, info_directory_path):
|
||
"""
|
||
加载额外描述文件
|
||
|
||
参数:
|
||
info_directory_path: 信息目录路径
|
||
"""
|
||
## 读取路径下的所有文件
|
||
for path in info_directory_path:
|
||
# file_extend_path = os.path.join(self.img_base_dir, info_directory, "description.txt")
|
||
try:
|
||
with open(path, "r") as f:
|
||
self.add_description += f.read()
|
||
except:
|
||
self.add_description = ""
|
||
|
||
def _create_temp_client(self):
|
||
"""
|
||
创建临时OpenAI客户端
|
||
|
||
返回:
|
||
OpenAI客户端实例
|
||
"""
|
||
try:
|
||
import gc
|
||
# 强制垃圾回收
|
||
gc.collect()
|
||
|
||
# 创建新的客户端实例
|
||
print(f"创建临时OpenAI客户端，API URL: {self.api_base_url}")
|
||
client = OpenAI(
|
||
base_url=self.api_base_url,
|
||
api_key=self.api_key
|
||
)
|
||
return client
|
||
except Exception as e:
|
||
print(f"创建OpenAI客户端失败: {str(e)}")
|
||
return None
|
||
|
||
def _close_client(self, client):
|
||
"""
|
||
关闭并清理OpenAI客户端
|
||
|
||
参数:
|
||
client: 需要关闭的客户端实例
|
||
"""
|
||
try:
|
||
# OpenAI客户端可能没有显式的close方法
|
||
# 将引用设为None，让Python垃圾回收处理
|
||
client = None
|
||
import gc
|
||
gc.collect()
|
||
print("OpenAI客户端资源已释放")
|
||
except Exception as e:
|
||
print(f"关闭客户端失败: {str(e)}")
|
||
|
||
|
||
def split_content(self, content):
|
||
"""
|
||
分割结果, 返回去除
|
||
```json
|
||
```的json内容
|
||
|
||
参数:
|
||
content: 需要分割的内容
|
||
|
||
返回:
|
||
分割后的json内容
|
||
"""
|
||
return json.loads(content.split("```json")[1].split("```")[0])
|
||
|
||
def generate_posters(self, poster_num, tweet_content, system_prompt=None, max_retries=3):
|
||
"""
|
||
生成海报内容
|
||
|
||
参数:
|
||
poster_num: 海报数量
|
||
tweet_content: 推文内容
|
||
system_prompt: 系统提示，默认为None则使用预设提示
|
||
max_retries: 最大重试次数
|
||
|
||
返回:
|
||
生成的海报内容
|
||
"""
|
||
if system_prompt is None:
|
||
system_prompt = """
|
||
你是一名资深海报设计师，有丰富的爆款海报设计经验，你现在要为旅游景点做宣传，在小红书上发布大量宣传海报。你的主要工作目标有2个：
|
||
1、你要根据我给你的图片描述和笔记推文内容，设计图文匹配的海报。
|
||
2、为海报设计文案，文案的<第一个小标题>和<第二个小标题>之间你需要检查是否逻辑关系合理，你将通过先去生成<第二个小标题>关于景区亮点的部分，再去综合判断<第一个小标题>应该如何搭配组合更符合两个小标题的逻辑再生成<第一个小标题>。
|
||
|
||
其中，生成三类标题文案的通用性要求如下：
|
||
1、生成的<大标题>字数必须小于8个字符
|
||
2、生成的<第一个小标题>字数和<第二个小标题>字数，两者都必须小8个字符
|
||
3、标题和文案都应符合中国社会主义核心价值观
|
||
|
||
接下来先开始生成<大标题>部分，由于海报是用来宣传旅游景点，生成的海报<大标题>必须使用以下8种格式之一：
|
||
①地名＋景点名（例如福建厦门鼓浪屿/厦门鼓浪屿）；
|
||
②地名+景点名+plog；
|
||
③拿捏+地名+景点名；
|
||
④地名+景点名+攻略；
|
||
⑤速通+地名+景点名
|
||
⑥推荐！+地名+景点名
|
||
⑦勇闯！+地名+景点名
|
||
⑧收藏！+地名+景点名
|
||
你需要随机挑选一种格式生成对应景点的文案，但是格式除了上面8种不可以有其他任何格式；同时尽量保证每一种格式出现的频率均衡。
|
||
接下来先去生成<第二个小标题>，<第二个小标题>文案的创作必须遵循以下原则：
|
||
请根据笔记内容和图片识别，用极简的文字概括这篇笔记和图片中景点的特色亮点，其中你可以参考以下词汇进行创作，这段文案字数控制6-8字符以内；
|
||
|
||
特色亮点可能会出现的词汇不完全举例：非遗、古建、绝佳山水、祈福圣地、研学圣地、解压天堂、中国小瑞士、秘境竹筏游等等类型词汇
|
||
|
||
接下来再去生成<第一个小标题>，<第一个小标题>文案的创作必须遵循以下原则：
|
||
这部分文案创作公式有5种，分别为：
|
||
①<受众人群画像>+<痛点词>
|
||
②<受众人群画像>
|
||
③<痛点词>
|
||
④<受众人群画像>+ | +<痛点词>
|
||
⑤<痛点词>+ | +<受众人群画像>
|
||
请你根据实际笔记内容，结合这部分文案创作公式，需要结合<受众人群画像>和<痛点词>时，必须根据<第二个小标题>的景点特征和所对应的完整笔记推文内容主旨，特征挑选对应<受众人群画像>和<痛点词>。
|
||
|
||
我给你提供受众人群画像库和痛点词库如下：
|
||
1、受众人群画像库：情侣党、亲子游、合家游、银发族、亲子研学、学生党、打工人、周边游、本地人、穷游党、性价比、户外人、美食党、出片
|
||
2、痛点词库：3天2夜、必去、看了都哭了、不能错过、一定要来、问爆了、超全攻略、必打卡、强推、懒人攻略、必游榜、小众打卡、狂喜等等。
|
||
|
||
你需要为每个请求至少生成{poster_num}个海报设计。请使用JSON格式输出结果，结构如下：
|
||
|
||
```json
|
||
[
|
||
{
|
||
"index": 1,
|
||
"main_title": "主标题内容",
|
||
"texts": ["第一个小标题", "第二个小标题"]
|
||
},
|
||
{
|
||
"index": 2,
|
||
"main_title": "主标题内容",
|
||
"texts": ["第一个小标题", "第二个小标题"]
|
||
}
|
||
// ... 更多海报
|
||
]
|
||
```
|
||
|
||
确保生成的数量与用户要求的数量一致。只生成上述JSON格式内容，不要有其他任何额外内容。
|
||
"""
|
||
|
||
user_content = f"""
|
||
海报数量：{poster_num};
|
||
景区介绍：{self.add_description};
|
||
推文内容：{tweet_content};
|
||
"""
|
||
|
||
# 最终响应内容
|
||
full_response = ""
|
||
|
||
# 创建临时客户端
|
||
temp_client = self._create_temp_client()
|
||
|
||
# 如果创建客户端失败，直接使用备用方案
|
||
if temp_client is None:
|
||
print("创建OpenAI客户端失败，使用备用方案生成内容")
|
||
return 404
|
||
else:
|
||
# 添加重试机制
|
||
for retry in range(max_retries):
|
||
try:
|
||
print(f"尝试连接API (尝试 {retry+1}/{max_retries})...")
|
||
|
||
# 计算退避时间（指数退避策略）：0, 2, 4, 8, 16...秒
|
||
if retry > 0:
|
||
backoff_time = min(2 ** (retry - 1) * 2, 30) # 最大等待30秒
|
||
print(f"等待 {backoff_time} 秒后重试...")
|
||
time.sleep(backoff_time)
|
||
|
||
# 设置超时时间随重试次数递增
|
||
timeout = 30 + (retry * 30) # 30, 60, 90, ...秒
|
||
|
||
chat_response = temp_client.chat.completions.create(
|
||
model=self.model_name,
|
||
messages=[
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": user_content}
|
||
],
|
||
stream=True,
|
||
temperature=self.temperature,
|
||
top_p=self.top_p,
|
||
presence_penalty=self.presence_penalty,
|
||
timeout=timeout # 设置请求超时时间
|
||
)
|
||
|
||
# 获取响应内容
|
||
for chunk in chat_response:
|
||
if chunk.choices and len(chunk.choices) > 0 and chunk.choices[0].delta.content is not None:
|
||
content = chunk.choices[0].delta.content
|
||
full_response += content
|
||
print(content, end="", flush=True)
|
||
if chunk.choices and len(chunk.choices) > 0 and chunk.choices[0].finish_reason == "stop":
|
||
break
|
||
|
||
print("\n") # 输出完成后换行
|
||
|
||
# 成功获取响应，跳出重试循环
|
||
break
|
||
|
||
except Exception as e:
|
||
error_msg = str(e)
|
||
print(f"API连接错误 (尝试 {retry+1}/{max_retries}): {error_msg}")
|
||
|
||
# 如果已经达到最大重试次数
|
||
if retry + 1 >= max_retries:
|
||
print("已达到最大重试次数，使用备用方案...")
|
||
# 生成备用内容（简单模板）
|
||
full_response = self._generate_fallback_content(poster_num)
|
||
else:
|
||
print(f"将在稍后重试，还剩 {max_retries - retry - 1} 次重试机会")
|
||
|
||
# 关闭临时客户端
|
||
self._close_client(temp_client)
|
||
|
||
# 生成时间戳
|
||
return full_response
|
||
|
||
def save_result(self, full_response):
|
||
"""
|
||
保存生成结果到文件
|
||
|
||
参数:
|
||
full_response: 生成的完整响应内容
|
||
|
||
返回:
|
||
结果文件路径
|
||
"""
|
||
# 生成时间戳
|
||
print(full_response)
|
||
date_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||
full_response = self.split_content(full_response)
|
||
# 创建结果文件路径
|
||
result_path = os.path.join(self.output_dir, f"{date_time}.json")
|
||
os.makedirs(os.path.dirname(result_path), exist_ok=True)
|
||
|
||
# 保存结果到文件
|
||
with open(result_path, "w", encoding="utf-8") as f:
|
||
json.dump(full_response, f, ensure_ascii=False)
|
||
|
||
print(f"结果已保存到: {result_path}")
|
||
return result_path
|
||
|
||
def run(self, info_directory, poster_num, tweet_content):
|
||
"""
|
||
运行完整的海报生成流程
|
||
|
||
参数:
|
||
info_directory: 信息目录
|
||
poster_num: 海报数量
|
||
tweet_content: 推文内容
|
||
|
||
返回:
|
||
结果保存路径
|
||
"""
|
||
## 读取资料文件
|
||
self.load_infomation(info_directory)
|
||
|
||
## 生成海报内容
|
||
full_response = self.generate_posters(poster_num, tweet_content)
|
||
|
||
if self.output_dir:
|
||
## 保存结果
|
||
return self.save_result(full_response)
|
||
else:
|
||
return full_response
|
||
|
||
def set_temperature(self, temperature):
|
||
self.temperature = temperature
|
||
|
||
def set_top_p(self, top_p):
|
||
self.top_p = top_p
|
||
|
||
def set_presence_penalty(self, presence_penalty):
|
||
self.presence_penalty = presence_penalty
|
||
|
||
def set_model_para(self, temperature, top_p, presence_penalty):
|
||
self.temperature = temperature
|
||
self.top_p = top_p
|
||
self.presence_penalty = presence_penalty
|
||
|
||
# def main():
|
||
# # 配置参数
|
||
# info_directory = [
|
||
# "/root/autodl-tmp/sanming_img/相机/甘露寺/description.txt"
|
||
# ] # 信息目录
|
||
# poster_num = 4 # 海报数量
|
||
|
||
# # 推文内容
|
||
# tweet_content = """<title>
|
||
# 🌿清明遛娃天花板！悬空古寺+非遗探秘
|
||
# </title>
|
||
# <content>
|
||
# 清明假期带娃哪里玩？泰宁甘露寺藏着明代建筑奇迹！一柱擎天的悬空阁楼+状元祈福传说，让孩子边玩边涨知识✨
|
||
|
||
# 🎒行程亮点：
|
||
# ✅ 安全科普第一站：讲解"一柱插地"千年不倒的秘密，用乐高积木模型让孩子理解力学原理
|
||
# ✅ 文化沉浸体验：穿汉服听"叶状元还愿建寺"故事，触摸3.38米粗的"状元柱"许愿
|
||
# ✅ 自然探索路线：连接金湖栈道徒步，观察丹霞地貌与古建筑的巧妙融合
|
||
|
||
# 📌实用攻略：
|
||
# 📍位置：福建省三明市泰宁县金湖西路（导航搜"甘露岩寺"）
|
||
# 🕒最佳时段：上午10点前抵达避开人流，下午可衔接参观明清园（80元/人）
|
||
# ⚠️注意事项：悬空栈道设置儿童安全绳租赁点，建议穿防滑鞋
|
||
|
||
# 💡亲子彩蛋：
|
||
# 1️⃣ 在"右鼓左钟"景观区玩声音实验，敲击不同岩石听回声差异
|
||
# 2️⃣ 领取任务卡完成"寻找建筑中的T形拱"小游戏，集章兑换非遗木雕书签
|
||
# 3️⃣ 结合清明节俗，用竹简模板书写祈福语系在古松枝头
|
||
|
||
# 周边推荐：游览完可直奔尚书第明代古民居群，对比不同时期建筑特色，晚餐推荐尝泰宁特色"灯盏糕"，亲子套票更划算！
|
||
|
||
# 清明带着孩子来场穿越850年的建筑探险，把课本里的力学知识变成触手可及的历史课堂！🌸
|
||
|
||
# #清明节周边游 #亲子科普游 #福建遛娃 #泰宁旅行攻略 #建筑启蒙
|
||
# </content>
|
||
# """
|
||
|
||
# # 创建海报生成器
|
||
# generator = ContentGenerator()
|
||
|
||
# # 运行生成流程
|
||
# generator.run(info_directory, poster_num, tweet_content)
|
||
|
||
|
||
# if __name__ == "__main__":
|
||
# main()
|