136 lines
4.4 KiB
Python
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'
|