Spaces:
Runtime error
Runtime error
| """ | |
| 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() | |