mirror of https://github.com/open-mmlab/mmocr.git
Merge pull request #5 from yuexy/feature/crnn_and_robustscanner
[feature]: add CRNN and RobustScannerpull/2/head
Hongbin Sun
GitHub
commit
50287450ab
|
|
@ -1,5 +1,5 @@
|
|||
<div align="center">
|
||||
<img src="resources/mmocr-logo.jpg" width="500px"/>
|
||||
<img src="resources/mmocr-logo.png" width="500px"/>
|
||||
</div>
|
||||
|
||||
## Introduction
|
||||
|
|
|
|||
|
|
@ -0,0 +1,37 @@
|
|||
# An end-to-end trainable neural network for image-based sequence recognition and its application to scene text recognition
|
||||
|
||||
## Introduction
|
||||
|
||||
[ALGORITHM]
|
||||
|
||||
```latex
|
||||
@article{shi2016end,
|
||||
title={An end-to-end trainable neural network for image-based sequence recognition and its application to scene text recognition},
|
||||
author={Shi, Baoguang and Bai, Xiang and Yao, Cong},
|
||||
journal={IEEE transactions on pattern analysis and machine intelligence},
|
||||
year={2016}
|
||||
}
|
||||
```
|
||||
|
||||
## Results and Models
|
||||
|
||||
### Train Dataset
|
||||
|
||||
| trainset | instance_num | repeat_num | note |
|
||||
| :------: | :----------: | :--------: | :---: |
|
||||
| Syn90k | 8919273 | 1 | synth |
|
||||
|
||||
### Test Dataset
|
||||
|
||||
| testset | instance_num | note |
|
||||
| :-----: | :----------: | :-----: |
|
||||
| IIIT5K | 3000 | regular |
|
||||
| SVT | 647 | regular |
|
||||
| IC13 | 1015 | regular |
|
||||
|
||||
## Results and models
|
||||
|
||||
| methods | | Regular Text | | | | Irregular Text | | download |
|
||||
| :-----: | :----: | :----------: | :--: | :-: | :--: | :------------: | :--: | :---------------------------------------------------------: |
|
||||
| methods | IIIT5K | SVT | IC13 | | IC15 | SVTP | CT80 |
|
||||
| CRNN | 80.5 | 81.5 | 86.5 | | - | - | - | [config](https://download.openmmlab.com/mmocr/textrecog/crnn/crnn_academic_dataset.py) [log]() [model](https) |
|
||||
|
|
@ -0,0 +1,155 @@
|
|||
_base_ = []
|
||||
checkpoint_config = dict(interval=1)
|
||||
# yapf:disable
|
||||
log_config = dict(
|
||||
interval=1,
|
||||
hooks=[
|
||||
dict(type='TextLoggerHook')
|
||||
|
||||
])
|
||||
# yapf:enable
|
||||
dist_params = dict(backend='nccl')
|
||||
log_level = 'INFO'
|
||||
load_from = None
|
||||
resume_from = None
|
||||
workflow = [('train', 1)]
|
||||
|
||||
# model
|
||||
label_convertor = dict(
|
||||
type='CTCConvertor', dict_type='DICT36', with_unknown=False, lower=True)
|
||||
|
||||
model = dict(
|
||||
type='CRNNNet',
|
||||
preprocessor=None,
|
||||
backbone=dict(type='VeryDeepVgg', leakyRelu=False, input_channels=1),
|
||||
encoder=None,
|
||||
decoder=dict(type='CRNNDecoder', in_channels=512, rnn_flag=True),
|
||||
loss=dict(type='CTCLoss'),
|
||||
label_convertor=label_convertor,
|
||||
pretrained=None)
|
||||
|
||||
train_cfg = None
|
||||
test_cfg = None
|
||||
|
||||
# optimizer
|
||||
optimizer = dict(type='Adadelta', lr=1.0)
|
||||
optimizer_config = dict(grad_clip=None)
|
||||
# learning policy
|
||||
lr_config = dict(policy='step', step=[])
|
||||
total_epochs = 5
|
||||
|
||||
# data
|
||||
img_norm_cfg = dict(mean=[0.5], std=[0.5])
|
||||
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='ResizeOCR',
|
||||
height=32,
|
||||
min_width=100,
|
||||
max_width=100,
|
||||
keep_aspect_ratio=False),
|
||||
dict(type='ToTensorOCR'),
|
||||
dict(type='NormalizeOCR', **img_norm_cfg),
|
||||
dict(
|
||||
type='Collect',
|
||||
keys=['img'],
|
||||
meta_keys=[
|
||||
'filename', 'ori_shape', 'img_shape', 'text', 'valid_ratio'
|
||||
]),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='ResizeOCR',
|
||||
height=32,
|
||||
min_width=4,
|
||||
max_width=None,
|
||||
keep_aspect_ratio=True),
|
||||
dict(type='ToTensorOCR'),
|
||||
dict(type='NormalizeOCR', **img_norm_cfg),
|
||||
dict(
|
||||
type='Collect',
|
||||
keys=['img'],
|
||||
meta_keys=['filename', 'ori_shape', 'img_shape', 'valid_ratio']),
|
||||
]
|
||||
|
||||
dataset_type = 'OCRDataset'
|
||||
|
||||
train_img_prefix = 'data/mixture/mnt/ramdisk/max/90kDICT32px'
|
||||
train_ann_file = 'data/mixture/mnt/ramdisk/max/90kDICT32px/label.txt'
|
||||
|
||||
train1 = dict(
|
||||
type=dataset_type,
|
||||
img_prefix=train_img_prefix,
|
||||
ann_file=train_ann_file,
|
||||
loader=dict(
|
||||
type='HardDiskLoader',
|
||||
repeat=1,
|
||||
parser=dict(
|
||||
type='LineStrParser',
|
||||
keys=['filename', 'text'],
|
||||
keys_idx=[0, 1],
|
||||
separator=' ')),
|
||||
pipeline=train_pipeline,
|
||||
test_mode=False)
|
||||
|
||||
test1 = dict(
|
||||
type=dataset_type,
|
||||
img_prefix=train_img_prefix,
|
||||
ann_file=train_ann_file,
|
||||
loader=dict(
|
||||
type='HardDiskLoader',
|
||||
repeat=1,
|
||||
parser=dict(
|
||||
type='LineStrParser',
|
||||
keys=['filename', 'text'],
|
||||
keys_idx=[0, 1],
|
||||
separator=' ')),
|
||||
pipeline=test_pipeline,
|
||||
test_mode=True)
|
||||
|
||||
test_img_prefix = 'data/mixture/'
|
||||
ic13_path = 'testset/icdar_2013/Challenge2_Test_Task3_Images/'
|
||||
test_img_prefix1 = test_img_prefix + ic13_path
|
||||
test_img_prefix2 = test_img_prefix + 'testset/IIIT5K/'
|
||||
test_img_prefix3 = test_img_prefix + 'testset/svt/'
|
||||
|
||||
test_ann_prefix = 'data/mixture/'
|
||||
test_ann_file1 = test_ann_prefix + 'testset/icdar_2013/test_label_1015.txt'
|
||||
test_ann_file2 = test_ann_prefix + 'testset/IIIT5K/label.txt'
|
||||
test_ann_file3 = test_ann_prefix + 'testset/svt/test_list.txt'
|
||||
|
||||
test1 = dict(
|
||||
type=dataset_type,
|
||||
img_prefix=test_img_prefix1,
|
||||
ann_file=test_ann_file1,
|
||||
loader=dict(
|
||||
type='HardDiskLoader',
|
||||
repeat=1,
|
||||
parser=dict(
|
||||
type='LineStrParser',
|
||||
keys=['filename', 'text'],
|
||||
keys_idx=[0, 1],
|
||||
separator=' ')),
|
||||
pipeline=test_pipeline,
|
||||
test_mode=True)
|
||||
|
||||
test2 = {key: value for key, value in test1.items()}
|
||||
test2['img_prefix'] = test_img_prefix2
|
||||
test2['ann_file'] = test_ann_file2
|
||||
|
||||
test3 = {key: value for key, value in test1.items()}
|
||||
test3['img_prefix'] = test_img_prefix3
|
||||
test3['ann_file'] = test_ann_file3
|
||||
|
||||
data = dict(
|
||||
samples_per_gpu=64,
|
||||
workers_per_gpu=4,
|
||||
train=dict(type='ConcatDataset', datasets=[train1]),
|
||||
val=dict(type='ConcatDataset', datasets=[test1, test2, test3]),
|
||||
test=dict(type='ConcatDataset', datasets=[test1, test2, test3]))
|
||||
|
||||
evaluation = dict(interval=1, metric='acc')
|
||||
|
||||
cudnn_benchmark = True
|
||||
|
|
@ -0,0 +1,6 @@
|
|||
_base_ = [
|
||||
'../../_base_/schedules/schedule_adadelta_8e.py',
|
||||
'../../_base_/default_runtime.py',
|
||||
'../../_base_/recog_datasets/toy_dataset.py',
|
||||
'../../_base_/recog_models/crnn.py'
|
||||
]
|
||||
|
|
@ -0,0 +1,12 @@
|
|||
_base_ = [
|
||||
'../../_base_/default_runtime.py',
|
||||
'../../_base_/recog_models/robust_scanner.py',
|
||||
'../../_base_/recog_datasets/toy_dataset.py'
|
||||
]
|
||||
|
||||
# optimizer
|
||||
optimizer = dict(type='Adam', lr=1e-3)
|
||||
optimizer_config = dict(grad_clip=None)
|
||||
# learning policy
|
||||
lr_config = dict(policy='step', step=[3, 4])
|
||||
total_epochs = 6
|
||||
|
|
@ -0,0 +1,198 @@
|
|||
_base_ = [
|
||||
'../../_base_/default_runtime.py',
|
||||
'../../_base_/recog_models/robust_scanner.py'
|
||||
]
|
||||
|
||||
# optimizer
|
||||
optimizer = dict(type='Adam', lr=1e-3)
|
||||
optimizer_config = dict(grad_clip=None)
|
||||
# learning policy
|
||||
lr_config = dict(policy='step', step=[3, 4])
|
||||
total_epochs = 5
|
||||
|
||||
img_norm_cfg = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
|
||||
train_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='ResizeOCR',
|
||||
height=48,
|
||||
min_width=48,
|
||||
max_width=160,
|
||||
keep_aspect_ratio=True,
|
||||
width_downsample_ratio=0.25),
|
||||
dict(type='ToTensorOCR'),
|
||||
dict(type='NormalizeOCR', **img_norm_cfg),
|
||||
dict(
|
||||
type='Collect',
|
||||
keys=['img'],
|
||||
meta_keys=[
|
||||
'filename', 'ori_shape', 'img_shape', 'text', 'valid_ratio'
|
||||
]),
|
||||
]
|
||||
test_pipeline = [
|
||||
dict(type='LoadImageFromFile'),
|
||||
dict(
|
||||
type='MultiRotateAugOCR',
|
||||
rotate_degrees=[0, 90, 270],
|
||||
transforms=[
|
||||
dict(
|
||||
type='ResizeOCR',
|
||||
height=48,
|
||||
min_width=48,
|
||||
max_width=160,
|
||||
keep_aspect_ratio=True,
|
||||
width_downsample_ratio=0.25),
|
||||
dict(type='ToTensorOCR'),
|
||||
dict(type='NormalizeOCR', **img_norm_cfg),
|
||||
dict(
|
||||
type='Collect',
|
||||
keys=['img'],
|
||||
meta_keys=[
|
||||
'filename', 'ori_shape', 'img_shape', 'valid_ratio'
|
||||
]),
|
||||
])
|
||||
]
|
||||
|
||||
dataset_type = 'OCRDataset'
|
||||
|
||||
prefix = 'data/mixture/'
|
||||
|
||||
train_img_prefix1 = prefix + 'icdar_2011/Challenge1_Training_Task3_Images_GT'
|
||||
train_img_prefix2 = prefix + 'icdar_2013/recog_train_data/' + \
|
||||
'Challenge2_Training_Task3_Images_GT'
|
||||
train_img_prefix3 = prefix + 'icdar_2015/ch4_training_word_images_gt'
|
||||
train_img_prefix4 = prefix + 'coco_text/train_words'
|
||||
train_img_prefix5 = prefix + 'III5K'
|
||||
train_img_prefix6 = prefix + 'SynthText_Add/SynthText_Add'
|
||||
train_img_prefix7 = prefix + 'SynthText/synthtext/SynthText_patch_horizontal'
|
||||
train_img_prefix8 = prefix + 'mnt/ramdisk/max/90kDICT32px'
|
||||
|
||||
train_ann_file1 = prefix + 'icdar_2011/training_label_fix.txt',
|
||||
train_ann_file2 = prefix + 'icdar_2013/recog_train_data/train_label.txt',
|
||||
train_ann_file3 = prefix + 'icdar_2015/training_label_fix.txt',
|
||||
train_ann_file4 = prefix + 'coco_text/train_label.txt',
|
||||
train_ann_file5 = prefix + 'III5K/train_label.txt',
|
||||
train_ann_file6 = prefix + 'SynthText_Add/SynthText_Add/' + \
|
||||
'annotationlist/label.lmdb',
|
||||
train_ann_file7 = prefix + 'SynthText/synthtext/shuffle_labels.lmdb',
|
||||
train_ann_file8 = prefix + 'mnt/ramdisk/max/90kDICT32px/shuffle_labels.lmdb'
|
||||
|
||||
train1 = dict(
|
||||
type=dataset_type,
|
||||
img_prefix=train_img_prefix1,
|
||||
ann_file=train_ann_file1,
|
||||
loader=dict(
|
||||
type='HardDiskLoader',
|
||||
repeat=20,
|
||||
parser=dict(
|
||||
type='LineStrParser',
|
||||
keys=['filename', 'text'],
|
||||
keys_idx=[0, 1],
|
||||
separator=' ')),
|
||||
pipeline=train_pipeline,
|
||||
test_mode=False)
|
||||
|
||||
train2 = {key: value for key, value in train1.items()}
|
||||
train2['img_prefix'] = train_img_prefix2
|
||||
train2['ann_file'] = train_ann_file2
|
||||
|
||||
train3 = {key: value for key, value in train1.items()}
|
||||
train3['img_prefix'] = train_img_prefix3
|
||||
train3['ann_file'] = train_ann_file3
|
||||
|
||||
train4 = {key: value for key, value in train1.items()}
|
||||
train4['img_prefix'] = train_img_prefix4
|
||||
train4['ann_file'] = train_ann_file4
|
||||
|
||||
train5 = {key: value for key, value in train1.items()}
|
||||
train5['img_prefix'] = train_img_prefix5
|
||||
train5['ann_file'] = train_ann_file5
|
||||
|
||||
train6 = dict(
|
||||
type=dataset_type,
|
||||
img_prefix=train_img_prefix6,
|
||||
ann_file=train_ann_file6,
|
||||
loader=dict(
|
||||
type='LmdbLoader',
|
||||
repeat=1,
|
||||
parser=dict(
|
||||
type='LineStrParser',
|
||||
keys=['filename', 'text'],
|
||||
keys_idx=[0, 1],
|
||||
separator=' ')),
|
||||
pipeline=train_pipeline,
|
||||
test_mode=False)
|
||||
|
||||
train7 = {key: value for key, value in train6.items()}
|
||||
train7['img_prefix'] = train_img_prefix7
|
||||
train7['ann_file'] = train_ann_file7
|
||||
|
||||
train8 = {key: value for key, value in train6.items()}
|
||||
train8['img_prefix'] = train_img_prefix8
|
||||
train8['ann_file'] = train_ann_file8
|
||||
|
||||
test_img_prefix1 = prefix + 'testset/IIIT5K/'
|
||||
test_img_prefix2 = prefix + 'testset/svt/'
|
||||
test_img_prefix3 = prefix + 'testset/icdar_2013/Challenge2_Test_Task3_Images/'
|
||||
test_img_prefix4 = prefix + 'testset/icdar_2015/ch4_test_word_images_gt'
|
||||
test_img_prefix5 = prefix + 'testset/svtp/'
|
||||
test_img_prefix6 = prefix + 'testset/ct80/'
|
||||
|
||||
test_ann_file1 = prefix + 'testset/IIIT5K/label.txt'
|
||||
test_ann_file2 = prefix + 'testset/svt/test_list.txt'
|
||||
test_ann_file3 = prefix + 'testset/icdar_2013/test_label_1015.txt'
|
||||
test_ann_file4 = prefix + 'testset/icdar_2015/test_label.txt'
|
||||
test_ann_file5 = prefix + 'testset/svtp/imagelist.txt'
|
||||
test_ann_file6 = prefix + 'testset/ct80/imagelist.txt'
|
||||
|
||||
test1 = dict(
|
||||
type=dataset_type,
|
||||
img_prefix=test_img_prefix1,
|
||||
ann_file=test_ann_file1,
|
||||
loader=dict(
|
||||
type='HardDiskLoader',
|
||||
repeat=1,
|
||||
parser=dict(
|
||||
type='LineStrParser',
|
||||
keys=['filename', 'text'],
|
||||
keys_idx=[0, 1],
|
||||
separator=' ')),
|
||||
pipeline=test_pipeline,
|
||||
test_mode=True)
|
||||
|
||||
test2 = {key: value for key, value in test1.items()}
|
||||
test2['img_prefix'] = test_img_prefix2
|
||||
test2['ann_file'] = test_ann_file2
|
||||
|
||||
test3 = {key: value for key, value in test1.items()}
|
||||
test3['img_prefix'] = test_img_prefix3
|
||||
test3['ann_file'] = test_ann_file3
|
||||
|
||||
test4 = {key: value for key, value in test1.items()}
|
||||
test4['img_prefix'] = test_img_prefix4
|
||||
test4['ann_file'] = test_ann_file4
|
||||
|
||||
test5 = {key: value for key, value in test1.items()}
|
||||
test5['img_prefix'] = test_img_prefix5
|
||||
test5['ann_file'] = test_ann_file5
|
||||
|
||||
test6 = {key: value for key, value in test1.items()}
|
||||
test6['img_prefix'] = test_img_prefix6
|
||||
test6['ann_file'] = test_ann_file6
|
||||
|
||||
data = dict(
|
||||
samples_per_gpu=64,
|
||||
workers_per_gpu=2,
|
||||
train=dict(
|
||||
type='ConcatDataset',
|
||||
datasets=[
|
||||
train1, train2, train3, train4, train5, train6, train7, train8
|
||||
]),
|
||||
val=dict(
|
||||
type='ConcatDataset',
|
||||
datasets=[test1, test2, test3, test4, test5, test6]),
|
||||
test=dict(
|
||||
type='ConcatDataset',
|
||||
datasets=[test1, test2, test3, test4, test5, test6]))
|
||||
|
||||
evaluation = dict(interval=1, metric='acc')
|
||||
|
|
@ -0,0 +1,216 @@
|
|||
import numpy as np
|
||||
|
||||
import mmocr.utils as utils
|
||||
from . import utils as eval_utils
|
||||
|
||||
|
||||
def compute_recall_precision(gt_polys, pred_polys):
|
||||
"""Compute the recall and the precision matrices between gt and predicted
|
||||
polygons.
|
||||
|
||||
Args:
|
||||
gt_polys (list[Polygon]): List of gt polygons.
|
||||
pred_polys (list[Polygon]): List of predicted polygons.
|
||||
|
||||
Returns:
|
||||
recall (ndarray): Recall matrix of size gt_num x det_num.
|
||||
precision (ndarray): Precision matrix of size gt_num x det_num.
|
||||
"""
|
||||
assert isinstance(gt_polys, list)
|
||||
assert isinstance(pred_polys, list)
|
||||
|
||||
gt_num = len(gt_polys)
|
||||
det_num = len(pred_polys)
|
||||
sz = [gt_num, det_num]
|
||||
|
||||
recall = np.zeros(sz)
|
||||
precision = np.zeros(sz)
|
||||
# compute area recall and precision for each (gt, det) pair
|
||||
# in one img
|
||||
for gt_id in range(gt_num):
|
||||
for pred_id in range(det_num):
|
||||
gt = gt_polys[gt_id]
|
||||
det = pred_polys[pred_id]
|
||||
|
||||
inter_area, _ = eval_utils.poly_intersection(det, gt)
|
||||
gt_area = gt.area()
|
||||
det_area = det.area()
|
||||
if gt_area != 0:
|
||||
recall[gt_id, pred_id] = inter_area / gt_area
|
||||
if det_area != 0:
|
||||
precision[gt_id, pred_id] = inter_area / det_area
|
||||
|
||||
return recall, precision
|
||||
|
||||
|
||||
def eval_hmean_ic13(det_boxes,
|
||||
gt_boxes,
|
||||
gt_ignored_boxes,
|
||||
precision_thr=0.4,
|
||||
recall_thr=0.8,
|
||||
center_dist_thr=1.0,
|
||||
one2one_score=1.,
|
||||
one2many_score=0.8,
|
||||
many2one_score=1.):
|
||||
"""Evalute hmean of text detection using the icdar2013 standard.
|
||||
|
||||
Args:
|
||||
det_boxes (list[list[list[float]]]): List of arrays of shape (n, 2k).
|
||||
Each element is the det_boxes for one img. k>=4.
|
||||
gt_boxes (list[list[list[float]]]): List of arrays of shape (m, 2k).
|
||||
Each element is the gt_boxes for one img. k>=4.
|
||||
gt_ignored_boxes (list[list[list[float]]]): List of arrays of
|
||||
(l, 2k). Each element is the ignored gt_boxes for one img. k>=4.
|
||||
precision_thr (float): Precision threshold of the iou of one
|
||||
(gt_box, det_box) pair.
|
||||
recall_thr (float): Recall threshold of the iou of one
|
||||
(gt_box, det_box) pair.
|
||||
center_dist_thr (float): Distance threshold of one (gt_box, det_box)
|
||||
center point pair.
|
||||
one2one_score (float): Reward when one gt matches one det_box.
|
||||
one2many_score (float): Reward when one gt matches many det_boxes.
|
||||
many2one_score (float): Reward when many gts match one det_box.
|
||||
|
||||
Returns:
|
||||
hmean (tuple[dict]): Tuple of dicts which encodes the hmean for
|
||||
the dataset and all images.
|
||||
"""
|
||||
assert utils.is_3dlist(det_boxes)
|
||||
assert utils.is_3dlist(gt_boxes)
|
||||
assert utils.is_3dlist(gt_ignored_boxes)
|
||||
|
||||
assert 0 <= precision_thr <= 1
|
||||
assert 0 <= recall_thr <= 1
|
||||
assert center_dist_thr > 0
|
||||
assert 0 <= one2one_score <= 1
|
||||
assert 0 <= one2many_score <= 1
|
||||
assert 0 <= many2one_score <= 1
|
||||
|
||||
img_num = len(det_boxes)
|
||||
assert img_num == len(gt_boxes)
|
||||
assert img_num == len(gt_ignored_boxes)
|
||||
|
||||
dataset_gt_num = 0
|
||||
dataset_pred_num = 0
|
||||
dataset_hit_recall = 0.0
|
||||
dataset_hit_prec = 0.0
|
||||
|
||||
img_results = []
|
||||
|
||||
for i in range(img_num):
|
||||
gt = gt_boxes[i]
|
||||
gt_ignored = gt_ignored_boxes[i]
|
||||
pred = det_boxes[i]
|
||||
|
||||
gt_num = len(gt)
|
||||
ignored_num = len(gt_ignored)
|
||||
pred_num = len(pred)
|
||||
|
||||
accum_recall = 0.
|
||||
accum_precision = 0.
|
||||
|
||||
gt_points = gt + gt_ignored
|
||||
gt_polys = [eval_utils.points2polygon(p) for p in gt_points]
|
||||
gt_ignored_index = [gt_num + i for i in range(len(gt_ignored))]
|
||||
gt_num = len(gt_polys)
|
||||
|
||||
pred_polys, pred_points, pred_ignored_index = eval_utils.ignore_pred(
|
||||
pred, gt_ignored_index, gt_polys, precision_thr)
|
||||
|
||||
if pred_num > 0 and gt_num > 0:
|
||||
|
||||
gt_hit = np.zeros(gt_num, np.int8).tolist()
|
||||
pred_hit = np.zeros(pred_num, np.int8).tolist()
|
||||
|
||||
# compute area recall and precision for each (gt, pred) pair
|
||||
# in one img.
|
||||
recall_mat, precision_mat = compute_recall_precision(
|
||||
gt_polys, pred_polys)
|
||||
|
||||
# match one gt to one pred box.
|
||||
for gt_id in range(gt_num):
|
||||
for pred_id in range(pred_num):
|
||||
if gt_hit[gt_id] != 0 or pred_hit[
|
||||
pred_id] != 0 or gt_id in gt_ignored_index \
|
||||
or pred_id in pred_ignored_index:
|
||||
continue
|
||||
match = eval_utils.one2one_match_ic13(
|
||||
gt_id, pred_id, recall_mat, precision_mat, recall_thr,
|
||||
precision_thr)
|
||||
|
||||
if match:
|
||||
gt_point = np.array(gt_points[gt_id])
|
||||
det_point = np.array(pred_points[pred_id])
|
||||
|
||||
norm_dist = eval_utils.box_center_distance(
|
||||
det_point, gt_point)
|
||||
norm_dist /= eval_utils.box_diag(
|
||||
det_point) + eval_utils.box_diag(gt_point)
|
||||
norm_dist *= 2.0
|
||||
|
||||
if norm_dist < center_dist_thr:
|
||||
gt_hit[gt_id] = 1
|
||||
pred_hit[pred_id] = 1
|
||||
accum_recall += one2one_score
|
||||
accum_precision += one2one_score
|
||||
|
||||
# match one gt to many det boxes.
|
||||
for gt_id in range(gt_num):
|
||||
if gt_id in gt_ignored_index:
|
||||
continue
|
||||
match, match_det_set = eval_utils.one2many_match_ic13(
|
||||
gt_id, recall_mat, precision_mat, recall_thr,
|
||||
precision_thr, gt_hit, pred_hit, pred_ignored_index)
|
||||
|
||||
if match:
|
||||
gt_hit[gt_id] = 1
|
||||
accum_recall += one2many_score
|
||||
accum_precision += one2many_score * len(match_det_set)
|
||||
for pred_id in match_det_set:
|
||||
pred_hit[pred_id] = 1
|
||||
|
||||
# match many gt to one det box. One pair of (det,gt) are matched
|
||||
# successfully if their recall, precision, normalized distance
|
||||
# meet some thresholds.
|
||||
for pred_id in range(pred_num):
|
||||
if pred_id in pred_ignored_index:
|
||||
continue
|
||||
|
||||
match, match_gt_set = eval_utils.many2one_match_ic13(
|
||||
pred_id, recall_mat, precision_mat, recall_thr,
|
||||
precision_thr, gt_hit, pred_hit, gt_ignored_index)
|
||||
|
||||
if match:
|
||||
pred_hit[pred_id] = 1
|
||||
accum_recall += many2one_score * len(match_gt_set)
|
||||
accum_precision += many2one_score
|
||||
for gt_id in match_gt_set:
|
||||
gt_hit[gt_id] = 1
|
||||
|
||||
gt_care_number = gt_num - ignored_num
|
||||
pred_care_number = pred_num - len(pred_ignored_index)
|
||||
|
||||
r, p, h = eval_utils.compute_hmean(accum_recall, accum_precision,
|
||||
gt_care_number, pred_care_number)
|
||||
|
||||
img_results.append({'recall': r, 'precision': p, 'hmean': h})
|
||||
|
||||
dataset_gt_num += gt_care_number
|
||||
dataset_pred_num += pred_care_number
|
||||
dataset_hit_recall += accum_recall
|
||||
dataset_hit_prec += accum_precision
|
||||
|
||||
total_r, total_p, total_h = eval_utils.compute_hmean(
|
||||
dataset_hit_recall, dataset_hit_prec, dataset_gt_num, dataset_pred_num)
|
||||
|
||||
dataset_results = {
|
||||
'num_gts': dataset_gt_num,
|
||||
'num_dets': dataset_pred_num,
|
||||
'num_recall': dataset_hit_recall,
|
||||
'num_precision': dataset_hit_prec,
|
||||
'recall': total_r,
|
||||
'precision': total_p,
|
||||
'hmean': total_h
|
||||
}
|
||||
|
||||
return dataset_results, img_results
|
||||
|
|
@ -0,0 +1,116 @@
|
|||
import numpy as np
|
||||
|
||||
import mmocr.utils as utils
|
||||
from . import utils as eval_utils
|
||||
|
||||
|
||||
def eval_hmean_iou(pred_boxes,
|
||||
gt_boxes,
|
||||
gt_ignored_boxes,
|
||||
iou_thr=0.5,
|
||||
precision_thr=0.5):
|
||||
"""Evalute hmean of text detection using IOU standard.
|
||||
|
||||
Args:
|
||||
pred_boxes (list[list[list[float]]]): Text boxes for an img list. Each
|
||||
box has 2k (>=8) values.
|
||||
gt_boxes (list[list[list[float]]]): Ground truth text boxes for an img
|
||||
list. Each box has 2k (>=8) values.
|
||||
gt_ignored_boxes (list[list[list[float]]]): Ignored ground truth text
|
||||
boxes for an img list. Each box has 2k (>=8) values.
|
||||
iou_thr (float): Iou threshold when one (gt_box, det_box) pair is
|
||||
matched.
|
||||
precision_thr (float): Precision threshold when one (gt_box, det_box)
|
||||
pair is matched.
|
||||
|
||||
Returns:
|
||||
hmean (tuple[dict]): Tuple of dicts indicates the hmean for the dataset
|
||||
and all images.
|
||||
"""
|
||||
assert utils.is_3dlist(pred_boxes)
|
||||
assert utils.is_3dlist(gt_boxes)
|
||||
assert utils.is_3dlist(gt_ignored_boxes)
|
||||
assert 0 <= iou_thr <= 1
|
||||
assert 0 <= precision_thr <= 1
|
||||
|
||||
img_num = len(pred_boxes)
|
||||
assert img_num == len(gt_boxes)
|
||||
assert img_num == len(gt_ignored_boxes)
|
||||
|
||||
dataset_gt_num = 0
|
||||
dataset_pred_num = 0
|
||||
dataset_hit_num = 0
|
||||
|
||||
img_results = []
|
||||
|
||||
for i in range(img_num):
|
||||
gt = gt_boxes[i]
|
||||
gt_ignored = gt_ignored_boxes[i]
|
||||
pred = pred_boxes[i]
|
||||
|
||||
gt_num = len(gt)
|
||||
gt_ignored_num = len(gt_ignored)
|
||||
pred_num = len(pred)
|
||||
|
||||
hit_num = 0
|
||||
|
||||
# get gt polygons.
|
||||
gt_all = gt + gt_ignored
|
||||
gt_polys = [eval_utils.points2polygon(p) for p in gt_all]
|
||||
gt_ignored_index = [gt_num + i for i in range(len(gt_ignored))]
|
||||
gt_num = len(gt_polys)
|
||||
pred_polys, _, pred_ignored_index = eval_utils.ignore_pred(
|
||||
pred, gt_ignored_index, gt_polys, precision_thr)
|
||||
|
||||
# match.
|
||||
if gt_num > 0 and pred_num > 0:
|
||||
sz = [gt_num, pred_num]
|
||||
iou_mat = np.zeros(sz)
|
||||
|
||||
gt_hit = np.zeros(gt_num, np.int8)
|
||||
pred_hit = np.zeros(pred_num, np.int8)
|
||||
|
||||
for gt_id in range(gt_num):
|
||||
for pred_id in range(pred_num):
|
||||
gt_pol = gt_polys[gt_id]
|
||||
det_pol = pred_polys[pred_id]
|
||||
|
||||
iou_mat[gt_id,
|
||||
pred_id] = eval_utils.poly_iou(det_pol, gt_pol)
|
||||
|
||||
for gt_id in range(gt_num):
|
||||
for pred_id in range(pred_num):
|
||||
if gt_hit[gt_id] != 0 or pred_hit[
|
||||
pred_id] != 0 or gt_id in gt_ignored_index \
|
||||
or pred_id in pred_ignored_index:
|
||||
continue
|
||||
if iou_mat[gt_id, pred_id] > iou_thr:
|
||||
gt_hit[gt_id] = 1
|
||||
pred_hit[pred_id] = 1
|
||||
hit_num += 1
|
||||
|
||||
gt_care_number = gt_num - gt_ignored_num
|
||||
pred_care_number = pred_num - len(pred_ignored_index)
|
||||
|
||||
r, p, h = eval_utils.compute_hmean(hit_num, hit_num, gt_care_number,
|
||||
pred_care_number)
|
||||
|
||||
img_results.append({'recall': r, 'precision': p, 'hmean': h})
|
||||
|
||||
dataset_hit_num += hit_num
|
||||
dataset_gt_num += gt_care_number
|
||||
dataset_pred_num += pred_care_number
|
||||
|
||||
dataset_r, dataset_p, dataset_h = eval_utils.compute_hmean(
|
||||
dataset_hit_num, dataset_hit_num, dataset_gt_num, dataset_pred_num)
|
||||
|
||||
dataset_results = {
|
||||
'num_gts': dataset_gt_num,
|
||||
'num_dets': dataset_pred_num,
|
||||
'num_match': dataset_hit_num,
|
||||
'recall': dataset_r,
|
||||
'precision': dataset_p,
|
||||
'hmean': dataset_h
|
||||
}
|
||||
|
||||
return dataset_results, img_results
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
import torch.nn as nn
|
||||
from mmcv.cnn import uniform_init, xavier_init
|
||||
|
||||
from mmdet.models.builder import BACKBONES
|
||||
|
||||
|
||||
@BACKBONES.register_module()
|
||||
class VeryDeepVgg(nn.Module):
|
||||
|
||||
def __init__(self, leakyRelu=True, input_channels=3):
|
||||
super().__init__()
|
||||
|
||||
ks = [3, 3, 3, 3, 3, 3, 2]
|
||||
ps = [1, 1, 1, 1, 1, 1, 0]
|
||||
ss = [1, 1, 1, 1, 1, 1, 1]
|
||||
nm = [64, 128, 256, 256, 512, 512, 512]
|
||||
|
||||
self.channels = nm
|
||||
|
||||
cnn = nn.Sequential()
|
||||
|
||||
def convRelu(i, batchNormalization=False):
|
||||
nIn = input_channels if i == 0 else nm[i - 1]
|
||||
nOut = nm[i]
|
||||
cnn.add_module('conv{0}'.format(i),
|
||||
nn.Conv2d(nIn, nOut, ks[i], ss[i], ps[i]))
|
||||
if batchNormalization:
|
||||
cnn.add_module('batchnorm{0}'.format(i), nn.BatchNorm2d(nOut))
|
||||
if leakyRelu:
|
||||
cnn.add_module('relu{0}'.format(i),
|
||||
nn.LeakyReLU(0.2, inplace=True))
|
||||
else:
|
||||
cnn.add_module('relu{0}'.format(i), nn.ReLU(True))
|
||||
|
||||
convRelu(0)
|
||||
cnn.add_module('pooling{0}'.format(0), nn.MaxPool2d(2, 2)) # 64x16x64
|
||||
convRelu(1)
|
||||
cnn.add_module('pooling{0}'.format(1), nn.MaxPool2d(2, 2)) # 128x8x32
|
||||
convRelu(2, True)
|
||||
convRelu(3)
|
||||
cnn.add_module('pooling{0}'.format(2),
|
||||
nn.MaxPool2d((2, 2), (2, 1), (0, 1))) # 256x4x16
|
||||
convRelu(4, True)
|
||||
convRelu(5)
|
||||
cnn.add_module('pooling{0}'.format(3),
|
||||
nn.MaxPool2d((2, 2), (2, 1), (0, 1))) # 512x2x16
|
||||
convRelu(6, True) # 512x1x16
|
||||
|
||||
self.cnn = cnn
|
||||
|
||||
def init_weights(self, pretrained=None):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
xavier_init(m)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
uniform_init(m)
|
||||
|
||||
def out_channels(self):
|
||||
return self.channels[-1]
|
||||
|
||||
def forward(self, x):
|
||||
output = self.cnn(x)
|
||||
|
||||
return output
|
||||
|
|
@ -0,0 +1,138 @@
|
|||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from mmocr.models.builder import DECODERS
|
||||
from mmocr.models.textrecog.layers import (DotProductAttentionLayer,
|
||||
PositionAwareLayer)
|
||||
from .base_decoder import BaseDecoder
|
||||
|
||||
|
||||
@DECODERS.register_module()
|
||||
class PositionAttentionDecoder(BaseDecoder):
|
||||
|
||||
def __init__(self,
|
||||
num_classes=None,
|
||||
rnn_layers=2,
|
||||
dim_input=512,
|
||||
dim_model=128,
|
||||
max_seq_len=40,
|
||||
mask=True,
|
||||
return_feature=False,
|
||||
encode_value=False):
|
||||
super().__init__()
|
||||
|
||||
self.num_classes = num_classes
|
||||
self.dim_input = dim_input
|
||||
self.dim_model = dim_model
|
||||
self.max_seq_len = max_seq_len
|
||||
self.return_feature = return_feature
|
||||
self.encode_value = encode_value
|
||||
self.mask = mask
|
||||
|
||||
self.embedding = nn.Embedding(self.max_seq_len + 1, self.dim_model)
|
||||
|
||||
self.position_aware_module = PositionAwareLayer(
|
||||
self.dim_model, rnn_layers)
|
||||
|
||||
self.attention_layer = DotProductAttentionLayer()
|
||||
|
||||
self.prediction = None
|
||||
if not self.return_feature:
|
||||
pred_num_classes = num_classes - 1
|
||||
self.prediction = nn.Linear(
|
||||
dim_model if encode_value else dim_input, pred_num_classes)
|
||||
|
||||
def init_weights(self):
|
||||
pass
|
||||
|
||||
def _get_position_index(self, length, batch_size, device=None):
|
||||
position_index = torch.arange(0, length, device=device)
|
||||
position_index = position_index.repeat([batch_size, 1])
|
||||
position_index = position_index.long()
|
||||
return position_index
|
||||
|
||||
def forward_train(self, feat, out_enc, targets_dict, img_metas):
|
||||
valid_ratios = [
|
||||
img_meta.get('valid_ratio', 1.0) for img_meta in img_metas
|
||||
] if self.mask else None
|
||||
|
||||
targets = targets_dict['padded_targets'].to(feat.device)
|
||||
|
||||
#
|
||||
n, c_enc, h, w = out_enc.size()
|
||||
assert c_enc == self.dim_model
|
||||
_, c_feat, _, _ = feat.size()
|
||||
assert c_feat == self.dim_input
|
||||
_, len_q = targets.size()
|
||||
assert len_q <= self.max_seq_len
|
||||
|
||||
position_index = self._get_position_index(len_q, n, feat.device)
|
||||
|
||||
position_out_enc = self.position_aware_module(out_enc)
|
||||
|
||||
query = self.embedding(position_index)
|
||||
query = query.permute(0, 2, 1).contiguous()
|
||||
key = position_out_enc.view(n, c_enc, h * w)
|
||||
if self.encode_value:
|
||||
value = out_enc.view(n, c_enc, h * w)
|
||||
else:
|
||||
value = feat.view(n, c_feat, h * w)
|
||||
|
||||
mask = None
|
||||
if valid_ratios is not None:
|
||||
mask = query.new_zeros((n, h, w))
|
||||
for i, valid_ratio in enumerate(valid_ratios):
|
||||
valid_width = min(w, math.ceil(w * valid_ratio))
|
||||
mask[i, :, valid_width:] = 1
|
||||
mask = mask.bool()
|
||||
mask = mask.view(n, h * w)
|
||||
|
||||
attn_out = self.attention_layer(query, key, value, mask)
|
||||
attn_out = attn_out.permute(0, 2, 1).contiguous() # [n, len_q, dim_v]
|
||||
|
||||
if self.return_feature:
|
||||
return attn_out
|
||||
|
||||
return self.prediction(attn_out)
|
||||
|
||||
def forward_test(self, feat, out_enc, img_metas):
|
||||
valid_ratios = [
|
||||
img_meta.get('valid_ratio', 1.0) for img_meta in img_metas
|
||||
] if self.mask else None
|
||||
|
||||
seq_len = self.max_seq_len
|
||||
n, c_enc, h, w = out_enc.size()
|
||||
assert c_enc == self.dim_model
|
||||
_, c_feat, _, _ = feat.size()
|
||||
assert c_feat == self.dim_input
|
||||
|
||||
position_index = self._get_position_index(seq_len, n, feat.device)
|
||||
|
||||
position_out_enc = self.position_aware_module(out_enc)
|
||||
|
||||
query = self.embedding(position_index)
|
||||
query = query.permute(0, 2, 1).contiguous()
|
||||
key = position_out_enc.view(n, c_enc, h * w)
|
||||
if self.encode_value:
|
||||
value = out_enc.view(n, c_enc, h * w)
|
||||
else:
|
||||
value = feat.view(n, c_feat, h * w)
|
||||
|
||||
mask = None
|
||||
if valid_ratios is not None:
|
||||
mask = query.new_zeros((n, h, w))
|
||||
for i, valid_ratio in enumerate(valid_ratios):
|
||||
valid_width = min(w, math.ceil(w * valid_ratio))
|
||||
mask[i, :, valid_width:] = 1
|
||||
mask = mask.bool()
|
||||
mask = mask.view(n, h * w)
|
||||
|
||||
attn_out = self.attention_layer(query, key, value, mask)
|
||||
attn_out = attn_out.permute(0, 2, 1).contiguous()
|
||||
|
||||
if self.return_feature:
|
||||
return attn_out
|
||||
|
||||
return self.prediction(attn_out)
|
||||
|
|
@ -0,0 +1,107 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from mmocr.models.builder import DECODERS, build_decoder
|
||||
from mmocr.models.textrecog.layers import RobustScannerFusionLayer
|
||||
from .base_decoder import BaseDecoder
|
||||
|
||||
|
||||
@DECODERS.register_module()
|
||||
class RobustScannerDecoder(BaseDecoder):
|
||||
|
||||
def __init__(self,
|
||||
num_classes=None,
|
||||
dim_input=512,
|
||||
dim_model=128,
|
||||
max_seq_len=40,
|
||||
start_idx=0,
|
||||
mask=True,
|
||||
padding_idx=None,
|
||||
encode_value=False,
|
||||
hybrid_decoder=None,
|
||||
position_decoder=None):
|
||||
super().__init__()
|
||||
self.num_classes = num_classes
|
||||
self.dim_input = dim_input
|
||||
self.dim_model = dim_model
|
||||
self.max_seq_len = max_seq_len
|
||||
self.encode_value = encode_value
|
||||
self.start_idx = start_idx
|
||||
self.padding_idx = padding_idx
|
||||
self.mask = mask
|
||||
|
||||
# init hybrid decoder
|
||||
hybrid_decoder.update(num_classes=self.num_classes)
|
||||
hybrid_decoder.update(dim_input=self.dim_input)
|
||||
hybrid_decoder.update(dim_model=self.dim_model)
|
||||
hybrid_decoder.update(start_idx=self.start_idx)
|
||||
hybrid_decoder.update(padding_idx=self.padding_idx)
|
||||
hybrid_decoder.update(max_seq_len=self.max_seq_len)
|
||||
hybrid_decoder.update(mask=self.mask)
|
||||
hybrid_decoder.update(encode_value=self.encode_value)
|
||||
hybrid_decoder.update(return_feature=True)
|
||||
|
||||
self.hybrid_decoder = build_decoder(hybrid_decoder)
|
||||
|
||||
# init position decoder
|
||||
position_decoder.update(num_classes=self.num_classes)
|
||||
position_decoder.update(dim_input=self.dim_input)
|
||||
position_decoder.update(dim_model=self.dim_model)
|
||||
position_decoder.update(max_seq_len=self.max_seq_len)
|
||||
position_decoder.update(mask=self.mask)
|
||||
position_decoder.update(encode_value=self.encode_value)
|
||||
position_decoder.update(return_feature=True)
|
||||
|
||||
self.position_decoder = build_decoder(position_decoder)
|
||||
|
||||
self.fusion_module = RobustScannerFusionLayer(
|
||||
self.dim_model if encode_value else dim_input)
|
||||
|
||||
pred_num_classes = num_classes - 1
|
||||
self.prediction = nn.Linear(dim_model if encode_value else dim_input,
|
||||
pred_num_classes)
|
||||
|
||||
def init_weights(self):
|
||||
pass
|
||||
|
||||
def forward_train(self, feat, out_enc, targets_dict, img_metas):
|
||||
hybrid_glimpse = self.hybrid_decoder.forward_train(
|
||||
feat, out_enc, targets_dict, img_metas)
|
||||
position_glimpse = self.position_decoder.forward_train(
|
||||
feat, out_enc, targets_dict, img_metas)
|
||||
|
||||
fusion_out = self.fusion_module(hybrid_glimpse, position_glimpse)
|
||||
|
||||
out = self.prediction(fusion_out)
|
||||
|
||||
return out
|
||||
|
||||
def forward_test(self, feat, out_enc, img_metas):
|
||||
seq_len = self.max_seq_len
|
||||
batch_size = feat.size(0)
|
||||
|
||||
decode_sequence = (feat.new_ones(
|
||||
(batch_size, seq_len)) * self.start_idx).long()
|
||||
|
||||
position_glimpse = self.position_decoder.forward_test(
|
||||
feat, out_enc, img_metas)
|
||||
|
||||
outputs = []
|
||||
for i in range(seq_len):
|
||||
hybrid_glimpse_step = self.hybrid_decoder.forward_test_step(
|
||||
feat, out_enc, decode_sequence, i, img_metas)
|
||||
|
||||
fusion_out = self.fusion_module(hybrid_glimpse_step,
|
||||
position_glimpse[:, i, :])
|
||||
|
||||
char_out = self.prediction(fusion_out)
|
||||
char_out = F.softmax(char_out, -1)
|
||||
outputs.append(char_out)
|
||||
_, max_idx = torch.max(char_out, dim=1, keepdim=False)
|
||||
if i < seq_len - 1:
|
||||
decode_sequence[:, i + 1] = max_idx
|
||||
|
||||
outputs = torch.stack(outputs, 1)
|
||||
|
||||
return outputs
|
||||
|
|
@ -0,0 +1,165 @@
|
|||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from mmocr.models.builder import DECODERS
|
||||
from mmocr.models.textrecog.layers import DotProductAttentionLayer
|
||||
from .base_decoder import BaseDecoder
|
||||
|
||||
|
||||
@DECODERS.register_module()
|
||||
class SequenceAttentionDecoder(BaseDecoder):
|
||||
|
||||
def __init__(self,
|
||||
num_classes=None,
|
||||
rnn_layers=2,
|
||||
dim_input=512,
|
||||
dim_model=128,
|
||||
max_seq_len=40,
|
||||
start_idx=0,
|
||||
mask=True,
|
||||
padding_idx=None,
|
||||
dropout_ratio=0,
|
||||
return_feature=False,
|
||||
encode_value=False):
|
||||
super().__init__()
|
||||
|
||||
self.num_classes = num_classes
|
||||
self.dim_input = dim_input
|
||||
self.dim_model = dim_model
|
||||
self.return_feature = return_feature
|
||||
self.encode_value = encode_value
|
||||
self.max_seq_len = max_seq_len
|
||||
self.start_idx = start_idx
|
||||
self.mask = mask
|
||||
|
||||
self.embedding = nn.Embedding(
|
||||
self.num_classes, self.dim_model, padding_idx=padding_idx)
|
||||
|
||||
self.sequence_layer = nn.LSTM(
|
||||
input_size=dim_model,
|
||||
hidden_size=dim_model,
|
||||
num_layers=rnn_layers,
|
||||
batch_first=True,
|
||||
dropout=dropout_ratio)
|
||||
|
||||
self.attention_layer = DotProductAttentionLayer()
|
||||
|
||||
self.prediction = None
|
||||
if not self.return_feature:
|
||||
pred_num_classes = num_classes - 1
|
||||
self.prediction = nn.Linear(
|
||||
dim_model if encode_value else dim_input, pred_num_classes)
|
||||
|
||||
def init_weights(self):
|
||||
pass
|
||||
|
||||
def forward_train(self, feat, out_enc, targets_dict, img_metas):
|
||||
valid_ratios = [
|
||||
img_meta.get('valid_ratio', 1.0) for img_meta in img_metas
|
||||
] if self.mask else None
|
||||
|
||||
targets = targets_dict['padded_targets'].to(feat.device)
|
||||
tgt_embedding = self.embedding(targets)
|
||||
|
||||
n, c_enc, h, w = out_enc.size()
|
||||
assert c_enc == self.dim_model
|
||||
_, c_feat, _, _ = feat.size()
|
||||
assert c_feat == self.dim_input
|
||||
_, len_q, c_q = tgt_embedding.size()
|
||||
assert c_q == self.dim_model
|
||||
assert len_q <= self.max_seq_len
|
||||
|
||||
query, _ = self.sequence_layer(tgt_embedding)
|
||||
query = query.permute(0, 2, 1).contiguous()
|
||||
key = out_enc.view(n, c_enc, h * w)
|
||||
if self.encode_value:
|
||||
value = key
|
||||
else:
|
||||
value = feat.view(n, c_feat, h * w)
|
||||
|
||||
mask = None
|
||||
if valid_ratios is not None:
|
||||
mask = query.new_zeros((n, h, w))
|
||||
for i, valid_ratio in enumerate(valid_ratios):
|
||||
valid_width = min(w, math.ceil(w * valid_ratio))
|
||||
mask[i, :, valid_width:] = 1
|
||||
mask = mask.bool()
|
||||
mask = mask.view(n, h * w)
|
||||
|
||||
attn_out = self.attention_layer(query, key, value, mask)
|
||||
attn_out = attn_out.permute(0, 2, 1).contiguous()
|
||||
|
||||
if self.return_feature:
|
||||
return attn_out
|
||||
|
||||
out = self.prediction(attn_out)
|
||||
|
||||
return out
|
||||
|
||||
def forward_test(self, feat, out_enc, img_metas):
|
||||
seq_len = self.max_seq_len
|
||||
batch_size = feat.size(0)
|
||||
|
||||
decode_sequence = (feat.new_ones(
|
||||
(batch_size, seq_len)) * self.start_idx).long()
|
||||
|
||||
outputs = []
|
||||
for i in range(seq_len):
|
||||
step_out = self.forward_test_step(feat, out_enc, decode_sequence,
|
||||
i, img_metas)
|
||||
outputs.append(step_out)
|
||||
_, max_idx = torch.max(step_out, dim=1, keepdim=False)
|
||||
if i < seq_len - 1:
|
||||
decode_sequence[:, i + 1] = max_idx
|
||||
|
||||
outputs = torch.stack(outputs, 1)
|
||||
|
||||
return outputs
|
||||
|
||||
def forward_test_step(self, feat, out_enc, decode_sequence, current_step,
|
||||
img_metas):
|
||||
valid_ratios = [
|
||||
img_meta.get('valid_ratio', 1.0) for img_meta in img_metas
|
||||
] if self.mask else None
|
||||
|
||||
embed = self.embedding(decode_sequence)
|
||||
|
||||
n, c_enc, h, w = out_enc.size()
|
||||
assert c_enc == self.dim_model
|
||||
_, c_feat, _, _ = feat.size()
|
||||
assert c_feat == self.dim_input
|
||||
_, _, c_q = embed.size()
|
||||
assert c_q == self.dim_model
|
||||
|
||||
query, _ = self.sequence_layer(embed)
|
||||
query = query.permute(0, 2, 1).contiguous()
|
||||
key = out_enc.view(n, c_enc, h * w)
|
||||
if self.encode_value:
|
||||
value = key
|
||||
else:
|
||||
value = feat.view(n, c_feat, h * w)
|
||||
|
||||
mask = None
|
||||
if valid_ratios is not None:
|
||||
mask = query.new_zeros((n, h, w))
|
||||
for i, valid_ratio in enumerate(valid_ratios):
|
||||
valid_width = min(w, math.ceil(w * valid_ratio))
|
||||
mask[i, :, valid_width:] = 1
|
||||
mask = mask.bool()
|
||||
mask = mask.view(n, h * w)
|
||||
|
||||
# [n, c, l]
|
||||
attn_out = self.attention_layer(query, key, value, mask)
|
||||
|
||||
out = attn_out[:, :, current_step]
|
||||
|
||||
if self.return_feature:
|
||||
return out
|
||||
|
||||
out = self.prediction(out)
|
||||
out = F.softmax(out, dim=-1)
|
||||
|
||||
return out
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
import torch.nn as nn
|
||||
from mmcv.cnn import xavier_init
|
||||
|
||||
from mmocr.models.builder import ENCODERS
|
||||
from .base_encoder import BaseEncoder
|
||||
|
||||
|
||||
@ENCODERS.register_module()
|
||||
class ChannelReductionEncoder(BaseEncoder):
|
||||
|
||||
def __init__(self, in_channels, out_channels):
|
||||
super().__init__()
|
||||
|
||||
self.layer = nn.Conv2d(
|
||||
in_channels, out_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def init_weights(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
xavier_init(m)
|
||||
|
||||
def forward(self, feat, img_metas=None):
|
||||
return self.layer(feat)
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class DotProductAttentionLayer(nn.Module):
|
||||
|
||||
def __init__(self, dim_model=None):
|
||||
super().__init__()
|
||||
|
||||
self.scale = dim_model**-0.5 if dim_model is not None else 1.
|
||||
|
||||
def forward(self, query, key, value, mask=None):
|
||||
n, seq_len = mask.size()
|
||||
logits = torch.matmul(query.permute(0, 2, 1), key) * self.scale
|
||||
|
||||
if mask is not None:
|
||||
mask = mask.view(n, 1, seq_len)
|
||||
logits = logits.masked_fill(mask, float('-inf'))
|
||||
|
||||
weights = F.softmax(logits, dim=2)
|
||||
|
||||
glimpse = torch.matmul(weights, value.transpose(1, 2))
|
||||
|
||||
glimpse = glimpse.permute(0, 2, 1).contiguous()
|
||||
|
||||
return glimpse
|
||||
|
|
@ -0,0 +1,35 @@
|
|||
import torch.nn as nn
|
||||
|
||||
|
||||
class PositionAwareLayer(nn.Module):
|
||||
|
||||
def __init__(self, dim_model, rnn_layers=2):
|
||||
super().__init__()
|
||||
|
||||
self.dim_model = dim_model
|
||||
|
||||
self.rnn = nn.LSTM(
|
||||
input_size=dim_model,
|
||||
hidden_size=dim_model,
|
||||
num_layers=rnn_layers,
|
||||
batch_first=True)
|
||||
|
||||
self.mixer = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
dim_model, dim_model, kernel_size=3, stride=1, padding=1),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(
|
||||
dim_model, dim_model, kernel_size=3, stride=1, padding=1))
|
||||
|
||||
def forward(self, img_feature):
|
||||
n, c, h, w = img_feature.size()
|
||||
|
||||
rnn_input = img_feature.permute(0, 2, 3, 1).contiguous()
|
||||
rnn_input = rnn_input.view(n * h, w, c)
|
||||
rnn_output, _ = self.rnn(rnn_input)
|
||||
rnn_output = rnn_output.view(n, h, w, c)
|
||||
rnn_output = rnn_output.permute(0, 3, 1, 2).contiguous()
|
||||
|
||||
out = self.mixer(rnn_output)
|
||||
|
||||
return out
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class RobustScannerFusionLayer(nn.Module):
|
||||
|
||||
def __init__(self, dim_model, dim=-1):
|
||||
super().__init__()
|
||||
|
||||
self.dim_model = dim_model
|
||||
self.dim = dim
|
||||
|
||||
self.linear_layer = nn.Linear(dim_model * 2, dim_model * 2)
|
||||
self.glu_layer = nn.GLU(dim=dim)
|
||||
|
||||
def forward(self, x0, x1):
|
||||
assert x0.size() == x1.size()
|
||||
fusion_input = torch.cat([x0, x1], self.dim)
|
||||
output = self.linear_layer(fusion_input)
|
||||
output = self.glu_layer(output)
|
||||
|
||||
return output
|
||||
|
|
@ -0,0 +1,10 @@
|
|||
from mmdet.models.builder import DETECTORS
|
||||
from .encode_decode_recognizer import EncodeDecodeRecognizer
|
||||
|
||||
|
||||
@DETECTORS.register_module()
|
||||
class RobustScanner(EncodeDecodeRecognizer):
|
||||
"""Implementation of `RobustScanner.
|
||||
|
||||
<https://arxiv.org/pdf/2007.07542.pdf>
|
||||
"""
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 11 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 31 KiB |
Loading…
Reference in New Issue