| 10 | |
| 11 | |
| 12 | def save_track(results, out_root, video_to_images, video_names, data_split='val'): |
| 13 | assert out_root is not None |
| 14 | out_dir = os.path.join(out_root, data_split) |
| 15 | if not os.path.exists(out_dir): |
| 16 | os.mkdir(out_dir) |
| 17 | |
| 18 | # save json. |
| 19 | # json_path = os.path.join(out_dir, "track_results.json") |
| 20 | # with open(json_path, "w") as f: |
| 21 | # f.write(json.dumps(results)) |
| 22 | # f.flush() |
| 23 | |
| 24 | # save it in standard format. |
| 25 | track_dir = os.path.join(out_dir, "tracks") |
| 26 | if not os.path.exists(track_dir): |
| 27 | os.mkdir(track_dir) |
| 28 | for video_id in video_to_images.keys(): |
| 29 | video_infos = video_to_images[video_id] |
| 30 | video_name = video_names[video_id] |
| 31 | file_path = os.path.join(track_dir, "{}.txt".format(video_name)) |
| 32 | f = open(file_path, "w") |
| 33 | tracks = defaultdict(list) |
| 34 | for video_info in video_infos: |
| 35 | image_id, frame_id = video_info["image_id"], video_info["frame_id"] |
| 36 | result = results[image_id] |
| 37 | for item in result: |
| 38 | if not ("tracking_id" in item): |
| 39 | raise NotImplementedError |
| 40 | tracking_id = item["tracking_id"] |
| 41 | bbox = item["bbox"] |
| 42 | bbox = [bbox[0], bbox[1], bbox[2], bbox[3], item['score'], item['active']] |
| 43 | tracks[tracking_id].append([frame_id] + bbox) |
| 44 | |
| 45 | rename_track_id = 0 |
| 46 | for track_id in sorted(tracks): |
| 47 | rename_track_id += 1 |
| 48 | for t in tracks[track_id]: |
| 49 | if t[6] > 0: |
| 50 | f.write("{},{},{:.2f},{:.2f},{:.2f},{:.2f},-1,-1,-1,-1\n".format( |
| 51 | t[0], rename_track_id, t[1], t[2], t[3] - t[1], t[4] - t[2])) |
| 52 | f.close() |