Spaces:
Runtime error
Runtime error
File size: 1,902 Bytes
7734d5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
"""
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()
|