wehub-resource-sync
281 lines
13 KiB
Python
281 lines
13 KiB
Python
import os
|
|
from fairseq import search
|
|
|
|
from fairseq import scoring, utils, metrics
|
|
from fairseq.data import Dictionary, encoders
|
|
from fairseq.tasks import LegacyFairseqTask, register_task
|
|
from fairseq.tasks.fairseq_task import FairseqTask
|
|
|
|
try:
|
|
from .data import SROIETextRecognitionDataset, Receipt53KDataset, SyntheticTextRecognitionDataset
|
|
from .data_aug import build_data_aug, OptForDataAugment, DataAugment
|
|
except ImportError:
|
|
from data import SROIETextRecognitionDataset, Receipt53KDataset, SyntheticTextRecognitionDataset
|
|
from data_aug import build_data_aug, OptForDataAugment, DataAugment
|
|
|
|
import logging
|
|
import torch
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@register_task('text_recognition')
|
|
class TextRecognitionTask(LegacyFairseqTask):
|
|
|
|
@staticmethod
|
|
def add_args(parser):
|
|
parser.add_argument('data', metavar='DIR',
|
|
help='the path to the data dir')
|
|
|
|
parser.add_argument('--reset-dictionary', action='store_true',
|
|
help='if reset dictionary and related parameters')
|
|
parser.add_argument('--adapt-dictionary', action='store_true',
|
|
help='if adapt dictionary and related parameters')
|
|
parser.add_argument('--adapt-encoder-pos-embed', action='store_true',
|
|
help='if adapt encoder pos embed')
|
|
|
|
parser.add_argument('--add-empty-sample', action='store_true',
|
|
help='add empty samples to the dataset (for multilingual dataset).')
|
|
parser.add_argument('--preprocess', default='ResizeNormalize', type=str,
|
|
help='the image preprocess methods (ResizeNormalize|DeiT)')
|
|
parser.add_argument('--decoder-pretrained', default=None, type=str,
|
|
help='seted to load the RoBERTa parameters to the decoder.')
|
|
parser.add_argument('--decoder-pretrained-url', default=None, type=str,
|
|
help='the ckpt url for decoder pretraining (only unilm for now)')
|
|
parser.add_argument('--dict-path-or-url', default=None, type=str,
|
|
help='the local path or url for dictionary file')
|
|
parser.add_argument('--input-size', type=int, nargs='+', help='images input size', required=True)
|
|
parser.add_argument('--data-type', type=str, default='SROIE',
|
|
help='the dataset type used for the task (SROIE or Receipt53K)')
|
|
|
|
# Augmentation parameters
|
|
parser.add_argument('--color-jitter', type=float, default=0.4, metavar='PCT',
|
|
help='Color jitter factor (default: 0.4)')
|
|
parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1', metavar='NAME',
|
|
help='Use AutoAugment policy. "v0" or "original". " + \
|
|
"(default: rand-m9-mstd0.5-inc1)'),
|
|
parser.add_argument('--smoothing', type=float, default=0.1, help='Label smoothing (default: 0.1)')
|
|
parser.add_argument('--train-interpolation', type=str, default='bicubic',
|
|
help='Training interpolation (random, bilinear, bicubic default: "bicubic")')
|
|
|
|
parser.add_argument('--repeated-aug', action='store_true')
|
|
parser.add_argument('--no-repeated-aug', action='store_false', dest='repeated_aug')
|
|
parser.set_defaults(repeated_aug=True)
|
|
|
|
# * Random Erase params
|
|
parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT',
|
|
help='Random erase prob (default: 0.25)')
|
|
parser.add_argument('--remode', type=str, default='pixel',
|
|
help='Random erase mode (default: "pixel")')
|
|
parser.add_argument('--recount', type=int, default=1,
|
|
help='Random erase count (default: 1)')
|
|
parser.add_argument('--resplit', action='store_true', default=False,
|
|
help='Do not random erase first (clean) augmentation split')
|
|
|
|
@classmethod
|
|
def setup_task(cls, args, **kwargs):
|
|
import urllib.request
|
|
import io
|
|
|
|
if getattr(args, "dict_path_or_url", None) is not None:
|
|
if args.dict_path_or_url.startswith('http'):
|
|
logger.info('Load dictionary from {}'.format(args.dict_path_or_url))
|
|
dict_content = urllib.request.urlopen(args.dict_path_or_url).read().decode()
|
|
dict_file_like = io.StringIO(dict_content)
|
|
target_dict = Dictionary.load(dict_file_like)
|
|
else:
|
|
target_dict = Dictionary.load(args.dict_path_or_url)
|
|
elif getattr(args, "decoder_pretrained", None) is not None:
|
|
if args.decoder_pretrained == 'unilm':
|
|
url = 'https://layoutlm.blob.core.windows.net/trocr/dictionaries/unilm3.dict.txt'
|
|
logger.info('Load unilm dictionary from {}'.format(url))
|
|
dict_content = urllib.request.urlopen(url).read().decode()
|
|
dict_file_like = io.StringIO(dict_content)
|
|
target_dict = Dictionary.load(dict_file_like)
|
|
elif args.decoder_pretrained.startswith('roberta'):
|
|
url = 'https://layoutlm.blob.core.windows.net/trocr/dictionaries/gpt2_with_mask.dict.txt'
|
|
logger.info('Load gpt2 dictionary from {}'.format(url))
|
|
dict_content = urllib.request.urlopen(url).read().decode()
|
|
dict_file_like = io.StringIO(dict_content)
|
|
target_dict = Dictionary.load(dict_file_like)
|
|
else:
|
|
raise ValueError('Unknown decoder_pretrained: {}'.format(args.decoder_pretrained))
|
|
else:
|
|
raise ValueError('Either dict_path_or_url or decoder_pretrained should be set.')
|
|
|
|
logger.info('[label] load dictionary: {} types'.format(len(target_dict)))
|
|
|
|
return cls(args, target_dict)
|
|
|
|
def __init__(self, args, target_dict):
|
|
|
|
super().__init__(args)
|
|
self.args = args
|
|
self.data_dir = args.data
|
|
self.target_dict = target_dict
|
|
if 'LOCAL_RANK' in os.environ and os.environ['LOCAL_RANK'] != '0':
|
|
torch.distributed.barrier()
|
|
self.bpe = self.build_bpe(args)
|
|
if 'LOCAL_RANK' in os.environ and os.environ['LOCAL_RANK'] == '0':
|
|
torch.distributed.barrier()
|
|
|
|
def load_dataset(self, split, **kwargs):
|
|
|
|
input_size = self.args.input_size
|
|
if isinstance(input_size, list):
|
|
if len(input_size) == 1:
|
|
input_size = (input_size[0], input_size[0])
|
|
else:
|
|
input_size = tuple(input_size)
|
|
elif isinstance(input_size, int):
|
|
input_size = (input_size, input_size)
|
|
|
|
logger.info('The input size is {}, the height is {} and the width is {}'.format(input_size, input_size[0], input_size[1]))
|
|
|
|
if self.args.preprocess == 'DA2':
|
|
tfm = build_data_aug(input_size, mode=split)
|
|
elif self.args.preprocess == 'RandAugment':
|
|
opt = OptForDataAugment(eval= (split != 'train'), isrand_aug=True, imgW=input_size[1], imgH=input_size[0], intact_prob=0.5, augs_num=3, augs_mag=None)
|
|
tfm = DataAugment(opt)
|
|
else:
|
|
raise Exception('Undeined image preprocess method.')
|
|
|
|
# load the dataset
|
|
if self.args.data_type == 'SROIE':
|
|
root_dir = os.path.join(self.data_dir, split)
|
|
self.datasets[split] = SROIETextRecognitionDataset(root_dir, tfm, self.bpe, self.target_dict)
|
|
elif self.args.data_type == 'Receipt53K':
|
|
gt_path = os.path.join(self.data_dir, 'gt_{}.txt'.format(split))
|
|
self.datasets[split] = Receipt53KDataset(gt_path, tfm, self.bpe, self.target_dict)
|
|
elif self.args.data_type == 'STR':
|
|
gt_path = os.path.join(self.data_dir, 'gt_{}.txt'.format(split))
|
|
self.datasets[split] = SyntheticTextRecognitionDataset(gt_path, tfm, self.bpe, self.target_dict)
|
|
else:
|
|
raise Exception('Not defined dataset type: ' + self.args.data_type)
|
|
|
|
@property
|
|
def source_dictionary(self):
|
|
return None
|
|
|
|
@property
|
|
def target_dictionary(self):
|
|
return self.target_dict
|
|
|
|
def build_generator(
|
|
self, models, args, seq_gen_cls=None, extra_gen_cls_kwargs=None
|
|
):
|
|
|
|
if getattr(args, "score_reference", False):
|
|
from fairseq.sequence_scorer import SequenceScorer
|
|
|
|
return SequenceScorer(
|
|
self.target_dictionary,
|
|
compute_alignment=getattr(args, "print_alignment", False),
|
|
)
|
|
|
|
from fairseq.sequence_generator import (
|
|
SequenceGenerator,
|
|
SequenceGeneratorWithAlignment,
|
|
)
|
|
try:
|
|
from .generator import TextRecognitionGenerator
|
|
except:
|
|
from generator import TextRecognitionGenerator
|
|
|
|
try:
|
|
from fairseq.fb_sequence_generator import FBSequenceGenerator
|
|
except ModuleNotFoundError:
|
|
pass
|
|
|
|
# Choose search strategy. Defaults to Beam Search.
|
|
sampling = getattr(args, "sampling", False)
|
|
sampling_topk = getattr(args, "sampling_topk", -1)
|
|
sampling_topp = getattr(args, "sampling_topp", -1.0)
|
|
diverse_beam_groups = getattr(args, "diverse_beam_groups", -1)
|
|
diverse_beam_strength = getattr(args, "diverse_beam_strength", 0.5)
|
|
match_source_len = getattr(args, "match_source_len", False)
|
|
diversity_rate = getattr(args, "diversity_rate", -1)
|
|
constrained = getattr(args, "constraints", False)
|
|
prefix_allowed_tokens_fn = getattr(args, "prefix_allowed_tokens_fn", None)
|
|
if (
|
|
sum(
|
|
int(cond)
|
|
for cond in [
|
|
sampling,
|
|
diverse_beam_groups > 0,
|
|
match_source_len,
|
|
diversity_rate > 0,
|
|
]
|
|
)
|
|
> 1
|
|
):
|
|
raise ValueError("Provided Search parameters are mutually exclusive.")
|
|
assert sampling_topk < 0 or sampling, "--sampling-topk requires --sampling"
|
|
assert sampling_topp < 0 or sampling, "--sampling-topp requires --sampling"
|
|
|
|
if sampling:
|
|
search_strategy = search.Sampling(
|
|
self.target_dictionary, sampling_topk, sampling_topp
|
|
)
|
|
elif diverse_beam_groups > 0:
|
|
search_strategy = search.DiverseBeamSearch(
|
|
self.target_dictionary, diverse_beam_groups, diverse_beam_strength
|
|
)
|
|
elif match_source_len:
|
|
# this is useful for tagging applications where the output
|
|
# length should match the input length, so we hardcode the
|
|
# length constraints for simplicity
|
|
search_strategy = search.LengthConstrainedBeamSearch(
|
|
self.target_dictionary,
|
|
min_len_a=1,
|
|
min_len_b=0,
|
|
max_len_a=1,
|
|
max_len_b=0,
|
|
)
|
|
elif diversity_rate > -1:
|
|
search_strategy = search.DiverseSiblingsSearch(
|
|
self.target_dictionary, diversity_rate
|
|
)
|
|
elif constrained:
|
|
search_strategy = search.LexicallyConstrainedBeamSearch(
|
|
self.target_dictionary, args.constraints
|
|
)
|
|
elif prefix_allowed_tokens_fn:
|
|
search_strategy = search.PrefixConstrainedBeamSearch(
|
|
self.target_dictionary, prefix_allowed_tokens_fn
|
|
)
|
|
else:
|
|
search_strategy = search.BeamSearch(self.target_dictionary)
|
|
|
|
extra_gen_cls_kwargs = extra_gen_cls_kwargs or {}
|
|
if seq_gen_cls is None:
|
|
if getattr(args, "print_alignment", False):
|
|
seq_gen_cls = SequenceGeneratorWithAlignment
|
|
extra_gen_cls_kwargs["print_alignment"] = args.print_alignment
|
|
elif getattr(args, "fb_seq_gen", False):
|
|
seq_gen_cls = FBSequenceGenerator
|
|
else:
|
|
seq_gen_cls = TextRecognitionGenerator
|
|
|
|
return seq_gen_cls(
|
|
models,
|
|
self.target_dictionary,
|
|
beam_size=getattr(args, "beam", 5),
|
|
max_len_a=getattr(args, "max_len_a", 0),
|
|
max_len_b=getattr(args, "max_len_b", 200),
|
|
min_len=getattr(args, "min_len", 1),
|
|
normalize_scores=(not getattr(args, "unnormalized", False)),
|
|
len_penalty=getattr(args, "lenpen", 1),
|
|
unk_penalty=getattr(args, "unkpen", 0),
|
|
temperature=getattr(args, "temperature", 1.0),
|
|
match_source_len=getattr(args, "match_source_len", False),
|
|
no_repeat_ngram_size=getattr(args, "no_repeat_ngram_size", 0),
|
|
search_strategy=search_strategy,
|
|
**extra_gen_cls_kwargs,
|
|
)
|
|
|
|
def filter_indices_by_size(
|
|
self, indices, dataset, max_positions=None, ignore_invalid_inputs=False
|
|
):
|
|
return indices |