frame_labels_hdf5_to_lmdb.py 文件源码

def main():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        'labels_hdf5',
        help=('Maps video names to a binary matrix of shape (num_frames, '
              'num_labels).'))
    parser.add_argument('output_lmdb')

    args = parser.parse_args()

    map_size = 2e9

    lmdb_environment = lmdb.open(args.output_lmdb, map_size=int(map_size))
    with lmdb_environment.begin(write=True) as lmdb_transaction, h5py.File(
            args.labels_hdf5) as labels:
        for video_name, file_labels in tqdm(labels.items()):
            file_labels = np.asarray(file_labels)
            for frame_number, frame_labels in enumerate(file_labels):
                key = '{}-{}'.format(video_name, frame_number + 1)
                lmdb_transaction.put(key, frame_labels.tobytes())