Skip to content

Commit

Permalink
[MODEL] Add Union14M trained models (#1960)
Browse files Browse the repository at this point in the history
* [Update] Add ABINet

* [Update] Add NRTR

* [Update] Add SATRN

* [Update] Add configs
  • Loading branch information
Mountchicken committed Sep 3, 2023
1 parent 1dcd6fa commit e50c5fd
Show file tree
Hide file tree
Showing 10 changed files with 601 additions and 0 deletions.
16 changes: 16 additions & 0 deletions configs/textrecog/abinet/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,22 @@ Linguistic knowledge is of great benefit to scene text recognition. However, how
2. Facts about the pretrained model: MMOCR does not have a systematic pipeline to pretrain the language model (LM) yet, thus the weights of LM are converted from [the official pretrained model](https://github.com/FangShancheng/ABINet). The weights of ABINet-Vision are directly used as the vision model of ABINet.
```

We also provide ABINet trained on [Union14M](https://github.com/Mountchicken/Union14M)

- Evaluated on six common benchmarks

| methods | pretrained | | Regular Text | | | | Irregular Text | | download |
| :---------------------------------------------------------------: | :--------: | :----: | :----------: | :-------: | :-: | :-------: | :------------: | :----: | :----------------------------------------------------------------- |
| | | IIIT5K | SVT | IC13-1015 | | IC15-2077 | SVTP | CT80 | |
| [ABINet-Vision](configs/textrecog/abinet/abinet-vision_10e_union14m.py) | - | 0.9730 | 0.9645 | 0.9552 | | 0.8536 | 0.8977 | 0.9479 | [model](https://download.openmmlab.com/mmocr/textrecog/abinet/abinet_union14m-cbf19742.pth) |

- Evaluated on [Union14M-Benchmark](https://github.com/Mountchicken/Union14M)

| Methods | | Unsolved Challenges | | | | | Additional Challenges | | General | download |
| ------------------------------------------------------ | ----- | ------------------- | -------- | ----------- | --- | ------- | --------------------- | ---------- | ------- | ------------------------------------------------------- |
| | Curve | Multi-Oriented | Artistic | Contextless | | Salient | Multi-Words | Incomplete | General | |
| [ABINet-Vision](configs/textrecog/abinet/abinet-vision_10e_union14m.py) | 0.750 | 0.615 | 0.653 | 0.711 | | 0.729 | 0.591 | 0.026 | 0.794 | [model](https://download.openmmlab.com/mmocr/textrecog/abinet/abinet_union14m-cbf19742.pth) |

## Citation

```bibtex
Expand Down
118 changes: 118 additions & 0 deletions configs/textrecog/abinet/abinet-vision_10e_union14m.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
_base_ = [
'../_base_/datasets/union14m_train.py',
'../_base_/datasets/union14m_benchmark.py',
'../_base_/datasets/cute80.py',
'../_base_/datasets/iiit5k.py',
'../_base_/datasets/svt.py',
'../_base_/datasets/svtp.py',
'../_base_/datasets/icdar2013.py',
'../_base_/datasets/icdar2015.py',
'../_base_/default_runtime.py',
'../_base_/schedules/schedule_adamw_cos_10e.py',
'_base_abinet.py',
]

load_from = 'https://download.openmmlab.com/mmocr/textrecog/abinet/abinet_pretrain-45deac15.pth' # noqa

_base_.pop('model')
dictionary = dict(
type='Dictionary',
dict_file= # noqa
'{{ fileDirname }}/../../../dicts/english_digits_symbols_space.txt',
with_padding=True,
with_unknown=True,
same_start_end=True,
with_start=True,
with_end=True)

model = dict(
type='ABINet',
backbone=dict(type='ResNetABI'),
encoder=dict(
type='ABIEncoder',
n_layers=3,
n_head=8,
d_model=512,
d_inner=2048,
dropout=0.1,
max_len=8 * 32,
),
decoder=dict(
type='ABIFuser',
vision_decoder=dict(
type='ABIVisionDecoder',
in_channels=512,
num_channels=64,
attn_height=8,
attn_width=32,
attn_mode='nearest',
init_cfg=dict(type='Xavier', layer='Conv2d')),
module_loss=dict(type='ABIModuleLoss'),
postprocessor=dict(type='AttentionPostprocessor'),
dictionary=dictionary,
max_seq_len=26,
),
data_preprocessor=dict(
type='TextRecogDataPreprocessor',
mean=[123.675, 116.28, 103.53],
std=[58.395, 57.12, 57.375]))

# dataset settings
train_list = [
_base_.union14m_challenging, _base_.union14m_hard, _base_.union14m_medium,
_base_.union14m_normal, _base_.union14m_easy
]
val_list = [
_base_.cute80_textrecog_test, _base_.iiit5k_textrecog_test,
_base_.svt_textrecog_test, _base_.svtp_textrecog_test,
_base_.icdar2013_textrecog_test, _base_.icdar2015_textrecog_test
]
test_list = [
_base_.union14m_benchmark_artistic,
_base_.union14m_benchmark_multi_oriented,
_base_.union14m_benchmark_contextless,
_base_.union14m_benchmark_curve,
_base_.union14m_benchmark_incomplete,
_base_.union14m_benchmark_incomplete_ori,
_base_.union14m_benchmark_multi_words,
_base_.union14m_benchmark_salient,
_base_.union14m_benchmark_general,
]

train_dataset = dict(
type='ConcatDataset', datasets=train_list, pipeline=_base_.train_pipeline)
test_dataset = dict(
type='ConcatDataset', datasets=test_list, pipeline=_base_.test_pipeline)
val_dataset = dict(
type='ConcatDataset', datasets=val_list, pipeline=_base_.test_pipeline)

train_dataloader = dict(
batch_size=128,
num_workers=24,
persistent_workers=True,
sampler=dict(type='DefaultSampler', shuffle=True),
dataset=train_dataset)

test_dataloader = dict(
batch_size=128,
num_workers=4,
persistent_workers=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=test_dataset)

val_dataloader = dict(
batch_size=128,
num_workers=4,
persistent_workers=True,
pin_memory=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=val_dataset)

val_evaluator = dict(
dataset_prefixes=['CUTE80', 'IIIT5K', 'SVT', 'SVTP', 'IC13', 'IC15'])
test_evaluator = dict(dataset_prefixes=[
'artistic', 'multi-oriented', 'contextless', 'curve', 'incomplete',
'incomplete-ori', 'multi-words', 'salient', 'general'
])
16 changes: 16 additions & 0 deletions configs/textrecog/aster/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,22 @@ A challenging aspect of scene text recognition is to handle text with distortion
| [ASTER](/configs/textrecog/aster/aster_resnet45_6e_st_mj.py) | ResNet45 | 0.9357 | 0.8949 | 0.9281 | | 0.7665 | 0.8062 | 0.8507 | [model](https://download.openmmlab.com/mmocr/textrecog/aster/aster_resnet45_6e_st_mj/aster_resnet45_6e_st_mj-cc56eca4.pth) \| [log](https://download.openmmlab.com/mmocr/textrecog/aster/aster_resnet45_6e_st_mj/20221214_232605.log) |
| [ASTER-TTA](/configs/textrecog/aster/aster_resnet45_6e_st_mj.py) | ResNet45 | 0.9337 | 0.8949 | 0.9251 | | 0.7925 | 0.8109 | 0.8507 | |

We also provide ASTER trained on [Union14M](https://github.com/Mountchicken/Union14M)

- Evaluated on six common benchmarks

| Methods | pretrained | | Regular Text | | | | Irregular Text | | download |
| :------------------------------------------------------------: | :--------: | :----: | :----------: | :-------: | :-: | :-------: | :------------: | :----: | :-------------------------------------------------------------------- |
| | | IIIT5K | SVT | IC13-1015 | | IC15-2077 | SVTP | CT80 | |
| [ASTER](configs/textrecog/aster/aster_resnet45_6e_union14m.py) | - | 0.9437 | 0.8903 | 0.9360 | | 0.7857 | 0.8093 | 0.9097 | [model](https://download.openmmlab.com/mmocr/textrecog/aster/aster_union14m/aster_union14m-230eb471.pth) |

- Evaluated on [Union14M-Benchmark](https://github.com/Mountchicken/Union14M)

| Methods | | Unsolved Challenges | | | | | Additional Challenges | | General | download |
| ------------------------------------------------------ | ----- | ------------------- | -------- | ----------- | --- | ------- | --------------------- | ---------- | ------- | ------------------------------------------------------- |
| | Curve | Multi-Oriented | Artistic | Contextless | | Salient | Multi-Words | Incomplete | General | |
| [ASTER](configs/textrecog/aster/aster_resnet45_6e_union14m.py) | 0.384 | 0.130 | 0.418 | 0.529 | | 0.319 | 0.498 | 0.013 | 0.667 | [model](https://download.openmmlab.com/mmocr/textrecog/aster/aster_union14m/aster_union14m-230eb471.pth) |

## Citation

```bibtex
Expand Down
91 changes: 91 additions & 0 deletions configs/textrecog/aster/aster_resnet45_6e_union14m.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
# training schedule for 1x
_base_ = [
'_base_aster.py',
'../_base_/datasets/union14m_train.py',
'../_base_/datasets/union14m_benchmark.py',
'../_base_/datasets/cute80.py',
'../_base_/datasets/iiit5k.py',
'../_base_/datasets/svt.py',
'../_base_/datasets/svtp.py',
'../_base_/datasets/icdar2013.py',
'../_base_/datasets/icdar2015.py',
'../_base_/default_runtime.py',
'../_base_/schedules/schedule_adamw_cos_6e.py',
]

dictionary = dict(
type='Dictionary',
dict_file= # noqa
'{{ fileDirname }}/../../../dicts/english_digits_symbols_space.txt',
with_padding=True,
with_unknown=True,
same_start_end=True,
with_start=True,
with_end=True)

# dataset settings
train_list = [
_base_.union14m_challenging, _base_.union14m_hard, _base_.union14m_medium,
_base_.union14m_normal, _base_.union14m_easy
]
val_list = [
_base_.cute80_textrecog_test, _base_.iiit5k_textrecog_test,
_base_.svt_textrecog_test, _base_.svtp_textrecog_test,
_base_.icdar2013_textrecog_test, _base_.icdar2015_textrecog_test
]
test_list = [
_base_.union14m_benchmark_artistic,
_base_.union14m_benchmark_multi_oriented,
_base_.union14m_benchmark_contextless,
_base_.union14m_benchmark_curve,
_base_.union14m_benchmark_incomplete,
_base_.union14m_benchmark_incomplete_ori,
_base_.union14m_benchmark_multi_words,
_base_.union14m_benchmark_salient,
_base_.union14m_benchmark_general,
]

default_hooks = dict(logger=dict(type='LoggerHook', interval=50))

auto_scale_lr = dict(base_batch_size=512)

train_dataset = dict(
type='ConcatDataset', datasets=train_list, pipeline=_base_.train_pipeline)
test_dataset = dict(
type='ConcatDataset', datasets=test_list, pipeline=_base_.test_pipeline)
val_dataset = dict(
type='ConcatDataset', datasets=val_list, pipeline=_base_.test_pipeline)

train_dataloader = dict(
batch_size=512,
num_workers=12,
persistent_workers=True,
pin_memory=True,
sampler=dict(type='DefaultSampler', shuffle=True),
dataset=train_dataset)

test_dataloader = dict(
batch_size=128,
num_workers=4,
persistent_workers=True,
pin_memory=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=test_dataset)

val_dataloader = dict(
batch_size=128,
num_workers=4,
persistent_workers=True,
pin_memory=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=val_dataset)

val_evaluator = dict(
dataset_prefixes=['CUTE80', 'IIIT5K', 'SVT', 'SVTP', 'IC13', 'IC15'])

test_evaluator = dict(dataset_prefixes=[
'artistic', 'multi-oriented', 'contextless', 'curve', 'incomplete',
'incomplete-ori', 'multi-words', 'salient', 'general'
])
16 changes: 16 additions & 0 deletions configs/textrecog/nrtr/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,22 @@ Scene text recognition has attracted a great many researches due to its importan
| [NRTR](/configs/textrecog/nrtr/nrtr_resnet31-1by16-1by8_6e_st_mj.py) | R31-1/16-1/8 | 0.9470 | 0.8918 | 0.9399 | | 0.7376 | 0.7969 | 0.8854 | [model](https://download.openmmlab.com/mmocr/textrecog/nrtr/nrtr_resnet31-1by16-1by8_6e_st_mj/nrtr_resnet31-1by16-1by8_6e_st_mj_20220920_143358-43767036.pth) \| [log](https://download.openmmlab.com/mmocr/textrecog/nrtr/nrtr_resnet31-1by16-1by8_6e_st_mj/20220920_143358.log) |
| [NRTR-TTA](/configs/textrecog/nrtr/nrtr_resnet31-1by16-1by8_6e_st_mj.py) | R31-1/16-1/8 | 0.9423 | 0.8903 | 0.9360 | | 0.7641 | 0.8016 | 0.8854 | |

We also provide NRTR trained on [Union14M](https://github.com/Mountchicken/Union14M)

- Evaluated on six common benchmarks

| Methods | Backbone | | Regular Text | | | | Irregular Text | | download |
| :--------------------------------------------------------------: | :---------: | :----: | :----------: | :-------: | :-: | :-------: | :------------: | :----: | :----------------------------------------------------------------- |
| | | IIIT5K | SVT | IC13-1015 | | IC15-2077 | SVTP | CT80 | |
| [NRTR](configs/textrecog/nrtr/nrtr_resnet31-1by8-1by4_union14m.py) | R31-1/8-1/4 | 0.9673 | 0.9320 | 0.9557 | | 0.8074 | 0.8357 | 0.9201 | [model](https://download.openmmlab.com/mmocr/textrecog/nrtr/nrtr_union14m/nrtr_union14m-606b6cba.pth) |

- Evaluated on [Union14M-Benchmark](https://github.com/Mountchicken/Union14M)

| Methods | Backbone | | Unsolved Challenges | | | | | Additional Challenges | | General | download |
| ------------------------------------------------ | :---------: | ----- | ------------------- | -------- | ----------- | --- | ------- | --------------------- | ---------- | ------- | -------------------------------------------------- |
| | | Curve | Multi-Oriented | Artistic | Contextless | | Salient | Multi-Words | Incomplete | General | |
| [NRTR](configs/textrecog/aster/aster_resnet45_6e_union14m.py) | R31-1/8-1/4 | 0.493 | 0.406 | 0.543 | 0.696 | | 0.429 | 0.755 | 0.015 | 0.752 | [model](https://download.openmmlab.com/mmocr/textrecog/nrtr/nrtr_union14m/nrtr_union14m-606b6cba.pth) |

## Citation

```bibtex
Expand Down
114 changes: 114 additions & 0 deletions configs/textrecog/nrtr/nrtr_resnet31-1by8-1by4_union14m.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
_base_ = [
'../_base_/datasets/union14m_train.py',
'../_base_/datasets/union14m_benchmark.py',
'../_base_/datasets/cute80.py',
'../_base_/datasets/iiit5k.py',
'../_base_/datasets/svt.py',
'../_base_/datasets/svtp.py',
'../_base_/datasets/icdar2013.py',
'../_base_/datasets/icdar2015.py',
'../_base_/default_runtime.py',
'../_base_/schedules/schedule_adam_base.py',
'_base_nrtr_resnet31.py',
]

# optimizer settings
train_cfg = dict(max_epochs=6)
# learning policy
param_scheduler = [
dict(type='MultiStepLR', milestones=[3, 4], end=6),
]

_base_.pop('model')
dictionary = dict(
type='Dictionary',
dict_file= # noqa
'{{ fileDirname }}/../../../dicts/english_digits_symbols_space.txt',
with_padding=True,
with_unknown=True,
same_start_end=True,
with_start=True,
with_end=True)

model = dict(
type='NRTR',
backbone=dict(
type='ResNet31OCR',
layers=[1, 2, 5, 3],
channels=[32, 64, 128, 256, 512, 512],
stage4_pool_cfg=dict(kernel_size=(2, 1), stride=(2, 1)),
last_stage_pool=False),
encoder=dict(type='NRTREncoder'),
decoder=dict(
type='NRTRDecoder',
module_loss=dict(
type='CEModuleLoss', ignore_first_char=True, flatten=True),
postprocessor=dict(type='AttentionPostprocessor'),
dictionary=dictionary,
max_seq_len=30,
),
data_preprocessor=dict(
type='TextRecogDataPreprocessor',
mean=[123.675, 116.28, 103.53],
std=[58.395, 57.12, 57.375]))

# dataset settings
train_list = [
_base_.union14m_challenging, _base_.union14m_hard, _base_.union14m_medium,
_base_.union14m_normal, _base_.union14m_easy
]
val_list = [
_base_.cute80_textrecog_test, _base_.iiit5k_textrecog_test,
_base_.svt_textrecog_test, _base_.svtp_textrecog_test,
_base_.icdar2013_textrecog_test, _base_.icdar2015_textrecog_test
]
test_list = [
_base_.union14m_benchmark_artistic,
_base_.union14m_benchmark_multi_oriented,
_base_.union14m_benchmark_contextless,
_base_.union14m_benchmark_curve,
_base_.union14m_benchmark_incomplete,
_base_.union14m_benchmark_incomplete_ori,
_base_.union14m_benchmark_multi_words,
_base_.union14m_benchmark_salient,
_base_.union14m_benchmark_general,
]

train_dataset = dict(
type='ConcatDataset', datasets=train_list, pipeline=_base_.train_pipeline)
test_dataset = dict(
type='ConcatDataset', datasets=test_list, pipeline=_base_.test_pipeline)
val_dataset = dict(
type='ConcatDataset', datasets=val_list, pipeline=_base_.test_pipeline)

train_dataloader = dict(
batch_size=128,
num_workers=24,
persistent_workers=True,
sampler=dict(type='DefaultSampler', shuffle=True),
dataset=train_dataset)

test_dataloader = dict(
batch_size=128,
num_workers=4,
persistent_workers=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=test_dataset)

val_dataloader = dict(
batch_size=128,
num_workers=4,
persistent_workers=True,
pin_memory=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=val_dataset)

val_evaluator = dict(
dataset_prefixes=['CUTE80', 'IIIT5K', 'SVT', 'SVTP', 'IC13', 'IC15'])

test_evaluator = dict(dataset_prefixes=[
'artistic', 'multi-oriented', 'contextless', 'curve', 'incomplete',
'incomplete-ori', 'multi-words', 'salient', 'general'
])

0 comments on commit e50c5fd

Please sign in to comment.