|
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152 |
- """
- 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()
|
