|
|
|
|
|
import argparse
|
|
|
|
|
|
import numpy as np
|
|
|
from mmengine import dump, load, track_iter_progress
|
|
|
|
|
|
|
|
|
def parse_args():
|
|
|
parser = argparse.ArgumentParser()
|
|
|
parser.add_argument(
|
|
|
'--det_test_result',
|
|
|
default='data/multisports/annotations/ms_det_proposals.pkl')
|
|
|
parser.add_argument(
|
|
|
'--stad_gt',
|
|
|
help='spatio-temporal action detection ground truth file',
|
|
|
default='data/multisports/annotations/multisports_GT.pkl')
|
|
|
parser.add_argument(
|
|
|
'--out_result',
|
|
|
default='data/multisports/annotations/multisports_proposals.pkl')
|
|
|
args = parser.parse_args()
|
|
|
return args
|
|
|
|
|
|
|
|
|
def dump_det_result(args):
|
|
|
print('loading test result...')
|
|
|
det_result = load(args.det_test_result)
|
|
|
stad_gt = load(args.stad_gt)
|
|
|
train_list = stad_gt['train_videos'][0]
|
|
|
val_list = stad_gt['test_videos'][0]
|
|
|
train_bbox_result = {}
|
|
|
val_bbox_result = {}
|
|
|
for sample in track_iter_progress(det_result):
|
|
|
bboxes = sample['pred_instances']['bboxes']
|
|
|
scores = sample['pred_instances']['scores']
|
|
|
h, w = sample['ori_shape']
|
|
|
bboxes[:, ::2] /= w
|
|
|
bboxes[:, 1::2] /= h
|
|
|
img_path = sample['img_path']
|
|
|
frm_key_list = img_path.split('.jpg')[0].split('/')
|
|
|
frm_key = ','.join([
|
|
|
f'{frm_key_list[-3]}/{frm_key_list[-2]}.mp4',
|
|
|
f'{int(frm_key_list[-1]):04d}'
|
|
|
])
|
|
|
bbox = np.concatenate([bboxes, scores[:, None]], axis=1)
|
|
|
|
|
|
vid_key = '/'.join(frm_key_list[-3:-1])
|
|
|
if vid_key in train_list:
|
|
|
train_bbox_result[frm_key] = bbox
|
|
|
elif vid_key in val_list:
|
|
|
val_bbox_result[frm_key] = bbox
|
|
|
else:
|
|
|
raise KeyError(vid_key)
|
|
|
dump(train_bbox_result, args.out_result[:-4] + '_train.pkl')
|
|
|
dump(val_bbox_result, args.out_result[:-4] + '_val.pkl')
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
args = parse_args()
|
|
|
dump_det_result(args)
|
|
|
|
