|
|
|
|
|
import copy |
|
|
import json |
|
|
import os |
|
|
import os.path as osp |
|
|
import tempfile |
|
|
|
|
|
import mmcv |
|
|
import numpy as np |
|
|
import pytest |
|
|
import torch |
|
|
from mmcv import Config |
|
|
from mmcv.parallel import MMDataParallel |
|
|
|
|
|
from mmocr.apis.test import single_gpu_test |
|
|
from mmocr.datasets import build_dataloader, build_dataset |
|
|
from mmocr.models import build_detector |
|
|
from mmocr.utils import check_argument, list_to_file, revert_sync_batchnorm |
|
|
|
|
|
|
|
|
def build_model(cfg): |
|
|
model = build_detector(cfg.model, test_cfg=cfg.get('test_cfg')) |
|
|
model = revert_sync_batchnorm(model) |
|
|
model = MMDataParallel(model) |
|
|
|
|
|
return model |
|
|
|
|
|
|
|
|
def generate_sample_dataloader(cfg, curr_dir, img_prefix='', ann_file=''): |
|
|
must_keys = ['img_norm_cfg', 'ori_filename', 'img_shape', 'ori_shape'] |
|
|
test_pipeline = cfg.data.test.pipeline |
|
|
for key in must_keys: |
|
|
if test_pipeline[1].type == 'MultiRotateAugOCR': |
|
|
collect_pipeline = test_pipeline[1]['transforms'][-1] |
|
|
else: |
|
|
collect_pipeline = test_pipeline[-1] |
|
|
if 'meta_keys' not in collect_pipeline: |
|
|
continue |
|
|
collect_pipeline['meta_keys'].append(key) |
|
|
|
|
|
img_prefix = osp.join(curr_dir, img_prefix) |
|
|
ann_file = osp.join(curr_dir, ann_file) |
|
|
test = copy.deepcopy(cfg.data.test.datasets[0]) |
|
|
test.img_prefix = img_prefix |
|
|
test.ann_file = ann_file |
|
|
cfg.data.workers_per_gpu = 0 |
|
|
cfg.data.test.datasets = [test] |
|
|
dataset = build_dataset(cfg.data.test) |
|
|
|
|
|
loader_cfg = { |
|
|
**dict((k, cfg.data[k]) for k in [ |
|
|
'workers_per_gpu', 'samples_per_gpu' |
|
|
] if k in cfg.data) |
|
|
} |
|
|
test_loader_cfg = { |
|
|
**loader_cfg, |
|
|
**dict(shuffle=False, drop_last=False), |
|
|
**cfg.data.get('test_dataloader', {}) |
|
|
} |
|
|
|
|
|
data_loader = build_dataloader(dataset, **test_loader_cfg) |
|
|
|
|
|
return data_loader |
|
|
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda') |
|
|
@pytest.mark.parametrize('cfg_file', [ |
|
|
'../configs/textrecog/sar/sar_r31_parallel_decoder_academic.py', |
|
|
'../configs/textrecog/crnn/crnn_academic_dataset.py', |
|
|
'../configs/textrecog/seg/seg_r31_1by16_fpnocr_academic.py' |
|
|
]) |
|
|
def test_single_gpu_test_recog(cfg_file): |
|
|
curr_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) |
|
|
config_file = os.path.join(curr_dir, cfg_file) |
|
|
cfg = Config.fromfile(config_file) |
|
|
|
|
|
model = build_model(cfg) |
|
|
img_prefix = 'data/ocr_toy_dataset/imgs' |
|
|
ann_file = 'data/ocr_toy_dataset/label.txt' |
|
|
data_loader = generate_sample_dataloader(cfg, curr_dir, img_prefix, |
|
|
ann_file) |
|
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdirname: |
|
|
out_dir = osp.join(tmpdirname, 'tmp') |
|
|
results = single_gpu_test(model, data_loader, out_dir=out_dir) |
|
|
assert check_argument.is_type_list(results, dict) |
|
|
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda') |
|
|
@pytest.mark.parametrize( |
|
|
'cfg_file', |
|
|
['../configs/textdet/psenet/psenet_r50_fpnf_600e_icdar2017.py']) |
|
|
def test_single_gpu_test_det(cfg_file): |
|
|
curr_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) |
|
|
config_file = os.path.join(curr_dir, cfg_file) |
|
|
cfg = Config.fromfile(config_file) |
|
|
|
|
|
model = build_model(cfg) |
|
|
img_prefix = 'data/toy_dataset/imgs' |
|
|
ann_file = 'data/toy_dataset/instances_test.json' |
|
|
data_loader = generate_sample_dataloader(cfg, curr_dir, img_prefix, |
|
|
ann_file) |
|
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdirname: |
|
|
out_dir = osp.join(tmpdirname, 'tmp') |
|
|
results = single_gpu_test(model, data_loader, out_dir=out_dir) |
|
|
assert check_argument.is_type_list(results, dict) |
|
|
|
|
|
|
|
|
def gene_sdmgr_model_dataloader(cfg, dirname, curr_dir, empty_img=False): |
|
|
json_obj = { |
|
|
'file_name': |
|
|
'1.jpg', |
|
|
'height': |
|
|
348, |
|
|
'width': |
|
|
348, |
|
|
'annotations': [{ |
|
|
'box': [114.0, 19.0, 230.0, 19.0, 230.0, 1.0, 114.0, 1.0], |
|
|
'text': |
|
|
'CHOEUN', |
|
|
'label': |
|
|
1 |
|
|
}] |
|
|
} |
|
|
ann_file = osp.join(dirname, 'test.txt') |
|
|
list_to_file(ann_file, [json.dumps(json_obj, ensure_ascii=False)]) |
|
|
|
|
|
if not empty_img: |
|
|
img = np.ones((348, 348, 3), dtype=np.uint8) |
|
|
img_file = osp.join(dirname, '1.jpg') |
|
|
mmcv.imwrite(img, img_file) |
|
|
|
|
|
test = copy.deepcopy(cfg.data.test) |
|
|
test.ann_file = ann_file |
|
|
test.img_prefix = dirname |
|
|
test.dict_file = osp.join(curr_dir, 'data/kie_toy_dataset/dict.txt') |
|
|
cfg.data.workers_per_gpu = 1 |
|
|
cfg.data.test = test |
|
|
cfg.model.class_list = osp.join(curr_dir, |
|
|
'data/kie_toy_dataset/class_list.txt') |
|
|
|
|
|
dataset = build_dataset(cfg.data.test) |
|
|
|
|
|
loader_cfg = { |
|
|
**dict((k, cfg.data[k]) for k in [ |
|
|
'workers_per_gpu', 'samples_per_gpu' |
|
|
] if k in cfg.data) |
|
|
} |
|
|
test_loader_cfg = { |
|
|
**loader_cfg, |
|
|
**dict(shuffle=False, drop_last=False), |
|
|
**cfg.data.get('test_dataloader', {}) |
|
|
} |
|
|
|
|
|
data_loader = build_dataloader(dataset, **test_loader_cfg) |
|
|
model = build_model(cfg) |
|
|
|
|
|
return model, data_loader |
|
|
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda') |
|
|
@pytest.mark.parametrize( |
|
|
'cfg_file', ['../configs/kie/sdmgr/sdmgr_unet16_60e_wildreceipt.py']) |
|
|
def test_single_gpu_test_kie(cfg_file): |
|
|
curr_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) |
|
|
config_file = os.path.join(curr_dir, cfg_file) |
|
|
cfg = Config.fromfile(config_file) |
|
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdirname: |
|
|
out_dir = osp.join(tmpdirname, 'tmp') |
|
|
model, data_loader = gene_sdmgr_model_dataloader( |
|
|
cfg, out_dir, curr_dir) |
|
|
results = single_gpu_test( |
|
|
model, data_loader, out_dir=out_dir, is_kie=True) |
|
|
assert check_argument.is_type_list(results, dict) |
|
|
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda') |
|
|
@pytest.mark.parametrize( |
|
|
'cfg_file', ['../configs/kie/sdmgr/sdmgr_novisual_60e_wildreceipt.py']) |
|
|
def test_single_gpu_test_kie_novisual(cfg_file): |
|
|
curr_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) |
|
|
config_file = os.path.join(curr_dir, cfg_file) |
|
|
cfg = Config.fromfile(config_file) |
|
|
meta_keys = list(cfg.data.test.pipeline[-1]['meta_keys']) |
|
|
must_keys = ['img_norm_cfg', 'ori_filename', 'img_shape'] |
|
|
for key in must_keys: |
|
|
meta_keys.append(key) |
|
|
|
|
|
cfg.data.test.pipeline[-1]['meta_keys'] = tuple(meta_keys) |
|
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdirname: |
|
|
out_dir = osp.join(tmpdirname, 'tmp') |
|
|
model, data_loader = gene_sdmgr_model_dataloader( |
|
|
cfg, out_dir, curr_dir, empty_img=True) |
|
|
results = single_gpu_test( |
|
|
model, data_loader, out_dir=out_dir, is_kie=True) |
|
|
assert check_argument.is_type_list(results, dict) |
|
|
|
|
|
model, data_loader = gene_sdmgr_model_dataloader( |
|
|
cfg, out_dir, curr_dir) |
|
|
results = single_gpu_test( |
|
|
model, data_loader, out_dir=out_dir, is_kie=True) |
|
|
assert check_argument.is_type_list(results, dict) |
|
|
|
