47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
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"
|
|
) |