|
|
| import argparse
|
| import multiprocessing
|
| import os
|
| import os.path as osp
|
|
|
| import numpy as np
|
| import scipy.interpolate
|
| from mmengine import dump, load
|
|
|
| args = None
|
|
|
|
|
| def parse_args():
|
| parser = argparse.ArgumentParser(description='ANet Feature Prepare')
|
| parser.add_argument('--rgb', default='', help='rgb feature root')
|
| parser.add_argument('--flow', default='', help='flow feature root')
|
| parser.add_argument('--dest', default='', help='dest root')
|
| parser.add_argument('--output-format', default='csv')
|
| args = parser.parse_args()
|
| return args
|
|
|
|
|
| def pool_feature(data, num_proposals=100, num_sample_bins=3, pool_type='mean'):
|
| """Pool features with arbitrary temporal length.
|
|
|
| Args:
|
| data (list[np.ndarray] | np.ndarray): Features of an untrimmed video,
|
| with arbitrary temporal length.
|
| num_proposals (int): The temporal dim of pooled feature. Default: 100.
|
| num_sample_bins (int): How many points to sample to get the feature
|
| vector at one timestamp. Default: 3.
|
| pool_type (str): Type of pooling to pool features. Choices are
|
| ['mean', 'max']. Default: 'mean'.
|
|
|
| Returns:
|
| np.ndarray: The pooled feature with shape num_proposals x feature_dim.
|
| """
|
| if len(data) == 1:
|
| return np.concatenate([data] * num_proposals)
|
| x_range = list(range(len(data)))
|
| f = scipy.interpolate.interp1d(x_range, data, axis=0)
|
| eps = 1e-4
|
| start, end = eps, len(data) - 1 - eps
|
| anchor_size = (end - start) / num_proposals
|
| ptr = start
|
| feature = []
|
| for _ in range(num_proposals):
|
| x_new = [
|
| ptr + i / num_sample_bins * anchor_size
|
| for i in range(num_sample_bins)
|
| ]
|
| y_new = f(x_new)
|
| if pool_type == 'mean':
|
| y_new = np.mean(y_new, axis=0)
|
| elif pool_type == 'max':
|
| y_new = np.max(y_new, axis=0)
|
| else:
|
| raise NotImplementedError('Unsupported pool type')
|
| feature.append(y_new)
|
| ptr += anchor_size
|
| feature = np.stack(feature)
|
| return feature
|
|
|
|
|
| def merge_feat(name):
|
|
|
| rgb_feat = load(osp.join(args.rgb, name))
|
| flow_feat = load(osp.join(args.flow, name))
|
| rgb_feat = pool_feature(rgb_feat)
|
| flow_feat = pool_feature(flow_feat)
|
| feat = np.concatenate([rgb_feat, flow_feat], axis=-1)
|
| if not osp.exists(args.dest):
|
| os.system(f'mkdir -p {args.dest}')
|
| if args.output_format == 'pkl':
|
| dump(feat, osp.join(args.dest, name))
|
| elif args.output_format == 'csv':
|
| feat = feat.tolist()
|
| lines = []
|
| line0 = ','.join([f'f{i}' for i in range(400)])
|
| lines.append(line0)
|
| for line in feat:
|
| lines.append(','.join([f'{x:.4f}' for x in line]))
|
| with open(osp.join(args.dest, name.replace('.pkl', '.csv')), 'w') as f:
|
| f.write('\n'.join(lines))
|
|
|
|
|
| def main():
|
| global args
|
| args = parse_args()
|
| rgb_feat = [file for file in os.listdir(args.rgb) if file.endswith('.pkl')]
|
| flow_feat = [
|
| file for file in os.listdir(args.flow) if file.endswith('.pkl')
|
| ]
|
| assert set(rgb_feat) == set(flow_feat)
|
|
|
|
|
| pool = multiprocessing.Pool(32)
|
| pool.map(merge_feat, rgb_feat)
|
|
|
|
|
| if __name__ == '__main__':
|
| main()
|
|
|
