mmdeploy/tests/test_mmocr/test_mmocr_export.py

136 lines
4.4 KiB
Python

import mmcv
import numpy as np
import pytest
from mmdeploy.apis.utils import build_dataloader, build_dataset, create_input
from mmdeploy.utils.constants import Codebase, Task
@pytest.mark.parametrize('task', [Task.TEXT_DETECTION, Task.TEXT_RECOGNITION])
def test_create_input(task):
if task == Task.TEXT_DETECTION:
test = dict(
type='IcdarDataset',
pipeline=[{
'type': 'LoadImageFromFile',
'color_type': 'color_ignore_orientation'
}, {
'type':
'MultiScaleFlipAug',
'img_scale': (128, 64),
'flip':
False,
'transforms': [
{
'type': 'Resize',
'img_scale': (256, 128),
'keep_ratio': True
},
{
'type': 'Normalize',
'mean': [123.675, 116.28, 103.53],
'std': [58.395, 57.12, 57.375],
'to_rgb': True
},
{
'type': 'Pad',
'size_divisor': 32
},
{
'type': 'ImageToTensor',
'keys': ['img']
},
{
'type': 'Collect',
'keys': ['img']
},
]
}])
imgs = [np.random.rand(128, 64, 3).astype(np.uint8)]
elif task == Task.TEXT_RECOGNITION:
test = dict(
type='UniformConcatDataset',
pipeline=[
{
'type': 'LoadImageFromFile',
'color_type': 'grayscale'
},
{
'type': 'ResizeOCR',
'height': 32,
'min_width': 32,
'max_width': None,
'keep_aspect_ratio': True
},
{
'type': 'Normalize',
'mean': [127],
'std': [127]
},
{
'type': 'DefaultFormatBundle'
},
{
'type': 'Collect',
'keys': ['img'],
'meta_keys': ['filename', 'resize_shape', 'valid_ratio']
},
])
imgs = [np.random.random((32, 32, 3)).astype(np.uint8)]
data = dict(test=test)
model_cfg = mmcv.Config(dict(data=data))
inputs = create_input(
Codebase.MMOCR,
task,
model_cfg,
imgs,
input_shape=imgs[0].shape[0:2],
device='cpu')
assert inputs is not None, 'Failed to create input'
@pytest.mark.parametrize('task', [Task.TEXT_DETECTION, Task.TEXT_RECOGNITION])
def test_build_dataset(task):
import tempfile
import os
ann_file, ann_path = tempfile.mkstemp()
if task == Task.TEXT_DETECTION:
data = dict(
test={
'type': 'IcdarDataset',
'ann_file':
'tests/test_mmocr/data/icdar2015/instances_test.json',
'img_prefix': 'tests/test_mmocr/data/icdar2015/imgs',
'pipeline': [
{
'type': 'LoadImageFromFile'
},
]
})
elif task == Task.TEXT_RECOGNITION:
data = dict(
test={
'type': 'OCRDataset',
'ann_file': ann_path,
'img_prefix': '',
'loader': {
'type': 'HardDiskLoader',
'repeat': 1,
'parser': {
'type': 'LineStrParser',
'keys': ['filename', 'text'],
'keys_idx': [0, 1],
'separator': ' '
}
},
'pipeline': [],
'test_mode': True
})
dataset_cfg = mmcv.Config(dict(data=data))
dataset = build_dataset(
Codebase.MMOCR, dataset_cfg=dataset_cfg, dataset_type='test')
assert dataset is not None, 'Failed to build dataset'
dataloader = build_dataloader(Codebase.MMOCR, dataset, 1, 1)
os.close(ann_file)
assert dataloader is not None, 'Failed to build dataloader'