mirror of https://github.com/open-mmlab/mmocr.git
209 lines
7.1 KiB
Python
209 lines
7.1 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
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
|
|
from mmocr.models import build_detector
|
|
from mmocr.registry import DATASETS
|
|
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 = DATASETS.build(cfg.data.test)
|
|
|
|
loader_cfg = {
|
|
**{
|
|
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 = DATASETS.build(cfg.data.test)
|
|
|
|
loader_cfg = {
|
|
**{
|
|
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)
|