File size: 2,276 Bytes
fe6a6e3 |
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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import json
import argparse
def parse_args():
parser = argparse.ArgumentParser("D2 model converter")
parser.add_argument("--results", default="", type=str, help="Path to the GLEE output dir ")
parser.add_argument("--refer", default="", type=str, help="Path to the BURST annotation val dir")
return parser.parse_args()
def main():
args = parse_args()
ori_anno = json.load(open(args.results,'rb'))
reference_anno = json.load(open(args.refer,'rb'))
num_tracks = 0
num_miss_video = 0
id_mapping = {}
for i, cate_info in enumerate(reference_anno['categories']):
new_id = i
old_id = cate_info['id']
id_mapping.update({new_id:old_id})
ref_sequences_dict = {}
for ref in reference_anno['sequences']:
ref_sequences_dict[ref['id']] = ref
# ids = [v['category_id'] for v in ori_anno]
sequences_dict = {}
for seg in ori_anno:
vid = seg['video_id']
if vid not in sequences_dict.keys():
# import pdb;pdb.set_trace()
sequences_dict[vid] = {
'id': vid,
'width': ref_sequences_dict[vid]['width'],
'height': ref_sequences_dict[vid]['height'],
'seq_name': ref_sequences_dict[vid]['seq_name'],
'dataset': ref_sequences_dict[vid]['dataset'],
'annotated_image_paths': ref_sequences_dict[vid]['annotated_image_paths'],
'fps': ref_sequences_dict[vid]['fps'],
'segmentations': [{} for i in range(len(seg['segmentations']))],
'track_category_ids': {},
}
track_id = str(len(sequences_dict[vid]['track_category_ids']) + 1)
for frame, rles in enumerate(seg['segmentations']):
sequences_dict[vid]['segmentations'][frame][track_id] = {'rle': rles['counts'], 'score':seg['score']}
# import pdb;pdb.set_trace()
sequences_dict[vid]['track_category_ids'][track_id] = id_mapping[seg['category_id']]
results = {'sequences':[]}
for k,v in sequences_dict.items():
results['sequences'].append(v)
with open('converted_tao_results.json', 'w') as f:
json.dump(results, f)
if __name__ == "__main__":
main()
|