53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
"""
|
|
Copyright (c) https://github.com/xingyizhou/CenterTrack
|
|
Modified by Peize Sun, Rufeng Zhang
|
|
"""
|
|
# coding: utf-8
|
|
import os
|
|
import json
|
|
import logging
|
|
from collections import defaultdict
|
|
|
|
|
|
def save_track(results, out_root, video_to_images, video_names, data_split='val'):
|
|
assert out_root is not None
|
|
out_dir = os.path.join(out_root, data_split)
|
|
if not os.path.exists(out_dir):
|
|
os.mkdir(out_dir)
|
|
|
|
# save json.
|
|
# json_path = os.path.join(out_dir, "track_results.json")
|
|
# with open(json_path, "w") as f:
|
|
# f.write(json.dumps(results))
|
|
# f.flush()
|
|
|
|
# save it in standard format.
|
|
track_dir = os.path.join(out_dir, "tracks")
|
|
if not os.path.exists(track_dir):
|
|
os.mkdir(track_dir)
|
|
for video_id in video_to_images.keys():
|
|
video_infos = video_to_images[video_id]
|
|
video_name = video_names[video_id]
|
|
file_path = os.path.join(track_dir, "{}.txt".format(video_name))
|
|
f = open(file_path, "w")
|
|
tracks = defaultdict(list)
|
|
for video_info in video_infos:
|
|
image_id, frame_id = video_info["image_id"], video_info["frame_id"]
|
|
result = results[image_id]
|
|
for item in result:
|
|
if not ("tracking_id" in item):
|
|
raise NotImplementedError
|
|
tracking_id = item["tracking_id"]
|
|
bbox = item["bbox"]
|
|
bbox = [bbox[0], bbox[1], bbox[2], bbox[3], item['score'], item['active']]
|
|
tracks[tracking_id].append([frame_id] + bbox)
|
|
|
|
rename_track_id = 0
|
|
for track_id in sorted(tracks):
|
|
rename_track_id += 1
|
|
for t in tracks[track_id]:
|
|
if t[6] > 0:
|
|
f.write("{},{},{:.2f},{:.2f},{:.2f},{:.2f},-1,-1,-1,-1\n".format(
|
|
t[0], rename_track_id, t[1], t[2], t[3] - t[1], t[4] - t[2]))
|
|
f.close()
|