| |
| |
| |
| |
|
|
|
|
| import pathlib |
| import json |
| import lmdb |
|
|
| from xmlparser import XMLRawDataset, ListRawDataset |
|
|
| import argparse |
|
|
|
|
| class Env: |
| def __init__(self, output_path, interval_writeCache=1000): |
| self.output_path = output_path |
| self.env = lmdb.open(str(output_path), map_size=1099511627776) |
| self.cache = dict() |
| self.n = 0 |
| self.interval = interval_writeCache |
|
|
| def finish_line(self): |
| self.n += 1 |
| if self.n % 1000 == 0: |
| self.writeCache() |
|
|
| def writeCache(self): |
| with self.env.begin(write=True) as txn: |
| for k, v in self.cache.items(): |
| txn.put(k, v) |
| self.cache = {} |
| print(f'Written {self.n} lines @ {self.output_path}') |
|
|
|
|
| def createDataset(input_path, output_path, db_type='xml', dry_run=False): |
| p = pathlib.Path(output_path) |
| p.mkdir(parents=True, exist_ok=True) |
|
|
| if db_type == 'xml': |
| generator = XMLRawDataset.from_list(input_path, image_type=XMLRawDataset.IMAGE_TYPE_ENCODED) |
| elif db_type == 'list': |
| generator = ListRawDataset(input_path, image_type=XMLRawDataset.IMAGE_TYPE_ENCODED) |
| if dry_run: |
| return |
|
|
| |
| env = Env(output_path[0]) |
| env.cache['dbtype'.encode()] = 'xml'.encode() |
|
|
| for il, (g, line) in enumerate(generator): |
| env.cache[f'{env.n:09d}-direction'.encode()] = line.get('direction').encode() |
| env.cache[f'{env.n:09d}-label'.encode()] = line.get('label').encode() |
| env.cache[f'{env.n:09d}-cattrs'.encode()] = json.dumps(line.get('cattrs')).encode() |
| env.cache[f'{env.n:09d}-image'.encode()] = g |
| env.finish_line() |
|
|
| env.cache['n_line'.encode()] = str(env.n).encode() |
| env.writeCache() |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--input_path', nargs='+', required=True) |
| parser.add_argument('--output_path', required=True) |
| parser.add_argument('--db_type', default='xml', choices=['xml', 'list']) |
| parser.add_argument('--dry-run', action='store_true') |
| opt = parser.parse_args() |
| createDataset(opt.input_path, opt.output_path, opt.db_type, dry_run=opt.dry_run) |
|
|
