366 lines
11 KiB
Python
Executable File
366 lines
11 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""一次性格式转换：TRM4 flat tree.json -> TRM5 TreeIndex JSON。
|
||
|
||
用法: python tools/convert_flat_to_treeindex.py <src_dir> <dst_dir>
|
||
|
||
遍历 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 <src_dir> <dst_dir>",
|
||
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()
|