openocr-python 0.0.2__py3-none-any.whl

openocr/tools/data/simple_dataset.py

@@ -0,0 +1,263 @@
+import json
+import math
+import os
+import random
+import traceback
+
+import cv2
+import numpy as np
+from torch.utils.data import Dataset
+
+from openrec.preprocess import create_operators, transform
+
+
+class SimpleDataSet(Dataset):
+
+    def __init__(self, config, mode, logger, seed=None, epoch=0):
+        super(SimpleDataSet, self).__init__()
+        self.logger = logger
+        self.mode = mode.lower()
+
+        global_config = config['Global']
+        dataset_config = config[mode]['dataset']
+        loader_config = config[mode]['loader']
+
+        self.delimiter = dataset_config.get('delimiter', '\t')
+        label_file_list = dataset_config.pop('label_file_list')
+        data_source_num = len(label_file_list)
+        ratio_list = dataset_config.get('ratio_list', 1.0)
+        if isinstance(ratio_list, (float, int)):
+            ratio_list = [float(ratio_list)] * int(data_source_num)
+
+        assert len(
+            ratio_list
+        ) == data_source_num, 'The length of ratio_list should be the same as the file_list.'
+        self.data_dir = dataset_config['data_dir']
+        self.do_shuffle = loader_config['shuffle']
+        self.seed = seed
+        logger.info(f'Initialize indexs of datasets: {label_file_list}')
+        self.data_lines = self.get_image_info_list(label_file_list, ratio_list)
+        self.data_idx_order_list = list(range(len(self.data_lines)))
+        if self.mode == 'train' and self.do_shuffle:
+            self.shuffle_data_random()
+
+        self.set_epoch_as_seed(self.seed, dataset_config)
+
+        self.ops = create_operators(dataset_config['transforms'],
+                                    global_config)
+        self.ext_op_transform_idx = dataset_config.get('ext_op_transform_idx',
+                                                       2)
+        self.need_reset = True in [x < 1 for x in ratio_list]
+
+    def set_epoch_as_seed(self, seed, dataset_config):
+        if self.mode == 'train':
+            try:
+                border_map_id = [
+                    index for index, dictionary in enumerate(
+                        dataset_config['transforms'])
+                    if 'MakeBorderMap' in dictionary
+                ][0]
+                shrink_map_id = [
+                    index for index, dictionary in enumerate(
+                        dataset_config['transforms'])
+                    if 'MakeShrinkMap' in dictionary
+                ][0]
+                dataset_config['transforms'][border_map_id]['MakeBorderMap'][
+                    'epoch'] = seed if seed is not None else 0
+                dataset_config['transforms'][shrink_map_id]['MakeShrinkMap'][
+                    'epoch'] = seed if seed is not None else 0
+            except Exception:
+                return
+
+    def get_image_info_list(self, file_list, ratio_list):
+        if isinstance(file_list, str):
+            file_list = [file_list]
+        data_lines = []
+        for idx, file in enumerate(file_list):
+            with open(file, 'rb') as f:
+                lines = f.readlines()
+                if self.mode == 'train' or ratio_list[idx] < 1.0:
+                    random.seed(self.seed)
+                    lines = random.sample(lines,
+                                          round(len(lines) * ratio_list[idx]))
+                data_lines.extend(lines)
+        return data_lines
+
+    def shuffle_data_random(self):
+        random.seed(self.seed)
+        random.shuffle(self.data_lines)
+        return
+
+    def _try_parse_filename_list(self, file_name):
+        # multiple images -> one gt label
+        if len(file_name) > 0 and file_name[0] == '[':
+            try:
+                info = json.loads(file_name)
+                file_name = random.choice(info)
+            except:
+                pass
+        return file_name
+
+    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:
+            file_idx = self.data_idx_order_list[np.random.randint(
+                self.__len__())]
+            data_line = self.data_lines[file_idx]
+            data_line = data_line.decode('utf-8')
+            substr = data_line.strip('\n').split(self.delimiter)
+            file_name = substr[0]
+            file_name = self._try_parse_filename_list(file_name)
+            label = substr[1]
+            img_path = os.path.join(self.data_dir, file_name)
+            data = {'img_path': img_path, 'label': label}
+            if not os.path.exists(img_path):
+                continue
+            with open(data['img_path'], 'rb') as f:
+                img = f.read()
+                data['image'] = img
+            data = transform(data, load_data_ops)
+
+            if data is None:
+                continue
+            if 'polys' in data.keys():
+                if data['polys'].shape[1] != 4:
+                    continue
+            ext_data.append(data)
+        return ext_data
+
+    def __getitem__(self, idx):
+        file_idx = self.data_idx_order_list[idx]
+        data_line = self.data_lines[file_idx]
+        try:
+            data_line = data_line.decode('utf-8')
+            substr = data_line.strip('\n').split(self.delimiter)
+            file_name = substr[0]
+            file_name = self._try_parse_filename_list(file_name)
+            label = substr[1]
+            img_path = os.path.join(self.data_dir, file_name)
+            data = {'img_path': img_path, 'label': label}
+
+            if not os.path.exists(img_path):
+                raise Exception('{} does not exist!'.format(img_path))
+            with open(data['img_path'], 'rb') as f:
+                img = f.read()
+                data['image'] = img
+            data['ext_data'] = self.get_ext_data()
+            outs = transform(data, self.ops)
+        except:
+            self.logger.error(
+                'When parsing line {}, error happened with msg: {}'.format(
+                    data_line, traceback.format_exc()))
+            outs = None
+        if outs is None:
+            # during evaluation, we should fix the idx to get same results for many times of evaluation.
+            rnd_idx = np.random.randint(self.__len__(
+            )) if self.mode == 'train' else (idx + 1) % self.__len__()
+            return self.__getitem__(rnd_idx)
+        return outs
+
+    def __len__(self):
+        return len(self.data_idx_order_list)
+
+
+class MultiScaleDataSet(SimpleDataSet):
+
+    def __init__(self, config, mode, logger, seed=None):
+        super(MultiScaleDataSet, self).__init__(config, mode, logger, seed)
+        self.ds_width = config[mode]['dataset'].get('ds_width', False)
+        if self.ds_width:
+            self.wh_aware()
+
+    def wh_aware(self):
+        data_line_new = []
+        wh_ratio = []
+        for lins in self.data_lines:
+            data_line_new.append(lins)
+            lins = lins.decode('utf-8')
+            name, label, w, h = lins.strip('\n').split(self.delimiter)
+            wh_ratio.append(float(w) / float(h))
+
+        self.data_lines = data_line_new
+        self.wh_ratio = np.array(wh_ratio)
+        self.wh_ratio_sort = np.argsort(self.wh_ratio)
+        self.data_idx_order_list = list(range(len(self.data_lines)))
+
+    def resize_norm_img(self, data, imgW, imgH, padding=True):
+        img = data['image']
+        h = img.shape[0]
+        w = img.shape[1]
+        if not padding:
+            resized_image = cv2.resize(img, (imgW, imgH),
+                                       interpolation=cv2.INTER_LINEAR)
+            resized_w = imgW
+        else:
+            ratio = w / float(h)
+            if math.ceil(imgH * ratio) > imgW:
+                resized_w = imgW
+            else:
+                resized_w = int(math.ceil(imgH * ratio))
+            resized_image = cv2.resize(img, (resized_w, imgH))
+        resized_image = resized_image.astype('float32')
+
+        resized_image = resized_image.transpose((2, 0, 1)) / 255
+        resized_image -= 0.5
+        resized_image /= 0.5
+        padding_im = np.zeros((3, imgH, imgW), dtype=np.float32)
+        padding_im[:, :, :resized_w] = resized_image
+        valid_ratio = min(1.0, float(resized_w / imgW))
+        data['image'] = padding_im
+        data['valid_ratio'] = valid_ratio
+        return data
+
+    def __getitem__(self, properties):
+        # properites is a tuple, contains (width, height, index)
+        img_height = properties[1]
+        idx = properties[2]
+        if self.ds_width and properties[3] is not None:
+            wh_ratio = properties[3]
+            img_width = img_height * (1 if int(round(wh_ratio)) == 0 else int(
+                round(wh_ratio)))
+            file_idx = self.wh_ratio_sort[idx]
+        else:
+            file_idx = self.data_idx_order_list[idx]
+            img_width = properties[0]
+            wh_ratio = None
+
+        data_line = self.data_lines[file_idx]
+        try:
+            data_line = data_line.decode('utf-8')
+            substr = data_line.strip('\n').split(self.delimiter)
+            file_name = substr[0]
+            file_name = self._try_parse_filename_list(file_name)
+            label = substr[1]
+            img_path = os.path.join(self.data_dir, file_name)
+            data = {'img_path': img_path, 'label': label}
+            if not os.path.exists(img_path):
+                raise Exception('{} does not exist!'.format(img_path))
+            with open(data['img_path'], 'rb') as f:
+                img = f.read()
+                data['image'] = img
+            data['ext_data'] = self.get_ext_data()
+            outs = transform(data, self.ops[:-1])
+            if outs is not None:
+                outs = self.resize_norm_img(outs, img_width, img_height)
+                outs = transform(outs, self.ops[-1:])
+        except:
+            self.logger.error(
+                'When parsing line {}, error happened with msg: {}'.format(
+                    data_line, traceback.format_exc()))
+            outs = None
+        if outs is None:
+            # during evaluation, we should fix the idx to get same results for many times of evaluation.
+            rnd_idx = np.random.randint(self.__len__(
+            )) if self.mode == 'train' else (idx + 1) % self.__len__()
+            return self.__getitem__([img_width, img_height, rnd_idx, wh_ratio])
+        return outs

openocr/tools/data/strlmdb_dataset.py

@@ -0,0 +1,143 @@
+import os
+
+import cv2
+import lmdb
+import numpy as np
+from torch.utils.data import Dataset
+
+from openrec.preprocess import create_operators, transform
+
+
+class STRLMDBDataSet(Dataset):
+
+    def __init__(self, config, mode, logger, seed=None, epoch=1, gpu_i=0):
+        super(STRLMDBDataSet, 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']
+        data_dir = '../training_aug_lmdb_noerror/ep' + str(epoch)
+        self.do_shuffle = loader_config['shuffle']
+
+        self.lmdb_sets = self.load_hierarchical_lmdb_dataset(data_dir)
+        logger.info('Initialize indexs of datasets:%s' % 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)
+
+        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/engine/__init__.py

@@ -0,0 +1,5 @@
+from . import config, trainer
+from .config import *
+from .trainer import *
+
+__all__ = config.__all__ + trainer.__all__

openocr/tools/engine/config.py

@@ -0,0 +1,158 @@
+import os
+from argparse import ArgumentParser, RawDescriptionHelpFormatter
+from collections.abc import Mapping
+
+import yaml
+
+__all__ = ['Config']
+
+
+class ArgsParser(ArgumentParser):
+
+    def __init__(self):
+        super(ArgsParser,
+              self).__init__(formatter_class=RawDescriptionHelpFormatter)
+        self.add_argument('-o',
+                          '--opt',
+                          nargs='*',
+                          help='set configuration options')
+        self.add_argument('--local_rank')
+
+    def parse_args(self, argv=None):
+        args = super(ArgsParser, self).parse_args(argv)
+        assert args.config is not None, 'Please specify --config=configure_file_path.'
+        args.opt = self._parse_opt(args.opt)
+        return args
+
+    def _parse_opt(self, opts):
+        config = {}
+        if not opts:
+            return config
+        for s in opts:
+            s = s.strip()
+            k, v = s.split('=', 1)
+            if '.' not in k:
+                config[k] = yaml.load(v, Loader=yaml.Loader)
+            else:
+                keys = k.split('.')
+                if keys[0] not in config:
+                    config[keys[0]] = {}
+                cur = config[keys[0]]
+                for idx, key in enumerate(keys[1:]):
+                    if idx == len(keys) - 2:
+                        cur[key] = yaml.load(v, Loader=yaml.Loader)
+                    else:
+                        cur[key] = {}
+                        cur = cur[key]
+        return config
+
+
+class AttrDict(dict):
+    """Single level attribute dict, NOT recursive."""
+
+    def __init__(self, **kwargs):
+        super(AttrDict, self).__init__()
+        super(AttrDict, self).update(kwargs)
+
+    def __getattr__(self, key):
+        if key in self:
+            return self[key]
+        raise AttributeError("object has no attribute '{}'".format(key))
+
+
+def _merge_dict(config, merge_dct):
+    """Recursive dict merge. Inspired by :meth:``dict.update()``, instead of
+    updating only top-level keys, dict_merge recurses down into dicts nested to
+    an arbitrary depth, updating keys. The ``merge_dct`` is merged into
+    ``dct``.
+
+    Args:
+        config: dict onto which the merge is executed
+        merge_dct: dct merged into config
+
+    Returns: dct
+    """
+    for key, value in merge_dct.items():
+        sub_keys = key.split('.')
+        key = sub_keys[0]
+        if key in config and len(sub_keys) > 1:
+            _merge_dict(config[key], {'.'.join(sub_keys[1:]): value})
+        elif key in config and isinstance(config[key], dict) and isinstance(
+                value, Mapping):
+            _merge_dict(config[key], value)
+        else:
+            config[key] = value
+    return config
+
+
+def print_dict(cfg, print_func=print, delimiter=0):
+    """Recursively visualize a dict and indenting acrrording by the
+    relationship of keys."""
+    for k, v in sorted(cfg.items()):
+        if isinstance(v, dict):
+            print_func('{}{} : '.format(delimiter * ' ', str(k)))
+            print_dict(v, print_func, delimiter + 4)
+        elif isinstance(v, list) and len(v) >= 1 and isinstance(v[0], dict):
+            print_func('{}{} : '.format(delimiter * ' ', str(k)))
+            for value in v:
+                print_dict(value, print_func, delimiter + 4)
+        else:
+            print_func('{}{} : {}'.format(delimiter * ' ', k, v))
+
+
+class Config(object):
+
+    def __init__(self, config_path, BASE_KEY='_BASE_'):
+        self.BASE_KEY = BASE_KEY
+        self.cfg = self._load_config_with_base(config_path)
+
+    def _load_config_with_base(self, file_path):
+        """Load config from file.
+
+        Args:
+            file_path (str): Path of the config file to be loaded.
+
+        Returns: global config
+        """
+        _, ext = os.path.splitext(file_path)
+        assert ext in ['.yml', '.yaml'], 'only support yaml files for now'
+
+        with open(file_path) as f:
+            file_cfg = yaml.load(f, Loader=yaml.Loader)
+
+        # NOTE: cfgs outside have higher priority than cfgs in _BASE_
+        if self.BASE_KEY in file_cfg:
+            all_base_cfg = AttrDict()
+            base_ymls = list(file_cfg[self.BASE_KEY])
+            for base_yml in base_ymls:
+                if base_yml.startswith('~'):
+                    base_yml = os.path.expanduser(base_yml)
+                if not base_yml.startswith('/'):
+                    base_yml = os.path.join(os.path.dirname(file_path),
+                                            base_yml)
+
+                with open(base_yml) as f:
+                    base_cfg = self._load_config_with_base(base_yml)
+                    all_base_cfg = _merge_dict(all_base_cfg, base_cfg)
+
+            del file_cfg[self.BASE_KEY]
+            file_cfg = _merge_dict(all_base_cfg, file_cfg)
+        file_cfg['filename'] = os.path.splitext(
+            os.path.split(file_path)[-1])[0]
+        return file_cfg
+
+    def merge_dict(self, args):
+        self.cfg = _merge_dict(self.cfg, args)
+
+    def print_cfg(self, print_func=print):
+        """Recursively visualize a dict and indenting acrrording by the
+        relationship of keys."""
+        print_func('----------- Config -----------')
+        print_dict(self.cfg, print_func)
+        print_func('---------------------------------------------')
+
+    def save(self, p, cfg=None):
+        if cfg is None:
+            cfg = self.cfg
+        with open(p, 'w') as f:
+            yaml.dump(dict(cfg), f, default_flow_style=False, sort_keys=False)