File size: 3,052 Bytes
186701e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
# Copyright (c) OpenMMLab. All rights reserved.
import os
from tempfile import NamedTemporaryFile, TemporaryDirectory
import numpy as np
import pytest
import torch
from mmengine import Config
try:
import importlib
importlib.import_module('mmdeploy')
except ImportError:
pytest.skip('mmdeploy is not installed.', allow_module_level=True)
import mmdeploy.backend.onnxruntime as ort_apis
from mmdeploy.apis import build_task_processor
from mmdeploy.codebase import import_codebase
from mmdeploy.utils import load_config
from mmdeploy.utils.config_utils import register_codebase
from mmdeploy.utils.test import SwitchBackendWrapper
try:
codebase = register_codebase('mmyolo')
import_codebase(codebase, ['mmyolo.deploy'])
except ImportError:
pytest.skip('mmyolo is not installed.', allow_module_level=True)
model_cfg_path = 'tests/test_deploy/data/model.py'
model_cfg = load_config(model_cfg_path)[0]
model_cfg.test_dataloader.dataset.data_root = \
'tests/data'
model_cfg.test_dataloader.dataset.ann_file = 'coco_sample.json'
model_cfg.test_evaluator.ann_file = \
'tests/coco_sample.json'
deploy_cfg = Config(
dict(
backend_config=dict(type='onnxruntime'),
codebase_config=dict(
type='mmyolo',
task='ObjectDetection',
post_processing=dict(
score_threshold=0.05,
confidence_threshold=0.005, # for YOLOv3
iou_threshold=0.5,
max_output_boxes_per_class=200,
pre_top_k=5000,
keep_top_k=100,
background_label_id=-1,
),
module=['mmyolo.deploy']),
onnx_config=dict(
type='onnx',
export_params=True,
keep_initializers_as_inputs=False,
opset_version=11,
input_shape=None,
input_names=['input'],
output_names=['dets', 'labels'])))
onnx_file = NamedTemporaryFile(suffix='.onnx').name
task_processor = None
img_shape = (32, 32)
img = np.random.rand(*img_shape, 3)
@pytest.fixture(autouse=True)
def init_task_processor():
global task_processor
task_processor = build_task_processor(model_cfg, deploy_cfg, 'cpu')
@pytest.fixture
def backend_model():
from mmdeploy.backend.onnxruntime import ORTWrapper
ort_apis.__dict__.update({'ORTWrapper': ORTWrapper})
wrapper = SwitchBackendWrapper(ORTWrapper)
wrapper.set(
outputs={
'dets': torch.rand(1, 10, 5).sort(2).values,
'labels': torch.randint(0, 10, (1, 10))
})
yield task_processor.build_backend_model([''])
wrapper.recover()
def test_visualize(backend_model):
img_path = 'tests/data/color.jpg'
input_dict, _ = task_processor.create_input(
img_path, input_shape=img_shape)
results = backend_model.test_step(input_dict)[0]
with TemporaryDirectory() as dir:
filename = dir + 'tmp.jpg'
task_processor.visualize(img, results, filename, 'window')
assert os.path.exists(filename)
|