Spaces:
Runtime error
Runtime error
| # 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 | |
| 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) | |