fix(tools): handle None card in flat tree conversion
This commit is contained in:
Executable
+365
@@ -0,0 +1,365 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user