Spaces:
Sleeping
Sleeping
# 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) | |
def init_task_processor(): | |
global task_processor | |
task_processor = build_task_processor(model_cfg, deploy_cfg, 'cpu') | |
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) | |