openocr-python 0.0.2__py3-none-any.whl

openocr/openrec/preprocess/visionlan_label_encode.py

@@ -0,0 +1,67 @@
+from random import sample
+
+import numpy as np
+
+from .ctc_label_encode import BaseRecLabelEncode
+
+
+class VisionLANLabelEncode(BaseRecLabelEncode):
+    """Convert between text-label and text-index."""
+
+    def __init__(self,
+                 max_text_length,
+                 character_dict_path=None,
+                 use_space_char=False,
+                 **kwargs):
+        super(VisionLANLabelEncode,
+              self).__init__(max_text_length, character_dict_path,
+                             use_space_char)
+        self.dict = {}
+        for i, char in enumerate(self.character):
+            self.dict[char] = i
+
+    def __call__(self, data):
+        text = data['label']  # original string
+        # generate occluded text
+        len_str = len(text)
+        if len_str <= 0:
+            return None
+        change_num = 1
+        order = list(range(len_str))
+        change_id = sample(order, change_num)[0]
+        label_sub = text[change_id]
+        if change_id == (len_str - 1):
+            label_res = text[:change_id]
+        elif change_id == 0:
+            label_res = text[1:]
+        else:
+            label_res = text[:change_id] + text[change_id + 1:]
+
+        data['label_res'] = label_res  # remaining string
+        data['label_sub'] = label_sub  # occluded character
+        data['label_id'] = change_id  # character index
+        # encode label
+        text = self.encode(text)
+        if text is None:
+            return None
+        text = [i + 1 for i in text]
+        data['length'] = np.array(len(text))
+        text = text + [0] * (self.max_text_len + 1 - len(text))
+        data['label'] = np.array(text)
+        label_res = self.encode(label_res)
+        label_sub = self.encode(label_sub)
+        if label_res is None:
+            label_res = []
+        else:
+            label_res = [i + 1 for i in label_res]
+        if label_sub is None:
+            label_sub = []
+        else:
+            label_sub = [i + 1 for i in label_sub]
+        data['length_res'] = np.array(len(label_res))
+        data['length_sub'] = np.array(len(label_sub))
+        label_res = label_res + [0] * (self.max_text_len - len(label_res))
+        label_sub = label_sub + [0] * (self.max_text_len - len(label_sub))
+        data['label_res'] = np.array(label_res)
+        data['label_sub'] = np.array(label_sub)
+        return data

openocr/tools/create_lmdb_dataset.py

@@ -0,0 +1,118 @@
+import os
+import lmdb
+import cv2
+from tqdm import tqdm
+import numpy as np
+import io
+from PIL import Image
+""" a modified version of CRNN torch repository https://github.com/bgshih/crnn/blob/master/tool/create_dataset.py """
+
+
+def get_datalist(data_dir, data_path, max_len):
+    """
+    获取训练和验证的数据list
+    :param data_dir: 数据集根目录
+    :param data_path: 训练的dataset文件列表，每个文件内以如下格式存储 ‘path/to/img\tlabel’
+    :return:
+    """
+    train_data = []
+    if isinstance(data_path, list):
+        for p in data_path:
+            train_data.extend(get_datalist(data_dir, p, max_len))
+    else:
+        with open(data_path, 'r', encoding='utf-8') as f:
+            for line in tqdm(f.readlines(),
+                             desc=f'load data from {data_path}'):
+                line = (line.strip('\n').replace('.jpg ', '.jpg\t').replace(
+                    '.png ', '.png\t').split('\t'))
+                if len(line) > 1:
+                    img_path = os.path.join(data_dir, line[0].strip(' '))
+                    label = line[1]
+                    if len(label) > max_len:
+                        continue
+                    if os.path.exists(
+                            img_path) and os.path.getsize(img_path) > 0:
+                        train_data.append([str(img_path), label])
+    return train_data
+
+
+def checkImageIsValid(imageBin):
+    if imageBin is None:
+        return False
+    imageBuf = np.frombuffer(imageBin, dtype=np.uint8)
+    img = cv2.imdecode(imageBuf, cv2.IMREAD_GRAYSCALE)
+    imgH, imgW = img.shape[0], img.shape[1]
+    if imgH * imgW == 0:
+        return False
+    return True
+
+
+def writeCache(env, cache):
+    with env.begin(write=True) as txn:
+        for k, v in cache.items():
+            txn.put(k, v)
+
+
+def createDataset(data_list, outputPath, checkValid=True):
+    """
+    Create LMDB dataset for training and evaluation.
+    ARGS:
+        inputPath  : input folder path where starts imagePath
+        outputPath : LMDB output path
+        gtFile     : list of image path and label
+        checkValid : if true, check the validity of every image
+    """
+    os.makedirs(outputPath, exist_ok=True)
+    env = lmdb.open(outputPath, map_size=1099511627776)
+    cache = {}
+    cnt = 1
+    for imagePath, label in tqdm(data_list,
+                                 desc=f'make dataset, save to {outputPath}'):
+        with open(imagePath, 'rb') as f:
+            imageBin = f.read()
+            buf = io.BytesIO(imageBin)
+            w, h = Image.open(buf).size
+        if checkValid:
+            try:
+                if not checkImageIsValid(imageBin):
+                    print('%s is not a valid image' % imagePath)
+                    continue
+            except:
+                continue
+
+        imageKey = 'image-%09d'.encode() % cnt
+        labelKey = 'label-%09d'.encode() % cnt
+        whKey = 'wh-%09d'.encode() % cnt
+        cache[imageKey] = imageBin
+        cache[labelKey] = label.encode()
+        cache[whKey] = (str(w) + '_' + str(h)).encode()
+
+        if cnt % 1000 == 0:
+            writeCache(env, cache)
+            cache = {}
+        cnt += 1
+    nSamples = cnt - 1
+    cache['num-samples'.encode()] = str(nSamples).encode()
+    writeCache(env, cache)
+    print('Created dataset with %d samples' % nSamples)
+
+
+if __name__ == '__main__':
+    data_dir = './Union14M-L/'
+    label_file_list = [
+        './Union14M-L/train_annos/filter_jsonl_mmocr0.x/filter_train_challenging.jsonl.txt',
+        './Union14M-L/train_annos/filter_jsonl_mmocr0.x/filter_train_easy.jsonl.txt',
+        './Union14M-L/train_annos/filter_jsonl_mmocr0.x/filter_train_hard.jsonl.txt',
+        './Union14M-L/train_annos/filter_jsonl_mmocr0.x/filter_train_medium.jsonl.txt',
+        './Union14M-L/train_annos/filter_jsonl_mmocr0.x/filter_train_normal.jsonl.txt'
+    ]
+    save_path_root = './Union14M-L-LMDB-Filtered/'
+
+    for data_list in label_file_list:
+        save_path = save_path_root + data_list.split('/')[-1].split(
+            '.')[0] + '/'
+        os.makedirs(save_path, exist_ok=True)
+        print(save_path)
+        train_data_list = get_datalist(data_dir, data_list, 800)
+
+        createDataset(train_data_list, save_path)

openocr/tools/data/__init__.py

@@ -0,0 +1,94 @@
+import os
+import sys
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
+
+import copy
+
+from torch.utils.data import DataLoader, DistributedSampler
+
+from tools.data.lmdb_dataset import LMDBDataSet
+from tools.data.lmdb_dataset_test import LMDBDataSetTest
+from tools.data.multi_scale_sampler import MultiScaleSampler
+from tools.data.ratio_dataset import RatioDataSet
+from tools.data.ratio_dataset_test import RatioDataSetTest
+from tools.data.ratio_dataset_tvresize_test import RatioDataSetTVResizeTest
+from tools.data.ratio_dataset_tvresize import RatioDataSetTVResize
+from tools.data.ratio_sampler import RatioSampler
+from tools.data.simple_dataset import MultiScaleDataSet, SimpleDataSet
+from tools.data.strlmdb_dataset import STRLMDBDataSet
+
+__all__ = [
+    'build_dataloader',
+    'transform',
+    'create_operators',
+]
+
+
+def build_dataloader(config, mode, logger, seed=None, epoch=3):
+    config = copy.deepcopy(config)
+
+    support_dict = [
+        'SimpleDataSet', 'LMDBDataSet', 'MultiScaleDataSet', 'STRLMDBDataSet',
+        'LMDBDataSetTest', 'RatioDataSet', 'RatioDataSetTest',
+        'RatioDataSetTVResize', 'RatioDataSetTVResizeTest'
+    ]
+    module_name = config[mode]['dataset']['name']
+    assert module_name in support_dict, Exception(
+        'DataSet only support {}/{}'.format(support_dict, module_name))
+    assert mode in ['Train', 'Eval',
+                    'Test'], 'Mode should be Train, Eval or Test.'
+
+    dataset = eval(module_name)(config, mode, logger, seed, epoch=epoch)
+    loader_config = config[mode]['loader']
+    batch_size = loader_config['batch_size_per_card']
+    drop_last = loader_config['drop_last']
+    shuffle = loader_config['shuffle']
+    num_workers = loader_config['num_workers']
+    if 'pin_memory' in loader_config.keys():
+        pin_memory = loader_config['use_shared_memory']
+    else:
+        pin_memory = False
+
+    sampler = None
+    batch_sampler = None
+    if 'sampler' in config[mode]:
+        config_sampler = config[mode]['sampler']
+        sampler_name = config_sampler.pop('name')
+        batch_sampler = eval(sampler_name)(dataset, **config_sampler)
+    elif config['Global']['distributed'] and mode == 'Train':
+        sampler = DistributedSampler(dataset=dataset, shuffle=shuffle)
+
+    if 'collate_fn' in loader_config:
+        from . import collate_fn
+
+        collate_fn = getattr(collate_fn, loader_config['collate_fn'])()
+    else:
+        collate_fn = None
+    if batch_sampler is None:
+        data_loader = DataLoader(
+            dataset=dataset,
+            sampler=sampler,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            collate_fn=collate_fn,
+            batch_size=batch_size,
+            drop_last=drop_last,
+        )
+    else:
+        data_loader = DataLoader(
+            dataset=dataset,
+            batch_sampler=batch_sampler,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            collate_fn=collate_fn,
+        )
+    if len(data_loader) == 0:
+        logger.error(
+            f'No Images in {mode.lower()} dataloader, please ensure\n'
+            '\t1. The images num in the train label_file_list should be larger than or equal with batch size.\n'
+            '\t2. The annotation file and path in the configuration file are provided normally.\n'
+            '\t3. The BatchSize is large than images.')
+        sys.exit()
+    return data_loader

openocr/tools/data/collate_fn.py

@@ -0,0 +1,100 @@
+import numbers
+from collections import defaultdict
+
+import numpy as np
+import torch
+
+
+class DictCollator(object):
+    """data batch."""
+
+    def __call__(self, batch):
+        data_dict = defaultdict(list)
+        to_tensor_keys = []
+        for sample in batch:
+            for k, v in sample.items():
+                if isinstance(v, (np.ndarray, torch.Tensor, numbers.Number)):
+                    if k not in to_tensor_keys:
+                        to_tensor_keys.append(k)
+                data_dict[k].append(v)
+        for k in to_tensor_keys:
+            data_dict[k] = torch.from_numpy(data_dict[k])
+        return data_dict
+
+
+class ListCollator(object):
+    """data batch."""
+
+    def __call__(self, batch):
+        data_dict = defaultdict(list)
+        to_tensor_idxs = []
+        for sample in batch:
+            for idx, v in enumerate(sample):
+                if isinstance(v, (np.ndarray, torch.Tensor, numbers.Number)):
+                    if idx not in to_tensor_idxs:
+                        to_tensor_idxs.append(idx)
+                data_dict[idx].append(v)
+        for idx in to_tensor_idxs:
+            data_dict[idx] = torch.from_numpy(data_dict[idx])
+        return list(data_dict.values())
+
+
+class SSLRotateCollate(object):
+    """
+    bach: [
+        [(4*3xH*W), (4,)]
+        [(4*3xH*W), (4,)]
+        ...
+    ]
+    """
+
+    def __call__(self, batch):
+        output = [np.concatenate(d, axis=0) for d in zip(*batch)]
+        return output
+
+
+class DyMaskCollator(object):
+    """
+    batch: [
+        image [batch_size, channel, maxHinbatch, maxWinbatch]
+        image_mask [batch_size, channel, maxHinbatch, maxWinbatch]
+        label [batch_size, maxLabelLen]
+        label_mask [batch_size, maxLabelLen]
+        ...
+    ]
+    """
+
+    def __call__(self, batch):
+        max_width, max_height, max_length = 0, 0, 0
+        bs, channel = len(batch), batch[0][0].shape[0]
+        proper_items = []
+        for item in batch:
+            if item[0].shape[1] * max_width > 1600 * 320 or item[0].shape[
+                    2] * max_height > 1600 * 320:
+                continue
+            max_height = item[0].shape[
+                1] if item[0].shape[1] > max_height else max_height
+            max_width = item[0].shape[
+                2] if item[0].shape[2] > max_width else max_width
+            max_length = len(
+                item[1]) if len(item[1]) > max_length else max_length
+            proper_items.append(item)
+
+        images, image_masks = np.zeros(
+            (len(proper_items), channel, max_height, max_width),
+            dtype='float32'), np.zeros(
+                (len(proper_items), 1, max_height, max_width), dtype='float32')
+        labels, label_masks = np.zeros((len(proper_items), max_length),
+                                       dtype='int64'), np.zeros(
+                                           (len(proper_items), max_length),
+                                           dtype='int64')
+
+        for i in range(len(proper_items)):
+            _, h, w = proper_items[i][0].shape
+            images[i][:, :h, :w] = proper_items[i][0]
+            image_masks[i][:, :h, :w] = 1
+            l = len(proper_items[i][1])
+            labels[i][:l] = proper_items[i][1]
+            label_masks[i][:l] = 1
+
+        return images, image_masks, labels, label_masks

openocr/tools/data/lmdb_dataset.py

@@ -0,0 +1,142 @@
+import os
+
+import cv2
+import lmdb
+import numpy as np
+from torch.utils.data import Dataset
+
+from openrec.preprocess import create_operators, transform
+
+
+class LMDBDataSet(Dataset):
+
+    def __init__(self, config, mode, logger, seed=None, epoch=1):
+        super(LMDBDataSet, self).__init__()
+
+        global_config = config['Global']
+        dataset_config = config[mode]['dataset']
+        loader_config = config[mode]['loader']
+        loader_config['batch_size_per_card']
+        data_dir = dataset_config['data_dir']
+        self.do_shuffle = loader_config['shuffle']
+
+        self.lmdb_sets = self.load_hierarchical_lmdb_dataset(data_dir)
+        logger.info(f'Initialize indexs of datasets: {data_dir}')
+        self.data_idx_order_list = self.dataset_traversal()
+        if self.do_shuffle:
+            np.random.shuffle(self.data_idx_order_list)
+        self.ops = create_operators(dataset_config['transforms'],
+                                    global_config)
+        self.ext_op_transform_idx = dataset_config.get('ext_op_transform_idx',
+                                                       1)
+
+        ratio_list = dataset_config.get('ratio_list', [1.0])
+        self.need_reset = True in [x < 1 for x in ratio_list]
+
+    def load_hierarchical_lmdb_dataset(self, data_dir):
+        lmdb_sets = {}
+        dataset_idx = 0
+        for dirpath, dirnames, filenames in os.walk(data_dir + '/'):
+            if not dirnames:
+                env = lmdb.open(
+                    dirpath,
+                    max_readers=32,
+                    readonly=True,
+                    lock=False,
+                    readahead=False,
+                    meminit=False,
+                )
+                txn = env.begin(write=False)
+                num_samples = int(txn.get('num-samples'.encode()))
+                lmdb_sets[dataset_idx] = {
+                    'dirpath': dirpath,
+                    'env': env,
+                    'txn': txn,
+                    'num_samples': num_samples,
+                }
+                dataset_idx += 1
+        return lmdb_sets
+
+    def dataset_traversal(self):
+        lmdb_num = len(self.lmdb_sets)
+        total_sample_num = 0
+        for lno in range(lmdb_num):
+            total_sample_num += self.lmdb_sets[lno]['num_samples']
+        data_idx_order_list = np.zeros((total_sample_num, 2))
+        beg_idx = 0
+        for lno in range(lmdb_num):
+            tmp_sample_num = self.lmdb_sets[lno]['num_samples']
+            end_idx = beg_idx + tmp_sample_num
+            data_idx_order_list[beg_idx:end_idx, 0] = lno
+            data_idx_order_list[beg_idx:end_idx,
+                                1] = list(range(tmp_sample_num))
+            data_idx_order_list[beg_idx:end_idx, 1] += 1
+            beg_idx = beg_idx + tmp_sample_num
+        return data_idx_order_list
+
+    def get_img_data(self, value):
+        """get_img_data."""
+        if not value:
+            return None
+        imgdata = np.frombuffer(value, dtype='uint8')
+        if imgdata is None:
+            return None
+        imgori = cv2.imdecode(imgdata, 1)
+        if imgori is None:
+            return None
+        return imgori
+
+    def get_ext_data(self):
+        ext_data_num = 0
+        for op in self.ops:
+            if hasattr(op, 'ext_data_num'):
+                ext_data_num = getattr(op, 'ext_data_num')
+                break
+        load_data_ops = self.ops[:self.ext_op_transform_idx]
+        ext_data = []
+
+        while len(ext_data) < ext_data_num:
+            lmdb_idx, file_idx = self.data_idx_order_list[np.random.randint(
+                len(self))]
+            lmdb_idx = int(lmdb_idx)
+            file_idx = int(file_idx)
+            sample_info = self.get_lmdb_sample_info(
+                self.lmdb_sets[lmdb_idx]['txn'], file_idx)
+            if sample_info is None:
+                continue
+            img, label = sample_info
+            data = {'image': img, 'label': label}
+            data = transform(data, load_data_ops)
+            if data is None:
+                continue
+            ext_data.append(data)
+        return ext_data
+
+    def get_lmdb_sample_info(self, txn, index):
+        label_key = 'label-%09d'.encode() % index
+        label = txn.get(label_key)
+        if label is None:
+            return None
+        label = label.decode('utf-8')
+        img_key = 'image-%09d'.encode() % index
+        imgbuf = txn.get(img_key)
+        return imgbuf, label
+
+    def __getitem__(self, idx):
+        lmdb_idx, file_idx = self.data_idx_order_list[idx]
+        lmdb_idx = int(lmdb_idx)
+        file_idx = int(file_idx)
+        sample_info = self.get_lmdb_sample_info(
+            self.lmdb_sets[lmdb_idx]['txn'], file_idx)
+        if sample_info is None:
+            return self.__getitem__(np.random.randint(self.__len__()))
+        img, label = sample_info
+        data = {'image': img, 'label': label}
+        data['ext_data'] = self.get_ext_data()
+        outs = transform(data, self.ops)
+        if outs is None:
+            return self.__getitem__(np.random.randint(self.__len__()))
+        return outs
+
+    def __len__(self):
+        return self.data_idx_order_list.shape[0]

openocr/tools/data/lmdb_dataset_test.py

@@ -0,0 +1,166 @@
+import io
+import re
+import unicodedata
+
+import lmdb
+from PIL import Image
+from torch.utils.data import Dataset
+
+from openrec.preprocess import create_operators, transform
+
+
+class CharsetAdapter:
+    """Transforms labels according to the target charset."""
+
+    def __init__(self, target_charset) -> None:
+        super().__init__()
+        self.lowercase_only = target_charset == target_charset.lower()
+        self.uppercase_only = target_charset == target_charset.upper()
+        self.unsupported = re.compile(f'[^{re.escape(target_charset)}]')
+
+    def __call__(self, label):
+        if self.lowercase_only:
+            label = label.lower()
+        elif self.uppercase_only:
+            label = label.upper()
+        # Remove unsupported characters
+        label = self.unsupported.sub('', label)
+        return label
+
+
+class LMDBDataSetTest(Dataset):
+    """Dataset interface to an LMDB database.
+
+    It supports both labelled and unlabelled datasets. For unlabelled datasets,
+    the image index itself is returned as the label. Unicode characters are
+    normalized by default. Case-sensitivity is inferred from the charset.
+    Labels are transformed according to the charset.
+    """
+
+    def __init__(self,
+                 config,
+                 mode,
+                 logger,
+                 seed=None,
+                 epoch=1,
+                 gpu_i=0,
+                 max_label_len: int = 25,
+                 min_image_dim: int = 0,
+                 remove_whitespace: bool = True,
+                 normalize_unicode: bool = True,
+                 unlabelled: bool = False,
+                 transform=None):
+        dataset_config = config[mode]['dataset']
+        global_config = config['Global']
+        max_label_len = global_config['max_text_length']
+        self.root = dataset_config['data_dir']
+        self._env = None
+        self.unlabelled = unlabelled
+        self.transform = transform
+        self.labels = []
+        self.filtered_index_list = []
+        self.min_image_dim = min_image_dim
+        self.filter_label = dataset_config.get('filter_label',
+                                               True)  #'data_dir']filter_label
+        character_dict_path = global_config.get('character_dict_path', None)
+        use_space_char = global_config.get('use_space_char', False)
+        if character_dict_path is None:
+            char_test = '0123456789abcdefghijklmnopqrstuvwxyz'
+        else:
+            char_test = ''
+            with open(character_dict_path, 'rb') as fin:
+                lines = fin.readlines()
+                for line in lines:
+                    line = line.decode('utf-8').strip('\n').strip('\r\n')
+                    char_test += line
+            if use_space_char:
+                char_test += ' '
+        self.ops = create_operators(dataset_config['transforms'],
+                                    global_config)
+        self.num_samples = self._preprocess_labels(char_test,
+                                                   remove_whitespace,
+                                                   normalize_unicode,
+                                                   max_label_len,
+                                                   min_image_dim)
+
+    def __del__(self):
+        if self._env is not None:
+            self._env.close()
+            self._env = None
+
+    def _create_env(self):
+        return lmdb.open(self.root,
+                         max_readers=1,
+                         readonly=True,
+                         create=False,
+                         readahead=False,
+                         meminit=False,
+                         lock=False)
+
+    @property
+    def env(self):
+        if self._env is None:
+            self._env = self._create_env()
+        return self._env
+
+    def _preprocess_labels(self, charset, remove_whitespace, normalize_unicode,
+                           max_label_len, min_image_dim):
+        charset_adapter = CharsetAdapter(charset)
+        with self._create_env() as env, env.begin() as txn:
+            num_samples = int(txn.get('num-samples'.encode()))
+            if self.unlabelled:
+                return num_samples
+            for index in range(num_samples):
+                index += 1  # lmdb starts with 1
+                label_key = f'label-{index:09d}'.encode()
+                label = txn.get(label_key).decode()
+                # Normally, whitespace is removed from the labels.
+                if remove_whitespace:
+                    label = ''.join(label.split())
+                # Normalize unicode composites (if any) and convert to compatible ASCII characters
+                if self.filter_label:
+                    # if normalize_unicode:
+                    label = unicodedata.normalize('NFKD', label).encode(
+                        'ascii', 'ignore').decode()
+                # Filter by length before removing unsupported characters. The original label might be too long.
+                if len(label) > max_label_len:
+                    continue
+
+                if self.filter_label:
+                    label = charset_adapter(label)
+                # We filter out samples which don't contain any supported characters
+                if not label:
+                    continue
+                # Filter images that are too small.
+                if min_image_dim > 0:
+                    img_key = f'image-{index:09d}'.encode()
+                    img = txn.get(img_key)
+                    data = {'image': img, 'label': label}
+                    outs = transform(data, self.ops)
+                    if outs is None:
+                        continue
+                    buf = io.BytesIO(img)
+                    w, h = Image.open(buf).size
+                    if w < self.min_image_dim or h < self.min_image_dim:
+                        continue
+                self.labels.append(label)
+                self.filtered_index_list.append(index)
+        return len(self.labels)
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, index):
+        if self.unlabelled:
+            label = index
+        else:
+            label = self.labels[index]
+            index = self.filtered_index_list[index]
+
+        img_key = f'image-{index:09d}'.encode()
+        with self.env.begin() as txn:
+            img = txn.get(img_key)
+        data = {'image': img, 'label': label}
+        outs = transform(data, self.ops)
+
+        return outs