import json
from collections import defaultdict


def f_detections_to_objects(detections_path: str, output_path: str):
    """
    将输入文件中的yolo检测数据转换为object数据并保存到输出文件中
    :param detections_path: 输入文件路径
    :param output_path: 输出文件路径
    :return: None
    """
    with open(detections_path, 'r', encoding='utf-8') as f:
        detections_data = json.load(f)

    class_set = set()
    track_info = defaultdict(lambda: {"class_id": None, "start_frame": float('inf'), "end_frame": 0})

    for frame_str, detections in detections_data.items():
        frame = int(frame_str)
        for det in detections:
            track_id = det["track_id"]
            class_id = det["class_id"]
            class_str = det["class_str"]

            class_set.add(class_str)

            if track_info[track_id]["class_id"] is None:
                track_info[track_id]["class_id"] = class_id

            if frame < track_info[track_id]["start_frame"]:
                track_info[track_id]["start_frame"] = frame
            if frame > track_info[track_id]["end_frame"]:
                track_info[track_id]["end_frame"] = frame

    result = {
        "class_list": list(class_set),
        "track_id_list": dict(track_info)
    }

    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=2)

if __name__ == "__main__":
    f_detections_to_objects(
        "output/santai5/frame_detections.json",
        "output/santai5/objects.json"
    )