forked from open-mmlab/mmocr
/
lmdb_util.py
47 lines (41 loc) · 1.58 KB
/
lmdb_util.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import shutil
import sys
import time
from pathlib import Path
import lmdb
from mmocr.utils import list_from_file
def lmdb_converter(img_list_file, output, batch_size=1000, coding='utf-8'):
# read img_list_file
lines = list_from_file(img_list_file)
# create lmdb database
if Path(output).is_dir():
while True:
print('%s already exist, delete or not? [Y/n]' % output)
Yn = input().strip()
if Yn in ['Y', 'y']:
shutil.rmtree(output)
break
if Yn in ['N', 'n']:
return
print('create database %s' % output)
Path(output).mkdir(parents=True, exist_ok=False)
env = lmdb.open(output, map_size=1099511627776)
# build lmdb
beg_time = time.strftime('%H:%M:%S')
for beg_index in range(0, len(lines), batch_size):
end_index = min(beg_index + batch_size, len(lines))
sys.stdout.write('\r[%s-%s], processing [%d-%d] / %d' %
(beg_time, time.strftime('%H:%M:%S'), beg_index,
end_index, len(lines)))
sys.stdout.flush()
batch = [(str(index).encode(coding), lines[index].encode(coding))
for index in range(beg_index, end_index)]
with env.begin(write=True) as txn:
cursor = txn.cursor()
cursor.putmulti(batch, dupdata=False, overwrite=True)
sys.stdout.write('\n')
with env.begin(write=True) as txn:
key = 'total_number'.encode(coding)
value = str(len(lines)).encode(coding)
txn.put(key, value)
print('done', flush=True)