yunqueai/baijiahao_text_to_image
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

添加Gemini生图功能到article_auto_image_matching.py - 未匹配成功时自动生成图片

Browse Source
This commit is contained in:
shengyudong
2026-02-02 18:16:56 +08:00
parent b9f65e2f9b
commit 3edd58dd4a
1 changed files with 315 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

319
article_auto_image_matching.py
View File

@@ -364,9 +364,305 @@ class ArticleImageMatcher:
self.log_to_database('ERROR', error_msg, traceback.format_exc())
return False
def generate_image_with_gemini(self, prompt: str, article_tags: List[str], article_id: int) -> Optional[str]:
"""
使用Gemini生成图片并上传到服务器
Args:
prompt: 图片生成提示词
article_tags: 文章标签列表，用于查询department和keywords
article_id: 文章ID，用于关联图片
Returns:
上传后的图片URL，失败返回None
"""
try:
# 导入必要的库
from google import genai
from google.genai.client import HttpOptions
client = genai.Client(
http_options=HttpOptions(base_url="https://work.poloapi.com"),
api_key="sk-V4tPnDgzFPa7nxWrvKnNJsW8ZcBXXPuGmjfgvPVRnwpHoeob"
)
logger.info(f"正在调用Gemini API生成图片，提示词: {prompt[:50]}...")
# 生成内容
response = client.models.generate_content(
model="gemini-3-pro-image-preview",
contents=[prompt],
)
# 检查是否有候选答案
if not response.candidates:
raise Exception("Gemini API未返回任何候选答案")
# 处理响应
candidate = response.candidates[0]
if not candidate.content or not candidate.content.parts:
raise Exception("Gemini API返回的候选答案中没有内容部分")
for part in candidate.content.parts:
if hasattr(part, 'inline_data') and part.inline_data is not None:
image_data = part.inline_data
if image_data.data is not None:
# 生成唯一的文件名（基于时间戳）
timestamp_ms = int(time.time() * 1000)
image_filename = f"{timestamp_ms}.png"
today_date = datetime.now().strftime("%Y%m%d")
image_url_path = f"{today_date}/{image_filename}"
temp_filename = f"temp_generated_image_{timestamp_ms}.png"
# 保存图片数据到临时文件
with open(temp_filename, 'wb') as f:
f.write(image_data.data)
logger.info(f"Gemini生成图片成功: {temp_filename}")
# 先将图片信息插入数据库
image_info = self.insert_generated_image_to_db(image_filename, image_url_path, article_tags)
if not image_info:
raise Exception("插入图片信息到数据库失败")
logger.info(f"图片信息已插入数据库，tag_image_id: {image_info['tag_image_id']}, image_id: {image_info['image_id']}")
# 使用tag_image_id上传图片到服务器
uploaded_url = self.upload_image_to_server(temp_filename, image_info['tag_image_id'])
# 将文章与图片的关联信息插入ai_article_images表
article_image_id = self.insert_article_image_relation_for_generated(
article_id=article_id,
image_id=image_info['image_id'],
image_url=image_info['image_url'],
image_thumb_url=image_info['image_thumb_url'],
tag_image_id=image_info['tag_image_id'],
keywords_id=image_info['keywords_id'],
keywords_name=image_info['keywords_name'],
department_id=image_info['department_id'],
department_name=image_info['department_name'],
image_source=0 # 默认值
)
if article_image_id:
logger.info(f"文章图片关联信息已创建，ai_article_images.id: {article_image_id}")
# 删除临时文件
os.remove(temp_filename)
logger.info(f"图片已上传到服务器: {uploaded_url}")
return uploaded_url
raise Exception("Gemini API未返回有效的图片数据")
except ImportError:
logger.error("错误：未安装google-genai库，请运行 'pip install google-genai' 进行安装")
return None
except Exception as e:
error_msg = f"Gemini生成图片异常: {e}"
logger.error(error_msg)
self.log_to_database('ERROR', error_msg, traceback.format_exc())
return None
def insert_generated_image_to_db(self, image_name: str, image_url: str, article_tags: List[str]) -> Optional[Dict]:
"""
将Gemini生成的图片信息插入数据库
Args:
image_name: 图片文件名
image_url: 图片URL路径
article_tags: 文章标签列表
Returns:
包含插入信息的字典
"""
try:
connection = self.db_manager.get_connection()
try:
with connection.cursor(pymysql.cursors.DictCursor) as cursor:
# 根据文章标签查询ai_image_tags表
if article_tags:
query = """
SELECT department_name, keywords_name, department_id, keywords_id, tag_id
FROM ai_image_tags
WHERE tag_name = %s
LIMIT 1
"""
cursor.execute(query, (article_tags[0],))
tag_info = cursor.fetchone()
if tag_info:
department = tag_info['department_name']
keywords = tag_info['keywords_name']
department_id = tag_info['department_id']
keywords_id = tag_info['keywords_id']
tag_id = tag_info['tag_id']
tag_name = article_tags[0]
else:
department = "AI生成"
keywords = "AI图片"
department_id = 1
keywords_id = 1
tag_id = 1
tag_name = article_tags[0] if article_tags else "AI生成"
else:
department = "AI生成"
keywords = "AI图片"
department_id = 1
keywords_id = 1
tag_id = 1
tag_name = "AI生成"
# 插入ai_images表
insert_image_query = """
INSERT INTO ai_images
(image_name, image_url, image_thumb_url, department, keywords, image_type, upload_user_id, status)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
"""
cursor.execute(insert_image_query, (
image_name, image_url, '', department, keywords, 'medical', 1, 'active'
))
image_id = cursor.lastrowid
logger.info(f"图片信息已插入ai_images表，image_id: {image_id}")
# 插入ai_image_tags表
insert_tag_query = """
INSERT INTO ai_image_tags
(image_id, image_name, image_url, image_thumb_url, tag_id, tag_name,
keywords_id, keywords_name, department_id, department_name,
image_source, created_user_id, image_attached_article_count)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
"""
cursor.execute(insert_tag_query, (
image_id, image_name, image_url, '', tag_id, tag_name,
keywords_id, keywords, department_id, department,
3, 1, 0 # image_source: 3表示AI生成
))
tag_image_id = cursor.lastrowid
logger.info(f"图片标签信息已插入ai_image_tags表，tag_image_id: {tag_image_id}")
connection.commit()
return {
'tag_image_id': tag_image_id,
'image_id': image_id,
'image_url': image_url,
'image_thumb_url': '',
'keywords_id': keywords_id,
'keywords_name': keywords,
'department_id': department_id,
'department_name': department
}
finally:
connection.close()
except Exception as e:
logger.error(f"插入图片信息到数据库失败: {e}")
return None
def upload_image_to_server(self, image_path: str, tag_image_id: int) -> str:
"""
上传图片到服务器
Args:
image_path: 本地图片路径
tag_image_id: 图片标签ID
Returns:
服务器上的图片URL
"""
base_url = "http://47.99.184.230:8324"
jwt_token = self.login_and_get_jwt_token(base_url)
if not jwt_token:
raise Exception("获取JWT token失败，无法上传图片")
upload_url = f"{base_url}/api/images/upload"
headers = {'Authorization': f'Bearer {jwt_token}'}
with open(image_path, 'rb') as image_file:
files = {'file': image_file}
data = {'tag_image_id': tag_image_id}
response = requests.post(upload_url, headers=headers, files=files, data=data)
logger.info(f"图片上传响应状态码: {response.status_code}")
if response.status_code == 200:
result = response.json()
if result.get('code') == 200:
return result['data']['http_image_url']
else:
raise Exception(f"图片上传失败: {result.get('message', '未知错误')}")
else:
raise Exception(f"图片上传请求失败，状态码: {response.status_code}")
def login_and_get_jwt_token(self, base_url: str) -> Optional[str]:
"""登录获取JWT token"""
login_url = f"{base_url}/api/auth/login"
login_data = {"username": "user010", "password": "@5^2W6R7"}
try:
response = requests.post(login_url, json=login_data, headers={'Content-Type': 'application/json'})
if response.status_code == 200:
result = response.json()
if result.get('code') == 200:
logger.info("JWT token获取成功")
return result['data']['token']
logger.error(f"登录失败: {response.status_code}")
return None
except Exception as e:
logger.error(f"登录异常: {e}")
return None
def insert_article_image_relation_for_generated(self, article_id: int, image_id: int, image_url: str,
image_thumb_url: str, tag_image_id: int, keywords_id: int,
keywords_name: str, department_id: int, department_name: str,
image_source: int = 0) -> Optional[int]:
"""
将文章与生成图片的关联信息插入ai_article_images表
"""
try:
connection = self.db_manager.get_connection()
try:
with connection.cursor(pymysql.cursors.DictCursor) as cursor:
# 查询当前文章下已有图片的最大sort_order
query_max_sort = """
SELECT COALESCE(MAX(sort_order), 0) as max_sort_order
FROM ai_article_images
WHERE article_id = %s
"""
cursor.execute(query_max_sort, (article_id,))
result = cursor.fetchone()
max_sort_order = result['max_sort_order'] if result else 0
new_sort_order = max_sort_order + 1
# 插入ai_article_images表
insert_query = """
INSERT INTO ai_article_images
(article_id, image_id, image_url, image_thumb_url, image_tag_id, sort_order,
keywords_id, keywords_name, department_id, department_name, image_source)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
"""
cursor.execute(insert_query, (
article_id, image_id, image_url, image_thumb_url, tag_image_id, new_sort_order,
keywords_id, keywords_name, department_id, department_name, image_source
))
article_image_id = cursor.lastrowid
logger.info(f"文章图片关联信息已插入ai_article_images表，id: {article_image_id}")
connection.commit()
return article_image_id
finally:
connection.close()
except Exception as e:
logger.error(f"插入文章图片关联信息失败: {e}")
return None
def match_article_with_images(self, article_data: Dict, available_images: List[Dict]) -> bool:
"""
为单篇文章匹配图片
为单篇文章匹配图片，未成功匹配时调用Gemini生图
Args:
article_data: 文章数据
@@ -377,6 +673,7 @@ class ArticleImageMatcher:
"""
article_id = article_data['article_id']
article_title = article_data.get('title', '')
article_content = article_data.get('content', '')
coze_tag = article_data.get('coze_tag', '')
try:
@@ -419,9 +716,23 @@ class ArticleImageMatcher:
else:
return False
else:
logger.info(f"文章 {article_id} 未找到合适的匹配图片")
self.log_to_database('WARNING', f"文章未找到匹配图片", f"文章ID: {article_id}")
return False
# 未找到合适的匹配图片，使用Gemini生成图片
logger.info(f"文章 {article_id} 未找到合适的匹配图片，调用Gemini生成图片")
self.log_to_database('WARNING', f"文章未找到匹配图片，尝试生成图片", f"文章ID: {article_id}")
# 构建生成提示词
prompt = f"与'{article_title}'相关的插图，标签: {', '.join(article_tags)}"
generated_image_url = self.generate_image_with_gemini(prompt, article_tags, article_id)
if generated_image_url:
logger.info(f"文章 {article_id} 成功生成图片: {generated_image_url}")
self.log_to_database('INFO', f"文章生成图片成功",
f"文章ID: {article_id}, 图片URL: {generated_image_url}")
return True
else:
logger.error(f"文章 {article_id} 生成图片失败")
self.log_to_database('ERROR', f"文章生成图片失败", f"文章ID: {article_id}")
return False
except Exception as e:
error_msg = f"匹配文章图片异常 - 文章ID: {article_id}, 错误: {e}"