mmocr.datasets.recog_lmdb_dataset 源代码
# Copyright (c) OpenMMLab. All rights reserved.
import json
import os.path as osp
import warnings
from typing import Callable, List, Optional, Sequence, Union
from mmengine.dataset import BaseDataset
from mmengine.utils import is_abs
from mmocr.registry import DATASETS, TASK_UTILS
[文档]@DATASETS.register_module()
class RecogLMDBDataset(BaseDataset):
r"""RecogLMDBDataset for text recognition.
The annotation format should be in lmdb format. We support two lmdb
formats, one is the lmdb file with only labels generated by txt2lmdb
(deprecated), and another one is the lmdb file generated by recog2lmdb.
The former format stores string in `filename text` format directly in lmdb,
while the latter uses `image_key` as well as `label_key` for querying.
Args:
ann_file (str): Annotation file path. Defaults to ''.
parse_cfg (dict, optional): Config of parser for parsing annotations.
Use ``LineJsonParser`` when the annotation file is in jsonl format
with keys of ``filename`` and ``text``. The keys in parse_cfg
should be consistent with the keys in jsonl annotations. The first
key in parse_cfg should be the key of the path in jsonl
annotations. The second key in parse_cfg should be the key of the
text in jsonl Use ``LineStrParser`` when the annotation file is in
txt format. Defaults to
``dict(type='LineJsonParser', keys=['filename', 'text'])``.
metainfo (dict, optional): Meta information for dataset, such as class
information. Defaults to None.
data_root (str): The root directory for ``data_prefix`` and
``ann_file``. Defaults to ''.
data_prefix (dict): Prefix for training data. Defaults to
``dict(img_path='')``.
filter_cfg (dict, optional): Config for filter data. Defaults to None.
indices (int or Sequence[int], optional): Support using first few
data in annotation file to facilitate training/testing on a smaller
dataset. Defaults to None which means using all ``data_infos``.
serialize_data (bool, optional): Whether to hold memory using
serialized objects, when enabled, data loader workers can use
shared RAM from master process instead of making a copy. Defaults
to True.
pipeline (list, optional): Processing pipeline. Defaults to [].
test_mode (bool, optional): ``test_mode=True`` means in test phase.
Defaults to False.
lazy_init (bool, optional): Whether to load annotation during
instantiation. In some cases, such as visualization, only the meta
information of the dataset is needed, which is not necessary to
load annotation file. ``RecogLMDBDataset`` can skip load
annotations to save time by set ``lazy_init=False``.
Defaults to False.
max_refetch (int, optional): If ``RecogLMDBdataset.prepare_data`` get a
None img. The maximum extra number of cycles to get a valid
image. Defaults to 1000.
"""
def __init__(self,
ann_file: str = '',
parser_cfg: Optional[dict] = dict(
type='LineJsonParser', keys=['filename', 'text']),
metainfo: Optional[dict] = None,
data_root: Optional[str] = '',
data_prefix: dict = dict(img_path=''),
filter_cfg: Optional[dict] = None,
indices: Optional[Union[int, Sequence[int]]] = None,
serialize_data: bool = True,
pipeline: List[Union[dict, Callable]] = [],
test_mode: bool = False,
lazy_init: bool = False,
max_refetch: int = 1000) -> None:
if parser_cfg['type'] != 'LineJsonParser':
raise ValueError('We only support using LineJsonParser '
'to parse lmdb file. Please use LineJsonParser '
'in the dataset config')
self.parser = TASK_UTILS.build(parser_cfg)
self.ann_file = ann_file
self.deprecated_format = False
env = self._get_env(root=data_root)
with env.begin(write=False) as txn:
try:
self.total_number = int(
txn.get(b'num-samples').decode('utf-8'))
except AttributeError:
warnings.warn(
'DeprecationWarning: The lmdb dataset generated with '
'txt2lmdb will be deprecate, please use the latest '
'tools/data/utils/recog2lmdb to generate lmdb dataset. '
'See https://mmocr.readthedocs.io/en/latest/tools.html#'
'convert-text-recognition-dataset-to-lmdb-format for '
'details.', UserWarning)
self.total_number = int(
txn.get(b'total_number').decode('utf-8'))
self.deprecated_format = True
# The lmdb file may contain only the label, or it may contain both
# the label and the image, so we use image_key here for probing.
image_key = f'image-{1:09d}'
if txn.get(image_key.encode('utf-8')) is None:
self.label_only = True
else:
self.label_only = False
super().__init__(
ann_file=ann_file,
metainfo=metainfo,
data_root=data_root,
data_prefix=data_prefix,
filter_cfg=filter_cfg,
indices=indices,
serialize_data=serialize_data,
pipeline=pipeline,
test_mode=test_mode,
lazy_init=lazy_init,
max_refetch=max_refetch)
[文档] def load_data_list(self) -> List[dict]:
"""Load annotations from an annotation file named as ``self.ann_file``
Returns:
List[dict]: A list of annotation.
"""
if not hasattr(self, 'env'):
self.env = self._get_env()
data_list = []
with self.env.begin(write=False) as txn:
for i in range(self.total_number):
if self.deprecated_format:
line = txn.get(str(i).encode('utf-8')).decode('utf-8')
filename, text = line.strip('/n').split(' ')
line = json.dumps(
dict(filename=filename, text=text), ensure_ascii=False)
else:
i = i + 1
label_key = f'label-{i:09d}'
if self.label_only:
line = txn.get(
label_key.encode('utf-8')).decode('utf-8')
else:
img_key = f'image-{i:09d}'
text = txn.get(
label_key.encode('utf-8')).decode('utf-8')
line = json.dumps(
dict(filename=img_key, text=text),
ensure_ascii=False)
data_list.append(self.parse_data_info(line))
return data_list
[文档] def parse_data_info(self, raw_anno_info: str) -> Union[dict, List[dict]]:
"""Parse raw annotation to target format.
Args:
raw_anno_info (str): One raw data information loaded
from ``ann_file``.
Returns:
(dict): Parsed annotation.
"""
data_info = {}
parsed_anno = self.parser(raw_anno_info)
img_path = osp.join(self.data_prefix['img_path'],
parsed_anno[self.parser.keys[0]])
data_info['img_path'] = img_path
data_info['instances'] = [dict(text=parsed_anno[self.parser.keys[1]])]
return data_info
def _get_env(self, root=''):
"""Get lmdb environment from self.ann_file.
Returns:
Lmdb environment.
"""
try:
import lmdb
except ImportError:
raise ImportError(
'Please install lmdb to enable RecogLMDBDataset.')
lmdb_path = self.ann_file if is_abs(self.ann_file) else osp.join(
root, self.ann_file)
return lmdb.open(
lmdb_path,
max_readers=1,
readonly=True,
lock=False,
readahead=False,
meminit=False,
)
[文档] def close(self):
"""Close lmdb environment."""
if hasattr(self, 'env'):
self.env.close()
del self.env