#!/usr/bin/env python # -*- coding: utf-8 -*- import os import json import shutil import csv import traceback import re import argparse from datetime import datetime import sqlite3 import logging import base64 # 配置日志 logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[logging.StreamHandler()] ) logger = logging.getLogger(__name__) # 内置数据库记录功能 def init_database(db_path): """初始化数据库，创建表结构""" try: conn = sqlite3.connect(db_path) conn.execute("PRAGMA foreign_keys = OFF") # 禁用外键约束 cursor = conn.cursor() # 创建内容表 cursor.execute(""" CREATE TABLE IF NOT EXISTS contents ( id INTEGER PRIMARY KEY AUTOINCREMENT, entry_id TEXT NOT NULL UNIQUE, output_txt_path TEXT, poster_path TEXT, article_json_path TEXT, product TEXT, object TEXT, date TEXT, logic TEXT, judge_status INTEGER, is_distributed INTEGER DEFAULT 0, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) # 创建索引 cursor.execute("CREATE INDEX IF NOT EXISTS idx_contents_entry_id ON contents(entry_id)") conn.commit() logger.info("数据库初始化成功") return conn except sqlite3.Error as e: logger.error(f"初始化数据库失败: {e}") return None def record_to_database( db_path, entry_id, output_txt_path=None, poster_path=None, article_json_path=None, product=None, object=None, date=None, logic=None, judge_status=None, is_distributed=0 ): """将内容记录到数据库""" try: # 检查数据库是否存在，如果不存在则初始化 if not os.path.exists(db_path): logger.info(f"数据库文件不存在: {db_path}，将自动创建") conn = init_database(db_path) if not conn: return False else: try: conn = sqlite3.connect(db_path) conn.execute("PRAGMA foreign_keys = OFF") # 禁用外键约束 except sqlite3.Error as e: logger.error(f"连接数据库失败: {e}") return False try: cursor = conn.cursor() # 准备数据 data = ( entry_id, output_txt_path or '', poster_path or '', article_json_path or '', product or '', object or '', date or '', logic or '', judge_status if judge_status is not None else None, is_distributed ) # 插入或更新内容 cursor.execute(""" INSERT OR REPLACE INTO contents (entry_id, output_txt_path, poster_path, article_json_path, product, object, date, logic, judge_status, is_distributed) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, data) conn.commit() logger.info(f"已将内容 {entry_id} 记录到数据库") return True except Exception as e: logger.error(f"记录内容到数据库失败: {e}") try: conn.rollback() except: pass return False finally: try: conn.close() except: pass except Exception as e: logger.error(f"记录提取内容时发生错误: {e}") return False def convert_json_to_txt_content(json_path, prefer_original=False): """ 读取 JSON 文件，提取标题、内容和标签，移除 Markdown 格式， 并返回格式化文本。 根据JSON文件中的状态字段决定使用什么内容： - 如果judged=True，使用审核后内容 - 如果judged=False，使用原始内容 所有情况都优先使用base64编码的字段，因为这些字段能正确保留特殊字符和换行符 Args: json_path: JSON文件路径 prefer_original: 参数保留但不再使用 """ print(f" - 正在读取 JSON: {json_path}") if not os.path.exists(json_path): print(f" - 警告: JSON 文件不存在: {json_path}") return None, f"文件未找到: {json_path}" try: with open(json_path, 'r', encoding='utf-8') as f_json: data = json.load(f_json) # 提取状态字段 judged = data.get('judged', False) print(f" - 文件状态: judged={judged}") # 初始化变量 title = None content = None original_title = None original_content = None tags = None original_tags = None # =================解码所有可能的base64字段================= # 解码标题和内容字段 if "title_base64" in data: try: title = base64.b64decode(data["title_base64"]).decode('utf-8') print(f" - 成功从base64解码审核后标题") except Exception as e: print(f" - 警告: 标题base64解码失败: {e}") if "content_base64" in data: try: content = base64.b64decode(data["content_base64"]).decode('utf-8') print(f" - 成功从base64解码审核后内容") except Exception as e: print(f" - 警告: 内容base64解码失败: {e}") # 解码原始标题和内容字段 if "original_title_base64" in data: try: original_title = base64.b64decode(data["original_title_base64"]).decode('utf-8') print(f" - 成功从base64解码原始标题") except Exception as e: print(f" - 警告: 原始标题base64解码失败: {e}") if "original_content_base64" in data: try: original_content = base64.b64decode(data["original_content_base64"]).decode('utf-8') print(f" - 成功从base64解码原始内容") except Exception as e: print(f" - 警告: 原始内容base64解码失败: {e}") # 解码标签字段 if "tags_base64" in data: try: tags = base64.b64decode(data["tags_base64"]).decode('utf-8') print(f" - 成功从base64解码审核后标签") except Exception as e: print(f" - 警告: 标签base64解码失败: {e}") if "original_tags_base64" in data: try: original_tags = base64.b64decode(data["original_tags_base64"]).decode('utf-8') print(f" - 成功从base64解码原始标签") except Exception as e: print(f" - 警告: 原始标签base64解码失败: {e}") # =================回退到非base64字段================= # 如果base64解码失败，尝试使用普通字段 if title is None and "title" in data: title = data["title"] print(f" - 使用普通字段标题") if content is None and "content" in data: content = data["content"] print(f" - 使用普通字段内容") if original_title is None and "original_title" in data: original_title = data["original_title"] print(f" - 使用普通字段原始标题") if original_content is None and "original_content" in data: original_content = data["original_content"] print(f" - 使用普通字段原始内容") if tags is None and "tags" in data: tags = data["tags"] print(f" - 使用普通字段标签") elif tags is None and "tag" in data: tags = data["tag"] print(f" - 使用普通tag字段作为标签") if original_tags is None and "original_tags" in data: original_tags = data["original_tags"] print(f" - 使用普通字段原始标签") # =================根据状态字段决定使用哪些内容================= final_title = None final_content = None final_tags = None # 简化逻辑：如果已审核，使用审核后内容；否则使用原始内容 if judged: print(f" - 使用审核后内容 (judged=True)") final_title = title final_content = content final_tags = tags else: print(f" - 使用原始内容 (judged=False)") final_title = original_title final_content = original_content final_tags = original_tags # 确保所有字段都有值 final_title = final_title or "未找到标题" final_content = final_content or "未找到内容" final_tags = final_tags or "未找到标签" # 移除Markdown格式，但保留换行符 content_no_format = re.sub(r'\*\*(.*?)\*\*', r'\1', final_content) # 组合输出文本，保留内容的所有换行符 result = final_title + "\n\n" + content_no_format if final_tags and final_tags != "未找到标签": result += "\n\n" + final_tags print(f" - 内容处理完成，最终文本长度: {len(result)} 字符") return result, None except json.JSONDecodeError: print(f" - 错误: JSON 格式无效: {json_path}") return None, f"无效的 JSON 格式: {json_path}" except Exception as e: print(f" - 错误: 处理 JSON 时出错: {e}") traceback.print_exc() return None, f"处理 JSON 时出错: {e}" def process_txt_content(txt_path): """ 直接读取TXT文件内容，移除Markdown格式，并返回处理后的文本。 Args: txt_path: TXT文件路径 Returns: tuple: (处理后的内容，错误信息) """ print(f" - 正在读取TXT: {txt_path}") if not os.path.exists(txt_path): print(f" - 警告: TXT文件不存在: {txt_path}") return None, f"文件未找到: {txt_path}" try: # 读取TXT文件内容 with open(txt_path, 'r', encoding='utf-8') as f_txt: content = f_txt.read() # 移除Markdown格式，但保留换行符 # 处理粗体 content_no_format = re.sub(r'\*\*(.*?)\*\*', r'\1', content) # 处理斜体 content_no_format = re.sub(r'\*(.*?)\*', r'\1', content_no_format) # 处理链接 [文本](链接) content_no_format = re.sub(r'\[(.*?)\]\(.*?\)', r'\1', content_no_format) # 处理标题 # 文本 content_no_format = re.sub(r'^#+ (.*?)$', r'\1', content_no_format, flags=re.MULTILINE) return content_no_format, None except Exception as e: print(f" - 错误: 处理TXT时出错: {e}") return None, f"处理TXT时出错: {e}" def load_topic_data(source_dir, run_id): """ 加载选题数据 Args: source_dir: 源目录路径 run_id: 运行ID Returns: dict: 以topic_index为键的选题数据字典 """ topic_file_path = os.path.join(source_dir, f"tweet_topic_{run_id}.json") topic_data = {} if os.path.exists(topic_file_path): try: with open(topic_file_path, 'r', encoding='utf-8') as f: topics = json.load(f) # 将选题数据转换为以index为键的字典 for topic in topics: index = topic.get("index") if index: topic_data[index] = topic print(f"成功加载选题数据，共{len(topic_data)}条") except Exception as e: print(f"加载选题数据时出错: {e}") else: print(f"警告: 未找到选题文件: {topic_file_path}") return topic_data def process_result_directory(source_dir, output_dir, run_id=None, prefer_original=False, db_path=None): """ 处理指定的结果目录，提取内容并渲染到输出目录。 Args: source_dir: 源目录路径，包含i_j子目录 output_dir: 输出目录路径 run_id: 可选的运行ID，如果不提供则使用源目录名 prefer_original: 是否优先使用原始内容，无视judge_success结果 db_path: 数据库路径，若不提供则使用默认路径 """ if not os.path.isdir(source_dir): print(f"错误: 源目录不存在: {source_dir}") return # 设置默认数据库路径 if db_path is None: db_path = '/root/autodl-tmp/TravelContentCreator/distribution.db' # 数据库是否启用 db_enabled = True # 创建输出目录 os.makedirs(output_dir, exist_ok=True) print(f"确保输出目录存在: {output_dir}") # 提取run_id if not run_id: run_id = os.path.basename(source_dir) # 加载选题数据 topic_data = load_topic_data(source_dir, run_id) # 创建CSV清单，添加选题相关字段 csv_path = os.path.join(output_dir, f"manifest_{run_id}.csv") csv_data = [ [ "EntryID", "TopicIndex", "VariantIndex", "Date", "Logic", "Object", "Product", "ProductLogic", "Style", "StyleLogic", "TargetAudience", "TargetAudienceLogic", "SourcePath", "ArticleJsonPath", "OutputTxtPath", "PosterPath", "AdditionalImagesCount", "Status", "Details", "JudgeStatus", "ContentSource", "RecordedInDB", "IsDistributed" ] ] # 查找所有i_j目录 entry_pattern = re.compile(r"^(\d+)_(\d+)$") entries = [] for item in os.listdir(source_dir): item_path = os.path.join(source_dir, item) match = entry_pattern.match(item) if os.path.isdir(item_path) and match: entries.append(item) if not entries: print(f"警告: 在源目录中未找到任何i_j格式的子目录") return print(f"找到 {len(entries)} 个条目目录") # 处理每个条目 for entry in sorted(entries): entry_path = os.path.join(source_dir, entry) output_entry_path = os.path.join(output_dir, entry) print(f"\n处理条目: {entry}") # 解析topic_index和variant_index match = entry_pattern.match(entry) topic_index = match.group(1) variant_index = match.group(2) # 获取该话题的选题信息 topic_info = topic_data.get(topic_index, {}) # 创建记录 record = { "EntryID": entry, "TopicIndex": topic_index, "VariantIndex": variant_index, "Date": topic_info.get("date", ""), "Logic": topic_info.get("logic", ""), "Object": topic_info.get("object", ""), "Product": topic_info.get("product", ""), "ProductLogic": topic_info.get("product_logic", ""), "Style": topic_info.get("style", ""), "StyleLogic": topic_info.get("style_logic", ""), "TargetAudience": topic_info.get("target_audience", ""), "TargetAudienceLogic": topic_info.get("target_audience_logic", ""), "SourcePath": entry_path, "ArticleJsonPath": "", "OutputTxtPath": "", "PosterPath": "", "AdditionalImagesCount": 0, "Status": "Processing", "Details": "", "JudgeStatus": "", "ContentSource": "unknown", "RecordedInDB": "No", "IsDistributed": "No" } # 创建输出条目目录 try: os.makedirs(output_entry_path, exist_ok=True) except Exception as e: record["Status"] = "Failed" record["Details"] = f"创建输出目录失败: {e}" csv_data.append([record[col] for col in csv_data[0]]) print(f" - 错误: {record['Details']}") continue # 1. 处理article内容 - 优先使用JSON文件 output_txt_path = os.path.join(output_entry_path, "article.txt") record["OutputTxtPath"] = output_txt_path # 读取article.json json_path = os.path.join(entry_path, "article.json") record["ArticleJsonPath"] = json_path content_processed = False # 优先从JSON提取内容 if os.path.exists(json_path): try: # 从JSON文件提取审核状态 with open(json_path, 'r', encoding='utf-8') as f_json: article_data = json.load(f_json) # 提取judge_success状态 if "judge_success" in article_data: record["JudgeStatus"] = str(article_data["judge_success"]) elif "judged" in article_data: record["JudgeStatus"] = "已审核" if article_data["judged"] else "未审核" # 使用convert_json_to_txt_content函数处理JSON文件 processed_content, error = convert_json_to_txt_content(json_path, prefer_original) if error: print(f" - 警告: 从JSON提取内容失败: {error}") else: try: with open(output_txt_path, 'w', encoding='utf-8') as f_txt: f_txt.write(processed_content) print(f" - 成功从JSON提取并写入内容到: {output_txt_path}") record["ContentSource"] = "json_file" content_processed = True except Exception as e: print(f" - 警告: 写入从JSON提取的内容时出错: {e}") except Exception as e: print(f" - 警告: 处理JSON文件时出错: {e}") # 如果从JSON提取内容失败，尝试使用现有的TXT文件 if not content_processed: input_txt_path = os.path.join(entry_path, "article.txt") if os.path.exists(input_txt_path): processed_content, error = process_txt_content(input_txt_path) if error: record["Status"] = "Partial" record["Details"] += f"文章处理失败: {error}; " print(f" - 错误: {record['Details']}") else: try: with open(output_txt_path, 'w', encoding='utf-8') as f_txt: f_txt.write(processed_content) print(f" - 成功写入处理后的文本文件: {output_txt_path}") record["ContentSource"] = "txt_file" content_processed = True except Exception as e: record["Status"] = "Partial" record["Details"] += f"写入文本文件失败: {e}; " print(f" - 错误: {record['Details']}") else: record["Status"] = "Partial" record["Details"] += "无法从JSON或TXT获取内容; " print(f" - 警告: {record['Details']}") # 2. 处理海报图片 poster_dir = os.path.join(entry_path, "poster") poster_jpg_path = os.path.join(poster_dir, "poster.jpg") output_poster_path = os.path.join(output_entry_path, "poster.jpg") record["PosterPath"] = output_poster_path if os.path.exists(poster_jpg_path): try: shutil.copy2(poster_jpg_path, output_poster_path) print(f" - 成功复制海报图片: {output_poster_path}") except Exception as e: record["Status"] = "Partial" record["Details"] += f"复制海报图片失败: {e}; " print(f" - 错误: {record['Details']}") else: record["Status"] = "Partial" record["Details"] += "海报图片不存在; " print(f" - 警告: {record['Details']}") # 3. 处理额外图片 image_dir = os.path.join(entry_path, "image") output_image_dir = os.path.join(output_entry_path, "additional_images") if os.path.exists(image_dir) and os.path.isdir(image_dir): try: os.makedirs(output_image_dir, exist_ok=True) image_count = 0 for filename in os.listdir(image_dir): if filename.startswith("additional_") and filename.endswith(".jpg"): source_file = os.path.join(image_dir, filename) dest_file = os.path.join(output_image_dir, filename) # 复制图片 shutil.copy2(source_file, dest_file) image_count += 1 record["AdditionalImagesCount"] = image_count print(f" - 复制了 {image_count} 张额外图片到: {output_image_dir}") except Exception as e: record["Status"] = "Partial" record["Details"] += f"处理额外图片时出错: {e}; " print(f" - 错误: {record['Details']}") else: record["AdditionalImagesCount"] = 0 print(f" - 没有找到额外图片目录") # 更新状态 if record["Status"] == "Processing": record["Status"] = "Success" record["Details"] = "处理成功完成" # 4. 将内容记录到数据库 if db_enabled: try: # 准备judge_status值 if record["JudgeStatus"] == "True": judge_status = 1 elif record["JudgeStatus"] == "False": judge_status = 0 else: judge_status = None # 调用数据库记录函数 success = record_to_database( db_path, entry_id=record["EntryID"], output_txt_path=record["OutputTxtPath"], poster_path=record["PosterPath"], article_json_path=record["ArticleJsonPath"], product=record["Product"], object=record["Object"], date=record["Date"], logic=record["Logic"], judge_status=judge_status, is_distributed=0 # 默认为未分发 ) if success: record["RecordedInDB"] = "Yes" print(f" - 成功将内容记录到数据库") else: record["RecordedInDB"] = "Failed" print(f" - 警告: 内容记录到数据库失败") except Exception as e: record["RecordedInDB"] = "Error" print(f" - 错误: 记录到数据库时发生异常: {e}") traceback.print_exc() # 打印详细的异常堆栈 else: record["RecordedInDB"] = "Disabled" print(f" - 信息: 数据库记录功能已禁用") # 添加记录到CSV数据 csv_data.append([record[col] for col in csv_data[0]]) # 写入CSV清单 try: print(f"\n正在写入清单CSV: {csv_path}") with open(csv_path, 'w', newline='', encoding='utf-8-sig') as f_csv: writer = csv.writer(f_csv) writer.writerows(csv_data) print(f"清单CSV生成成功") except Exception as e: print(f"写入CSV文件时出错: {e}") traceback.print_exc() print(f"\n处理完成. 共处理 {len(entries)} 个条目.") print(f"结果保存在: {output_dir}") def main(): parser = argparse.ArgumentParser(description="从TravelContentCreator结果目录提取内容并渲染到指定目录") parser.add_argument("--source", type=str, help="源目录路径") parser.add_argument("--output", type=str, help="输出目录路径") parser.add_argument("--run-id", type=str, help="自定义运行ID") parser.add_argument("--prefer-original", action="store_true", help="优先使用原始内容，忽略审核结果") parser.add_argument("--db-path", type=str, help="数据库路径，若不提供则使用默认路径") parser.add_argument("--disable-db", action="store_true", help="禁用数据库记录功能") args = parser.parse_args() # 默认值设置 source = args.source if args.source else "/root/autodl-tmp/TravelContentCreator/result/2025-05-20_15-37-25" output = args.output if args.output else "/root/autodl-tmp/TravelContentCreator/output/2025-05-20_15-37-25" run_id = args.run_id if args.run_id else os.path.basename(source) prefer_original = args.prefer_original db_path = args.db_path if args.db_path else '/root/autodl-tmp/TravelContentCreator/distribution.db' print("-" * 60) print(f"开始提取和渲染流程") print(f"源目录: {source}") print(f"输出目录: {output}") print(f"运行ID: {run_id}") if prefer_original: print("内容模式: 优先使用原始内容") else: print("内容模式: 根据审核结果选择内容") if args.disable_db: print("数据库记录: 已禁用") else: print(f"数据库记录: 已启用 (路径: {db_path})") print("-" * 60) process_result_directory(source, output, run_id, prefer_original, db_path) print("\n脚本执行完毕.") if __name__ == "__main__": main()