diff --git a/tools/convert_flat_to_treeindex.py b/tools/convert_flat_to_treeindex.py new file mode 100755 index 0000000..549ab82 --- /dev/null +++ b/tools/convert_flat_to_treeindex.py @@ -0,0 +1,365 @@ +#!/usr/bin/env python3 +"""一次性格式转换:TRM4 flat tree.json -> TRM5 TreeIndex JSON。 + +用法: python tools/convert_flat_to_treeindex.py + +遍历 src_dir 下每个 video_id 子目录中的 tree.json(TRM4 flat 格式), +转换为 TRM5 TreeIndex 嵌套格式并写入 dst_dir 对应子目录。 + +app/core/adapters 不 import 此脚本。迁移完成后归档至 tools/archived/。 +""" + +from __future__ import annotations + +import json +import sys +from datetime import datetime +from pathlib import Path +from typing import Any + +# --------------------------------------------------------------------------- +# Card 字段默认值(处理 TRM4 可能缺失的字段) +# --------------------------------------------------------------------------- + +_L3_CARD_DEFAULTS: dict[str, Any] = { + "frame_summary": "", + "visible_entities": [], + "ongoing_actions": [], + "visible_text": [], + "spatial_layout": "", + "visual_attributes": {}, +} + +_L2_CARD_DEFAULTS: dict[str, Any] = { + "event_description": "", + "entities": [], + "actions": [], + "action_subjects": [], + "visible_text": [], + "spatial_relations": "", + "state_changes": None, +} + +_L1_CARD_DEFAULTS: dict[str, Any] = { + "scene_summary": "", + "main_setting": "", + "key_entities": [], + "main_actions": [], + "topic_keywords": [], + "visible_text": [], + "temporal_flow": "", +} + + +# --------------------------------------------------------------------------- +# Card 构建辅助 +# --------------------------------------------------------------------------- + + +def _build_card(raw_card: dict[str, Any], defaults: dict[str, Any]) -> dict[str, Any]: + """从 TRM4 原始 card 字典构建 TRM5 card,缺失字段用默认值填充。 + + 参数: + raw_card: TRM4 tree.json 中节点的 card 字典。 + defaults: 该层级的默认值字典。 + + 返回: + 仅包含目标字段的 card 字典(字段集合与 defaults 一致)。 + """ + return {key: raw_card.get(key, default) for key, default in defaults.items()} + + +# --------------------------------------------------------------------------- +# 节点转换 +# --------------------------------------------------------------------------- + + +def _convert_l3( + node: dict[str, Any], + video_id: str, +) -> dict[str, Any]: + """将 TRM4 flat L3 节点转换为 TRM5 嵌套 L3 节点。 + + 参数: + node: TRM4 flat 格式的 L3 节点字典。 + video_id: 视频 ID,用于计算 frame_path 的相对路径后缀。 + + 返回: + TRM5 格式的 L3 节点字典。 + """ + node_id: str = node["node_id"] + raw_card = node.get("card") or {} + card = _build_card(raw_card, _L3_CARD_DEFAULTS) + + # frame_path: frames/{suffix}.jpg,suffix = node_id 去掉 video_id 前缀 + 下划线 + prefix = f"{video_id}_" + suffix = node_id[len(prefix) :] if node_id.startswith(prefix) else node_id + frame_path = f"frames/{suffix}.jpg" + + return { + "id": node_id, + "card": card, + "timestamp": node.get("frame_timestamp"), + "frame_path": frame_path, + "subtitle": node.get("subtitle"), + } + + +def _convert_l2( + node: dict[str, Any], + l3_children: list[dict[str, Any]], +) -> dict[str, Any]: + """将 TRM4 flat L2 节点转换为 TRM5 嵌套 L2 节点。 + + 参数: + node: TRM4 flat 格式的 L2 节点字典。 + l3_children: 已转换的 L3 子节点列表(按 time_range 排序)。 + + 返回: + TRM5 格式的 L2 节点字典。 + """ + raw_card = node.get("card") or {} + card = _build_card(raw_card, _L2_CARD_DEFAULTS) + time_range = node.get("time_range") + + return { + "id": node["node_id"], + "card": card, + "time_range": time_range, + "children": l3_children, + } + + +def _convert_l1( + node: dict[str, Any], + l2_children: list[dict[str, Any]], +) -> dict[str, Any]: + """将 TRM4 flat L1 节点转换为 TRM5 嵌套 L1 节点。 + + 参数: + node: TRM4 flat 格式的 L1 节点字典。 + l2_children: 已转换的 L2 子节点列表(按 time_range 排序)。 + + 返回: + TRM5 格式的 L1 节点字典。 + """ + raw_card = node.get("card") or {} + card = _build_card(raw_card, _L1_CARD_DEFAULTS) + time_range = node.get("time_range") + + return { + "id": node["node_id"], + "card": card, + "time_range": time_range, + "children": l2_children, + } + + +# --------------------------------------------------------------------------- +# 排序辅助 +# --------------------------------------------------------------------------- + + +def _sort_key_time_range(node: dict[str, Any]) -> float: + """按 time_range 的起始时间排序。 + + 参数: + node: TRM4 节点字典。 + + 返回: + 起始时间(float),无 time_range 时返回 0.0。 + """ + tr = node.get("time_range") + if tr and len(tr) >= 1: + return float(tr[0]) + return 0.0 + + +def _sort_key_timestamp(converted: dict[str, Any]) -> float: + """按 timestamp 排序(L3 转换后的字典)。 + + 参数: + converted: 已转换的 TRM5 L3 节点字典。 + + 返回: + timestamp(float),无值时返回 0.0。 + """ + ts = converted.get("timestamp") + return float(ts) if ts is not None else 0.0 + + +# --------------------------------------------------------------------------- +# 单棵树转换 +# --------------------------------------------------------------------------- + + +def convert_single_tree(flat_data: dict[str, Any], source_path: str) -> dict[str, Any]: + """将单个 TRM4 flat tree.json 转换为 TRM5 TreeIndex 字典。 + + 参数: + flat_data: TRM4 flat tree.json 解析后的字典。 + source_path: 原始数据路径(写入 metadata.source_path)。 + + 返回: + TRM5 TreeIndex 格式的字典(可直接 json.dump 或传入 TreeIndex.from_dict)。 + + 异常: + ValueError: 无法从 flat_data 中提取 video_id。 + """ + video_id = flat_data.get("video_id") or flat_data.get("videoID") + if not video_id: + raise ValueError("flat tree.json 中缺少 video_id / videoID 字段") + + nodes: dict[str, dict[str, Any]] = flat_data.get("nodes", {}) + + # Phase 1: 按层级分组 + l1_nodes: list[dict[str, Any]] = [] + l2_nodes: list[dict[str, Any]] = [] + l3_nodes: list[dict[str, Any]] = [] + + for node in nodes.values(): + level = node.get("level") + if level == 1: + l1_nodes.append(node) + elif level == 2: + l2_nodes.append(node) + elif level == 3: + l3_nodes.append(node) + + # Phase 2: 构建 parent -> children 映射 + # L3 按 parent_id 分组 + l3_by_parent: dict[str, list[dict[str, Any]]] = {} + for n in l3_nodes: + pid = n.get("parent_id", "") + l3_by_parent.setdefault(pid, []).append(n) + + # L2 按 parent_id 分组 + l2_by_parent: dict[str, list[dict[str, Any]]] = {} + for n in l2_nodes: + pid = n.get("parent_id", "") + l2_by_parent.setdefault(pid, []).append(n) + + # Phase 3: 自底向上构建嵌套结构 + # 转换 L2 -> 附带转换后的 L3 children + converted_l2_by_id: dict[str, dict[str, Any]] = {} + for l2 in l2_nodes: + l2_id = l2["node_id"] + raw_l3_children = l3_by_parent.get(l2_id, []) + # 先转换 L3,再按 timestamp 排序 + converted_l3 = [_convert_l3(n, video_id) for n in raw_l3_children] + converted_l3.sort(key=_sort_key_timestamp) + converted_l2_by_id[l2_id] = _convert_l2(l2, converted_l3) + + # 转换 L1 -> 附带转换后的 L2 children + roots: list[dict[str, Any]] = [] + l1_nodes.sort(key=_sort_key_time_range) + + for l1 in l1_nodes: + l1_id = l1["node_id"] + raw_l2_children = l2_by_parent.get(l1_id, []) + raw_l2_children.sort(key=_sort_key_time_range) + l2_children = [converted_l2_by_id[n["node_id"]] for n in raw_l2_children] + roots.append(_convert_l1(l1, l2_children)) + + # Phase 4: 构建 TreeIndex 字典 + return { + "metadata": { + "source_path": source_path, + "modality": "video", + "created_at": datetime.now().isoformat(), + }, + "roots": roots, + } + + +# --------------------------------------------------------------------------- +# 批量转换入口 +# --------------------------------------------------------------------------- + + +def convert_directory(src_dir: str, dst_dir: str) -> tuple[int, int]: + """批量转换目录下所有 TRM4 tree.json 到 TRM5 TreeIndex 格式。 + + 参数: + src_dir: 源目录(TRM4 store/videos/),其下每个子目录含 tree.json。 + dst_dir: 目标目录(TRM5 store/videos/),保持同名子目录结构。 + + 返回: + (成功数, 失败数) 元组。 + """ + src_path = Path(src_dir) + dst_path = Path(dst_dir) + + if not src_path.is_dir(): + print(f"错误: 源目录不存在: {src_dir}", file=sys.stderr) + sys.exit(1) + + success_count = 0 + fail_count = 0 + + tree_files = sorted(src_path.glob("*/tree.json")) + total = len(tree_files) + print(f"发现 {total} 个 tree.json 待转换") + + for idx, tree_file in enumerate(tree_files, 1): + video_id = tree_file.parent.name + out_dir = dst_path / video_id + out_file = out_dir / "tree.json" + + try: + with open(tree_file, encoding="utf-8") as f: + flat_data = json.load(f) + + result = convert_single_tree(flat_data, source_path=str(tree_file)) + + out_dir.mkdir(parents=True, exist_ok=True) + with open(out_file, "w", encoding="utf-8") as f: + json.dump(result, f, ensure_ascii=False, indent=2) + + # 统计节点数 + n_l1 = len(result["roots"]) + n_l2 = sum(len(r["children"]) for r in result["roots"]) + n_l3 = sum(len(l2["children"]) for r in result["roots"] for l2 in r["children"]) + print(f"[{idx}/{total}] {video_id}: L1={n_l1}, L2={n_l2}, L3={n_l3}") + success_count += 1 + + except Exception as e: + print(f"[{idx}/{total}] {video_id}: 失败 - {e}", file=sys.stderr) + fail_count += 1 + + return success_count, fail_count + + +# --------------------------------------------------------------------------- +# CLI 入口 +# --------------------------------------------------------------------------- + + +def main() -> None: + """CLI 入口:解析参数并执行批量转换。""" + if len(sys.argv) != 3: + print( + "用法: python tools/convert_flat_to_treeindex.py ", + file=sys.stderr, + ) + print(" src_dir: TRM4 store/videos/ 目录(含 video_id/tree.json)", file=sys.stderr) + print(" dst_dir: TRM5 store/videos/ 目标目录", file=sys.stderr) + sys.exit(1) + + src_dir = sys.argv[1] + dst_dir = sys.argv[2] + + print(f"源目录: {src_dir}") + print(f"目标目录: {dst_dir}") + print() + + success, fail = convert_directory(src_dir, dst_dir) + + print() + print(f"转换完成: 成功 {success}, 失败 {fail}") + if fail > 0: + sys.exit(1) + + +if __name__ == "__main__": + main()