Upload convert_ytvis2tao.py

annotations/convert_ytvis2tao.py

@@ -0,0 +1,82 @@
+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()
+
+
+