| import argparse |
| import os |
| import os.path as osp |
| import lmdb |
| import pyarrow as pa |
| import json |
| from tqdm import tqdm |
| import warnings |
| warnings.filterwarnings("ignore") |
|
|
|
|
| def loads_pyarrow(buf): |
| """ |
| Args: |
| buf: the output of `dumps`. |
| """ |
| return pa.deserialize(buf) |
|
|
|
|
| def raw_reader(path): |
| with open(path, 'rb') as f: |
| bin_data = f.read() |
| return bin_data |
|
|
|
|
| def dumps_pyarrow(obj): |
| """ |
| Serialize an object. |
| Returns: |
| Implementation-dependent bytes-like object |
| """ |
| return pa.serialize(obj).to_buffer() |
|
|
|
|
| def folder2lmdb(json_data, img_dir, mask_dir, output_dir, split, write_frequency=1000): |
| lmdb_path = osp.join(output_dir, "%s.lmdb" % split) |
| isdir = os.path.isdir(lmdb_path) |
|
|
| print("Generate LMDB to %s" % lmdb_path) |
| db = lmdb.open(lmdb_path, subdir=isdir, map_size=1099511627776 * 2, readonly=False, meminit=False, map_async=True) |
|
|
| txn = db.begin(write=True) |
| tbar = tqdm(json_data) |
| for idx, item in enumerate(tbar): |
| img = raw_reader(osp.join(img_dir, item['img_name'])) |
| mask = raw_reader(osp.join(mask_dir, f"{item['segment_id']}.png")) |
| data = { |
| 'img': img, |
| 'mask': mask, |
| 'cat': item['cat'], |
| 'seg_id': item['segment_id'], |
| 'img_name': item['img_name'], |
| 'num_sents': len(item['sentences']), |
| 'sents': [i['sent'] for i in item['sentences']] |
| } |
| txn.put(u'{}'.format(idx).encode('ascii'), dumps_pyarrow(data)) |
| if idx % write_frequency == 0: |
| |
| txn.commit() |
| txn = db.begin(write=True) |
|
|
| |
| txn.commit() |
| keys = [u'{}'.format(k).encode('ascii') for k in range(idx + 1)] |
| with db.begin(write=True) as txn: |
| txn.put(b'__keys__', dumps_pyarrow(keys)) |
| txn.put(b'__len__', dumps_pyarrow(len(keys))) |
|
|
| print("Flushing database ...") |
| db.sync() |
| db.close() |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser(description='COCO Folder to LMDB.') |
| parser.add_argument('-j', '--json-dir', type=str, |
| default='', |
| help='the name of json file.') |
| parser.add_argument('-i', '--img-dir', type=str, |
| default='refcoco+', |
| help='the folder of images.') |
| parser.add_argument('-m', '--mask-dir', type=str, |
| default='refcoco+', |
| help='the folder of masks.') |
| parser.add_argument('-o', '--output-dir', type=str, |
| default='refcoco+', |
| help='the folder of output lmdb file.') |
| parser.add_argument('-s', '--split', type=str, |
| default='train', |
| help='the split type.') |
| args = parser.parse_args() |
| return args |
|
|
|
|
| if __name__ == '__main__': |
| args = parse_args() |
| args.split = osp.basename(args.json_dir).split(".")[0] |
| os.makedirs(args.output_dir, exist_ok=True) |
|
|
| with open(args.json_dir, 'rb') as f: |
| json_data = json.load(f) |
|
|
| folder2lmdb(json_data, args.img_dir, args.mask_dir, args.output_dir, args.split) |
|
|
