(json_data, img_dir, mask_dir, output_dir, split, write_frequency=1000)
| 33 | |
| 34 | |
| 35 | def folder2lmdb(json_data, img_dir, mask_dir, output_dir, split, write_frequency=1000): |
| 36 | lmdb_path = osp.join(output_dir, "%s.lmdb" % split) |
| 37 | isdir = os.path.isdir(lmdb_path) |
| 38 | |
| 39 | print("Generate LMDB to %s" % lmdb_path) |
| 40 | db = lmdb.open(lmdb_path, subdir=isdir, |
| 41 | map_size=1099511627776 * 2, readonly=False, |
| 42 | meminit=False, map_async=True) |
| 43 | |
| 44 | txn = db.begin(write=True) |
| 45 | tbar = tqdm(json_data) |
| 46 | for idx, item in enumerate(tbar): |
| 47 | img = raw_reader(osp.join(img_dir, item['img_name'])) |
| 48 | mask = raw_reader(osp.join(mask_dir, f"{item['segment_id']}.png")) |
| 49 | data = {'img': img, 'mask': mask, 'cat': item['cat'], |
| 50 | 'seg_id': item['segment_id'], 'img_name': item['img_name'], |
| 51 | 'num_sents': item['sentences_num'], 'sents': [i['sent'] for i in item['sentences']]} |
| 52 | txn.put(u'{}'.format(idx).encode('ascii'), dumps_pyarrow(data)) |
| 53 | if idx % write_frequency == 0: |
| 54 | # print("[%d/%d]" % (idx, len(data_loader))) |
| 55 | txn.commit() |
| 56 | txn = db.begin(write=True) |
| 57 | |
| 58 | # finish iterating through dataset |
| 59 | txn.commit() |
| 60 | keys = [u'{}'.format(k).encode('ascii') for k in range(idx + 1)] |
| 61 | with db.begin(write=True) as txn: |
| 62 | txn.put(b'__keys__', dumps_pyarrow(keys)) |
| 63 | txn.put(b'__len__', dumps_pyarrow(len(keys))) |
| 64 | |
| 65 | print("Flushing database ...") |
| 66 | db.sync() |
| 67 | db.close() |
| 68 | |
| 69 | |
| 70 | def parse_args(): |
no test coverage detected