File size: 2,647 Bytes
d3dbf03 |
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 |
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
from collections import defaultdict
from mmengine import dump, load
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
'stad_anno', help='spatiotemporal action detection anno path')
parser.add_argument('det_path', help='output detection anno path')
args = parser.parse_args()
return args
def generate_mmdet_coco_anno(args):
ori_anno = load(args.stad_anno)
train_videos = ori_anno['train_videos']
val_videos = ori_anno['test_videos']
videos = {'train': train_videos, 'val': val_videos}
for split in ['train', 'val']:
img_id = 0
bbox_id = 0
img_list = []
anno_list = []
for vid in videos[split][0]:
vid_tubes = ori_anno['gttubes'][vid]
height, width = ori_anno['resolution'][vid]
frm2bbox = defaultdict(list)
for label_idx, tube_list in vid_tubes.items():
for tube in tube_list:
for frm_anno in tube:
frm_idx, bbox = frm_anno[0], frm_anno[1:]
frm2bbox[frm_idx].append({'label': 0, 'bbox': bbox})
for frm_idx, frm_bboxes in frm2bbox.items():
img_path = f'{vid}/{int(frm_idx):05d}.jpg'
img_instance = {
'file_name': img_path,
'height': height,
'width': width,
'id': img_id
}
img_list.append(img_instance)
for bbox_info in frm_bboxes:
label = bbox_info['label']
x1, y1, x2, y2 = bbox_info['bbox']
bbox = [x1, y1, x2 - x1, y2 - y1]
anno_instance = {
'area': bbox[2] * bbox[3],
'image_id': img_id,
'bbox': bbox,
'category_id': label,
'iscrowd': 0,
'id': bbox_id
}
anno_list.append(anno_instance)
bbox_id += 1
img_id += 1
total_anno = {
'images': img_list,
'annotations': anno_list,
'categories': [{
'id': 0,
'name': 'person'
}],
}
dump(total_anno, args.det_path[:-5] + f'_{split}' + args.det_path[-5:])
if __name__ == '__main__':
args = parse_args()
generate_mmdet_coco_anno(args)
|