From 3edd58dd4aaa2cd1c78df9b3bf04fd0975c606f6 Mon Sep 17 00:00:00 2001 From: shengyudong Date: Mon, 2 Feb 2026 18:16:56 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0Gemini=E7=94=9F=E5=9B=BE?= =?UTF-8?q?=E5=8A=9F=E8=83=BD=E5=88=B0article=5Fauto=5Fimage=5Fmatching.py?= =?UTF-8?q?=20-=20=E6=9C=AA=E5=8C=B9=E9=85=8D=E6=88=90=E5=8A=9F=E6=97=B6?= =?UTF-8?q?=E8=87=AA=E5=8A=A8=E7=94=9F=E6=88=90=E5=9B=BE=E7=89=87?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- article_auto_image_matching.py | 319 ++++++++++++++++++++++++++++++++- 1 file changed, 315 insertions(+), 4 deletions(-) diff --git a/article_auto_image_matching.py b/article_auto_image_matching.py index 643c6ec..4794a8f 100644 --- a/article_auto_image_matching.py +++ b/article_auto_image_matching.py @@ -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}"