mirror of
https://github.com/PaddlePaddle/PaddleOCR.git
synced 2026-09-21 13:23:18 +08:00
301 lines
10 KiB
Python
301 lines
10 KiB
Python
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
import copy
|
||
import numpy as np
|
||
import os
|
||
import lmdb
|
||
import random
|
||
import signal
|
||
import paddle
|
||
from paddle.io import Dataset, DataLoader, DistributedBatchSampler, BatchSampler
|
||
|
||
from .imaug import transform, create_operators
|
||
from ppocr.utils.logging import get_logger
|
||
|
||
|
||
def term_mp(sig_num, frame):
|
||
""" kill all child processes
|
||
"""
|
||
pid = os.getpid()
|
||
pgid = os.getpgid(os.getpid())
|
||
print("main proc {} exit, kill process group " "{}".format(pid, pgid))
|
||
os.killpg(pgid, signal.SIGKILL)
|
||
|
||
|
||
signal.signal(signal.SIGINT, term_mp)
|
||
signal.signal(signal.SIGTERM, term_mp)
|
||
|
||
|
||
class ModeException(Exception):
|
||
"""
|
||
ModeException
|
||
"""
|
||
|
||
def __init__(self, message='', mode=''):
|
||
message += "\nOnly the following 3 modes are supported: " \
|
||
"train, valid, test. Given mode is {}".format(mode)
|
||
super(ModeException, self).__init__(message)
|
||
|
||
|
||
class SampleNumException(Exception):
|
||
"""
|
||
SampleNumException
|
||
"""
|
||
|
||
def __init__(self, message='', sample_num=0, batch_size=1):
|
||
message += "\nError: The number of the whole data ({}) " \
|
||
"is smaller than the batch_size ({}), and drop_last " \
|
||
"is turnning on, so nothing will feed in program, " \
|
||
"Terminated now. Please reset batch_size to a smaller " \
|
||
"number or feed more data!".format(sample_num, batch_size)
|
||
super(SampleNumException, self).__init__(message)
|
||
|
||
|
||
def get_file_list(file_list, data_dir, delimiter='\t'):
|
||
"""
|
||
read label list from file and shuffle the list
|
||
|
||
Args:
|
||
params(dict):
|
||
"""
|
||
if isinstance(file_list, str):
|
||
file_list = [file_list]
|
||
data_source_list = []
|
||
for file in file_list:
|
||
with open(file) as f:
|
||
full_lines = [line.strip() for line in f]
|
||
for line in full_lines:
|
||
try:
|
||
img_path, label = line.split(delimiter)
|
||
except:
|
||
logger = get_logger()
|
||
logger.warning('label error in {}'.format(line))
|
||
img_path = os.path.join(data_dir, img_path)
|
||
data = {'img_path': img_path, 'label': label}
|
||
data_source_list.append(data)
|
||
return data_source_list
|
||
|
||
|
||
class LMDBDateSet(Dataset):
|
||
def __init__(self, config, global_config):
|
||
super(LMDBDateSet, self).__init__()
|
||
self.data_list = self.load_lmdb_dataset(
|
||
config['file_list'], global_config['max_text_length'])
|
||
random.shuffle(self.data_list)
|
||
|
||
self.ops = create_operators(config['transforms'], global_config)
|
||
|
||
# for rec
|
||
character = ''
|
||
for op in self.ops:
|
||
if hasattr(op, 'character'):
|
||
character = getattr(op, 'character')
|
||
|
||
self.info_dict = {'character': character}
|
||
|
||
def load_lmdb_dataset(self, data_dir, max_text_length):
|
||
self.env = lmdb.open(
|
||
data_dir,
|
||
max_readers=32,
|
||
readonly=True,
|
||
lock=False,
|
||
readahead=False,
|
||
meminit=False)
|
||
if not self.env:
|
||
print('cannot create lmdb from %s' % (data_dir))
|
||
exit(0)
|
||
|
||
filtered_index_list = []
|
||
with self.env.begin(write=False) as txn:
|
||
nSamples = int(txn.get('num-samples'.encode()))
|
||
self.nSamples = nSamples
|
||
for index in range(self.nSamples):
|
||
index += 1 # lmdb starts with 1
|
||
label_key = 'label-%09d'.encode() % index
|
||
label = txn.get(label_key).decode('utf-8')
|
||
if len(label) > max_text_length:
|
||
# print(f'The length of the label is longer than max_length: length
|
||
# {len(label)}, {label} in dataset {self.root}')
|
||
continue
|
||
|
||
# By default, images containing characters which are not in opt.character are filtered.
|
||
# You can add [UNK] token to `opt.character` in utils.py instead of this filtering.
|
||
filtered_index_list.append(index)
|
||
return filtered_index_list
|
||
|
||
def print_lmdb_sets_info(self, lmdb_sets):
|
||
lmdb_info_strs = []
|
||
for dataset_idx in range(len(lmdb_sets)):
|
||
tmp_str = " %s:%d," % (lmdb_sets[dataset_idx]['dirpath'],
|
||
lmdb_sets[dataset_idx]['num_samples'])
|
||
lmdb_info_strs.append(tmp_str)
|
||
lmdb_info_strs = ''.join(lmdb_info_strs)
|
||
logger = get_logger()
|
||
logger.info("DataSummary:" + lmdb_info_strs)
|
||
return
|
||
|
||
def __getitem__(self, idx):
|
||
idx = self.data_list[idx]
|
||
with self.env.begin(write=False) as txn:
|
||
label_key = 'label-%09d'.encode() % idx
|
||
label = txn.get(label_key)
|
||
if label is not None:
|
||
label = label.decode('utf-8')
|
||
img_key = 'image-%09d'.encode() % idx
|
||
imgbuf = txn.get(img_key)
|
||
data = {'image': imgbuf, 'label': label}
|
||
outs = transform(data, self.ops)
|
||
else:
|
||
outs = None
|
||
if outs is None:
|
||
return self.__getitem__(np.random.randint(self.__len__()))
|
||
return outs
|
||
|
||
def __len__(self):
|
||
return len(self.data_list)
|
||
|
||
|
||
class SimpleDataSet(Dataset):
|
||
def __init__(self, config, global_config):
|
||
super(SimpleDataSet, self).__init__()
|
||
delimiter = config.get('delimiter', '\t')
|
||
self.data_list = get_file_list(config['file_list'], config['data_dir'],
|
||
delimiter)
|
||
random.shuffle(self.data_list)
|
||
|
||
self.ops = create_operators(config['transforms'], global_config)
|
||
|
||
# for rec
|
||
character = ''
|
||
for op in self.ops:
|
||
if hasattr(op, 'character'):
|
||
character = getattr(op, 'character')
|
||
|
||
self.info_dict = {'character': character}
|
||
|
||
def __getitem__(self, idx):
|
||
data = copy.deepcopy(self.data_list[idx])
|
||
with open(data['img_path'], 'rb') as f:
|
||
img = f.read()
|
||
data['image'] = img
|
||
outs = transform(data, self.ops)
|
||
if outs is None:
|
||
return self.__getitem__(np.random.randint(self.__len__()))
|
||
return outs
|
||
|
||
def __len__(self):
|
||
return len(self.data_list)
|
||
|
||
|
||
class BatchBalancedDataLoader(object):
|
||
def __init__(self,
|
||
dataset_list: list,
|
||
ratio_list: list,
|
||
distributed,
|
||
device,
|
||
loader_args: dict):
|
||
"""
|
||
对datasetlist里的dataset按照ratio_list里对应的比例组合，似的每个batch里的数据按按照比例采样的
|
||
:param dataset_list: 数据集列表
|
||
:param ratio_list: 比例列表
|
||
:param loader_args: dataloader的配置
|
||
"""
|
||
assert sum(ratio_list) == 1 and len(dataset_list) == len(ratio_list)
|
||
|
||
self.dataset_len = 0
|
||
self.data_loader_list = []
|
||
self.dataloader_iter_list = []
|
||
all_batch_size = loader_args.pop('batch_size')
|
||
batch_size_list = list(
|
||
map(int, [max(1.0, all_batch_size * x) for x in ratio_list]))
|
||
remain_num = all_batch_size - sum(batch_size_list)
|
||
batch_size_list[np.argmax(ratio_list)] += remain_num
|
||
|
||
for _dataset, _batch_size in zip(dataset_list, batch_size_list):
|
||
if distributed:
|
||
batch_sampler_class = DistributedBatchSampler
|
||
else:
|
||
batch_sampler_class = BatchSampler
|
||
batch_sampler = batch_sampler_class(
|
||
dataset=_dataset,
|
||
batch_size=_batch_size,
|
||
shuffle=loader_args['shuffle'],
|
||
drop_last=loader_args['drop_last'], )
|
||
_data_loader = DataLoader(
|
||
dataset=_dataset,
|
||
batch_sampler=batch_sampler,
|
||
places=device,
|
||
num_workers=loader_args['num_workers'],
|
||
return_list=True, )
|
||
self.data_loader_list.append(_data_loader)
|
||
self.dataloader_iter_list.append(iter(_data_loader))
|
||
self.dataset_len += len(_dataset)
|
||
|
||
def __iter__(self):
|
||
return self
|
||
|
||
def __len__(self):
|
||
return min([len(x) for x in self.data_loader_list])
|
||
|
||
def __next__(self):
|
||
batch = []
|
||
for i, data_loader_iter in enumerate(self.dataloader_iter_list):
|
||
try:
|
||
_batch_i = next(data_loader_iter)
|
||
batch.append(_batch_i)
|
||
except StopIteration:
|
||
self.dataloader_iter_list[i] = iter(self.data_loader_list[i])
|
||
_batch_i = next(self.dataloader_iter_list[i])
|
||
batch.append(_batch_i)
|
||
except ValueError:
|
||
pass
|
||
if len(batch) > 0:
|
||
batch_list = []
|
||
batch_item_size = len(batch[0])
|
||
for i in range(batch_item_size):
|
||
cur_item_list = [batch_i[i] for batch_i in batch]
|
||
batch_list.append(paddle.concat(cur_item_list, axis=0))
|
||
else:
|
||
batch_list = batch[0]
|
||
return batch_list
|
||
|
||
|
||
def fill_batch(batch):
|
||
"""
|
||
2020.09.08： The current paddle version only supports returning data with the same length.
|
||
Therefore, fill in the batches with inconsistent lengths.
|
||
this method is currently only useful for text detection
|
||
"""
|
||
keys = list(range(len(batch[0])))
|
||
v_max_len_dict = {}
|
||
for k in keys:
|
||
v_max_len_dict[k] = max([len(item[k]) for item in batch])
|
||
for item in batch:
|
||
length = []
|
||
for k in keys:
|
||
v = item[k]
|
||
length.append(len(v))
|
||
assert isinstance(v, np.ndarray)
|
||
if len(v) == v_max_len_dict[k]:
|
||
continue
|
||
try:
|
||
tmp_shape = [v_max_len_dict[k] - len(v)] + list(v[0].shape)
|
||
except:
|
||
a = 1
|
||
tmp_array = np.zeros(tmp_shape, dtype=v[0].dtype)
|
||
new_array = np.concatenate([v, tmp_array])
|
||
item[k] = new_array
|
||
item.append(length)
|
||
return batch
|