mirror of https://github.com/open-mmlab/mmocr.git
132 lines
3.8 KiB
Python
132 lines
3.8 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import copy
|
|
from os.path import dirname, exists, join
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
|
|
def _demo_mm_inputs(num_kernels=0, input_shape=(1, 3, 300, 300),
|
|
num_items=None): # yapf: disable
|
|
"""Create a superset of inputs needed to run test or train batches.
|
|
|
|
Args:
|
|
input_shape (tuple): Input batch dimensions.
|
|
|
|
num_items (None | list[int]): Specifies the number of boxes
|
|
for each batch item.
|
|
"""
|
|
|
|
(N, C, H, W) = input_shape
|
|
rng = np.random.RandomState(0)
|
|
imgs = rng.rand(*input_shape)
|
|
|
|
img_metas = [{
|
|
'img_shape': (H, W, C),
|
|
'ori_shape': (H, W, C),
|
|
'pad_shape': (H, W, C),
|
|
'filename': '<demo>.png',
|
|
} for _ in range(N)]
|
|
relations = [torch.randn(10, 10, 5) for _ in range(N)]
|
|
texts = [torch.ones(10, 16) for _ in range(N)]
|
|
gt_bboxes = [torch.Tensor([[2, 2, 4, 4]]).expand(10, 4) for _ in range(N)]
|
|
gt_labels = [torch.ones(10, 11).long() for _ in range(N)]
|
|
|
|
mm_inputs = {
|
|
'imgs': torch.FloatTensor(imgs).requires_grad_(True),
|
|
'img_metas': img_metas,
|
|
'relations': relations,
|
|
'texts': texts,
|
|
'gt_bboxes': gt_bboxes,
|
|
'gt_labels': gt_labels
|
|
}
|
|
return mm_inputs
|
|
|
|
|
|
def _get_config_directory():
|
|
"""Find the predefined detector config directory."""
|
|
try:
|
|
# Assume we are running in the source mmocr repo
|
|
repo_dpath = dirname(dirname(dirname(__file__)))
|
|
except NameError:
|
|
# For IPython development when this __file__ is not defined
|
|
import mmocr
|
|
repo_dpath = dirname(dirname(mmocr.__file__))
|
|
config_dpath = join(repo_dpath, 'configs')
|
|
if not exists(config_dpath):
|
|
raise Exception('Cannot find config path')
|
|
return config_dpath
|
|
|
|
|
|
def _get_config_module(fname):
|
|
"""Load a configuration as a python module."""
|
|
from mmcv import Config
|
|
config_dpath = _get_config_directory()
|
|
config_fpath = join(config_dpath, fname)
|
|
config_mod = Config.fromfile(config_fpath)
|
|
return config_mod
|
|
|
|
|
|
def _get_detector_cfg(fname):
|
|
"""Grab configs necessary to create a detector.
|
|
|
|
These are deep copied to allow for safe modification of parameters without
|
|
influencing other tests.
|
|
"""
|
|
config = _get_config_module(fname)
|
|
config.model.class_list = None
|
|
model = copy.deepcopy(config.model)
|
|
return model
|
|
|
|
|
|
@pytest.mark.parametrize('cfg_file', [
|
|
'kie/sdmgr/sdmgr_novisual_60e_wildreceipt.py',
|
|
'kie/sdmgr/sdmgr_unet16_60e_wildreceipt.py'
|
|
])
|
|
def test_sdmgr_pipeline(cfg_file):
|
|
model = _get_detector_cfg(cfg_file)
|
|
|
|
from mmocr.models import build_detector
|
|
detector = build_detector(model)
|
|
|
|
input_shape = (1, 3, 128, 128)
|
|
|
|
mm_inputs = _demo_mm_inputs(0, input_shape)
|
|
|
|
imgs = mm_inputs.pop('imgs')
|
|
img_metas = mm_inputs.pop('img_metas')
|
|
relations = mm_inputs.pop('relations')
|
|
texts = mm_inputs.pop('texts')
|
|
gt_bboxes = mm_inputs.pop('gt_bboxes')
|
|
gt_labels = mm_inputs.pop('gt_labels')
|
|
|
|
# Test forward train
|
|
losses = detector.forward(
|
|
imgs,
|
|
img_metas,
|
|
relations=relations,
|
|
texts=texts,
|
|
gt_bboxes=gt_bboxes,
|
|
gt_labels=gt_labels)
|
|
assert isinstance(losses, dict)
|
|
|
|
# Test forward test
|
|
with torch.no_grad():
|
|
batch_results = []
|
|
for idx in range(len(img_metas)):
|
|
result = detector.forward(
|
|
imgs[idx:idx + 1],
|
|
None,
|
|
return_loss=False,
|
|
relations=[relations[idx]],
|
|
texts=[texts[idx]],
|
|
gt_bboxes=[gt_bboxes[idx]])
|
|
batch_results.append(result)
|
|
|
|
# Test show_result
|
|
results = {'nodes': torch.randn(1, 3)}
|
|
boxes = [[1, 1, 2, 1, 2, 2, 1, 2]]
|
|
img = np.random.rand(5, 5, 3)
|
|
detector.show_result(img, results, boxes)
|