import json
import time
from lib.qwen_fun_vid import frame_all_to_obj_vid

# 测试数据路径
input_video_path = "runs/segment/predict/04-01.avi"
json_path = "output_2/04-01/frame_all.json"
output_video_root = "output_2/04-01/test_optimized"

# 读取标注数据
print("读取标注数据...")
with open(json_path, 'r', encoding='utf-8') as f:
    json_data = json.load(f)

print(f"标注数据包含 {len(json_data)} 帧")

# 统计track_id数量
track_ids = set()
for frame_id_str, detections in json_data.items():
    for det in detections:
        track_id = det.get("track_id", -1)
        track_ids.add(track_id)
print(f"共包含 {len(track_ids)} 个track_id")

# 测试优化后的函数
print("\n开始测试优化后的函数...")
start_time = time.time()
frame_all_to_obj_vid(json_data, input_video_path, output_video_root)
end_time = time.time()

print(f"\n测试完成，耗时: {end_time - start_time:.2f} 秒")
print(f"输出视频保存到: {output_video_root}")