diff --git a/backend_ops/ncnn/pyncnn_ext/CMakeLists.txt b/backend_ops/ncnn/pyncnn_ext/CMakeLists.txt index 87108ec247..a3e3733dfa 100755 --- a/backend_ops/ncnn/pyncnn_ext/CMakeLists.txt +++ b/backend_ops/ncnn/pyncnn_ext/CMakeLists.txt @@ -11,4 +11,4 @@ pybind11_add_module(ncnn_ext ncnn_ext.cpp) target_link_libraries(ncnn_ext PUBLIC ncnn ${SHARED_TARGET}) set_target_properties( ncnn_ext PROPERTIES LIBRARY_OUTPUT_DIRECTORY - ${CMAKE_SOURCE_DIR}/mmdeploy/apis/ncnn) + ${CMAKE_SOURCE_DIR}/mmdeploy/backend/ncnn) diff --git a/mmdeploy/__init__.py b/mmdeploy/__init__.py index a73f89927d..e96395d7cb 100644 --- a/mmdeploy/__init__.py +++ b/mmdeploy/__init__.py @@ -7,28 +7,3 @@ importlib.import_module('mmdeploy.mmcv') else: logging.debug('mmcv is not installed.') - -if importlib.util.find_spec('mmcls'): - importlib.import_module('mmdeploy.mmcls') -else: - logging.debug('mmcls is not installed.') - -if importlib.util.find_spec('mmdet'): - importlib.import_module('mmdeploy.mmdet') -else: - logging.debug('mmdet is not installed.') - -if importlib.util.find_spec('mmseg'): - importlib.import_module('mmdeploy.mmseg') -else: - logging.debug('mmseg is not installed.') - -if importlib.util.find_spec('mmocr'): - importlib.import_module('mmdeploy.mmocr') -else: - logging.debug('mmocr is not installed.') - -if importlib.util.find_spec('mmedit'): - importlib.import_module('mmdeploy.mmedit') -else: - logging.debug('mmedit is not installed.') diff --git a/mmdeploy/apis/__init__.py b/mmdeploy/apis/__init__.py index 36065f5285..f923a11fab 100644 --- a/mmdeploy/apis/__init__.py +++ b/mmdeploy/apis/__init__.py @@ -2,13 +2,11 @@ from .extract_model import extract_model from .inference import inference_model from .pytorch2onnx import torch2onnx, torch2onnx_impl -from .test import post_process_outputs, single_gpu_test -from .utils import (build_dataloader, build_dataset, get_tensor_from_input, - init_backend_model) +from .utils import build_task_processor, get_predefined_partition_cfg +from .visualize import visualize_model __all__ = [ - 'create_calib_table', 'torch2onnx_impl', 'torch2onnx', 'extract_model', - 'inference_model', 'init_backend_model', 'single_gpu_test', - 'post_process_outputs', 'build_dataset', 'get_tensor_from_input', - 'build_dataloader' + 'create_calib_table', 'extract_model', 'inference_model', 'torch2onnx', + 'torch2onnx_impl', 'build_task_processor', 'get_predefined_partition_cfg', + 'visualize_model' ] diff --git a/mmdeploy/apis/calibration.py b/mmdeploy/apis/calibration.py index 035520eb47..a435008ba3 100644 --- a/mmdeploy/apis/calibration.py +++ b/mmdeploy/apis/calibration.py @@ -7,9 +7,7 @@ from mmdeploy.core import (RewriterContext, patch_model, reset_mark_function_count) -from mmdeploy.utils import cfg_apply_marks, get_codebase, load_config -from .utils import (build_dataloader, build_dataset, get_tensor_from_input, - init_pytorch_model, run_inference) +from mmdeploy.utils import cfg_apply_marks, load_config def create_calib_table(calib_file: str, @@ -46,32 +44,33 @@ def create_calib_table(calib_file: str, # load dataset_cfg if necessary dataset_cfg = load_config(dataset_cfg)[0] - codebase = get_codebase(deploy_cfg) + from mmdeploy.apis.utils import build_task_processor + task_processor = build_task_processor(model_cfg, deploy_cfg, device) + apply_marks = cfg_apply_marks(deploy_cfg) backend = 'default' - model = init_pytorch_model( - codebase, model_cfg, model_checkpoint, device=device) - dataset = build_dataset(codebase, dataset_cfg, dataset_type) + model = task_processor.init_pytorch_model(model_checkpoint) + dataset = task_processor.build_dataset(dataset_cfg, dataset_type) # patch model patched_model = patch_model(model, cfg=deploy_cfg, backend=backend) - with h5py.File(calib_file, mode='w') as calib_file: - calib_data_group = calib_file.create_group('calib_data') + with h5py.File(calib_file, mode='w') as file: + calib_data_group = file.create_group('calib_data') if not apply_marks: # create end2end group input_data_group = calib_data_group.create_group('end2end') input_group = input_data_group.create_group('input') - dataloader = build_dataloader( - codebase, dataset, 1, 1, dist=False, shuffle=False) + dataloader = task_processor.build_dataloader( + dataset, 1, 1, dist=False, shuffle=False) patched_model = MMDataParallel(patched_model, device_ids=[device_id]) prog_bar = mmcv.ProgressBar(len(dataset)) for data_id, input_data in enumerate(dataloader): if not apply_marks: # save end2end data - input_tensor = get_tensor_from_input(codebase, input_data) + input_tensor = task_processor.get_tensor_from_input(input_data) input_ndarray = input_tensor.detach().cpu().numpy() input_group.create_dataset( str(data_id), @@ -84,10 +83,10 @@ def create_calib_table(calib_file: str, cfg=deploy_cfg, backend=backend, create_calib=True, - calib_file=calib_file, + calib_file=file, data_id=data_id): reset_mark_function_count() - _ = run_inference(codebase, input_data, patched_model) - calib_file.flush() + _ = task_processor.run_inference(patched_model, input_data) + file.flush() prog_bar.update() diff --git a/mmdeploy/apis/extract_model.py b/mmdeploy/apis/extract_model.py index a7896e665a..e705198498 100644 --- a/mmdeploy/apis/extract_model.py +++ b/mmdeploy/apis/extract_model.py @@ -16,7 +16,7 @@ def extract_model(model: Union[str, onnx.ModelProto], start_name_map: Optional[Dict[str, str]] = None, end_name_map: Optional[Dict[str, str]] = None, dynamic_axes: Optional[Dict[str, Dict[int, str]]] = None, - save_file: Optional[str] = None): + save_file: Optional[str] = None) -> onnx.ModelProto: """Extract sub-model from an ONNX model. The sub-model is defined by the names of the input and output tensors diff --git a/mmdeploy/apis/inference.py b/mmdeploy/apis/inference.py index 8040ed8cd9..09c1720df5 100644 --- a/mmdeploy/apis/inference.py +++ b/mmdeploy/apis/inference.py @@ -1,75 +1,40 @@ -from typing import Optional, Sequence, Union +from typing import Any, Sequence, Union import mmcv import numpy as np import torch -from mmdeploy.utils import (Backend, get_backend, get_codebase, - get_input_shape, get_task_type, load_config) -from .utils import (create_input, init_backend_model, init_pytorch_model, - run_inference, visualize) +from mmdeploy.utils import get_input_shape, load_config def inference_model(model_cfg: Union[str, mmcv.Config], deploy_cfg: Union[str, mmcv.Config], - model: Union[str, Sequence[str], torch.nn.Module], - img: Union[str, np.ndarray], - device: str, - backend: Optional[Backend] = None, - output_file: Optional[str] = None, - show_result: bool = False): + backend_files: Sequence[str], img: Union[str, np.ndarray], + device: str) -> Any: """Run inference with PyTorch or backend model and show results. Args: model_cfg (str | mmcv.Config): Model config file or Config object. deploy_cfg (str | mmcv.Config): Deployment config file or Config object. - model (str | list[str], torch.nn.Module): Input model or file(s). + backend_files (Sequence[str]): Input backend model file(s). img (str | np.ndarray): Input image file or numpy array for inference. device (str): A string specifying device type. - backend (Backend): Specifying backend type, defaults to `None`. - output_file (str): Output file to save visualized image, defaults to - `None`. Only valid if `show_result` is set to `False`. - show_result (bool): Whether to show plotted image in windows, defaults - to `False`. + + Returns: + Any: The inference results """ deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg) - codebase = get_codebase(deploy_cfg) - task = get_task_type(deploy_cfg) - input_shape = get_input_shape(deploy_cfg) - if backend is None: - backend = get_backend(deploy_cfg) - - if isinstance(model, str): - model = [model] + from mmdeploy.apis.utils import build_task_processor + task_processor = build_task_processor(model_cfg, deploy_cfg, device) - if isinstance(model, (list, tuple)): - assert len(model) > 0, 'Model should have at least one element.' - assert all([isinstance(m, str) for m in model]), 'All elements in the \ - list should be str' + model = task_processor.init_backend_model(backend_files) - if backend == Backend.PYTORCH: - model = init_pytorch_model(codebase, model_cfg, model[0], device) - else: - device_id = -1 if device == 'cpu' else 0 - model = init_backend_model( - model, - model_cfg=model_cfg, - deploy_cfg=deploy_cfg, - device_id=device_id) - - model_inputs, _ = create_input(codebase, task, model_cfg, img, input_shape, - device) + input_shape = get_input_shape(deploy_cfg) + model_inputs, _ = task_processor.create_input(img, input_shape) with torch.no_grad(): - result = run_inference(codebase, model_inputs, model) + result = task_processor.run_inference(model, model_inputs) - visualize( - codebase, - img, - result=result, - model=model, - output_file=output_file, - backend=backend, - show_result=show_result) + return result diff --git a/mmdeploy/apis/ncnn/__init__.py b/mmdeploy/apis/ncnn/__init__.py index 28ed404be8..9d66d06530 100644 --- a/mmdeploy/apis/ncnn/__init__.py +++ b/mmdeploy/apis/ncnn/__init__.py @@ -1,28 +1,8 @@ -import importlib -import os.path as osp - -from .init_plugins import get_onnx2ncnn_path, get_ops_path - -__all__ = ['get_ops_path', 'get_onnx2ncnn_path'] - - -def is_available(): - """Check whether ncnn with extension is installed. - - Returns: - bool: True if ncnn and its extension are installed. - """ - ncnn_ops_path = get_ops_path() - if not osp.exists(ncnn_ops_path): - return False - has_pyncnn = importlib.util.find_spec('ncnn') is not None - has_pyncnn_ext = importlib.util.find_spec( - 'mmdeploy.apis.ncnn.ncnn_ext') is not None - - return has_pyncnn and has_pyncnn_ext +from mmdeploy.backend.ncnn import is_available +__all__ = ['is_available'] if is_available(): - from .ncnn_utils import NCNNWrapper - - __all__ += ['NCNNWrapper'] + from mmdeploy.backend.ncnn.onnx2ncnn import (onnx2ncnn, + get_output_model_file) + __all__ += ['onnx2ncnn', 'get_output_model_file'] diff --git a/mmdeploy/apis/onnxruntime/__init__.py b/mmdeploy/apis/onnxruntime/__init__.py index 7eed665244..281f1e3c87 100644 --- a/mmdeploy/apis/onnxruntime/__init__.py +++ b/mmdeploy/apis/onnxruntime/__init__.py @@ -1,22 +1,3 @@ -import importlib -import os.path as osp +from mmdeploy.backend.onnxruntime import is_available -from .init_plugins import get_ops_path - - -def is_available(): - """Check whether onnxruntime and its custom ops are installed. - - Returns: - bool: True if onnxruntime package is installed and its - custom ops are compiled. - """ - onnxruntime_op_path = get_ops_path() - if not osp.exists(onnxruntime_op_path): - return False - return importlib.util.find_spec('onnxruntime') is not None - - -if is_available(): - from .onnxruntime_utils import ORTWrapper - __all__ = ['get_ops_path', 'ORTWrapper'] +__all__ = ['is_available'] diff --git a/mmdeploy/apis/openvino/__init__.py b/mmdeploy/apis/openvino/__init__.py index 89547526d2..5d2e28fb15 100644 --- a/mmdeploy/apis/openvino/__init__.py +++ b/mmdeploy/apis/openvino/__init__.py @@ -1,19 +1,11 @@ -import importlib - - -def is_available() -> bool: - """Checking if OpenVINO is installed. - - Returns: - bool: True if OpenVINO is installed. - """ - return importlib.util.find_spec('openvino') is not None +from mmdeploy.backend.openvino import is_available +__all__ = ['is_available'] if is_available(): - from .openvino_utils import OpenVINOWrapper, get_input_shape_from_cfg - from .onnx2openvino import (onnx2openvino, get_output_model_file) - __all__ = [ - 'OpenVINOWrapper', 'onnx2openvino', 'get_output_model_file', - 'get_input_shape_from_cfg' + from mmdeploy.backend.openvino.onnx2openvino \ + import onnx2openvino, get_output_model_file + from .utils import get_input_shape_from_cfg + __all__ += [ + 'onnx2openvino', 'get_output_model_file', 'get_input_shape_from_cfg' ] diff --git a/mmdeploy/apis/openvino/utils.py b/mmdeploy/apis/openvino/utils.py new file mode 100644 index 0000000000..4c66f23aab --- /dev/null +++ b/mmdeploy/apis/openvino/utils.py @@ -0,0 +1,22 @@ +from typing import List + +import mmcv + + +def get_input_shape_from_cfg(config: mmcv.Config) -> List[int]: + """Get the input shape from the model config for OpenVINO Model Optimizer. + + Args: + config (mmcv.Config): Model config. + Returns: + List[int]: The input shape in [1, 3, H, W] format from config + or [1, 3, 800, 1344]. + """ + shape = [] + test_pipeline = config.get('test_pipeline', None) + if test_pipeline is not None: + img_scale = test_pipeline[1]['img_scale'] + shape = [1, 3, img_scale[1], img_scale[0]] + else: + shape = [1, 3, 800, 1344] + return shape diff --git a/mmdeploy/apis/ppl/__init__.py b/mmdeploy/apis/ppl/__init__.py index 4b0d4e954b..3357a3c50a 100644 --- a/mmdeploy/apis/ppl/__init__.py +++ b/mmdeploy/apis/ppl/__init__.py @@ -1,16 +1,8 @@ -import importlib - - -def is_available(): - """Check whether ppl is installed. - - Returns: - bool: True if ppl package is installed. - """ - return importlib.util.find_spec('pyppl') is not None +from mmdeploy.backend.ppl import is_available +__all__ = ['is_available'] if is_available(): - from .ppl_utils import PPLWrapper, register_engines - from .onnx2ppl import onnx2ppl - __all__ = ['register_engines', 'PPLWrapper', 'onnx2ppl'] + from mmdeploy.backend.ppl import onnx2ppl + + __all__ += ['onnx2ppl'] diff --git a/mmdeploy/apis/pytorch2onnx.py b/mmdeploy/apis/pytorch2onnx.py index 91b648af94..527e3b56c8 100644 --- a/mmdeploy/apis/pytorch2onnx.py +++ b/mmdeploy/apis/pytorch2onnx.py @@ -5,9 +5,8 @@ import torch from mmdeploy.core import RewriterContext, patch_model -from mmdeploy.utils import (get_backend, get_codebase, get_input_shape, - get_onnx_config, get_task_type, load_config) -from .utils import create_input, init_pytorch_model +from mmdeploy.utils import (get_backend, get_input_shape, get_onnx_config, + load_config) def torch2onnx_impl(model: torch.nn.Module, input: torch.Tensor, @@ -74,14 +73,13 @@ def torch2onnx(img: Any, mmcv.mkdir_or_exist(osp.abspath(work_dir)) output_file = osp.join(work_dir, save_file) - codebase = get_codebase(deploy_cfg) - task = get_task_type(deploy_cfg) input_shape = get_input_shape(deploy_cfg) - torch_model = init_pytorch_model(codebase, model_cfg, model_checkpoint, - device) - data, model_inputs = create_input(codebase, task, model_cfg, img, - input_shape, device) + from mmdeploy.apis import build_task_processor + task_processor = build_task_processor(model_cfg, deploy_cfg, device) + + torch_model = task_processor.init_pytorch_model(model_checkpoint) + data, model_inputs = task_processor.create_input(img, input_shape) if not isinstance(model_inputs, torch.Tensor): model_inputs = model_inputs[0] diff --git a/mmdeploy/apis/tensorrt/__init__.py b/mmdeploy/apis/tensorrt/__init__.py index f217613559..cb224a475e 100644 --- a/mmdeploy/apis/tensorrt/__init__.py +++ b/mmdeploy/apis/tensorrt/__init__.py @@ -1,34 +1,8 @@ -# flake8: noqa -import importlib -import os.path as osp - -from .init_plugins import get_ops_path, load_tensorrt_plugin - - -def is_available(): - """Check whether TensorRT and plugins are installed. - - Returns: - bool: True if TensorRT and plugins are installed. - """ - tensorrt_op_path = get_ops_path() - if not osp.exists(tensorrt_op_path): - return False - - return importlib.util.find_spec('tensorrt') is not None - +from mmdeploy.backend.tensorrt import is_available __all__ = ['is_available'] if is_available(): - from .onnx2tensorrt import onnx2tensorrt - from .tensorrt_utils import (TRTWrapper, create_trt_engine, - load_trt_engine, save_trt_engine) - - # load tensorrt plugin lib - load_tensorrt_plugin() + from mmdeploy.backend.tensorrt.onnx2tensorrt import onnx2tensorrt - __all__ += [ - 'create_trt_engine', 'save_trt_engine', 'load_trt_engine', - 'TRTWrapper', 'onnx2tensorrt' - ] + __all__ += ['onnx2tensorrt'] diff --git a/mmdeploy/apis/test.py b/mmdeploy/apis/test.py deleted file mode 100644 index 7c77bb5dbd..0000000000 --- a/mmdeploy/apis/test.py +++ /dev/null @@ -1,167 +0,0 @@ -import warnings -from typing import Optional - -import mmcv -import numpy as np -from torch import nn -from torch.utils.data import DataLoader, Dataset - -from mmdeploy.utils import Codebase - - -def single_gpu_test(codebase: Codebase, - model: nn.Module, - data_loader: DataLoader, - show: bool = False, - out_dir: Optional[str] = None, - show_score_thr: float = 0.3): - """Run test with single gpu. - - Args: - codebase (Codebase): Specifying codebase type. - model (torch.nn.Module): Input model from nn.Module. - data_loader (DataLoader): PyTorch data loader. - show (bool): Specifying whether to show plotted results. Defaults - to `False`. - out_dir (str): A directory to save results, defaults to `None`. - show_score_thr (float): A threshold to show detection results, - defaults to `0.3`. - - Returns: - list: The prediction results. - """ - if codebase == Codebase.MMCLS: - from mmcls.apis import single_gpu_test - outputs = single_gpu_test(model, data_loader, show, out_dir) - elif codebase == Codebase.MMDET: - from mmdet.apis import single_gpu_test - outputs = single_gpu_test(model, data_loader, show, out_dir, - show_score_thr) - elif codebase == Codebase.MMSEG: - from mmseg.apis import single_gpu_test - outputs = single_gpu_test(model, data_loader, show, out_dir) - elif codebase == Codebase.MMOCR: - from mmdet.apis import single_gpu_test - outputs = single_gpu_test(model, data_loader, show, out_dir) - elif codebase == Codebase.MMEDIT: - from mmedit.apis import single_gpu_test - outputs = single_gpu_test(model, data_loader, show, out_dir) - else: - raise NotImplementedError(f'Unknown codebase type: {codebase.value}') - return outputs - - -def post_process_outputs(outputs: list, - dataset: Dataset, - model_cfg: mmcv.Config, - codebase: Codebase, - metrics: Optional[str] = None, - out: Optional[str] = None, - metric_options: Optional[dict] = None, - format_only: bool = False): - """Perform post-processing to predictions of model. - - Args: - outputs (list): A list of predictions of model inference. - dataset (Dataset): Input dataset to run test. - model_cfg (mmcv.Config): The model config. - codebase (Codebase): Specifying codebase type. - metrics (str): Evaluation metrics, which depends on - the codebase and the dataset, e.g., "bbox", "segm", "proposal" - for COCO, and "mAP", "recall" for PASCAL VOC in mmdet; "accuracy", - "precision", "recall", "f1_score", "support" for single label - dataset, and "mAP", "CP", "CR", "CF1", "OP", "OR", "OF1" for - multi-label dataset in mmcls. Defaults is `None`. - out (str): Output result file in pickle format, defaults to `None`. - metric_options (dict): Custom options for evaluation, will be kwargs - for dataset.evaluate() function. Defaults to `None`. - format_only (bool): Format the output results without perform - evaluation. It is useful when you want to format the result - to a specific format and submit it to the test server. Defaults - to `False`. - """ - if codebase == Codebase.MMCLS: - if metrics: - results = dataset.evaluate(outputs, metrics, metric_options) - for k, v in results.items(): - print(f'\n{k} : {v:.2f}') - else: - warnings.warn('Evaluation metrics are not specified.') - scores = np.vstack(outputs) - pred_score = np.max(scores, axis=1) - pred_label = np.argmax(scores, axis=1) - pred_class = [dataset.CLASSES[lb] for lb in pred_label] - results = { - 'pred_score': pred_score, - 'pred_label': pred_label, - 'pred_class': pred_class - } - if not out: - print('\nthe predicted result for the first element is ' - f'pred_score = {pred_score[0]:.2f}, ' - f'pred_label = {pred_label[0]} ' - f'and pred_class = {pred_class[0]}. ' - 'Specify --out to save all results to files.') - if out: - print(f'\nwriting results to {out}') - mmcv.dump(results, out) - - elif codebase == Codebase.MMDET: - if out: - print(f'\nwriting results to {out}') - mmcv.dump(outputs, out) - kwargs = {} if metric_options is None else metric_options - if format_only: - dataset.format_results(outputs, **kwargs) - if metrics: - eval_kwargs = model_cfg.get('evaluation', {}).copy() - # hard-code way to remove EvalHook args - for key in [ - 'interval', 'tmpdir', 'start', 'gpu_collect', 'save_best', - 'rule' - ]: - eval_kwargs.pop(key, None) - eval_kwargs.update(dict(metric=metrics, **kwargs)) - print(dataset.evaluate(outputs, **eval_kwargs)) - - elif codebase == Codebase.MMSEG: - if out: - print(f'\nwriting results to {out}') - mmcv.dump(outputs, out) - kwargs = {} if metric_options is None else metric_options - if format_only: - dataset.format_results(outputs, **kwargs) - if metrics: - dataset.evaluate(outputs, metrics, **kwargs) - - elif codebase == Codebase.MMOCR: - if out: - print(f'\nwriting results to {out}') - mmcv.dump(outputs, out) - kwargs = {} if metric_options is None else metric_options - if format_only: - dataset.format_results(outputs, **kwargs) - if metrics: - eval_kwargs = model_cfg.get('evaluation', {}).copy() - # hard-code way to remove EvalHook args - for key in [ - 'interval', 'tmpdir', 'start', 'gpu_collect', 'save_best', - 'rule' - ]: - eval_kwargs.pop(key, None) - eval_kwargs.update(dict(metric=metrics, **kwargs)) - print(dataset.evaluate(outputs, **eval_kwargs)) - - elif codebase == Codebase.MMEDIT: - if out: - print(f'\nwriting results to {out}') - mmcv.dump(outputs, out) - # The Dataset doesn't need metrics - print('\n') - # print metrics - stats = dataset.evaluate(outputs) - for stat in stats: - print('Eval-{}: {}'.format(stat, stats[stat])) - - else: - raise NotImplementedError(f'Unknown codebase type: {codebase.value}') diff --git a/mmdeploy/apis/utils.py b/mmdeploy/apis/utils.py index cc603bd8c0..bb30ef8649 100644 --- a/mmdeploy/apis/utils.py +++ b/mmdeploy/apis/utils.py @@ -1,367 +1,42 @@ -from typing import Any, Dict, Optional, Sequence, Union - import mmcv -import numpy as np -import torch -from torch.utils.data import Dataset - -from mmdeploy.utils import Backend, Codebase, Task, get_codebase, load_config - - -def init_pytorch_model(codebase: Codebase, - model_cfg: Union[str, mmcv.Config], - model_checkpoint: Optional[str] = None, - device: str = 'cuda:0', - cfg_options: Optional[Dict] = None): - """Initialize torch model. - - Args: - codebase (Codebase): Specifying codebase type. - model_cfg (str | mmcv.Config): Model config file or Config object. - model_checkpoint (str): The checkpoint file of torch model, defaults - to `None`. - device (str): A string specifying device type, defaults to 'cuda:0'. - cfg_options (dict): Optional config key-pair parameters. - - Returns: - nn.Module: An initialized torch model. - """ - if codebase == Codebase.MMCLS: - from mmcls.apis import init_model - model = init_model(model_cfg, model_checkpoint, device, cfg_options) - - elif codebase == Codebase.MMDET: - from mmdet.apis import init_detector - model = init_detector(model_cfg, model_checkpoint, device, cfg_options) - - elif codebase == Codebase.MMSEG: - from mmseg.apis import init_segmentor - from mmdeploy.mmseg.export import convert_syncbatchnorm - model = init_segmentor(model_cfg, model_checkpoint, device) - model = convert_syncbatchnorm(model) - - elif codebase == Codebase.MMOCR: - from mmocr.apis import init_detector - model = init_detector(model_cfg, model_checkpoint, device, cfg_options) - - elif codebase == Codebase.MMEDIT: - from mmedit.apis import init_model - model = init_model(model_cfg, model_checkpoint, device) - model.forward = model.forward_dummy - else: - raise NotImplementedError(f'Unknown codebase type: {codebase.value}') - - return model.eval() - - -def create_input(codebase: Codebase, - task: Task, - model_cfg: Union[str, mmcv.Config], - imgs: Union[str, np.ndarray], - input_shape: Sequence[int] = None, - device: str = 'cuda:0', - **kwargs): - """Create input for model. - - Args: - codebase (Codebase): Specifying codebase type. - task (Task): Specifying task type. - model_cfg (str | mmcv.Config): Model config file or loaded Config - object. - imgs (str | np.ndarray): Input image(s), accpeted data types are `str`, - `np.ndarray`. - input_shape (list[int]): Input shape of image in (width, height) - format, defaults to `None`. - device (str): A string specifying device type, defaults to 'cuda:0'. - - Returns: - tuple: (data, img), meta information for the input image and input - image tensor. - """ - model_cfg = load_config(model_cfg)[0] - - cfg = model_cfg.copy() - if codebase == Codebase.MMCLS: - from mmdeploy.mmcls.export import create_input - return create_input(task, cfg, imgs, input_shape, device, **kwargs) - - elif codebase == Codebase.MMDET: - from mmdeploy.mmdet.export import create_input - return create_input(task, cfg, imgs, input_shape, device, **kwargs) - - elif codebase == Codebase.MMOCR: - from mmdeploy.mmocr.export import create_input - return create_input(task, cfg, imgs, input_shape, device, **kwargs) - - elif codebase == Codebase.MMSEG: - from mmdeploy.mmseg.export import create_input - return create_input(task, cfg, imgs, input_shape, device, **kwargs) - - elif codebase == Codebase.MMEDIT: - from mmdeploy.mmedit.export import create_input - return create_input(task, cfg, imgs, input_shape, device, **kwargs) - else: - raise NotImplementedError(f'Unknown codebase type: {codebase.value}') +from mmdeploy.codebase import BaseTask, get_codebase_class +from mmdeploy.utils import get_codebase, get_task_type -def init_backend_model(model_files: Sequence[str], - model_cfg: Union[str, mmcv.Config], - deploy_cfg: Union[str, mmcv.Config], - device_id: int = 0, - **kwargs): - """Initialize backend model. +def build_task_processor(model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str) -> BaseTask: + """Build a task processor to manage the deploy pipeline. Args: - model_files (list[str]): Input model files. - model_cfg (str | mmcv.Config): Model config file or - loaded Config object. - deploy_cfg (str | mmcv.Config): Deployment config file or - loaded Config object. - device_id (int): An integer specifying device index. + model_cfg (str | mmcv.Config): Model config file. + deploy_cfg (str | mmcv.Config): Deployment config file. + device (str): A string specifying device type. Returns: - nn.Module: An initialized model. + BaseTask: A task processor. """ - deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg) + codebase_type = get_codebase(deploy_cfg) + codebase = get_codebase_class(codebase_type) + return codebase.build_task_processor(model_cfg, deploy_cfg, device) - codebase = get_codebase(deploy_cfg) - if codebase == Codebase.MMCLS: - from mmdeploy.mmcls.apis import build_classifier - return build_classifier( - model_files, model_cfg, deploy_cfg, device_id=device_id) - - elif codebase == Codebase.MMDET: - from mmdeploy.mmdet.apis import build_detector - return build_detector( - model_files, model_cfg, deploy_cfg, device_id=device_id) - - elif codebase == Codebase.MMSEG: - from mmdeploy.mmseg.apis import build_segmentor - return build_segmentor( - model_files, model_cfg, deploy_cfg, device_id=device_id) - - elif codebase == Codebase.MMOCR: - from mmdeploy.mmocr.apis import build_ocr_processor - return build_ocr_processor( - model_files, model_cfg, deploy_cfg, device_id=device_id) - - elif codebase == Codebase.MMEDIT: - from mmdeploy.mmedit.apis import build_editing_processor - return build_editing_processor(model_files, model_cfg, deploy_cfg, - device_id) - - else: - raise NotImplementedError(f'Unknown codebase type: {codebase.value}') - - -def run_inference(codebase: Codebase, model_inputs: dict, - model: torch.nn.Module): - """Run once inference for a model of nn.Module. - - Args: - codebase (Codebase): Specifying codebase type. - model_inputs (dict): A dict containing model inputs tensor and - meta info. - model (nn.Module): Input model. - - Returns: - list: The predictions of model inference. - """ - if codebase == Codebase.MMCLS: - return model(**model_inputs, return_loss=False)[0] - elif codebase == Codebase.MMDET: - return model(**model_inputs, return_loss=False, rescale=True)[0] - elif codebase == Codebase.MMSEG: - return model(**model_inputs, return_loss=False) - elif codebase == Codebase.MMOCR: - return model(**model_inputs, return_loss=False, rescale=True)[0] - elif codebase == Codebase.MMEDIT: - result = model(model_inputs['lq'])[0] - # TODO: (For mmedit codebase) - # The data type of pytorch backend is not consistent - if not isinstance(result, np.ndarray): - result = result.detach().cpu().numpy() - return result - else: - raise NotImplementedError(f'Unknown codebase type: {codebase.value}') - - -def visualize(codebase: Codebase, - image: Union[str, np.ndarray], - result: list, - model: torch.nn.Module, - output_file: str, - backend: Backend, - show_result: bool = False): - """Visualize predictions of a model. - - Args: - codebase (Codebase): Specifying codebase type. - image (str | np.ndarray): Input image to draw predictions on. - result (list): A list of predictions. - model (nn.Module): Input model. - output_file (str): Output file to save drawn image. - backend (Backend): Specifying backend type. - show_result (bool): Whether to show result in windows, defaults - to `False`. - """ - show_img = mmcv.imread(image) if isinstance(image, str) else image - output_file = None if show_result else output_file - - if codebase == Codebase.MMCLS: - from mmdeploy.mmcls.apis import show_result as show_result_mmcls - show_result_mmcls(model, show_img, result, output_file, backend, - show_result) - elif codebase == Codebase.MMDET: - from mmdeploy.mmdet.apis import show_result as show_result_mmdet - show_result_mmdet(model, show_img, result, output_file, backend, - show_result) - elif codebase == Codebase.MMSEG: - from mmdeploy.mmseg.apis import show_result as show_result_mmseg - show_result_mmseg(model, show_img, result, output_file, backend, - show_result) - elif codebase == Codebase.MMOCR: - from mmdeploy.mmocr.apis import show_result as show_result_mmocr - show_result_mmocr(model, show_img, result, output_file, backend, - show_result) - elif codebase == Codebase.MMEDIT: - from mmdeploy.mmedit.apis import show_result as show_result_mmedit - show_result_mmedit(result, output_file, backend, show_result) - - -def get_partition_cfg(codebase: Codebase, partition_type: str): - """Get a certain partition config. +def get_predefined_partition_cfg(deploy_cfg: mmcv.Config, partition_type: str): + """Get the predefined partition config. Notes: Currently only support mmdet codebase. Args: - codebase (Codebase): Specifying codebase type. + deploy_cfg (mmcv.Config): use deploy config to get the codebase and + task type. partition_type (str): A string specifying partition type. Returns: dict: A dictionary of partition config. """ - if codebase == Codebase.MMDET: - from mmdeploy.mmdet.export import get_partition_cfg \ - as get_partition_cfg_mmdet - return get_partition_cfg_mmdet(partition_type) - else: - raise NotImplementedError(f'Unknown codebase type: {codebase.value}') - - -def build_dataset(codebase: Codebase, - dataset_cfg: Union[str, mmcv.Config], - dataset_type: str = 'val', - **kwargs): - """Build dataset for different codebase. - - Args: - codebase (Codebase): Specifying codebase type. - dataset_cfg (str | mmcv.Config): Dataset config file or Config object. - dataset_type (str): Specifying dataset type, e.g.: 'train', 'test', - 'val', defaults to 'val'. - - Returns: - Dataset: The built dataset. - """ - if codebase == Codebase.MMCLS: - from mmdeploy.mmcls.export import build_dataset \ - as build_dataset_mmcls - return build_dataset_mmcls(dataset_cfg, dataset_type, **kwargs) - elif codebase == Codebase.MMDET: - from mmdeploy.mmdet.export import build_dataset \ - as build_dataset_mmdet - return build_dataset_mmdet(dataset_cfg, dataset_type, **kwargs) - elif codebase == Codebase.MMSEG: - from mmdeploy.mmseg.export import build_dataset as build_dataset_mmseg - return build_dataset_mmseg(dataset_cfg, dataset_type, **kwargs) - elif codebase == Codebase.MMEDIT: - from mmdeploy.mmedit.export import build_dataset \ - as build_dataset_mmedit - return build_dataset_mmedit(dataset_cfg, **kwargs) - elif codebase == Codebase.MMOCR: - from mmdeploy.mmocr.export import build_dataset as build_dataset_mmocr - return build_dataset_mmocr(dataset_cfg, dataset_type, **kwargs) - else: - raise NotImplementedError(f'Unknown codebase type: {codebase.value}') - - -def build_dataloader(codebase: Codebase, dataset: Dataset, - samples_per_gpu: int, workers_per_gpu: int, **kwargs): - """Build PyTorch dataloader. - - Args: - codebase (Codebase): Specifying codebase type. - dataset (Dataset): A PyTorch dataset. - samples_per_gpu (int): Number of training samples on each GPU, i.e., - batch size of each GPU. - workers_per_gpu (int): How many subprocesses to use for data loading - for each GPU. - - Returns: - DataLoader: A PyTorch dataloader. - """ - if codebase == Codebase.MMCLS: - from mmdeploy.mmcls.export import build_dataloader \ - as build_dataloader_mmcls - return build_dataloader_mmcls(dataset, samples_per_gpu, - workers_per_gpu, **kwargs) - elif codebase == Codebase.MMDET: - from mmdeploy.mmdet.export import build_dataloader \ - as build_dataloader_mmdet - return build_dataloader_mmdet(dataset, samples_per_gpu, - workers_per_gpu, **kwargs) - elif codebase == Codebase.MMSEG: - from mmdeploy.mmseg.export import build_dataloader \ - as build_dataloader_mmseg - return build_dataloader_mmseg(dataset, samples_per_gpu, - workers_per_gpu, **kwargs) - elif codebase == Codebase.MMEDIT: - from mmdeploy.mmedit.export import build_dataloader \ - as build_dataloader_mmedit - return build_dataloader_mmedit(dataset, samples_per_gpu, - workers_per_gpu, **kwargs) - elif codebase == Codebase.MMOCR: - from mmdeploy.mmocr.export import build_dataloader \ - as build_dataloader_mmocr - return build_dataloader_mmocr(dataset, samples_per_gpu, - workers_per_gpu, **kwargs) - else: - raise NotImplementedError(f'Unknown codebase type: {codebase.value}') - - -def get_tensor_from_input(codebase: Codebase, input_data: Dict[str, Any]): - """Get input tensor from input data. - - Args: - codebase (Codebase): Specifying codebase type. - input_data (dict): Input data containing meta info and image tensor. - - Returns: - torch.Tensor: An image in `Tensor`. - """ - if codebase == Codebase.MMCLS: - from mmdeploy.mmcls.export import get_tensor_from_input \ - as get_tensor_from_input_mmcls - return get_tensor_from_input_mmcls(input_data) - elif codebase == Codebase.MMDET: - from mmdeploy.mmdet.export import get_tensor_from_input \ - as get_tensor_from_input_mmdet - return get_tensor_from_input_mmdet(input_data) - elif codebase == Codebase.MMSEG: - from mmdeploy.mmseg.export import get_tensor_from_input \ - as get_tensor_from_input_mmseg - return get_tensor_from_input_mmseg(input_data) - elif codebase == Codebase.MMOCR: - from mmdeploy.mmocr.export import get_tensor_from_input \ - as get_tensor_from_input_mmocr - return get_tensor_from_input_mmocr(input_data) - elif codebase == Codebase.MMEDIT: - from mmdeploy.mmedit.export import get_tensor_from_input \ - as get_tensor_from_input_mmedit - return get_tensor_from_input_mmedit(input_data) - else: - raise NotImplementedError(f'Unknown codebase type: {codebase.value}') + codebase_type = get_codebase(deploy_cfg) + task = get_task_type(deploy_cfg) + codebase = get_codebase_class(codebase_type) + task_processor_class = codebase.get_task_class(task) + return task_processor_class.get_partition_cfg(partition_type) diff --git a/mmdeploy/apis/visualize.py b/mmdeploy/apis/visualize.py new file mode 100644 index 0000000000..ce43af3328 --- /dev/null +++ b/mmdeploy/apis/visualize.py @@ -0,0 +1,67 @@ +from typing import Optional, Sequence, Union + +import mmcv +import numpy as np +import torch + +from mmdeploy.codebase import BaseTask +from mmdeploy.utils import Backend, get_backend, get_input_shape, load_config + + +def visualize_model(model_cfg: Union[str, mmcv.Config], + deploy_cfg: Union[str, mmcv.Config], + model: Union[str, Sequence[str], BaseTask], + img: Union[str, np.ndarray], + device: str, + backend: Optional[Backend] = None, + output_file: Optional[str] = None, + show_result: bool = False): + """Run inference with PyTorch or backend model and show results. + + Args: + model_cfg (str | mmcv.Config): Model config file or Config object. + deploy_cfg (str | mmcv.Config): Deployment config file or Config + object. + model (str | list[str], BaseSubtask): Input model or file(s). + img (str | np.ndarray): Input image file or numpy array for inference. + device (str): A string specifying device type. + backend (Backend): Specifying backend type, defaults to `None`. + output_file (str): Output file to save visualized image, defaults to + `None`. Only valid if `show_result` is set to `False`. + show_result (bool): Whether to show plotted image in windows, defaults + to `False`. + """ + deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg) + + from mmdeploy.apis.utils import build_task_processor + task_processor = build_task_processor(model_cfg, deploy_cfg, device) + + input_shape = get_input_shape(deploy_cfg) + if backend is None: + backend = get_backend(deploy_cfg) + + if isinstance(model, str): + model = [model] + + if isinstance(model, (list, tuple)): + assert len(model) > 0, 'Model should have at least one element.' + assert all([isinstance(m, str) for m in model]), 'All elements in the \ + list should be str' + + if backend == Backend.PYTORCH: + model = task_processor.init_pytorch_model(model[0]) + else: + model = task_processor.init_backend_model(model) + + model_inputs, _ = task_processor.create_input(img, input_shape) + + with torch.no_grad(): + result = task_processor.run_inference(model, model_inputs)[0] + + task_processor.visualize( + image=img, + model=model, + result=result, + output_file=output_file, + window_name=backend.value, + show_result=show_result) diff --git a/mmdeploy/backend/__init__.py b/mmdeploy/backend/__init__.py new file mode 100644 index 0000000000..59675fc573 --- /dev/null +++ b/mmdeploy/backend/__init__.py @@ -0,0 +1,22 @@ +from mmdeploy.backend.ncnn import is_available as ncnn_available +from mmdeploy.backend.onnxruntime import is_available as ort_available +from mmdeploy.backend.openvino import is_available as openvino_available +from mmdeploy.backend.ppl import is_available as ppl_available +from mmdeploy.backend.tensorrt import is_available as trt_available + +__all__ = [] +if ncnn_available(): + from .ncnn import NCNNWrapper # noqa: F401,F403 + __all__.append('NCNNWrapper') +if ort_available(): + from .onnxruntime import ORTWrapper # noqa: F401,F403 + __all__.append('ORTWrapper') +if trt_available(): + from .tensorrt import TRTWrapper # noqa: F401,F403 + __all__.append('TRTWrapper') +if ppl_available(): + from .ppl import PPLWrapper # noqa: F401,F403 + __all__.append('PPLWrapper') +if openvino_available(): + from .openvino import OpenVINOWrapper # noqa: F401,F403 + __all__.append('OpenVINOWrapper') diff --git a/mmdeploy/backend/base/__init__.py b/mmdeploy/backend/base/__init__.py new file mode 100644 index 0000000000..9ac96fb31f --- /dev/null +++ b/mmdeploy/backend/base/__init__.py @@ -0,0 +1,8 @@ +from .backend_wrapper_registry import (BACKEND_WRAPPER, get_backend_file_count, + get_backend_wrapper_class) +from .base_wrapper import BaseWrapper + +__all__ = [ + 'BaseWrapper', 'BACKEND_WRAPPER', 'get_backend_wrapper_class', + 'get_backend_file_count' +] diff --git a/mmdeploy/backend/base/backend_wrapper_registry.py b/mmdeploy/backend/base/backend_wrapper_registry.py new file mode 100644 index 0000000000..0ca8798cef --- /dev/null +++ b/mmdeploy/backend/base/backend_wrapper_registry.py @@ -0,0 +1,27 @@ +from mmcv.utils import Registry + +from mmdeploy.utils.config_utils import Backend + + +def __build_backend_wrapper_class(backend: Backend, registry: Registry): + return registry.module_dict[backend.value] + + +BACKEND_WRAPPER = Registry('backend', __build_backend_wrapper_class) + + +def get_backend_wrapper_class(backend: Backend) -> type: + """Get the backend wrapper class from the registry. + + Args: + backend (Backend): The backend enum type. + + Returns: + type: The backend wrapper class + """ + return BACKEND_WRAPPER.build(backend) + + +def get_backend_file_count(backend: Backend): + backend_class = get_backend_wrapper_class(backend) + return backend_class.get_backend_file_count() diff --git a/mmdeploy/backend/base/base_wrapper.py b/mmdeploy/backend/base/base_wrapper.py new file mode 100644 index 0000000000..985315b8a8 --- /dev/null +++ b/mmdeploy/backend/base/base_wrapper.py @@ -0,0 +1,70 @@ +from abc import ABCMeta, abstractmethod +from typing import Dict, List, Sequence + +import torch + + +class BaseWrapper(torch.nn.Module, metaclass=ABCMeta): + """Abstract base class for backend wrappers. + + Args: + output_names (Sequence[str]): Names to model outputs in order, which is + useful when converting the output dict to a ordered list or converting + the output ordered list to a key-value dict. + """ + + def __init__(self, output_names: Sequence[str]): + super().__init__() + self._output_names = output_names + + @staticmethod + def get_backend_file_count() -> int: + """Return the count of backend file(s) + + Each backend has its own requirement on backend files (e.g., TensorRT + requires 1 .engine file and ncnn requires 2 files (.param, .bin)). This + interface allow developers to get the count of these required files. + + Returns: + int: The count of required backend file(s). + """ + return 1 + + @abstractmethod + def forward(self, inputs: Dict[str, + torch.Tensor]) -> Dict[str, torch.Tensor]: + """Run forward inference. + + Args: + inputs (Dict[str, torch.Tensor]): Key-value pairs of model inputs. + + Returns: + Dict[str, torch.Tensor]: Key-value pairs of model outputs. + """ + pass + + @property + def output_names(self): + """Return the output names.""" + return self._output_names + + @output_names.setter + def output_names(self, value): + """Set the output names.""" + self._output_names = value + + def output_to_list(self, output_dict: Dict[str, + torch.Tensor]) -> \ + List[torch.Tensor]: + """Convert the output dict of forward() to a tensor list. + + Args: + output_dict (Dict[str, torch.Tensor]): Key-value pairs of model + outputs. + + Returns: + List[torch.Tensor]: An output value list whose order is determined + by the ouput_names list. + """ + outputs = [output_dict[name] for name in self._output_names] + return outputs diff --git a/mmdeploy/backend/ncnn/__init__.py b/mmdeploy/backend/ncnn/__init__.py new file mode 100644 index 0000000000..ccfa0138b0 --- /dev/null +++ b/mmdeploy/backend/ncnn/__init__.py @@ -0,0 +1,26 @@ +import importlib +import os.path as osp + +from .init_plugins import get_ops_path + + +def is_available(): + """Check whether ncnn with extension is installed. + + Returns: + bool: True if ncnn and its extension are installed. + """ + ncnn_ops_path = get_ops_path() + if not osp.exists(ncnn_ops_path): + return False + has_pyncnn = importlib.util.find_spec('ncnn') is not None + has_pyncnn_ext = importlib.util.find_spec( + 'mmdeploy.backend.ncnn.ncnn_ext') is not None + + return has_pyncnn and has_pyncnn_ext + + +if is_available(): + from .wrapper import NCNNWrapper + + __all__ = ['NCNNWrapper'] diff --git a/mmdeploy/apis/ncnn/init_plugins.py b/mmdeploy/backend/ncnn/init_plugins.py similarity index 92% rename from mmdeploy/apis/ncnn/init_plugins.py rename to mmdeploy/backend/ncnn/init_plugins.py index d6b2615404..331a83e92f 100644 --- a/mmdeploy/apis/ncnn/init_plugins.py +++ b/mmdeploy/backend/ncnn/init_plugins.py @@ -2,7 +2,7 @@ import os -def get_ops_path(): +def get_ops_path() -> str: """Get NCNN custom ops library path. Returns: @@ -18,7 +18,7 @@ def get_ops_path(): return lib_path -def get_onnx2ncnn_path(): +def get_onnx2ncnn_path() -> str: """Get onnx2ncnn path. Returns: diff --git a/mmdeploy/backend/ncnn/onnx2ncnn.py b/mmdeploy/backend/ncnn/onnx2ncnn.py new file mode 100644 index 0000000000..b772a27718 --- /dev/null +++ b/mmdeploy/backend/ncnn/onnx2ncnn.py @@ -0,0 +1,40 @@ +from subprocess import call +from typing import List + +from .init_plugins import get_onnx2ncnn_path + + +def get_output_model_file(onnx_path: str, work_dir: str) -> List[str]: + """Returns the path to the .param, .bin file with export result. + + Args: + onnx_path (str): The path to the onnx model. + work_dir (str): The path to the directory for saving the results. + + Returns: + List[str]: The path to the files where the export result will be + located. + """ + save_param = onnx_path.replace('.onnx', '.param') + save_bin = onnx_path.replace('.onnx', '.bin') + + return [save_param, save_bin] + + +def onnx2ncnn(onnx_path: str, work_dir: str): + """Convert ONNX to ncnn. + + The inputs of ncnn include a model file and a weight file. We need to use + a executable program to convert the ".onnx" file to a ".param" file and + a ".bin" file. The output files will save to work_dir. + + Args: + onnx_path (str): The path of the onnx model. + work_dir (str): The path to the directory for saving the results. + """ + + onnx2ncnn_path = get_onnx2ncnn_path() + + save_param, save_bin = get_output_model_file(onnx_path, work_dir) + + call([onnx2ncnn_path, onnx_path, save_param, save_bin]) diff --git a/mmdeploy/apis/ncnn/ncnn_utils.py b/mmdeploy/backend/ncnn/wrapper.py similarity index 67% rename from mmdeploy/apis/ncnn/ncnn_utils.py rename to mmdeploy/backend/ncnn/wrapper.py index 4cf701cb00..5622a578bf 100644 --- a/mmdeploy/apis/ncnn/ncnn_utils.py +++ b/mmdeploy/backend/ncnn/wrapper.py @@ -1,24 +1,28 @@ import importlib -from typing import Dict, Iterable, Optional +from typing import Dict, Optional, Sequence import ncnn import numpy as np import torch +from mmdeploy.utils import Backend from mmdeploy.utils.timer import TimeCounter +from ..base import BACKEND_WRAPPER, BaseWrapper -class NCNNWrapper(torch.nn.Module): +@BACKEND_WRAPPER.register_module(Backend.NCNN.value) +class NCNNWrapper(BaseWrapper): """NCNN wrapper class for inference. Args: param_file (str): Path of a parameter file. bin_file (str): Path of a binary file. - output_names (list[str] | tuple[str]): Names to model outputs. Defaults - to `None`. + output_names (Sequence[str] | None): Names of model outputs in order. + Defaults to `None` and the wrapper will load the output names from + ncnn model. Examples: - >>> from mmdeploy.apis.ncnn import NCNNWrapper + >>> from mmdeploy.backend.ncnn import NCNNWrapper >>> import torch >>> >>> param_file = 'model.params' @@ -32,41 +36,36 @@ class NCNNWrapper(torch.nn.Module): def __init__(self, param_file: str, bin_file: str, - output_names: Optional[Iterable[str]] = None, + output_names: Optional[Sequence[str]] = None, **kwargs): - super(NCNNWrapper, self).__init__() net = ncnn.Net() - if importlib.util.find_spec('mmdeploy.apis.ncnn.ncnn_ext'): - from mmdeploy.apis.ncnn import ncnn_ext + if importlib.util.find_spec('mmdeploy.backend.ncnn.ncnn_ext'): + from mmdeploy.backend.ncnn import ncnn_ext ncnn_ext.register_mmdeploy_custom_layers(net) net.load_param(param_file) net.load_model(bin_file) self._net = net - self._output_names = output_names + if output_names is None: + assert hasattr(self._net, 'output_names') + output_names = self._net.output_names() - def set_output_names(self, output_names: Iterable[str]): - """Set names of the model outputs. + super().__init__(output_names) - Args: - output_names (list[str] | tuple[str]): Names to model outputs. - """ - self._output_names = output_names + @staticmethod + def get_backend_file_count() -> int: + """Return the count of backend file(s) - def get_output_names(self): - """Get names of the model outputs. + ncnn needs a .param file and a .bin file. So the count is 2. Returns: - list[str]: Names to model outputs. + int: The count of required backend file(s). """ - if self._output_names is not None: - return self._output_names - else: - assert hasattr(self._net, 'output_names') - return self._net.output_names() + return 2 - def forward(self, inputs: Dict[str, torch.Tensor]): + def forward(self, inputs: Dict[str, + torch.Tensor]) -> Dict[str, torch.Tensor]: """Run forward inference. Args: @@ -84,7 +83,7 @@ def forward(self, inputs: Dict[str, torch.Tensor]): 'NCNN only supports cpu device' # set output names - output_names = self.get_output_names() + output_names = self._output_names # create output dict outputs = dict([name, [None] * batch_size] for name in output_names) @@ -101,7 +100,8 @@ def forward(self, inputs: Dict[str, torch.Tensor]): ex.input(name, input_mat) # get outputs - result = self.ncnn_execute(extractor=ex, output_names=output_names) + result = self.__ncnn_execute( + extractor=ex, output_names=output_names) for name in output_names: outputs[name][batch_id] = torch.from_numpy( np.array(result[name])) @@ -113,8 +113,8 @@ def forward(self, inputs: Dict[str, torch.Tensor]): return outputs @TimeCounter.count_time() - def ncnn_execute(self, extractor: ncnn.Extractor, - output_names: Iterable[str]): + def __ncnn_execute(self, extractor: ncnn.Extractor, + output_names: Sequence[str]) -> Dict[str, ncnn.Mat]: """Run inference with NCNN. Args: diff --git a/mmdeploy/backend/onnxruntime/__init__.py b/mmdeploy/backend/onnxruntime/__init__.py new file mode 100644 index 0000000000..0ccaf386d5 --- /dev/null +++ b/mmdeploy/backend/onnxruntime/__init__.py @@ -0,0 +1,22 @@ +import importlib +import os.path as osp + +from .init_plugins import get_ops_path + + +def is_available(): + """Check whether onnxruntime and its custom ops are installed. + + Returns: + bool: True if onnxruntime package is installed and its + custom ops are compiled. + """ + onnxruntime_op_path = get_ops_path() + if not osp.exists(onnxruntime_op_path): + return False + return importlib.util.find_spec('onnxruntime') is not None + + +if is_available(): + from .wrapper import ORTWrapper + __all__ = ['ORTWrapper'] diff --git a/mmdeploy/apis/onnxruntime/init_plugins.py b/mmdeploy/backend/onnxruntime/init_plugins.py similarity index 93% rename from mmdeploy/apis/onnxruntime/init_plugins.py rename to mmdeploy/backend/onnxruntime/init_plugins.py index 54c1b72297..7f51e30b21 100644 --- a/mmdeploy/apis/onnxruntime/init_plugins.py +++ b/mmdeploy/backend/onnxruntime/init_plugins.py @@ -2,7 +2,7 @@ import os -def get_ops_path(): +def get_ops_path() -> str: """Get the library path of onnxruntime custom ops. Returns: diff --git a/mmdeploy/apis/onnxruntime/onnxruntime_utils.py b/mmdeploy/backend/onnxruntime/wrapper.py similarity index 59% rename from mmdeploy/apis/onnxruntime/onnxruntime_utils.py rename to mmdeploy/backend/onnxruntime/wrapper.py index 5b45287a79..84174bb569 100644 --- a/mmdeploy/apis/onnxruntime/onnxruntime_utils.py +++ b/mmdeploy/backend/onnxruntime/wrapper.py @@ -1,39 +1,43 @@ import logging import os.path as osp -from typing import Dict, Sequence +from typing import Dict, Optional, Sequence import numpy as np import onnxruntime as ort import torch +from mmdeploy.utils import Backend, parse_device_id from mmdeploy.utils.timer import TimeCounter +from ..base import BACKEND_WRAPPER, BaseWrapper from .init_plugins import get_ops_path -class ORTWrapper(torch.nn.Module): +@BACKEND_WRAPPER.register_module(Backend.ONNXRUNTIME.value) +class ORTWrapper(BaseWrapper): """ONNXRuntime wrapper for inference. - Args: - onnx_file (str): Input onnx model file. - device_id (int): The device id to input model. - output_names (list[str] | tuple[str]): Names to model outputs. - - Examples: - >>> from mmdeploy.apis.onnxruntime import ORTWrapper - >>> import torch - >>> - >>> onnx_file = 'model.onnx' - >>> model = ORTWrapper(onnx_file, -1) - >>> inputs = dict(input=torch.randn(1, 3, 224, 224, device='cpu')) - >>> outputs = model(inputs) - >>> print(outputs) + Args: + onnx_file (str): Input onnx model file. + device (str): The device to input model. + output_names (Sequence[str] | None): Names of model outputs in order. + Defaults to `None` and the wrapper will load the output names from + model. + + Examples: + >>> from mmdeploy.backend.onnxruntime import ORTWrapper + >>> import torch + >>> + >>> onnx_file = 'model.onnx' + >>> model = ORTWrapper(onnx_file, -1) + >>> inputs = dict(input=torch.randn(1, 3, 224, 224, device='cpu')) + >>> outputs = model(inputs) + >>> print(outputs) """ def __init__(self, onnx_file: str, - device_id: int, - output_names: Sequence[str] = None): - super(ORTWrapper, self).__init__() + device: str, + output_names: Optional[Sequence[str]] = None): # get the custom op path ort_custom_op_path = get_ops_path() session_options = ort.SessionOptions() @@ -41,13 +45,15 @@ def __init__(self, if osp.exists(ort_custom_op_path): session_options.register_custom_ops_library(ort_custom_op_path) logging.info(f'Successfully loaded onnxruntime custom ops from \ - {ort_custom_op_path}') + {ort_custom_op_path}') else: logging.warning(f'The library of onnxruntime custom ops does \ - not exist: {ort_custom_op_path}') + not exist: {ort_custom_op_path}') sess = ort.InferenceSession(onnx_file, session_options) + device_id = parse_device_id(device) + providers = ['CPUExecutionProvider'] options = [{}] is_cuda_available = ort.get_device() == 'GPU' @@ -59,18 +65,21 @@ def __init__(self, output_names = [_.name for _ in sess.get_outputs()] self.sess = sess self.io_binding = sess.io_binding() - self.output_names = output_names self.device_id = device_id self.is_cuda_available = is_cuda_available self.device_type = 'cuda' if is_cuda_available else 'cpu' - def forward(self, inputs: Dict[str, torch.Tensor]): + super().__init__(output_names) + + def forward(self, inputs: Dict[str, + torch.Tensor]) -> Dict[str, torch.Tensor]: """Run forward inference. Args: inputs (Dict[str, torch.Tensor]): The input name and tensor pairs. + Returns: - list[np.ndarray]: A list of output numpy array. + Dict[str, torch.Tensor]: The output name and tensor pairs. """ for name, input_tensor in inputs.items(): # set io binding for inputs/outputs @@ -84,16 +93,19 @@ def forward(self, inputs: Dict[str, torch.Tensor]): shape=input_tensor.shape, buffer_ptr=input_tensor.data_ptr()) - for name in self.output_names: + for name in self._output_names: self.io_binding.bind_output(name) # run session to get outputs - self.ort_execute(self.io_binding) - outputs = self.io_binding.copy_outputs_to_cpu() + self.__ort_execute(self.io_binding) + output_list = self.io_binding.copy_outputs_to_cpu() + outputs = {} + for output_name, numpy_tensor in zip(self._output_names, output_list): + outputs[output_name] = torch.from_numpy(numpy_tensor) return outputs @TimeCounter.count_time() - def ort_execute(self, io_binding: ort.IOBinding): + def __ort_execute(self, io_binding: ort.IOBinding): """Run inference with ONNXRuntime session. Args: diff --git a/mmdeploy/backend/openvino/__init__.py b/mmdeploy/backend/openvino/__init__.py new file mode 100644 index 0000000000..2c9f2c2903 --- /dev/null +++ b/mmdeploy/backend/openvino/__init__.py @@ -0,0 +1,16 @@ +import importlib + + +def is_available() -> bool: + """Checking if OpenVINO is installed. + + Returns: + bool: True if OpenVINO is installed. + """ + return importlib.util.find_spec('openvino') is not None + + +if is_available(): + from .wrapper import OpenVINOWrapper + from .onnx2openvino import get_output_model_file + __all__ = ['OpenVINOWrapper', 'get_output_model_file'] diff --git a/mmdeploy/apis/openvino/onnx2openvino.py b/mmdeploy/backend/openvino/onnx2openvino.py similarity index 100% rename from mmdeploy/apis/openvino/onnx2openvino.py rename to mmdeploy/backend/openvino/onnx2openvino.py diff --git a/mmdeploy/apis/openvino/openvino_utils.py b/mmdeploy/backend/openvino/wrapper.py similarity index 75% rename from mmdeploy/apis/openvino/openvino_utils.py rename to mmdeploy/backend/openvino/wrapper.py index 310ef25156..223c2ec9d9 100644 --- a/mmdeploy/apis/openvino/openvino_utils.py +++ b/mmdeploy/backend/openvino/wrapper.py @@ -1,40 +1,26 @@ import os.path as osp -from typing import Dict, List +from typing import Dict, Optional, Sequence -import mmcv import numpy as np import torch +from mmdeploy.utils import Backend from mmdeploy.utils.timer import TimeCounter +from ..base import BACKEND_WRAPPER, BaseWrapper -def get_input_shape_from_cfg(config: mmcv.Config) -> List[int]: - """Get the input shape from the model config for OpenVINO Model Optimizer. - - Args: - config (mmcv.Config): Model config. - Returns: - List[int]: The input shape in [1, 3, H, W] format from config - or [1, 3, 800, 1344]. - """ - shape = [] - test_pipeline = config.get('test_pipeline', None) - if test_pipeline is not None: - img_scale = test_pipeline[1]['img_scale'] - shape = [1, 3, img_scale[1], img_scale[0]] - else: - shape = [1, 3, 800, 1344] - return shape - - -class OpenVINOWrapper(torch.nn.Module): +@BACKEND_WRAPPER.register_module(Backend.OPENVINO.value) +class OpenVINOWrapper(BaseWrapper): """OpenVINO wrapper for inference in CPU. Args: ir_model_file (str): Input OpenVINO IR model file. + output_names (Sequence[str] | None): Names of model outputs in order. + Defaults to `None` and the wrapper will load the output names from + model. Examples: - >>> from mmdeploy.apis.openvino import OpenVINOWrapper + >>> from mmdeploy.backend.openvino import OpenVINOWrapper >>> import torch >>> >>> ir_model_file = 'model.xml' @@ -44,8 +30,11 @@ class OpenVINOWrapper(torch.nn.Module): >>> print(outputs) """ - def __init__(self, ir_model_file: str): - super(OpenVINOWrapper, self).__init__() + def __init__(self, + ir_model_file: str, + output_names: Optional[Sequence[str]] = None, + **kwargs): + from openvino.inference_engine import IECore self.ie = IECore() bin_path = osp.splitext(ir_model_file)[0] + '.bin' @@ -57,6 +46,12 @@ def __init__(self, ir_model_file: str): self.sess = self.ie.load_network( network=self.net, device_name=self.device.upper(), num_requests=1) + # TODO: Check if output_names can be read + if output_names is None: + output_names = [name for name in self.net.outputs] + + super().__init__(output_names) + def __update_device( self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: """Updates the device type to 'self.device' (cpu) for the input @@ -107,11 +102,13 @@ def forward(self, inputs: Dict[str, """ inputs = self.__update_device(inputs) self.__reshape(inputs) - outputs = self.openvino_execute(inputs) + outputs = self.__openvino_execute(inputs) + for output_name, numpy_tensor in outputs.items(): + outputs[output_name] = torch.from_numpy(numpy_tensor) return outputs @TimeCounter.count_time() - def openvino_execute( + def __openvino_execute( self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: """Run inference with OpenVINO IE. @@ -119,7 +116,7 @@ def openvino_execute( inputs (Dict[str, torch.Tensor]): The input name and tensor pairs. Returns: - Dict[str, torch.Tensor]: The output name and tensor pairs. + Dict[str, numpy.ndarray]: The output name and tensor pairs. """ outputs = self.sess.infer(inputs) return outputs diff --git a/mmdeploy/backend/ppl/__init__.py b/mmdeploy/backend/ppl/__init__.py new file mode 100644 index 0000000000..ddc0dab59b --- /dev/null +++ b/mmdeploy/backend/ppl/__init__.py @@ -0,0 +1,16 @@ +import importlib + + +def is_available(): + """Check whether ppl is installed. + + Returns: + bool: True if ppl package is installed. + """ + return importlib.util.find_spec('pyppl') is not None + + +if is_available(): + from .onnx2ppl import onnx2ppl + from .wrapper import PPLWrapper, register_engines + __all__ = ['onnx2ppl', 'register_engines', 'PPLWrapper'] diff --git a/mmdeploy/apis/ppl/onnx2ppl.py b/mmdeploy/backend/ppl/onnx2ppl.py similarity index 79% rename from mmdeploy/apis/ppl/onnx2ppl.py rename to mmdeploy/backend/ppl/onnx2ppl.py index 331f62cb22..5409bce995 100644 --- a/mmdeploy/apis/ppl/onnx2ppl.py +++ b/mmdeploy/backend/ppl/onnx2ppl.py @@ -1,25 +1,9 @@ from typing import Optional, Sequence -import torch from pyppl import nn as pplnn -from mmdeploy.apis.ppl import register_engines - - -def parse_cuda_device_id(device: str) -> int: - """Parse cuda device index from a string. - - Args: - device (str): The typical style of string specifying cuda device, - e.g.: 'cuda:0'. - - Returns: - int: The parsed device id, defaults to `0`. - """ - device_id = 0 - if len(device) >= 6: - device_id = torch.device(device).index - return device_id +from mmdeploy.utils.device import parse_cuda_device_id +from .wrapper import register_engines def onnx2ppl(algo_file: str, diff --git a/mmdeploy/apis/ppl/ppl_utils.py b/mmdeploy/backend/ppl/wrapper.py similarity index 76% rename from mmdeploy/apis/ppl/ppl_utils.py rename to mmdeploy/backend/ppl/wrapper.py index 9ad8fcfa8b..2b5317c909 100644 --- a/mmdeploy/apis/ppl/ppl_utils.py +++ b/mmdeploy/backend/ppl/wrapper.py @@ -1,13 +1,16 @@ import logging import sys -from typing import Dict, Sequence +from typing import Dict, List, Optional, Sequence import numpy as np +import onnx +import pyppl.common as pplcommon +import pyppl.nn as pplnn import torch -from pyppl import common as pplcommon -from pyppl import nn as pplnn +from mmdeploy.utils import Backend, parse_device_id from mmdeploy.utils.timer import TimeCounter +from ..base import BACKEND_WRAPPER, BaseWrapper def register_engines(device_id: int, @@ -15,7 +18,7 @@ def register_engines(device_id: int, quick_select: bool = False, input_shapes: Sequence[Sequence[int]] = None, export_algo_file: str = None, - import_algo_file: str = None): + import_algo_file: str = None) -> List[pplnn.Engine]: """Register engines for ppl runtime. Args: @@ -97,7 +100,8 @@ def register_engines(device_id: int, return engines -class PPLWrapper(torch.nn.Module): +@BACKEND_WRAPPER.register_module(Backend.PPL.value) +class PPLWrapper(BaseWrapper): """PPL wrapper for inference. Args: @@ -106,7 +110,7 @@ class PPLWrapper(torch.nn.Module): device_id (int): Device id to put model. Examples: - >>> from mmdeploy.apis.ppl import PPLWrapper + >>> from mmdeploy.backend.ppl import PPLWrapper >>> import torch >>> >>> onnx_file = 'model.onnx' @@ -116,8 +120,19 @@ class PPLWrapper(torch.nn.Module): >>> print(outputs) """ - def __init__(self, onnx_file: str, algo_file: str, device_id: int): - super(PPLWrapper, self).__init__() + def __init__(self, + onnx_file: str, + algo_file: str, + device: str, + output_names: Optional[Sequence[str]] = None, + **kwargs): + + # enable quick select by default to speed up pipeline + # TODO: open it to users after ppl supports saving serialized models + + # TODO: assert device is gpu + device_id = parse_device_id(device) + # enable quick select by default to speed up pipeline # TODO: disable_avx512 will be removed or open to users in config engines = register_engines( @@ -139,36 +154,39 @@ def __init__(self, onnx_file: str, algo_file: str, device_id: int): for i in range(runtime.GetInputCount()) } - def forward(self, inputs: Dict[str, torch.Tensor]): + if output_names is None: + model = onnx.load(onnx_file) + output_names = [node.name for node in model.graph.output] + + super().__init__(output_names) + + def forward(self, inputs: Dict[str, + torch.Tensor]) -> Dict[str, torch.Tensor]: """Run forward inference. Args: inputs (Dict[str, torch.Tensor]): Input name and tensor pairs. Return: - list[np.ndarray]: A list of output numpy array. + Dict[str, torch.Tensor]: The output name and tensor pairs. """ for name, input_tensor in inputs.items(): input_tensor = input_tensor.contiguous() self.inputs[name].ConvertFromHost(input_tensor.cpu().numpy()) - self.ppl_execute() - outputs = [] + self.__ppl_execute() + outputs = {} for i in range(self.runtime.GetOutputCount()): out_tensor = self.runtime.GetOutputTensor(i).ConvertToHost() - if out_tensor: - outputs.append(np.array(out_tensor, copy=False)) - else: - out_shape = self.runtime.GetOutputTensor( - i).GetShape().GetDims() - outputs.append(np.random.rand(*out_shape)) + name = self.output_names[i] + outputs[name] = torch.from_numpy(np.array(out_tensor, copy=False)) return outputs @TimeCounter.count_time() - def ppl_execute(self): + def __ppl_execute(self): """Run inference with PPL.""" status = self.runtime.Run() - assert status == pplcommon.RC_SUCCESS, 'Run() '\ - 'failed: ' + pplcommon.GetRetCodeStr(status) + assert status == pplcommon.RC_SUCCESS, 'Run() failed: ' + \ + pplcommon.GetRetCodeStr(status) status = self.runtime.Sync() - assert status == pplcommon.RC_SUCCESS, 'Sync() '\ - 'failed: ' + pplcommon.GetRetCodeStr(status) + assert status == pplcommon.RC_SUCCESS, 'Sync() failed: ' + \ + pplcommon.GetRetCodeStr(status) diff --git a/mmdeploy/backend/tensorrt/__init__.py b/mmdeploy/backend/tensorrt/__init__.py new file mode 100644 index 0000000000..e03048d729 --- /dev/null +++ b/mmdeploy/backend/tensorrt/__init__.py @@ -0,0 +1,30 @@ +# flake8: noqa +import importlib +import os.path as osp + +from .init_plugins import get_ops_path, load_tensorrt_plugin + + +def is_available(): + """Check whether TensorRT and plugins are installed. + + Returns: + bool: True if TensorRT and plugins are installed. + """ + tensorrt_op_path = get_ops_path() + if not osp.exists(tensorrt_op_path): + return False + + return importlib.util.find_spec('tensorrt') is not None + + +if is_available(): + from .utils import create_trt_engine, load_trt_engine, save_trt_engine + from .wrapper import TRTWrapper + + # load tensorrt plugin lib + load_tensorrt_plugin() + + __all__ = [ + 'create_trt_engine', 'save_trt_engine', 'load_trt_engine', 'TRTWrapper' + ] diff --git a/mmdeploy/apis/tensorrt/calib_utils.py b/mmdeploy/backend/tensorrt/calib_utils.py similarity index 96% rename from mmdeploy/apis/tensorrt/calib_utils.py rename to mmdeploy/backend/tensorrt/calib_utils.py index e2fad16028..3122730b93 100644 --- a/mmdeploy/apis/tensorrt/calib_utils.py +++ b/mmdeploy/backend/tensorrt/calib_utils.py @@ -59,7 +59,7 @@ def __del__(self): if hasattr(self, 'calib_file'): self.calib_file.close() - def get_batch(self, names: Sequence[str], **kwargs): + def get_batch(self, names: Sequence[str], **kwargs) -> list: """Get batch data.""" if self.count < self.dataset_length: @@ -95,7 +95,7 @@ def get_batch(self, names: Sequence[str], **kwargs): else: return None - def get_algorithm(self): + def get_algorithm(self) -> trt.CalibrationAlgoType: """Get Calibration algo type. Returns: @@ -103,7 +103,7 @@ def get_algorithm(self): """ return self.algorithm - def get_batch_size(self): + def get_batch_size(self) -> int: """Get batch size. Returns: diff --git a/mmdeploy/apis/tensorrt/init_plugins.py b/mmdeploy/backend/tensorrt/init_plugins.py similarity index 93% rename from mmdeploy/apis/tensorrt/init_plugins.py rename to mmdeploy/backend/tensorrt/init_plugins.py index fd06597d63..97b8efa599 100644 --- a/mmdeploy/apis/tensorrt/init_plugins.py +++ b/mmdeploy/backend/tensorrt/init_plugins.py @@ -4,7 +4,7 @@ import os -def get_ops_path(): +def get_ops_path() -> str: """Get path of the TensorRT plugin library. Returns: @@ -20,7 +20,7 @@ def get_ops_path(): return lib_path -def load_tensorrt_plugin(): +def load_tensorrt_plugin() -> bool: """Load TensorRT plugins library. Returns: diff --git a/mmdeploy/apis/tensorrt/onnx2tensorrt.py b/mmdeploy/backend/tensorrt/onnx2tensorrt.py similarity index 81% rename from mmdeploy/apis/tensorrt/onnx2tensorrt.py rename to mmdeploy/backend/tensorrt/onnx2tensorrt.py index 123150ffc2..b96e96cea9 100644 --- a/mmdeploy/apis/tensorrt/onnx2tensorrt.py +++ b/mmdeploy/backend/tensorrt/onnx2tensorrt.py @@ -6,24 +6,8 @@ import tensorrt as trt from mmdeploy.utils import (get_calib_filename, get_common_config, - get_model_inputs, load_config) -from .tensorrt_utils import create_trt_engine, save_trt_engine - - -def parse_device_id(device: str) -> int: - """Parse cuda device index from a string. - - Args: - device (str): The typical style of string specifying cuda device, - e.g.: 'cuda:0'. - - Returns: - int: The parsed device id, defaults to `0`. - """ - device_id = 0 - if len(device) >= 6: - device_id = int(device[5:]) - return device_id + get_model_inputs, load_config, parse_device_id) +from .utils import create_trt_engine, save_trt_engine def onnx2tensorrt(work_dir: str, diff --git a/mmdeploy/apis/tensorrt/tensorrt_utils.py b/mmdeploy/backend/tensorrt/utils.py similarity index 57% rename from mmdeploy/apis/tensorrt/tensorrt_utils.py rename to mmdeploy/backend/tensorrt/utils.py index dd4fa53970..4cd546460d 100644 --- a/mmdeploy/apis/tensorrt/tensorrt_utils.py +++ b/mmdeploy/backend/tensorrt/utils.py @@ -1,11 +1,10 @@ -from typing import Any, Dict, Sequence, Union +from typing import Dict, Sequence, Union import onnx import tensorrt as trt import torch from packaging import version -from mmdeploy.utils.timer import TimeCounter from .calib_utils import HDF5Calibrator @@ -17,7 +16,7 @@ def create_trt_engine(onnx_model: Union[str, onnx.ModelProto], int8_param: dict = None, max_workspace_size: int = 0, device_id: int = 0, - **kwargs): + **kwargs) -> trt.ICudaEngine: """Create a tensorrt engine from ONNX. Args: @@ -87,7 +86,8 @@ def create_trt_engine(onnx_model: Union[str, onnx.ModelProto], config.add_optimization_profile(profile) if fp16_mode: - builder.fp16_mode = fp16_mode + if version.parse(trt.__version__) < version.parse('8'): + builder.fp16_mode = fp16_mode config.set_flag(trt.BuilderFlag.FP16) if int8_mode: @@ -100,9 +100,9 @@ def create_trt_engine(onnx_model: Union[str, onnx.ModelProto], device_id=device_id, algorithm=int8_param.get( 'algorithm', trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2)) - - builder.int8_mode = int8_mode - builder.int8_calibrator = config.int8_calibrator + if version.parse(trt.__version__) < version.parse('8'): + builder.int8_mode = int8_mode + builder.int8_calibrator = config.int8_calibrator # create engine with torch.cuda.device(device): @@ -112,7 +112,7 @@ def create_trt_engine(onnx_model: Union[str, onnx.ModelProto], return engine -def save_trt_engine(engine: trt.ICudaEngine, path: str): +def save_trt_engine(engine: trt.ICudaEngine, path: str) -> None: """Serialize TensorRT engine to disk. Args: @@ -123,7 +123,7 @@ def save_trt_engine(engine: trt.ICudaEngine, path: str): f.write(bytearray(engine.serialize())) -def load_trt_engine(path: str): +def load_trt_engine(path: str) -> trt.ICudaEngine: """Deserialize TensorRT engine from disk. Args: @@ -139,7 +139,7 @@ def load_trt_engine(path: str): return engine -def torch_dtype_from_trt(dtype: trt.DataType): +def torch_dtype_from_trt(dtype: trt.DataType) -> torch.dtype: """Convert pytorch dtype to TensorRT dtype. Args: @@ -168,7 +168,6 @@ def torch_device_from_trt(device: trt.TensorLocation): Args: device (trt.TensorLocation): The device in tensorrt. - Returns: torch.device: The corresponding device in torch. """ @@ -178,108 +177,3 @@ def torch_device_from_trt(device: trt.TensorLocation): return torch.device('cpu') else: return TypeError(f'{device} is not supported by torch') - - -class TRTWrapper(torch.nn.Module): - """TensorRT engine wrapper for inference. - - Args: - engine (tensorrt.ICudaEngine): TensorRT engine to wrap. - - Note: - If the engine is converted from onnx model. The input_names and - output_names should be the same as onnx model. - - Examples: - >>> from mmdeploy.apis.tensorrt import TRTWrapper - >>> engine_file = 'resnet.engine' - >>> model = TRTWrapper(engine_file) - >>> inputs = dict(input=torch.randn(1, 3, 224, 224)) - >>> outputs = model(inputs) - >>> print(outputs) - """ - - def __init__(self, engine: Union[str, trt.ICudaEngine]): - super(TRTWrapper, self).__init__() - self.engine = engine - if isinstance(self.engine, str): - self.engine = load_trt_engine(engine) - - if not isinstance(self.engine, trt.ICudaEngine): - raise TypeError(f'`engine` should be str or trt.ICudaEngine, \ - but given: {type(self.engine)}') - - self._register_state_dict_hook(TRTWrapper._on_state_dict) - self.context = self.engine.create_execution_context() - - self._load_io_names() - - def _load_io_names(self): - """Load input/output names from engine.""" - names = [_ for _ in self.engine] - input_names = list(filter(self.engine.binding_is_input, names)) - output_names = list(set(names) - set(input_names)) - self.input_names = input_names - self.output_names = output_names - - def _on_state_dict(self, state_dict: Dict[str, Any], prefix: str): - """State dict hook - Args: - state_dict (Dict[str, Any]): A dict to save state information - such as the serialized engine, input/output names. - prefix (str): A string to be prefixed at the key of the - state dict. - """ - state_dict[prefix + 'engine'] = bytearray(self.engine.serialize()) - state_dict[prefix + 'input_names'] = self.input_names - state_dict[prefix + 'output_names'] = self.output_names - - def forward(self, inputs: Dict[str, torch.Tensor]): - """Run forward inference. - - Args: - inputs (Dict[str, torch.Tensor]): The input name and tensor pairs. - - Return: - Dict[str, torch.Tensor]: The output name and tensor pairs. - """ - assert self.input_names is not None - assert self.output_names is not None - bindings = [None] * (len(self.input_names) + len(self.output_names)) - - for input_name, input_tensor in inputs.items(): - idx = self.engine.get_binding_index(input_name) - - # All input tensors must be gpu variables - assert 'cuda' in input_tensor.device.type - - if input_tensor.dtype == torch.long: - input_tensor = input_tensor.int() - self.context.set_binding_shape(idx, tuple(input_tensor.shape)) - bindings[idx] = input_tensor.contiguous().data_ptr() - - # create output tensors - outputs = {} - for i, output_name in enumerate(self.output_names): - idx = self.engine.get_binding_index(output_name) - dtype = torch_dtype_from_trt(self.engine.get_binding_dtype(idx)) - shape = tuple(self.context.get_binding_shape(idx)) - - device = torch_device_from_trt(self.engine.get_location(idx)) - output = torch.empty(size=shape, dtype=dtype, device=device) - outputs[output_name] = output - bindings[idx] = output.data_ptr() - - self.trt_execute(bindings=bindings) - - return outputs - - @TimeCounter.count_time() - def trt_execute(self, bindings: Sequence[int]): - """Run inference with TensorRT. - - Args: - bindings (list[int]): A list of integer binding the input/output. - """ - self.context.execute_async_v2(bindings, - torch.cuda.current_stream().cuda_stream) diff --git a/mmdeploy/backend/tensorrt/wrapper.py b/mmdeploy/backend/tensorrt/wrapper.py new file mode 100644 index 0000000000..324ffa8f29 --- /dev/null +++ b/mmdeploy/backend/tensorrt/wrapper.py @@ -0,0 +1,123 @@ +from typing import Any, Dict, Optional, Sequence, Union + +import tensorrt as trt +import torch + +from mmdeploy.utils import Backend +from mmdeploy.utils.timer import TimeCounter +from ..base import BACKEND_WRAPPER, BaseWrapper +from .utils import load_trt_engine, torch_device_from_trt, torch_dtype_from_trt + + +@BACKEND_WRAPPER.register_module(Backend.TENSORRT.value) +class TRTWrapper(BaseWrapper): + """TensorRT engine wrapper for inference. + + Args: + engine (tensorrt.ICudaEngine): TensorRT engine to wrap. + output_names (Sequence[str] | None): Names of model outputs in order. + Defaults to `None` and the wrapper will load the output names from + model. + + Note: + If the engine is converted from onnx model. The input_names and + output_names should be the same as onnx model. + + Examples: + >>> from mmdeploy.backend.tensorrt import TRTWrapper + >>> engine_file = 'resnet.engine' + >>> model = TRTWrapper(engine_file) + >>> inputs = dict(input=torch.randn(1, 3, 224, 224)) + >>> outputs = model(inputs) + >>> print(outputs) + """ + + def __init__(self, + engine: Union[str, trt.ICudaEngine], + output_names: Optional[Sequence[str]] = None): + super().__init__(output_names) + self.engine = engine + if isinstance(self.engine, str): + self.engine = load_trt_engine(engine) + + if not isinstance(self.engine, trt.ICudaEngine): + raise TypeError(f'`engine` should be str or trt.ICudaEngine, \ + but given: {type(self.engine)}') + + self._register_state_dict_hook(TRTWrapper.__on_state_dict) + self.context = self.engine.create_execution_context() + + self.__load_io_names() + + def __load_io_names(self): + """Load input/output names from engine.""" + names = [_ for _ in self.engine] + input_names = list(filter(self.engine.binding_is_input, names)) + self._input_names = input_names + + if self._output_names is None: + output_names = list(set(names) - set(input_names)) + self._output_names = output_names + + def __on_state_dict(self, state_dict: Dict[str, Any], prefix: str): + """State dict hook + Args: + state_dict (Dict[str, Any]): A dict to save state information + such as the serialized engine, input/output names. + prefix (str): A string to be prefixed at the key of the + state dict. + """ + state_dict[prefix + 'engine'] = bytearray(self.engine.serialize()) + state_dict[prefix + 'input_names'] = self._input_names + state_dict[prefix + 'output_names'] = self._output_names + + def forward(self, inputs: Dict[str, + torch.Tensor]) -> Dict[str, torch.Tensor]: + """Run forward inference. + + Args: + inputs (Dict[str, torch.Tensor]): The input name and tensor pairs. + + Return: + Dict[str, torch.Tensor]: The output name and tensor pairs. + """ + assert self._input_names is not None + assert self._output_names is not None + bindings = [None] * (len(self._input_names) + len(self._output_names)) + + for input_name, input_tensor in inputs.items(): + idx = self.engine.get_binding_index(input_name) + + # All input tensors must be gpu variables + assert 'cuda' in input_tensor.device.type + + if input_tensor.dtype == torch.long: + input_tensor = input_tensor.int() + self.context.set_binding_shape(idx, tuple(input_tensor.shape)) + bindings[idx] = input_tensor.contiguous().data_ptr() + + # create output tensors + outputs = {} + for output_name in self._output_names: + idx = self.engine.get_binding_index(output_name) + dtype = torch_dtype_from_trt(self.engine.get_binding_dtype(idx)) + shape = tuple(self.context.get_binding_shape(idx)) + + device = torch_device_from_trt(self.engine.get_location(idx)) + output = torch.empty(size=shape, dtype=dtype, device=device) + outputs[output_name] = output + bindings[idx] = output.data_ptr() + + self.__trt_execute(bindings=bindings) + + return outputs + + @TimeCounter.count_time() + def __trt_execute(self, bindings: Sequence[int]): + """Run inference with TensorRT. + + Args: + bindings (list[int]): A list of integer binding the input/output. + """ + self.context.execute_async_v2(bindings, + torch.cuda.current_stream().cuda_stream) diff --git a/mmdeploy/codebase/__init__.py b/mmdeploy/codebase/__init__.py new file mode 100644 index 0000000000..f5710c8c70 --- /dev/null +++ b/mmdeploy/codebase/__init__.py @@ -0,0 +1,31 @@ +import importlib +import logging + +from .base import BaseTask, MMCodebase, get_codebase_class + +if importlib.util.find_spec('mmcls'): + importlib.import_module('mmdeploy.codebase.mmcls') +else: + logging.debug('mmcls is not installed.') + +if importlib.util.find_spec('mmdet'): + importlib.import_module('mmdeploy.codebase.mmdet') +else: + logging.debug('mmdet is not installed.') + +if importlib.util.find_spec('mmseg'): + importlib.import_module('mmdeploy.codebase.mmseg') +else: + logging.debug('mmseg is not installed.') + +if importlib.util.find_spec('mmocr'): + importlib.import_module('mmdeploy.codebase.mmocr') +else: + logging.debug('mmocr is not installed.') + +if importlib.util.find_spec('mmedit'): + importlib.import_module('mmdeploy.codebase.mmedit') +else: + logging.debug('mmedit is not installed.') + +__all__ = ['MMCodebase', 'BaseTask', 'get_codebase_class'] diff --git a/mmdeploy/codebase/base/__init__.py b/mmdeploy/codebase/base/__init__.py new file mode 100644 index 0000000000..f5bda10b70 --- /dev/null +++ b/mmdeploy/codebase/base/__init__.py @@ -0,0 +1,8 @@ +from .backend_model import BaseBackendModel +from .mmcodebase import CODEBASE, MMCodebase, get_codebase_class +from .task import BaseTask + +__all__ = [ + 'BaseBackendModel', 'BaseTask', 'MMCodebase', 'get_codebase_class', + 'CODEBASE' +] diff --git a/mmdeploy/codebase/base/backend_model.py b/mmdeploy/codebase/base/backend_model.py new file mode 100644 index 0000000000..fd93239480 --- /dev/null +++ b/mmdeploy/codebase/base/backend_model.py @@ -0,0 +1,78 @@ +from abc import ABCMeta, abstractmethod +from typing import Optional, Sequence + +import torch + +from mmdeploy.utils.constants import Backend + + +class BaseBackendModel(torch.nn.Module, metaclass=ABCMeta): + """A backend model wraps the details to initialize and run a backend + engine.""" + + def __init__(self): + super().__init__() + + @staticmethod + def _build_wrapper(backend: Backend, + backend_files: Sequence[str], + device: str, + output_names: Optional[Sequence[str]] = None): + """The default methods to build backend wrappers. + + Args: + backend (Backend): The backend enum type. + beckend_files (Sequence[str]): Paths to all required backend files( + e.g. '.onnx' for ONNX Runtime, '.param' and '.bin' for ncnn). + device (str): A string specifying device type. + output_names (Sequence[str] | None): Names of model outputs in + order. Defaults to `None` and the wrapper will load the output + names from the model. + """ + if backend == Backend.ONNXRUNTIME: + from mmdeploy.backend.onnxruntime import ORTWrapper + return ORTWrapper( + onnx_file=backend_files[0], + device=device, + output_names=output_names) + elif backend == Backend.TENSORRT: + from mmdeploy.backend.tensorrt import TRTWrapper + return TRTWrapper( + engine=backend_files[0], output_names=output_names) + elif backend == Backend.PPL: + from mmdeploy.backend.ppl import PPLWrapper + return PPLWrapper( + onnx_file=backend_files[0], + algo_file=backend_files[1], + device=device, + output_names=output_names) + elif backend == Backend.NCNN: + from mmdeploy.backend.ncnn import NCNNWrapper + return NCNNWrapper( + param_file=backend_files[0], + bin_file=backend_files[1], + output_names=output_names) + elif backend == Backend.OPENVINO: + from mmdeploy.backend.openvino import OpenVINOWrapper + return OpenVINOWrapper( + ir_model_file=backend_files[0], output_names=output_names) + else: + raise NotImplementedError(f'Unknown backend type: {backend.value}') + + @abstractmethod + def forward(self, *args, **kwargs): + """The forward interface that must be implemented. + + The arguments should align to forward() of the corresponding model of + OpenMMLab codebases + """ + pass + + @abstractmethod + def show_result(self, *args, **kwargs): + """The visualize interface that must be implemented. + + The arguments should align to show_result() of the corresponding model + of OpenMMLab codebases + """ + pass diff --git a/mmdeploy/codebase/base/mmcodebase.py b/mmdeploy/codebase/base/mmcodebase.py new file mode 100644 index 0000000000..552a12b6af --- /dev/null +++ b/mmdeploy/codebase/base/mmcodebase.py @@ -0,0 +1,122 @@ +from abc import ABCMeta, abstractmethod +from typing import Optional, Union + +import mmcv +import torch +from mmcv.utils.registry import Registry +from torch.utils.data import DataLoader, Dataset + +from mmdeploy.utils import Codebase, Task + + +class MMCodebase(metaclass=ABCMeta): + """Wrap the apis of OpenMMLab Codebase.""" + + task_registry: Registry = None + + def __init__() -> None: + pass + + @classmethod + def get_task_class(cls, task: Task) -> type: + """Get the task processors class according to the task type. + + Args: + task (Task): The task enumeration. + + Returns: + type: The task processor class. + """ + return cls.task_registry.module_dict[task.value] + + @staticmethod + @abstractmethod + def build_task_processor(model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str): + """The interface to build the task processors of the codebase. + + Args: + model_cfg (str | mmcv.Config): Model config file. + deploy_cfg (str | mmcv.Config): Deployment config file. + device (str): A string specifying device type. + + Returns: + BaseTask: A task processor. + """ + pass + + @staticmethod + @abstractmethod + def build_dataset(dataset_cfg: Union[str, mmcv.Config], + dataset_type: str = 'val', + **kwargs) -> Dataset: + """Build dataset for different codebase. + + Args: + dataset_cfg (str | mmcv.Config): Dataset config file or Config + object. + dataset_type (str): Specifying dataset type, e.g.: 'train', 'test', + 'val', defaults to 'val'. + + Returns: + Dataset: The built dataset. + """ + pass + + @staticmethod + @abstractmethod + def build_dataloader(dataset: Dataset, samples_per_gpu: int, + workers_per_gpu: int, **kwargs) -> DataLoader: + """Build PyTorch dataloader. + + Args: + dataset (Dataset): A PyTorch dataset. + samples_per_gpu (int): Number of training samples on each GPU, + i.e., batch size of each GPU. + workers_per_gpu (int): How many subprocesses to use for data + loading for each GPU. + + Returns: + DataLoader: A PyTorch dataloader. + """ + pass + + @staticmethod + @abstractmethod + def single_gpu_test(model: torch.nn.Module, + data_loader: DataLoader, + show: bool = False, + out_dir: Optional[str] = None, + **kwargs): + """Run test with single gpu. + + Args: + model (torch.nn.Module): Input model from nn.Module. + data_loader (DataLoader): PyTorch data loader. + show (bool): Specifying whether to show plotted results. Defaults + to `False`. + out_dir (str): A directory to save results, defaults to `None`. + + Returns: + list: The prediction results. + """ + pass + + +def __build_codebase_class(codebase: Codebase, registry: Registry): + return registry.module_dict[codebase.value] + + +CODEBASE = Registry('Codebases', build_func=__build_codebase_class) + + +def get_codebase_class(codebase: Codebase) -> type: + """Get the codebase class from the registry. + + Args: + codebase (Codebase): The codebase enum type. + + Returns: + type: The codebase class + """ + return CODEBASE.build(codebase) diff --git a/mmdeploy/codebase/base/task.py b/mmdeploy/codebase/base/task.py new file mode 100644 index 0000000000..373dfca653 --- /dev/null +++ b/mmdeploy/codebase/base/task.py @@ -0,0 +1,238 @@ +from abc import ABCMeta, abstractmethod +from typing import Any, Dict, Optional, Sequence, Tuple, Union + +import mmcv +import numpy as np +import torch +from torch.utils.data import DataLoader, Dataset + +from mmdeploy.utils import get_codebase + + +class BaseTask(metaclass=ABCMeta): + """Wrap the processing functions of a Computer Vision task. + + Args: + model_cfg (str | mmcv.Config): Model config file. + deploy_cfg (str | mmcv.Config): Deployment config file. + device (str): A string specifying device type. + """ + + def __init__(self, model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str): + + self.model_cfg = model_cfg + self.deploy_cfg = deploy_cfg + self.device = device + + codebase = get_codebase(deploy_cfg) + + from .mmcodebase import get_codebase_class + self.codebase_class = get_codebase_class(codebase) + + @abstractmethod + def init_backend_model(self, + model_files: Sequence[str] = None, + **kwargs) -> torch.nn.Module: + """Initialize backend model. + + Args: + model_files (Sequence[str]): Input model files. + + Returns: + nn.Module: An initialized backend model. + """ + pass + + @abstractmethod + def init_pytorch_model(self, + model_checkpoint: Optional[str] = None, + cfg_options: Optional[Dict] = None, + **kwargs) -> torch.nn.Module: + """Initialize torch model. + + Args: + model_checkpoint (str): The checkpoint file of torch model, + defaults to `None`. + cfg_options (dict): Optional config key-pair parameters. + + Returns: + nn.Module: An initialized torch model generated by other OpenMMLab + codebases. + """ + pass + + def build_dataset(self, + dataset_cfg: Union[str, mmcv.Config], + dataset_type: str = 'val', + **kwargs) -> Dataset: + """Build dataset for different codebase. + + Args: + dataset_cfg (str | mmcv.Config): Dataset config file or Config + object. + dataset_type (str): Specifying dataset type, e.g.: 'train', 'test', + 'val', defaults to 'val'. + + Returns: + Dataset: The built dataset. + """ + return self.codebase_class.build_dataset(dataset_cfg, dataset_type, + **kwargs) + + def build_dataloader(self, dataset: Dataset, samples_per_gpu: int, + workers_per_gpu: int, **kwargs) -> DataLoader: + """Build PyTorch dataloader. + + Args: + dataset (Dataset): A PyTorch dataset. + samples_per_gpu (int): Number of training samples on each GPU, + i.e., batch size of each GPU. + workers_per_gpu (int): How many subprocesses to use for data + loading for each GPU. + + Returns: + DataLoader: A PyTorch dataloader. + """ + return self.codebase_class.build_dataloader(dataset, samples_per_gpu, + workers_per_gpu, **kwargs) + + def single_gpu_test(self, + model: torch.nn.Module, + data_loader: DataLoader, + show: bool = False, + out_dir: Optional[str] = None, + **kwargs): + """Run test with single gpu. + + Args: + model (torch.nn.Module): Input model from nn.Module. + data_loader (DataLoader): PyTorch data loader. + show (bool): Specifying whether to show plotted results. Defaults + to `False`. + out_dir (str): A directory to save results, defaults to `None`. + + Returns: + list: The prediction results. + """ + return self.codebase_class.single_gpu_test(model, data_loader, show, + out_dir, **kwargs) + + @abstractmethod + def create_input(self, + imgs: Union[str, np.ndarray], + input_shape: Sequence[int] = None, + **kwargs) -> Tuple[Dict, torch.Tensor]: + """Create input for model. + + Args: + imgs (str | np.ndarray): Input image(s), accpeted data types are + `str`, `np.ndarray`. + input_shape (list[int]): Input shape of image in (width, height) + format, defaults to `None`. + + Returns: + tuple: (data, img), meta information for the input image and input + image tensor. + """ + pass + + @abstractmethod + def visualize(self, + model: torch.nn.Module, + image: Union[str, np.ndarray], + result: list, + output_file: str, + window_name: str = '', + show_result: bool = False, + **kwargs): + """Visualize predictions of a model. + + Args: + model (nn.Module): Input model. + image (str | np.ndarray): Input image to draw predictions on. + result (list): A list of predictions. + output_file (str): Output file to save drawn image. + backend (Backend): Specifying backend type. + window_name (str): The name of visualization window. Defaults to + an empty string. + show_result (bool): Whether to show result in windows, defaults + to `False`. + """ + pass + + @staticmethod + @abstractmethod + def run_inference(model, model_inputs: Dict[str, torch.Tensor]): + """Run inference once for a model of a OpenMMLab Codebase. + + Args: + model (nn.Module): Input model. + model_inputs (dict): A dict containing model inputs tensor and + meta info. + + Returns: + list: The predictions of model inference. + """ + pass + + @staticmethod + @abstractmethod + def get_partition_cfg(partition_type: str, **kwargs) -> Dict: + """Get a certain partition config. + + Args: + partition_type (str): A string specifying partition type. + + Returns: + dict: A dictionary of partition config. + """ + pass + + @staticmethod + @abstractmethod + def get_tensor_from_input(self, input_data: Dict[str, Any], + **kwargs) -> torch.Tensor: + """Get input tensor from input data. + + Args: + input_data (dict): Input data containing meta info and image + tensor. + Returns: + torch.Tensor: An image in `Tensor`. + """ + pass + + @staticmethod + @abstractmethod + def evaluate_outputs(model_cfg, + outputs: Sequence, + dataset: Dataset, + metrics: Optional[str] = None, + out: Optional[str] = None, + metric_options: Optional[dict] = None, + format_only: bool = False, + **kwargs): + """Perform post-processing to predictions of model. + + Args: + outputs (list): A list of predictions of model inference. + dataset (Dataset): Input dataset to run test. + model_cfg (mmcv.Config): The model config. + codebase (Codebase): Specifying codebase type. + metrics (str): Evaluation metrics, which depends on + the codebase and the dataset, e.g., "bbox", "segm", "proposal" + for COCO, and "mAP", "recall" for PASCAL VOC in mmdet; + "accuracy", "precision", "recall", "f1_score", "support" + for single label dataset, and "mAP", "CP", "CR", "CF1", + "OP", "OR", "OF1" for multi-label dataset in mmcls. + Defaults is `None`. + out (str): Output result file in pickle format, defaults to `None`. + metric_options (dict): Custom options for evaluation, will be + kwargs for dataset.evaluate() function. Defaults to `None`. + format_only (bool): Format the output results without perform + evaluation. It is useful when you want to format the result + to a specific format and submit it to the test server. Defaults + to `False`. + """ + pass diff --git a/mmdeploy/mmcls/__init__.py b/mmdeploy/codebase/mmcls/__init__.py similarity index 50% rename from mmdeploy/mmcls/__init__.py rename to mmdeploy/codebase/mmcls/__init__.py index d2b62e1cb6..33b69c74df 100644 --- a/mmdeploy/mmcls/__init__.py +++ b/mmdeploy/codebase/mmcls/__init__.py @@ -1,2 +1,2 @@ -from .export import * # noqa: F401,F403 +from .deploy import * # noqa: F401,F403 from .models import * # noqa: F401,F403 diff --git a/mmdeploy/codebase/mmcls/deploy/__init__.py b/mmdeploy/codebase/mmcls/deploy/__init__.py new file mode 100644 index 0000000000..42688e026e --- /dev/null +++ b/mmdeploy/codebase/mmcls/deploy/__init__.py @@ -0,0 +1,4 @@ +from .classification import Classification +from .mmclassification import MMClassification + +__all__ = ['MMClassification', 'Classification'] diff --git a/mmdeploy/codebase/mmcls/deploy/classification.py b/mmdeploy/codebase/mmcls/deploy/classification.py new file mode 100644 index 0000000000..0a3dba3a91 --- /dev/null +++ b/mmdeploy/codebase/mmcls/deploy/classification.py @@ -0,0 +1,234 @@ +import logging +from typing import Any, Dict, Optional, Sequence, Tuple, Union + +import mmcv +import numpy as np +import torch +from torch.utils.data import Dataset + +from mmdeploy.codebase.base import BaseTask +from mmdeploy.utils import Task +from .mmclassification import MMCLS_TASK + + +@MMCLS_TASK.register_module(Task.CLASSIFICATION.value) +class Classification(BaseTask): + """Classification task class. + + Args: + model_cfg (mmcv.Config): Original PyTorch model config file. + deploy_cfg (mmcv.Config): Deployment config file or loaded Config + object. + device (str): A string represents device type. + """ + + def __init__(self, model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str): + super(Classification, self).__init__(model_cfg, deploy_cfg, device) + + def init_backend_model(self, + model_files: Sequence[str] = None, + **kwargs) -> torch.nn.Module: + """Initialize backend model. + + Args: + model_files (Sequence[str]): Input model files. + + Returns: + nn.Module: An initialized backend model. + """ + from .classification_model import build_classification_model + + model = build_classification_model( + model_files, self.model_cfg, self.deploy_cfg, device=self.device) + + return model.eval() + + def init_pytorch_model(self, + model_checkpoint: Optional[str] = None, + cfg_options: Optional[Dict] = None, + **kwargs) -> torch.nn.Module: + """Initialize torch model. + + Args: + model_checkpoint (str): The checkpoint file of torch model, + Default: None. + cfg_options (dict): Optional config key-pair parameters. + + Returns: + nn.Module: An initialized torch model generated by OpenMMLab + codebases. + """ + from mmcls.apis import init_model + model = init_model(self.model_cfg, model_checkpoint, self.device, + cfg_options) + + return model.eval() + + def create_input(self, + imgs: Union[str, np.ndarray], + input_shape: Optional[Sequence[int]] = None) \ + -> Tuple[Dict, torch.Tensor]: + """Create input for classifier. + + Args: + imgs (Any): Input image(s), accepted data type are `str`, + `np.ndarray`, `torch.Tensor`. + input_shape (list[int]): A list of two integer in (width, height) + format specifying input shape. Default: None. + + Returns: + tuple: (data, img), meta information for the input image and input. + """ + import logging + + from mmcls.datasets.pipelines import Compose + from mmcv.parallel import collate, scatter + + cfg = self.model_cfg.copy() + if isinstance(imgs, str): + if cfg.data.test.pipeline[0]['type'] != 'LoadImageFromFile': + cfg.data.test.pipeline.insert(0, + dict(type='LoadImageFromFile')) + data = dict(img_info=dict(filename=imgs), img_prefix=None) + else: + if cfg.data.test.pipeline[0]['type'] == 'LoadImageFromFile': + cfg.data.test.pipeline.pop(0) + data = dict(img=imgs) + # check whether input_shape is valid + if input_shape is not None: + if 'crop_size' in cfg.data.test.pipeline[2]: + crop_size = cfg.data.test.pipeline[2]['crop_size'] + if tuple(input_shape) != (crop_size, crop_size): + logging.warning( + f'`input shape` should be equal to `crop_size`: {crop_size},\ + but given: {input_shape}') + test_pipeline = Compose(cfg.data.test.pipeline) + data = test_pipeline(data) + data = collate([data], samples_per_gpu=1) + data['img'] = [data['img']] + if self.device != 'cpu': + data = scatter(data, [self.device])[0] + return data, data['img'] + + def visualize(self, + model: torch.nn.Module, + image: Union[str, np.ndarray], + result: list, + output_file: str, + window_name: str = '', + show_result: bool = False): + """Visualize predictions of a model. + + Args: + model (nn.Module): Input model. + image (str | np.ndarray): Input image to draw predictions on. + result (list): A list of predictions. + output_file (str): Output file to save drawn image. + window_name (str): The name of visualization window. Defaults to + an empty string. + show_result (bool): Whether to show result in windows. + Default: False. + """ + show_img = mmcv.imread(image) if isinstance(image, str) else image + output_file = None if show_result else output_file + pred_score = np.max(result) + pred_label = np.argmax(result) + result = {'pred_label': pred_label, 'pred_score': float(pred_score)} + result['pred_class'] = model.CLASSES[result['pred_label']] + return model.show_result( + show_img, + result, + show=show_result, + win_name=window_name, + out_file=output_file) + + @staticmethod + def run_inference(model: torch.nn.Module, + model_inputs: Dict[str, torch.Tensor]) -> list: + """Run inference once for a classification model of mmcls. + + Args: + model (nn.Module): Input model. + model_inputs (dict): A dict containing model inputs tensor and + meta info. + + Returns: + list: The predictions of model inference. + """ + return model(**model_inputs, return_loss=False) + + @staticmethod + def get_partition_cfg(partition_type: str) -> Dict: + """Get a certain partition config. + + Args: + partition_type (str): A string specifying partition type. + + Returns: + dict: A dictionary of partition config. + """ + raise NotImplementedError('Not supported yet.') + + @staticmethod + def get_tensor_from_input(input_data: Dict[str, Any]) -> torch.Tensor: + """Get input tensor from input data. + + Args: + input_data (tuple): Input data containing meta info and image + tensor. + Returns: + torch.Tensor: An image in `Tensor`. + """ + return input_data['img'] + + @staticmethod + def evaluate_outputs(model_cfg: mmcv.Config, + outputs: list, + dataset: Dataset, + metrics: Optional[str] = None, + out: Optional[str] = None, + metric_options: Optional[dict] = None, + format_only: bool = False) -> None: + """Perform post-processing to predictions of model. + + Args: + model_cfg (mmcv.Config): The model config. + outputs (list): A list of predictions of model inference. + dataset (Dataset): Input dataset to run test. + metrics (str): Evaluation metrics, which depends on + the codebase and the dataset, e.g., "mAP" in mmcls. + out (str): Output result file in pickle format, Default: None. + metric_options (dict): Custom options for evaluation, will be + kwargs for dataset.evaluate() function. Default: None. + format_only (bool): Format the output results without perform + evaluation. It is useful when you want to format the result + to a specific format and submit it to the test server. + Default: False. + """ + import warnings + + if metrics: + results = dataset.evaluate(outputs, metrics, metric_options) + for k, v in results.items(): + logging.info(f'\n{k} : {v:.2f}') + else: + warnings.warn('Evaluation metrics are not specified.') + scores = np.vstack(outputs) + pred_score = np.max(scores, axis=1) + pred_label = np.argmax(scores, axis=1) + pred_class = [dataset.CLASSES[lb] for lb in pred_label] + results = { + 'pred_score': pred_score, + 'pred_label': pred_label, + 'pred_class': pred_class + } + if not out: + logging.info('\nthe predicted result for the first element is ' + f'pred_score = {pred_score[0]:.2f}, ' + f'pred_label = {pred_label[0]} ' + f'and pred_class = {pred_class[0]}. ' + 'Specify --out to save all results to files.') + if out: + logging.info(f'\nwriting results to {out}') + mmcv.dump(results, out) diff --git a/mmdeploy/codebase/mmcls/deploy/classification_model.py b/mmdeploy/codebase/mmcls/deploy/classification_model.py new file mode 100644 index 0000000000..062eac2f02 --- /dev/null +++ b/mmdeploy/codebase/mmcls/deploy/classification_model.py @@ -0,0 +1,160 @@ +from typing import List, Sequence, Union + +import mmcv +import numpy as np +import torch +from mmcls.datasets import DATASETS +from mmcls.models.classifiers.base import BaseClassifier + +from mmdeploy.codebase.base import BaseBackendModel +from mmdeploy.utils import Backend, get_backend, get_onnx_config, load_config + + +class End2EndModel(BaseBackendModel): + """End to end model for inference of classification. + + Args: + backend (Backend): The backend enum, specifying backend type. + backend_files (Sequence[str]): Paths to all required backend files(e.g. + '.onnx' for ONNX Runtime, '.param' and '.bin' for ncnn). + device (str): A string represents device type. + class_names (Sequence[str]): A list of string specifying class names. + deploy_cfg (str | mmcv.Config): Deployment config file or loaded Config + object. + """ + + def __init__( + self, + backend: Backend, + backend_files: Sequence[str], + device: str, + class_names: Sequence[str], + deploy_cfg: Union[str, mmcv.Config] = None, + ): + super(End2EndModel, self).__init__() + self.CLASSES = class_names + self.deploy_cfg = deploy_cfg + self._init_wrapper( + backend=backend, backend_files=backend_files, device=device) + + def _init_wrapper(self, backend: Backend, backend_files: Sequence[str], + device: str): + onnx_config = get_onnx_config(self.deploy_cfg) + output_names = onnx_config['output_names'] + self.wrapper = BaseBackendModel._build_wrapper( + backend=backend, + backend_files=backend_files, + device=device, + output_names=output_names) + + def forward(self, img: List[torch.Tensor], *args, **kwargs) -> list: + """Run forward inference. + + Args: + img (List[torch.Tensor]): A list contains input image(s) + in [N x C x H x W] format. + *args: Other arguments. + **kwargs: Other key-pair arguments. + + Returns: + list: A list contains predictions. + """ + + input_img = img[0].contiguous() + outputs = self.forward_test(input_img, *args, **kwargs) + + return list(outputs) + + def forward_test(self, imgs: torch.Tensor, *args, **kwargs) -> \ + List[np.ndarray]: + """The interface for forward test. + + Args: + imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. + + Returns: + List[np.ndarray]: A list of classification prediction. + """ + outputs = self.wrapper({'input': imgs}) + outputs = self.wrapper.output_to_list(outputs) + outputs = [out.detach().cpu().numpy() for out in outputs] + return outputs + + def show_result(self, + img: np.ndarray, + result: list, + win_name: str, + show: bool = True, + out_file: str = None): + """Show predictions of classification. + Args: + img: (np.ndarray): Input image to draw predictions. + result (list): A list of predictions. + win_name (str): The name of visualization window. + show (bool): Whether to show plotted image in windows. Defaults to + `True`. + out_file (str): Output image file to save drawn predictions. + + Returns: + np.ndarray: Drawn image, only if not `show` or `out_file`. + """ + return BaseClassifier.show_result( + self, img, result, show=show, win_name=win_name, out_file=out_file) + + +def get_classes_from_config(model_cfg: Union[str, mmcv.Config]): + """Get class name from config. + + Args: + model_cfg (str | mmcv.Config): Input model config file or + Config object. + + Returns: + list[str]: A list of string specifying names of different class. + """ + model_cfg = load_config(model_cfg)[0] + module_dict = DATASETS.module_dict + data_cfg = model_cfg.data + + if 'train' in data_cfg: + module = module_dict[data_cfg.train.type] + elif 'val' in data_cfg: + module = module_dict[data_cfg.val.type] + elif 'test' in data_cfg: + module = module_dict[data_cfg.test.type] + else: + raise RuntimeError(f'No dataset config found in: {model_cfg}') + + return module.CLASSES + + +def build_classification_model(model_files: Sequence[str], + model_cfg: Union[str, mmcv.Config], + deploy_cfg: Union[str, mmcv.Config], + device: str, **kwargs): + """Build classification model for different backend. + + Args: + model_files (Sequence[str]): Input model file(s). + model_cfg (str | mmcv.Config): Input model config file or Config + object. + deploy_cfg (str | mmcv.Config): Input deployment config file or + Config object. + device (str): Device to input model. + + Returns: + BaseBackendModel: Classifier for a configured backend. + """ + # load cfg if necessary + deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg) + + backend = get_backend(deploy_cfg) + class_names = get_classes_from_config(model_cfg) + backend_classifier = End2EndModel( + backend, + model_files, + device, + class_names, + deploy_cfg=deploy_cfg, + **kwargs) + return backend_classifier diff --git a/mmdeploy/codebase/mmcls/deploy/mmclassification.py b/mmdeploy/codebase/mmcls/deploy/mmclassification.py new file mode 100644 index 0000000000..8d55a01b12 --- /dev/null +++ b/mmdeploy/codebase/mmcls/deploy/mmclassification.py @@ -0,0 +1,140 @@ +from typing import List, Optional, Union + +import mmcv +import torch +from mmcv.utils import Registry +from torch.utils.data import DataLoader, Dataset + +from mmdeploy.codebase.base import CODEBASE, BaseTask, MMCodebase +from mmdeploy.utils import Codebase, get_task_type + + +def __build_mmcls_task(model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str, registry: Registry) -> BaseTask: + task = get_task_type(deploy_cfg) + return registry.module_dict[task.value](model_cfg, deploy_cfg, device) + + +MMCLS_TASK = Registry('mmcls_tasks', build_func=__build_mmcls_task) + + +@CODEBASE.register_module(Codebase.MMCLS.value) +class MMClassification(MMCodebase): + """mmclassification codebase class.""" + + task_registry = MMCLS_TASK + + def __init__(self): + super(MMClassification, self).__init__() + + @staticmethod + def build_task_processor(model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str) -> BaseTask: + """The interface to build the task processors of mmseg. + + Args: + model_cfg (mmcv.Config): Model config file. + deploy_cfg (mmcv.Config): Deployment config file. + device (str): A string specifying device type. + + Returns: + BaseTask: A task processor. + """ + return MMCLS_TASK.build(model_cfg, deploy_cfg, device) + + @staticmethod + def build_dataset(dataset_cfg: Union[str, mmcv.Config], + dataset_type: str = 'val', + **kwargs) -> Dataset: + """Build dataset for classification. + + Args: + dataset_cfg (str | mmcv.Config): The input dataset config. + dataset_type (str): A string represents dataset type, e.g.: 'train' + , 'test', 'val'. + Default: 'val'. + + Returns: + Dataset: A PyTorch dataset. + """ + + from mmcls.datasets import build_dataset as build_dataset_mmcls + + from mmdeploy.utils import load_config + + dataset_cfg = load_config(dataset_cfg)[0] + data = dataset_cfg.data + assert dataset_type in data + + dataset = build_dataset_mmcls(data[dataset_type]) + + return dataset + + def build_dataloader(dataset: Dataset, + samples_per_gpu: int, + workers_per_gpu: int, + num_gpus: int = 1, + dist: bool = False, + shuffle: bool = False, + round_up: bool = True, + seed: Optional[int] = None, + pin_memory: bool = True, + persistent_workers: bool = True, + **kwargs) -> DataLoader: + """Build dataloader for classifier. + + Args: + dataset (Dataset): Input dataset. + samples_per_gpu (int): Number of training samples on each GPU, + i.e., batch size of each GPU. + workers_per_gpu (int): How many subprocesses to use for data + loading for each GPU. + num_gpus (int): Number of GPUs. Only used in non-distributed + training. + dist (bool): Distributed training/test or not. Default: False. + shuffle (bool): Whether to shuffle the data at every epoch. + Default: False. + round_up (bool): Whether to round up the length of dataset by + adding extra samples to make it evenly divisible. + Default: True. + seed (int): An integer set to be seed. Default: None. + pin_memory (bool): Whether to use pin_memory in DataLoader. + Default: True. + persistent_workers (bool): If `True`, the data loader will not + shutdown the worker processes after a dataset has been + consumed once. This allows to maintain the workers Dataset + instances alive. The argument also has effect in + PyTorch>=1.7.0. Default: True. + kwargs: Any other keyword argument to be used to initialize + DataLoader. + + Returns: + DataLoader: A PyTorch dataloader. + """ + from mmcls.datasets import build_dataloader as build_dataloader_mmcls + return build_dataloader_mmcls(dataset, samples_per_gpu, + workers_per_gpu, num_gpus, dist, shuffle, + round_up, seed, pin_memory, + persistent_workers, **kwargs) + + @staticmethod + def single_gpu_test(model: torch.nn.Module, + data_loader: DataLoader, + show: bool = False, + out_dir: Optional[str] = None, + **kwargs) -> List: + """Run test with single gpu. + + Args: + model (torch.nn.Module): Input model from nn.Module. + data_loader (DataLoader): PyTorch data loader. + show (bool): Specifying whether to show plotted results. + Default: False. + out_dir (str): A directory to save results, Default: None. + + Returns: + list: The prediction results. + """ + from mmcls.apis import single_gpu_test + outputs = single_gpu_test(model, data_loader, show, out_dir, **kwargs) + return outputs diff --git a/mmdeploy/mmcls/models/__init__.py b/mmdeploy/codebase/mmcls/models/__init__.py similarity index 100% rename from mmdeploy/mmcls/models/__init__.py rename to mmdeploy/codebase/mmcls/models/__init__.py diff --git a/mmdeploy/mmcls/models/backbones/__init__.py b/mmdeploy/codebase/mmcls/models/backbones/__init__.py similarity index 100% rename from mmdeploy/mmcls/models/backbones/__init__.py rename to mmdeploy/codebase/mmcls/models/backbones/__init__.py diff --git a/mmdeploy/mmcls/models/backbones/shufflenet_v2.py b/mmdeploy/codebase/mmcls/models/backbones/shufflenet_v2.py similarity index 100% rename from mmdeploy/mmcls/models/backbones/shufflenet_v2.py rename to mmdeploy/codebase/mmcls/models/backbones/shufflenet_v2.py diff --git a/mmdeploy/mmcls/models/classifiers/__init__.py b/mmdeploy/codebase/mmcls/models/classifiers/__init__.py similarity index 100% rename from mmdeploy/mmcls/models/classifiers/__init__.py rename to mmdeploy/codebase/mmcls/models/classifiers/__init__.py diff --git a/mmdeploy/mmcls/models/classifiers/base.py b/mmdeploy/codebase/mmcls/models/classifiers/base.py similarity index 100% rename from mmdeploy/mmcls/models/classifiers/base.py rename to mmdeploy/codebase/mmcls/models/classifiers/base.py diff --git a/mmdeploy/mmcls/models/heads/__init__.py b/mmdeploy/codebase/mmcls/models/heads/__init__.py similarity index 100% rename from mmdeploy/mmcls/models/heads/__init__.py rename to mmdeploy/codebase/mmcls/models/heads/__init__.py diff --git a/mmdeploy/mmcls/models/heads/cls_head.py b/mmdeploy/codebase/mmcls/models/heads/cls_head.py similarity index 100% rename from mmdeploy/mmcls/models/heads/cls_head.py rename to mmdeploy/codebase/mmcls/models/heads/cls_head.py diff --git a/mmdeploy/mmcls/models/heads/multi_label_head.py b/mmdeploy/codebase/mmcls/models/heads/multi_label_head.py similarity index 100% rename from mmdeploy/mmcls/models/heads/multi_label_head.py rename to mmdeploy/codebase/mmcls/models/heads/multi_label_head.py diff --git a/mmdeploy/codebase/mmdet/__init__.py b/mmdeploy/codebase/mmdet/__init__.py new file mode 100644 index 0000000000..66d489e720 --- /dev/null +++ b/mmdeploy/codebase/mmdet/__init__.py @@ -0,0 +1,9 @@ +from .core import * # noqa: F401,F403 +from .deploy import (MMDetection, ObjectDetection, clip_bboxes, + get_post_processing_params, pad_with_value) +from .models import * # noqa: F401,F403 + +__all__ = [ + 'get_post_processing_params', 'clip_bboxes', 'pad_with_value', + 'MMDetection', 'ObjectDetection' +] diff --git a/mmdeploy/mmdet/core/__init__.py b/mmdeploy/codebase/mmdet/core/__init__.py similarity index 100% rename from mmdeploy/mmdet/core/__init__.py rename to mmdeploy/codebase/mmdet/core/__init__.py diff --git a/mmdeploy/codebase/mmdet/core/bbox/__init__.py b/mmdeploy/codebase/mmdet/core/bbox/__init__.py new file mode 100644 index 0000000000..1d9a90bde6 --- /dev/null +++ b/mmdeploy/codebase/mmdet/core/bbox/__init__.py @@ -0,0 +1,3 @@ +from .delta_xywh_bbox_coder import * # noqa: F401,F403 +from .tblr_bbox_coder import * # noqa: F401,F403 +from .transforms import * # noqa: F401,F403 diff --git a/mmdeploy/mmdet/core/bbox/coder/delta_xywh_bbox_coder.py b/mmdeploy/codebase/mmdet/core/bbox/delta_xywh_bbox_coder.py similarity index 94% rename from mmdeploy/mmdet/core/bbox/coder/delta_xywh_bbox_coder.py rename to mmdeploy/codebase/mmdet/core/bbox/delta_xywh_bbox_coder.py index 541d5a9a5c..0d989d6ad9 100644 --- a/mmdeploy/mmdet/core/bbox/coder/delta_xywh_bbox_coder.py +++ b/mmdeploy/codebase/mmdet/core/bbox/delta_xywh_bbox_coder.py @@ -95,7 +95,7 @@ def delta2bbox(ctx, y2 = gy + gh * 0.5 if clip_border and max_shape is not None: - from mmdeploy.mmdet.export import clip_bboxes + from mmdeploy.codebase.mmdet.deploy import clip_bboxes x1, y1, x2, y2 = clip_bboxes(x1, y1, x2, y2, max_shape) bboxes = torch.stack([x1, y1, x2, y2], dim=-1).view(deltas.size()) @@ -105,16 +105,16 @@ def delta2bbox(ctx, @FUNCTION_REWRITER.register_rewriter( func_name='mmdet.core.bbox.coder.delta_xywh_bbox_coder.delta2bbox', # noqa backend='ncnn') -def delta2bbox_ncnn(ctx, - rois, - deltas, - means=(0., 0., 0., 0.), - stds=(1., 1., 1., 1.), - max_shape=None, - wh_ratio_clip=16 / 1000, - clip_border=True, - add_ctr_clamp=False, - ctr_clamp=32): +def delta2bbox__ncnn(ctx, + rois, + deltas, + means=(0., 0., 0., 0.), + stds=(1., 1., 1., 1.), + max_shape=None, + wh_ratio_clip=16 / 1000, + clip_border=True, + add_ctr_clamp=False, + ctr_clamp=32): """Rewrite `delta2bbox` for ncnn backend. Batch dimension is not supported by ncnn, but supported by pytorch. @@ -216,7 +216,7 @@ def delta2bbox_ncnn(ctx, y2 = gy + gh * 0.5 if clip_border and max_shape is not None: - from mmdeploy.mmdet.export import clip_bboxes + from mmdeploy.codebase.mmdet.deploy import clip_bboxes x1, y1, x2, y2 = clip_bboxes(x1, y1, x2, y2, max_shape) bboxes = torch.stack([x1, y1, x2, y2], dim=-1).view(deltas.size()) diff --git a/mmdeploy/mmdet/core/bbox/coder/tblr_bbox_coder.py b/mmdeploy/codebase/mmdet/core/bbox/tblr_bbox_coder.py similarity index 93% rename from mmdeploy/mmdet/core/bbox/coder/tblr_bbox_coder.py rename to mmdeploy/codebase/mmdet/core/bbox/tblr_bbox_coder.py index 7c5ff5c71f..a568ce8ef2 100644 --- a/mmdeploy/mmdet/core/bbox/coder/tblr_bbox_coder.py +++ b/mmdeploy/codebase/mmdet/core/bbox/tblr_bbox_coder.py @@ -65,7 +65,7 @@ def tblr2bboxes(ctx, ymax = prior_centers[..., 1].unsqueeze(-1) + bottom if clip_border and max_shape is not None: - from mmdeploy.mmdet.export import clip_bboxes + from mmdeploy.codebase.mmdet.deploy import clip_bboxes xmin, ymin, xmax, ymax = clip_bboxes(xmin, ymin, xmax, ymax, max_shape) bboxes = torch.cat([xmin, ymin, xmax, ymax], dim=-1).view(priors.size()) @@ -75,13 +75,13 @@ def tblr2bboxes(ctx, @FUNCTION_REWRITER.register_rewriter( func_name='mmdet.core.bbox.coder.tblr_bbox_coder.tblr2bboxes', backend='ncnn') -def tblr2bboxes_ncnn(ctx, - priors, - tblr, - normalizer=4.0, - normalize_by_wh=True, - max_shape=None, - clip_border=True): +def tblr2bboxes__ncnn(ctx, + priors, + tblr, + normalizer=4.0, + normalize_by_wh=True, + max_shape=None, + clip_border=True): """Rewrite `tblr2bboxes` for ncnn backend. Batch dimension is not supported by ncnn, but supported by pytorch. @@ -137,7 +137,7 @@ def tblr2bboxes_ncnn(ctx, ymax = prior_centers[..., 1].unsqueeze(-1) + bottom if clip_border and max_shape is not None: - from mmdeploy.mmdet.export import clip_bboxes + from mmdeploy.codebase.mmdet.deploy import clip_bboxes xmin, ymin, xmax, ymax = clip_bboxes(xmin, ymin, xmax, ymax, max_shape) bboxes = torch.cat([xmin, ymin, xmax, ymax], dim=-1).view(priors.size()) diff --git a/mmdeploy/mmdet/core/bbox/transforms.py b/mmdeploy/codebase/mmdet/core/bbox/transforms.py similarity index 87% rename from mmdeploy/mmdet/core/bbox/transforms.py rename to mmdeploy/codebase/mmdet/core/bbox/transforms.py index a4713b31b0..7f0a0d016f 100644 --- a/mmdeploy/mmdet/core/bbox/transforms.py +++ b/mmdeploy/codebase/mmdet/core/bbox/transforms.py @@ -1,10 +1,12 @@ import torch -from mmdeploy.mmdet.export import clip_bboxes +from mmdeploy.codebase.mmdet.deploy import clip_bboxes def distance2bbox(points, distance, max_shape=None): - """Decode distance prediction to bounding box. + """Rewrite `mmdet.core.bbox.transforms.distance2bbox` + + Decode distance prediction to bounding box. Args: points (Tensor): Shape (B, N, 2) or (N, 2). diff --git a/mmdeploy/mmdet/core/post_processing/__init__.py b/mmdeploy/codebase/mmdet/core/post_processing/__init__.py similarity index 100% rename from mmdeploy/mmdet/core/post_processing/__init__.py rename to mmdeploy/codebase/mmdet/core/post_processing/__init__.py diff --git a/mmdeploy/mmdet/core/post_processing/bbox_nms.py b/mmdeploy/codebase/mmdet/core/post_processing/bbox_nms.py similarity index 97% rename from mmdeploy/mmdet/core/post_processing/bbox_nms.py rename to mmdeploy/codebase/mmdet/core/post_processing/bbox_nms.py index 8f4b7be375..70aa1afd92 100644 --- a/mmdeploy/mmdet/core/post_processing/bbox_nms.py +++ b/mmdeploy/codebase/mmdet/core/post_processing/bbox_nms.py @@ -129,7 +129,7 @@ def _multiclass_nms(boxes: Tensor, @FUNCTION_REWRITER.register_rewriter( - func_name='mmdeploy.mmdet.core.post_processing._multiclass_nms', + func_name='mmdeploy.codebase.mmdet.core.post_processing._multiclass_nms', backend='tensorrt') def multiclass_nms_static(ctx, boxes: Tensor, @@ -173,4 +173,5 @@ def multiclass_nms_static(ctx, @mark('multiclass_nms', inputs=['boxes', 'scores'], outputs=['dets', 'labels']) def multiclass_nms(*args, **kwargs): """Wrapper function for `_multiclass_nms`.""" - return mmdeploy.mmdet.core.post_processing._multiclass_nms(*args, **kwargs) + return mmdeploy.codebase.mmdet.core.post_processing._multiclass_nms( + *args, **kwargs) diff --git a/mmdeploy/codebase/mmdet/deploy/__init__.py b/mmdeploy/codebase/mmdet/deploy/__init__.py new file mode 100644 index 0000000000..c1563b7c3f --- /dev/null +++ b/mmdeploy/codebase/mmdet/deploy/__init__.py @@ -0,0 +1,8 @@ +from .mmdetection import MMDetection +from .object_detection import ObjectDetection +from .utils import clip_bboxes, get_post_processing_params, pad_with_value + +__all__ = [ + 'get_post_processing_params', 'clip_bboxes', 'pad_with_value', + 'MMDetection', 'ObjectDetection' +] diff --git a/mmdeploy/codebase/mmdet/deploy/mmdetection.py b/mmdeploy/codebase/mmdet/deploy/mmdetection.py new file mode 100644 index 0000000000..7d063dcc4e --- /dev/null +++ b/mmdeploy/codebase/mmdet/deploy/mmdetection.py @@ -0,0 +1,142 @@ +from typing import Optional, Union + +import mmcv +import torch +from mmcv.utils import Registry +from mmdet.datasets import replace_ImageToTensor +from torch.utils.data import DataLoader, Dataset + +from mmdeploy.utils import Codebase, get_task_type +from ...base import CODEBASE, BaseTask, MMCodebase + + +def __build_mmdet_task(model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str, registry: Registry) -> BaseTask: + task = get_task_type(deploy_cfg) + return registry.module_dict[task.value](model_cfg, deploy_cfg, device) + + +MMDET_TASK = Registry('mmdet_tasks', build_func=__build_mmdet_task) + + +@CODEBASE.register_module(Codebase.MMDET.value) +class MMDetection(MMCodebase): + + task_registry = MMDET_TASK + + def __init__(self) -> None: + super().__init__() + + @staticmethod + def build_task_processor(model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str): + """The interface to build the task processors of mmdet. + + Args: + model_cfg (str | mmcv.Config): Model config file. + deploy_cfg (str | mmcv.Config): Deployment config file. + device (str): A string specifying device type. + + Returns: + BaseTask: A task processor. + """ + return MMDET_TASK.build(model_cfg, deploy_cfg, device) + + @staticmethod + def build_dataset(dataset_cfg: Union[str, mmcv.Config], + dataset_type: str = 'val', + **kwargs) -> Dataset: + """Build dataset for detection. + + Args: + dataset_cfg (str | mmcv.Config): The input dataset config. + dataset_type (str): A string represents dataset type, e.g.: 'train' + , 'test', 'val'. Defaults to 'val'. + + Returns: + Dataset: A PyTorch dataset. + """ + from mmdet.datasets import build_dataset as build_dataset_mmdet + + assert dataset_type in dataset_cfg.data + data_cfg = dataset_cfg.data[dataset_type] + # in case the dataset is concatenated + if isinstance(data_cfg, dict): + data_cfg.test_mode = True + samples_per_gpu = data_cfg.get('samples_per_gpu', 1) + if samples_per_gpu > 1: + # Replace 'ImageToTensor' to 'DefaultFormatBundle' + data_cfg.pipeline = replace_ImageToTensor(data_cfg.pipeline) + elif isinstance(data_cfg, list): + for ds_cfg in data_cfg: + ds_cfg.test_mode = True + samples_per_gpu = max( + [ds_cfg.get('samples_per_gpu', 1) for ds_cfg in data_cfg]) + if samples_per_gpu > 1: + for ds_cfg in data_cfg: + ds_cfg.pipeline = replace_ImageToTensor(ds_cfg.pipeline) + dataset = build_dataset_mmdet(data_cfg) + + return dataset + + @staticmethod + def build_dataloader(dataset: Dataset, + samples_per_gpu: int, + workers_per_gpu: int, + num_gpus: int = 1, + dist: bool = False, + shuffle: bool = False, + seed: Optional[int] = None, + **kwargs) -> DataLoader: + """Build dataloader for detection. + + Args: + dataset (Dataset): Input dataset. + samples_per_gpu (int): Number of training samples on each GPU, i.e. + ,batch size of each GPU. + workers_per_gpu (int): How many subprocesses to use for data + loading for each GPU. + num_gpus (int): Number of GPUs. Only used in non-distributed + training. dist (bool): Distributed training/test or not. + Defaults to `False`.shuffle (bool): Whether to shuffle the + data at every epoch. + Defaults to `False`. + seed (int): An integer set to be seed. Default is `None`. + kwargs: Any other keyword argument to be used to initialize + DataLoader. + + Returns: + DataLoader: A PyTorch dataloader. + """ + from mmdet.datasets import build_dataloader as build_dataloader_mmdet + return build_dataloader_mmdet( + dataset, + samples_per_gpu, + workers_per_gpu, + num_gpus=num_gpus, + dist=dist, + shuffle=shuffle, + seed=seed, + **kwargs) + + @staticmethod + def single_gpu_test(model: torch.nn.Module, + data_loader: DataLoader, + show: bool = False, + out_dir: Optional[str] = None, + **kwargs): + """Run test with single gpu. + + Args: + model (torch.nn.Module): Input model from nn.Module. + data_loader (DataLoader): PyTorch data loader. + show (bool): Specifying whether to show plotted results. Defaults + to `False`. + out_dir (str): A directory to save results, defaults to `None`. + + Returns: + list: The prediction results. + """ + from mmdet.apis import single_gpu_test + outputs = single_gpu_test(model, data_loader, show, out_dir, **kwargs) + return outputs diff --git a/mmdeploy/mmdet/export/model_partition.py b/mmdeploy/codebase/mmdet/deploy/model_partition_cfg.py similarity index 80% rename from mmdeploy/mmdet/export/model_partition.py rename to mmdeploy/codebase/mmdet/deploy/model_partition_cfg.py index d53522d4cc..76b11d5bae 100644 --- a/mmdeploy/mmdet/export/model_partition.py +++ b/mmdeploy/codebase/mmdet/deploy/model_partition_cfg.py @@ -59,17 +59,3 @@ }, ) ]) - - -def get_partition_cfg(partition_type: str): - """Get a certain partition config for mmdet. - - Args: - partition_type (str): A string specifying partition type. - - Returns: - dict: A dictionary of partition config. - """ - assert (partition_type - in MMDET_PARTITION_CFG), f'Unknown partition_type {partition_type}' - return MMDET_PARTITION_CFG[partition_type] diff --git a/mmdeploy/codebase/mmdet/deploy/object_detection.py b/mmdeploy/codebase/mmdet/deploy/object_detection.py new file mode 100644 index 0000000000..e47e89e409 --- /dev/null +++ b/mmdeploy/codebase/mmdet/deploy/object_detection.py @@ -0,0 +1,235 @@ +import logging +from typing import Any, Dict, Optional, Sequence, Tuple, Union + +import mmcv +import numpy as np +import torch +from torch.utils.data import Dataset + +from mmdeploy.utils import Task +from ...base import BaseTask +from .mmdetection import MMDET_TASK + + +@MMDET_TASK.register_module(Task.OBJECT_DETECTION.value) +class ObjectDetection(BaseTask): + + def __init__(self, model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str) -> None: + super().__init__(model_cfg, deploy_cfg, device) + + def init_backend_model(self, + model_files: Optional[str] = None, + **kwargs) -> torch.nn.Module: + """Initialize backend model. + + Args: + model_files (Sequence[str]): Input model files. + + Returns: + nn.Module: An initialized backend model. + """ + from .object_detection_model import build_object_detection_model + model = build_object_detection_model( + model_files, self.model_cfg, self.deploy_cfg, device=self.device) + return model.eval() + + def init_pytorch_model(self, + model_checkpoint: Optional[str] = None, + cfg_options: Optional[Dict] = None, + **kwargs) -> torch.nn.Module: + """Initialize torch model. + + Args: + model_checkpoint (str): The checkpoint file of torch model, + defaults to `None`. + cfg_options (dict): Optional config key-pair parameters. + + Returns: + nn.Module: An initialized torch model generated by other OpenMMLab + codebases. + """ + from mmdet.apis import init_detector + model = init_detector(self.model_cfg, model_checkpoint, self.device, + cfg_options) + return model.eval() + + def create_input(self, + imgs: Union[str, np.ndarray], + input_shape: Sequence[int] = None) \ + -> Tuple[Dict, torch.Tensor]: + """Create input for detector. + + Args: + task (Task): Specifying task type. + imgs (Any): Input image(s), accpeted data type are `str`, + `np.ndarray`, `torch.Tensor`. + input_shape (list[int]): A list of two integer in (width, height) + format specifying input shape. Defaults to `None`. + + Returns: + tuple: (data, img), meta information for the input image and input. + """ + from mmdet.datasets import replace_ImageToTensor + from mmdet.datasets.pipelines import Compose + from mmcv.parallel import collate, scatter + + cfg = self.model_cfg.copy() + + if not isinstance(imgs, (list, tuple)): + imgs = [imgs] + + if isinstance(imgs[0], np.ndarray): + cfg = cfg.copy() + # set loading pipeline type + cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam' + # for static exporting + if input_shape is not None: + cfg.data.test.pipeline[1]['img_scale'] = tuple(input_shape) + transforms = cfg.data.test.pipeline[1]['transforms'] + for trans in transforms: + trans_type = trans['type'] + if trans_type == 'Resize': + trans['keep_ratio'] = False + elif trans_type == 'Pad': + trans['size_divisor'] = 1 + + cfg.data.test.pipeline = replace_ImageToTensor(cfg.data.test.pipeline) + test_pipeline = Compose(cfg.data.test.pipeline) + data_list = [] + for img in imgs: + # prepare data + if isinstance(img, np.ndarray): + # directly add img + data = dict(img=img) + else: + # add information into dict + data = dict(img_info=dict(filename=img), img_prefix=None) + # build the data pipeline + data = test_pipeline(data) + data_list.append(data) + + data = collate(data_list, samples_per_gpu=len(imgs)) + + data['img_metas'] = [ + img_metas.data[0] for img_metas in data['img_metas'] + ] + data['img'] = [img.data[0] for img in data['img']] + if self.device != 'cpu': + data = scatter(data, [self.device])[0] + + return data, data['img'] + + def visualize(self, + model, + image: Union[str, np.ndarray], + result: list, + output_file: str, + window_name: str, + show_result: bool = False, + score_thr=0.3): + """Visualize predictions of a model. + + Args: + model (nn.Module): Input model. + image (str | np.ndarray): Input image to draw predictions on. + result (list): A list of predictions. + output_file (str): Output file to save drawn image. + window_name (str): The name of visualization window. Defaults to + an empty string. + show_result (bool): Whether to show result in windows, defaults + to `False`. + score_thr (float): The score threshold to display the bbox. + Defaults to 0.3. + """ + show_img = mmcv.imread(image) if isinstance(image, str) else image + output_file = None if show_result else output_file + model.show_result( + show_img, + result=result, + win_name=window_name, + show=show_result, + out_file=output_file, + score_thr=score_thr) + + @staticmethod + def run_inference(model, model_inputs: Dict[str, torch.Tensor]): + """Run inference once for a object detection model of mmdet. + + Args: + model (nn.Module): Input model. + model_inputs (dict): A dict containing model inputs tensor and + meta info. + + Returns: + list: The predictions of model inference. + """ + return model(**model_inputs, return_loss=False, rescale=True) + + @staticmethod + def get_partition_cfg(partition_type: str) -> Dict: + """Get a certain partition config for mmdet. + + Args: + partition_type (str): A string specifying partition type. + + Returns: + dict: A dictionary of partition config. + """ + from .model_partition_cfg import MMDET_PARTITION_CFG + assert (partition_type in MMDET_PARTITION_CFG), \ + f'Unknown partition_type {partition_type}' + return MMDET_PARTITION_CFG[partition_type] + + @staticmethod + def get_tensor_from_input(input_data: Dict[str, Any]) -> torch.Tensor: + """Get input tensor from input data. + + Args: + input_data (dict): Input data containing meta info and image + tensor. + Returns: + torch.Tensor: An image in `Tensor`. + """ + return input_data['img'][0] + + @staticmethod + def evaluate_outputs(model_cfg, + outputs: Sequence, + dataset: Dataset, + metrics: Optional[str] = None, + out: Optional[str] = None, + metric_options: Optional[dict] = None, + format_only: bool = False): + """Perform post-processing to predictions of model. + + Args: + outputs (list): A list of predictions of model inference. + dataset (Dataset): Input dataset to run test. + metrics (str): Evaluation metrics, which depends on + the codebase and the dataset, e.g., "bbox", "segm", "proposal" + for COCO, and "mAP", "recall" for PASCAL VOC in mmdet. + out (str): Output result file in pickle format, defaults to `None`. + metric_options (dict): Custom options for evaluation, will be + kwargs for dataset.evaluate() function. Defaults to `None`. + format_only (bool): Format the output results without perform + evaluation. It is useful when you want to format the result + to a specific format and submit it to the test server. Defaults + to `False`. + """ + if out: + logging.info(f'\nwriting results to {out}') + mmcv.dump(outputs, out) + kwargs = {} if metric_options is None else metric_options + if format_only: + dataset.format_results(outputs, **kwargs) + if metrics: + eval_kwargs = model_cfg.get('evaluation', {}).copy() + # hard-code way to remove EvalHook args + for key in [ + 'interval', 'tmpdir', 'start', 'gpu_collect', 'save_best', + 'rule' + ]: + eval_kwargs.pop(key, None) + eval_kwargs.update(dict(metric=metrics, **kwargs)) + print(dataset.evaluate(outputs, **eval_kwargs)) diff --git a/mmdeploy/codebase/mmdet/deploy/object_detection_model.py b/mmdeploy/codebase/mmdet/deploy/object_detection_model.py new file mode 100644 index 0000000000..236d4a104f --- /dev/null +++ b/mmdeploy/codebase/mmdet/deploy/object_detection_model.py @@ -0,0 +1,579 @@ +from functools import partial +from typing import List, Sequence, Tuple, Union + +import mmcv +import numpy as np +import torch +import torch.nn.functional as F +from mmcv.utils import Registry +from mmdet.core import bbox2result +from mmdet.datasets import DATASETS +from mmdet.models import BaseDetector + +from mmdeploy.backend.base import get_backend_file_count +from mmdeploy.codebase.base import BaseBackendModel +from mmdeploy.codebase.mmdet import get_post_processing_params, multiclass_nms +from mmdeploy.utils import (Backend, get_backend, get_onnx_config, + get_partition_config, load_config) + + +def __build_backend_model(partition_name: str, backend: Backend, + backend_files: Sequence[str], device: str, + class_names: Sequence[str], + model_cfg: Union[str, mmcv.Config], + deploy_cfg: Union[str, mmcv.Config], + registry: Registry, **kwargs): + return registry.module_dict[partition_name]( + backend=backend, + backend_files=backend_files, + class_names=class_names, + device=device, + model_cfg=model_cfg, + deploy_cfg=deploy_cfg, + **kwargs) + + +# Use registry to store models with different partition methods +# If a model doesn't need to partition, we don't need this registry +__BACKEND_MODEl = mmcv.utils.Registry( + 'backend_detectors', build_func=__build_backend_model) + + +@__BACKEND_MODEl.register_module('end2end') +class End2EndModel(BaseBackendModel): + """End to end model for inference of detection. + + TODO: UPDATE this docstring + Args: + class_names (Sequence[str]): A list of string specifying class names. + device_id (int): An integer represents device index. + """ + + def __init__(self, backend: Backend, backend_files: Sequence[str], + device: str, class_names: Sequence[str], + deploy_cfg: Union[str, mmcv.Config], **kwargs): + super().__init__() + self.CLASSES = class_names + self.deploy_cfg = deploy_cfg + self._init_wrapper( + backend=backend, backend_files=backend_files, device=device) + + def _init_wrapper(self, backend, backend_files, device): + onnx_config = get_onnx_config(self.deploy_cfg) + output_names = onnx_config['output_names'] + self.wrapper = BaseBackendModel._build_wrapper( + backend=backend, + backend_files=backend_files, + device=device, + output_names=output_names) + + @staticmethod + def __clear_outputs( + test_outputs: List[Union[torch.Tensor, np.ndarray]] + ) -> List[Union[List[torch.Tensor], List[np.ndarray]]]: + """Removes additional outputs and detections with zero and negative + score. + + Args: + test_outputs (List[Union[torch.Tensor, np.ndarray]]): + outputs of forward_test. + + Returns: + List[Union[List[torch.Tensor], List[np.ndarray]]]: + outputs with without zero score object. + """ + batch_size = len(test_outputs[0]) + + num_outputs = len(test_outputs) + outputs = [[None for _ in range(batch_size)] + for _ in range(num_outputs)] + + for i in range(batch_size): + inds = test_outputs[0][i, :, 4] > 0.0 + for output_id in range(num_outputs): + outputs[output_id][i] = test_outputs[output_id][i, inds, ...] + return outputs + + @staticmethod + def __postprocessing_masks(det_bboxes: np.ndarray, + det_masks: np.ndarray, + img_w: int, + img_h: int, + mask_thr_binary: float = 0.5) -> np.ndarray: + """Additional processing of masks. Resizes masks from [num_det, 28, 28] + to [num_det, img_w, img_h]. Analog of the 'mmdeploy.codebase.mmdet. + models.roi_heads.fcn_mask_head._do_paste_mask' function. + + Args: + det_bboxes (np.ndarray): Bbox of shape [num_det, 5] + det_masks (np.ndarray): Masks of shape [num_det, 28, 28]. + img_w (int): Width of the original image. + img_h (int): Height of the original image. + mask_thr_binary (float): The threshold for the mask. + + Returns: + np.ndarray: masks of shape [N, num_det, img_w, img_h]. + """ + masks = det_masks + bboxes = det_bboxes + + if isinstance(masks, np.ndarray): + masks = torch.tensor(masks) + bboxes = torch.tensor(bboxes) + + result_masks = [] + for bbox, mask in zip(bboxes, masks): + + x0_int, y0_int = 0, 0 + x1_int, y1_int = img_w, img_h + + img_y = torch.arange(y0_int, y1_int, dtype=torch.float32) + 0.5 + img_x = torch.arange(x0_int, x1_int, dtype=torch.float32) + 0.5 + x0, y0, x1, y1 = bbox + + img_y = (img_y - y0) / (y1 - y0) * 2 - 1 + img_x = (img_x - x0) / (x1 - x0) * 2 - 1 + if torch.isinf(img_x).any(): + inds = torch.where(torch.isinf(img_x)) + img_x[inds] = 0 + if torch.isinf(img_y).any(): + inds = torch.where(torch.isinf(img_y)) + img_y[inds] = 0 + + gx = img_x[None, :].expand(img_y.size(0), img_x.size(0)) + gy = img_y[:, None].expand(img_y.size(0), img_x.size(0)) + grid = torch.stack([gx, gy], dim=2) + + img_masks = F.grid_sample( + mask.to(dtype=torch.float32)[None, None, :, :], + grid[None, :, :, :], + align_corners=False) + + mask = img_masks + mask = (mask >= mask_thr_binary).to(dtype=torch.bool) + result_masks.append(mask.numpy()) + result_masks = np.concatenate(result_masks, axis=1) + return result_masks.squeeze(0) + + def forward(self, img: Sequence[torch.Tensor], img_metas: Sequence[dict], + *args, **kwargs): + """Run forward inference. + + Args: + img (Sequence[torch.Tensor]): A list contains input image(s) + in [N x C x H x W] format. + img_metas (Sequence[dict]): A list of meta info for image(s). + *args: Other arguments. + **kwargs: Other key-pair arguments. + + Returns: + list: A list contains predictions. + """ + input_img = img[0].contiguous() + outputs = self.forward_test(input_img, img_metas, *args, **kwargs) + outputs = End2EndModel.__clear_outputs(outputs) + batch_dets, batch_labels = outputs[:2] + batch_masks = outputs[2] if len(outputs) == 3 else None + batch_size = input_img.shape[0] + img_metas = img_metas[0] + results = [] + rescale = kwargs.get('rescale', True) + for i in range(batch_size): + dets, labels = batch_dets[i], batch_labels[i] + if rescale: + scale_factor = img_metas[i]['scale_factor'] + + if isinstance(scale_factor, (list, tuple, np.ndarray)): + assert len(scale_factor) == 4 + scale_factor = np.array(scale_factor)[None, :] # [1,4] + dets[:, :4] /= scale_factor + + if 'border' in img_metas[i]: + # offset pixel of the top-left corners between original image + # and padded/enlarged image, 'border' is used when exporting + # CornerNet and CentripetalNet to onnx + x_off = img_metas[i]['border'][2] + y_off = img_metas[i]['border'][0] + dets[:, [0, 2]] -= x_off + dets[:, [1, 3]] -= y_off + dets[:, :4] *= (dets[:, :4] > 0).astype(dets.dtype) + + dets_results = bbox2result(dets, labels, len(self.CLASSES)) + + if batch_masks is not None: + masks = batch_masks[i] + img_h, img_w = img_metas[i]['img_shape'][:2] + ori_h, ori_w = img_metas[i]['ori_shape'][:2] + export_postprocess_mask = True + if self.deploy_cfg is not None: + + mmdet_deploy_cfg = get_post_processing_params( + self.deploy_cfg) + # this flag enable postprocess when export. + export_postprocess_mask = mmdet_deploy_cfg.get( + 'export_postprocess_mask', True) + if not export_postprocess_mask: + masks = End2EndModel.__postprocessing_masks( + dets[:, :4], masks, ori_w, ori_h) + else: + masks = masks[:, :img_h, :img_w] + # avoid to resize masks with zero dim + if rescale and masks.shape[0] != 0: + masks = masks.astype(np.float32) + masks = torch.from_numpy(masks) + masks = torch.nn.functional.interpolate( + masks.unsqueeze(0), size=(ori_h, ori_w)) + masks = masks.squeeze(0).detach().numpy() + if masks.dtype != np.bool: + masks = masks >= 0.5 + segms_results = [[] for _ in range(len(self.CLASSES))] + for j in range(len(dets)): + segms_results[labels[j]].append(masks[j]) + results.append((dets_results, segms_results)) + else: + results.append(dets_results) + return results + + def forward_test(self, imgs: torch.Tensor, *args, **kwargs) -> \ + Tuple[np.ndarray, np.ndarray]: + """The interface for forward test. + + Args: + imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. + + Returns: + tuple[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] + and class labels of shape [N, num_det]. + """ + outputs = self.wrapper({'input': imgs}) + outputs = self.wrapper.output_to_list(outputs) + outputs = [out.detach().cpu().numpy() for out in outputs] + return outputs + + def show_result(self, + img: np.ndarray, + result: list, + win_name: str, + show: bool = True, + score_thr: float = 0.3, + out_file=None): + return BaseDetector.show_result( + self, + img=img, + result=result, + score_thr=score_thr, + show=show, + win_name=win_name, + out_file=out_file) + + +@__BACKEND_MODEl.register_module('single_stage') +class PartitionSingleStageModel(End2EndModel): + """Partitioned single stage detection model. + + Args: + model_file (str): The path of input model file. + class_names (Sequence[str]): A list of string specifying class names. + model_cfg: (str | mmcv.Config): Input model config. + deploy_cfg: (str | mmcv.Config): Input deployment config. + device_id (int): An integer represents device index. + """ + + def __init__(self, backend: Backend, backend_files: Sequence[str], + device: str, class_names: Sequence[str], + model_cfg: Union[str, mmcv.Config], + deploy_cfg: Union[str, mmcv.Config], **kwargs): + super().__init__(backend, backend_files, device, class_names, + deploy_cfg, **kwargs) + # load cfg if necessary + model_cfg = load_config(model_cfg)[0] + self.model_cfg = model_cfg + + def _init_wrapper(self, backend, backend_files, device): + self.wrapper = BaseBackendModel._build_wrapper( + backend=backend, + backend_files=backend_files, + device=device, + output_names=['scores', 'boxes']) + + def partition0_postprocess(self, scores: torch.Tensor, + bboxes: torch.Tensor): + """Perform post-processing for partition 0. + + Args: + scores (Tensor): The detection scores of shape + [N, num_boxes, num_classes]. + bboxes (Tensor): The bounding boxes of shape [N, num_boxes, 4]. + + Returns: + tuple[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] and + class labels of shape [N, num_det]. + """ + cfg = self.model_cfg.model.test_cfg + deploy_cfg = self.deploy_cfg + + post_params = get_post_processing_params(deploy_cfg) + max_output_boxes_per_class = post_params.max_output_boxes_per_class + iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) + score_threshold = cfg.get('score_thr', post_params.score_threshold) + pre_top_k = -1 if post_params.pre_top_k >= bboxes.shape[1] \ + else post_params.pre_top_k + keep_top_k = cfg.get('max_per_img', post_params.keep_top_k) + ret = multiclass_nms( + bboxes, + scores, + max_output_boxes_per_class, + iou_threshold=iou_threshold, + score_threshold=score_threshold, + pre_top_k=pre_top_k, + keep_top_k=keep_top_k) + ret = [r.cpu() for r in ret] + return ret + + def forward_test(self, imgs: torch.Tensor, *args, **kwargs): + """Implement forward test. + + Args: + imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. + + Returns: + list[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] and + class labels of shape [N, num_det]. + """ + outputs = self.wrapper({'input': imgs}) + outputs = self.wrapper.output_to_list(outputs) + scores, bboxes = outputs[:2] + return self.partition0_postprocess(scores, bboxes) + + +@__BACKEND_MODEl.register_module('two_stage') +class PartitionTwoStageModel(End2EndModel): + """Partitioned two stage detection model. + + Args: + class_names (Sequence[str]): A list of string specifying class names. + model_cfg: (str | mmcv.Config): Input model config. + deploy_cfg: (str | mmcv.Config): Input deployment config. + device_id (int): An integer represents device index. + """ + + def __init__(self, backend: Backend, backend_files: Sequence[str], + device: str, class_names: Sequence[str], + model_cfg: Union[str, mmcv.Config], + deploy_cfg: Union[str, mmcv.Config], **kwargs): + + # load cfg if necessary + model_cfg = load_config(model_cfg)[0] + + self.model_cfg = model_cfg + + super().__init__(backend, backend_files, device, class_names, + deploy_cfg, **kwargs) + from mmdet.models.builder import build_head, build_roi_extractor + from ..models.roi_heads.bbox_head import bbox_head__get_bboxes + + self.bbox_roi_extractor = build_roi_extractor( + model_cfg.model.roi_head.bbox_roi_extractor) + self.bbox_head = build_head(model_cfg.model.roi_head.bbox_head) + + class Context: + pass + + ctx = Context() + ctx.cfg = self.deploy_cfg + self.bbox_head__get_bboxes = partial(bbox_head__get_bboxes, ctx) + + def _init_wrapper(self, backend, backend_files, device): + n = get_backend_file_count(backend) + num_feat = self.model_cfg['model']['neck']['num_outs'] + partition0_output_names = [ + 'feat/{}'.format(i) for i in range(num_feat) + ] + ['scores', 'boxes'] + + self.first_wrapper = BaseBackendModel._build_wrapper( + backend, backend_files[0:n], device, partition0_output_names) + + self.second_wrapper = BaseBackendModel._build_wrapper( + backend, backend_files[n:2 * n], device, + ['cls_score', 'bbox_pred']) + + def partition0_postprocess(self, x: Sequence[torch.Tensor], + scores: torch.Tensor, bboxes: torch.Tensor): + """Perform post-processing for partition 0. + + Args: + x (tuple[Tensor]): Feature maps of all scale levels. + scores (Tensor): The detection scores of shape + [N, num_boxes, num_classes]. + bboxes (Tensor): The bounding boxes of shape [N, num_boxes, 4]. + + Returns: + tuple(Tensor, Tensor): rois and bbox_feats. + """ + # rpn-nms + roi-extractor + cfg = self.model_cfg.model.test_cfg.rpn + deploy_cfg = self.deploy_cfg + + post_params = get_post_processing_params(deploy_cfg) + iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) + score_threshold = cfg.get('score_thr', post_params.score_threshold) + pre_top_k = -1 if post_params.pre_top_k >= bboxes.shape[1] \ + else post_params.pre_top_k + keep_top_k = cfg.get('max_per_img', post_params.keep_top_k) + # only one class in rpn + max_output_boxes_per_class = keep_top_k + proposals, _ = multiclass_nms( + bboxes, + scores, + max_output_boxes_per_class, + iou_threshold=iou_threshold, + score_threshold=score_threshold, + pre_top_k=pre_top_k, + keep_top_k=keep_top_k) + + rois = proposals + batch_index = torch.arange( + rois.shape[0], device=rois.device).float().view(-1, 1, 1).expand( + rois.size(0), rois.size(1), 1) + rois = torch.cat([batch_index, rois[..., :4]], dim=-1) + batch_size = rois.shape[0] + num_proposals_per_img = rois.shape[1] + + # Eliminate the batch dimension + rois = rois.view(-1, 5) + bbox_feats = self.bbox_roi_extractor( + x[:self.bbox_roi_extractor.num_inputs], rois) + + rois = rois.reshape(batch_size, num_proposals_per_img, rois.size(-1)) + return rois, bbox_feats + + def partition1_postprocess(self, rois: torch.Tensor, + cls_score: torch.Tensor, + bbox_pred: torch.Tensor, + img_metas: Sequence[dict]): + """Perform post-processing for partition 1. + Args: + rois (torch.Tensor): Input tensor of roi. + cls_score (torch.Tensor): Scores of all classes. + bbox_pred (torch.Tensor): Bounding box proposals. + img_metas (Sequence[dict]): A list of image(s) meta information. + + Returns: + tuple[Tensor, Tensor]: dets of shape [N, num_det, 5] and class + labels of shape [N, num_det]. + """ + batch_size = rois.shape[0] + num_proposals_per_img = rois.shape[1] + + cls_score = cls_score.reshape(batch_size, num_proposals_per_img, + cls_score.size(-1)) + + bbox_pred = bbox_pred.reshape(batch_size, num_proposals_per_img, + bbox_pred.size(-1)) + + rcnn_test_cfg = self.model_cfg.model.test_cfg.rcnn + return self.bbox_head__get_bboxes(self.bbox_head, rois, cls_score, + bbox_pred, + img_metas[0][0]['img_shape'], + rcnn_test_cfg) + + def forward_test(self, imgs: torch.Tensor, img_metas: Sequence[dict], + *args, **kwargs): + """Implement forward test. + + Args: + imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. + img_metas (Sequence[dict]): A list of image(s) meta information. + + Returns: + tuple[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] and + class labels of shape [N, num_det]. + """ + outputs = self.first_wrapper({'input': imgs}) + outputs = self.first_wrapper.output_to_list(outputs) + feats = outputs[:-2] + scores, bboxes = outputs[-2:] + + # partition0_postprocess + rois, bbox_feats = self.partition0_postprocess(feats, scores, bboxes) + + # partition1 forward + bbox_feats = bbox_feats.contiguous() + outputs = self.second_wrapper({'bbox_feats': bbox_feats}) + outputs = self.second_wrapper.output_to_list(outputs) + cls_score, bbox_pred = outputs[:2] + + # partition1_postprocess + outputs = self.partition1_postprocess(rois, cls_score, bbox_pred, + img_metas) + outputs = [out.detach().cpu() for out in outputs] + return outputs + + +def get_classes_from_config(model_cfg: Union[str, mmcv.Config], **kwargs): + """Get class name from config. + + Args: + model_cfg (str | mmcv.Config): Input model config file or + Config object. + + Returns: + list[str]: A list of string specifying names of different class. + """ + # load cfg if necessary + model_cfg = load_config(model_cfg)[0] + module_dict = DATASETS.module_dict + data_cfg = model_cfg.data + + if 'test' in data_cfg: + module = module_dict[data_cfg.test.type] + elif 'val' in data_cfg: + module = module_dict[data_cfg.val.type] + elif 'train' in data_cfg: + module = module_dict[data_cfg.train.type] + else: + raise RuntimeError(f'No dataset config found in: {model_cfg}') + + return module.CLASSES + + +def build_object_detection_model(model_files: Sequence[str], + model_cfg: Union[str, mmcv.Config], + deploy_cfg: Union[str, mmcv.Config], + device: str, **kwargs): + """Build object detection model for different backends. + + Args: + model_files (Sequence[str]): Input model file(s). + model_cfg (str | mmcv.Config): Input model config file or Config + object. + deploy_cfg (str | mmcv.Config): Input deployment config file or + Config object. + device (str): Device to input model + + Returns: + DeployBaseDetector: Detector for a configured backend. + """ + # load cfg if necessary + deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg) + + backend = get_backend(deploy_cfg) + class_names = get_classes_from_config(model_cfg) + + # Default Config is 'end2end' + partition_type = 'end2end' + partition_config = get_partition_config(deploy_cfg) + if partition_config is not None: + partition_type = partition_config.get('type', None) + + backend_detector = __BACKEND_MODEl.build( + partition_type, + backend=backend, + backend_files=model_files, + class_names=class_names, + device=device, + model_cfg=model_cfg, + deploy_cfg=deploy_cfg, + **kwargs) + + return backend_detector diff --git a/mmdeploy/codebase/mmdet/deploy/utils.py b/mmdeploy/codebase/mmdet/deploy/utils.py new file mode 100644 index 0000000000..d8981a3cf6 --- /dev/null +++ b/mmdeploy/codebase/mmdet/deploy/utils.py @@ -0,0 +1,99 @@ +from typing import Any, Optional, Sequence, Union + +import mmcv +import torch +from torch import Tensor + +from mmdeploy.utils import load_config + + +def get_post_processing_params(deploy_cfg: Union[str, mmcv.Config]): + """Get mmdet post-processing parameters from config. + + Args: + deploy_cfg (str | mmcv.Config): The path or content of config. + + Returns: + dict: A dict of parameters for mmdet. + """ + deploy_cfg = load_config(deploy_cfg)[0] + codebase_key = 'codebase_config' + assert codebase_key in deploy_cfg + codebase_config = deploy_cfg[codebase_key] + post_params = codebase_config.get('post_processing', None) + assert post_params is not None, 'Failed to get `post_processing`.' + return post_params + + +def clip_bboxes(x1: Tensor, y1: Tensor, x2: Tensor, y2: Tensor, + max_shape: Union[Tensor, Sequence[int]]): + """Clip bboxes for onnx. + + Since torch.clamp cannot have dynamic `min` and `max`, we scale the + boxes by 1/max_shape and clamp in the range [0, 1] if necessary. + + Args: + x1 (Tensor): The x1 for bounding boxes. + y1 (Tensor): The y1 for bounding boxes. + x2 (Tensor): The x2 for bounding boxes. + y2 (Tensor): The y2 for bounding boxes. + max_shape (Tensor | Sequence[int]): The (H,W) of original image. + Returns: + tuple(Tensor): The clipped x1, y1, x2, y2. + """ + assert len(max_shape) == 2, '`max_shape` should be [h, w]' + if isinstance(max_shape, torch.Tensor): + # scale by 1/max_shape + x1 = x1 / max_shape[1] + y1 = y1 / max_shape[0] + x2 = x2 / max_shape[1] + y2 = y2 / max_shape[0] + + # clamp [0, 1] + x1 = torch.clamp(x1, 0, 1) + y1 = torch.clamp(y1, 0, 1) + x2 = torch.clamp(x2, 0, 1) + y2 = torch.clamp(y2, 0, 1) + + # scale back + x1 = x1 * max_shape[1] + y1 = y1 * max_shape[0] + x2 = x2 * max_shape[1] + y2 = y2 * max_shape[0] + else: + x1 = torch.clamp(x1, 0, max_shape[1]) + y1 = torch.clamp(y1, 0, max_shape[0]) + x2 = torch.clamp(x2, 0, max_shape[1]) + y2 = torch.clamp(y2, 0, max_shape[0]) + return x1, y1, x2, y2 + + +def pad_with_value(x: Tensor, + pad_dim: int, + pad_size: int, + pad_value: Optional[Any] = None): + """Pad a tensor with a value along some dim. + + Args: + x (Tensor): Input tensor. + pad_dim (int): Along which dim to pad. + pad_size (int): To which size to pad. + pad_value (Any): Filled value for padding. Defaults to `None`. + + Returns: + Tensor: Padded tensor. + """ + num_dims = len(x.shape) + pad_slice = (slice(None, None, None), ) * num_dims + pad_slice = pad_slice[:pad_dim] + (slice(0, 1, + 1), ) + pad_slice[pad_dim + 1:] + repeat_size = [1] * num_dims + repeat_size[pad_dim] = pad_size + + x_pad = x.__getitem__(pad_slice) + if pad_value is not None: + x_pad = x_pad * 0 + pad_value + + x_pad = x_pad.repeat(*repeat_size) + x = torch.cat([x, x_pad], dim=pad_dim) + return x diff --git a/mmdeploy/mmdet/models/__init__.py b/mmdeploy/codebase/mmdet/models/__init__.py similarity index 100% rename from mmdeploy/mmdet/models/__init__.py rename to mmdeploy/codebase/mmdet/models/__init__.py diff --git a/mmdeploy/codebase/mmdet/models/dense_heads/__init__.py b/mmdeploy/codebase/mmdet/models/dense_heads/__init__.py new file mode 100644 index 0000000000..eb6c5937ee --- /dev/null +++ b/mmdeploy/codebase/mmdet/models/dense_heads/__init__.py @@ -0,0 +1,17 @@ +from .anchor_head import anchor_head__get_bboxes, anchor_head__get_bboxes__ncnn +from .atss_head import atss_head__get_bboxes +from .fcos_head import fcos_head__get_bboxes, fcos_head__get_bboxes__ncnn +from .fovea_head import fovea_head__get_bboxes +from .rpn_head import rpn_head__get_bboxes +from .vfnet_head import vfnet_head__get_bboxes +from .yolo_head import yolov3_head__get_bboxes, yolov3_head__get_bboxes__ncnn +from .yolox_head import yolox_head__get_bboxes + +__all__ = [ + 'anchor_head__get_bboxes', 'anchor_head__get_bboxes__ncnn', + 'atss_head__get_bboxes', 'fcos_head__get_bboxes', + 'fcos_head__get_bboxes__ncnn', 'fovea_head__get_bboxes', + 'rpn_head__get_bboxes', 'vfnet_head__get_bboxes', + 'yolov3_head__get_bboxes', 'yolov3_head__get_bboxes__ncnn', + 'yolox_head__get_bboxes' +] diff --git a/mmdeploy/mmdet/models/dense_heads/anchor_head.py b/mmdeploy/codebase/mmdet/models/dense_heads/anchor_head.py similarity index 95% rename from mmdeploy/mmdet/models/dense_heads/anchor_head.py rename to mmdeploy/codebase/mmdet/models/dense_heads/anchor_head.py index fdd33599d2..e394d4b50d 100644 --- a/mmdeploy/mmdet/models/dense_heads/anchor_head.py +++ b/mmdeploy/codebase/mmdet/models/dense_heads/anchor_head.py @@ -1,14 +1,13 @@ import torch +from mmdeploy.codebase.mmdet import (get_post_processing_params, + multiclass_nms, pad_with_value) from mmdeploy.core import FUNCTION_REWRITER -from mmdeploy.mmdet.core import multiclass_nms -from mmdeploy.mmdet.export import pad_with_value -from mmdeploy.utils import (Backend, get_backend, get_mmdet_params, - is_dynamic_shape) +from mmdeploy.utils import Backend, get_backend, is_dynamic_shape @FUNCTION_REWRITER.register_rewriter( - func_name='mmdet.models.AnchorHead.get_bboxes') + func_name='mmdet.models.dense_heads.AnchorHead.get_bboxes') def anchor_head__get_bboxes(ctx, self, cls_scores, @@ -129,7 +128,7 @@ def anchor_head__get_bboxes(ctx, if not with_nms: return batch_mlvl_bboxes, batch_mlvl_scores - post_params = get_mmdet_params(deploy_cfg) + post_params = get_post_processing_params(deploy_cfg) max_output_boxes_per_class = post_params.max_output_boxes_per_class iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) score_threshold = cfg.get('score_thr', post_params.score_threshold) @@ -146,7 +145,7 @@ def anchor_head__get_bboxes(ctx, @FUNCTION_REWRITER.register_rewriter( - func_name='mmdet.models.AnchorHead.get_bboxes', backend='ncnn') + func_name='mmdet.models.dense_heads.AnchorHead.get_bboxes', backend='ncnn') def anchor_head__get_bboxes__ncnn(ctx, self, cls_scores, @@ -257,7 +256,7 @@ def anchor_head__get_bboxes__ncnn(ctx, if not with_nms: return batch_mlvl_bboxes, batch_mlvl_scores - post_params = get_mmdet_params(deploy_cfg) + post_params = get_post_processing_params(deploy_cfg) max_output_boxes_per_class = post_params.max_output_boxes_per_class iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) score_threshold = cfg.get('score_thr', post_params.score_threshold) diff --git a/mmdeploy/mmdet/models/dense_heads/atss_head.py b/mmdeploy/codebase/mmdet/models/dense_heads/atss_head.py similarity index 93% rename from mmdeploy/mmdet/models/dense_heads/atss_head.py rename to mmdeploy/codebase/mmdet/models/dense_heads/atss_head.py index fea597c86c..ea2d4f5a7c 100644 --- a/mmdeploy/mmdet/models/dense_heads/atss_head.py +++ b/mmdeploy/codebase/mmdet/models/dense_heads/atss_head.py @@ -1,11 +1,11 @@ import torch +from mmdeploy.codebase.mmdet import get_post_processing_params, multiclass_nms from mmdeploy.core import FUNCTION_REWRITER -from mmdeploy.mmdet.core import multiclass_nms -from mmdeploy.utils import get_mmdet_params -@FUNCTION_REWRITER.register_rewriter('mmdet.models.ATSSHead.get_bboxes') +@FUNCTION_REWRITER.register_rewriter( + 'mmdet.models.dense_heads.ATSSHead.get_bboxes') def atss_head__get_bboxes(ctx, self, cls_scores, @@ -15,7 +15,7 @@ def atss_head__get_bboxes(ctx, cfg=None, rescale=False, with_nms=True): - """Rewrite `get_bboxes` from ATSSHead for default backend. + """Rewrite `get_bboxes` of `ATSSHead` for default backend. Rewrite this function to deploy model, transform network output for a batch into bbox predictions. @@ -96,7 +96,7 @@ def atss_head__get_bboxes(ctx, batch_mlvl_scores.shape) batch_mlvl_scores = batch_mlvl_scores * batch_mlvl_centerness deploy_cfg = ctx.cfg - post_params = get_mmdet_params(deploy_cfg) + post_params = get_post_processing_params(deploy_cfg) max_output_boxes_per_class = post_params.max_output_boxes_per_class iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) score_threshold = cfg.get('score_thr', post_params.score_threshold) diff --git a/mmdeploy/mmdet/models/dense_heads/fcos_head.py b/mmdeploy/codebase/mmdet/models/dense_heads/fcos_head.py similarity index 94% rename from mmdeploy/mmdet/models/dense_heads/fcos_head.py rename to mmdeploy/codebase/mmdet/models/dense_heads/fcos_head.py index c6584c8369..6f27e0df0e 100644 --- a/mmdeploy/mmdet/models/dense_heads/fcos_head.py +++ b/mmdeploy/codebase/mmdet/models/dense_heads/fcos_head.py @@ -1,14 +1,13 @@ import torch +from mmdeploy.codebase.mmdet import (distance2bbox, get_post_processing_params, + multiclass_nms, pad_with_value) from mmdeploy.core import FUNCTION_REWRITER -from mmdeploy.mmdet.core import distance2bbox, multiclass_nms -from mmdeploy.mmdet.export import pad_with_value -from mmdeploy.utils import (Backend, get_backend, get_mmdet_params, - is_dynamic_shape) +from mmdeploy.utils import Backend, get_backend, is_dynamic_shape @FUNCTION_REWRITER.register_rewriter( - func_name='mmdet.models.FCOSHead.get_bboxes') + func_name='mmdet.models.dense_heads.FCOSHead.get_bboxes') def fcos_head__get_bboxes(ctx, self, cls_scores, @@ -18,7 +17,7 @@ def fcos_head__get_bboxes(ctx, with_nms=True, cfg=None, **kwargs): - """Rewrite `get_bboxes` of FCOSHead for default backend. + """Rewrite `get_bboxes` of `FCOSHead` for default backend. Rewrite this function to support deployment of default backend and dynamic shape export. Transform network output for a batch into @@ -130,7 +129,7 @@ def fcos_head__get_bboxes(ctx, return batch_mlvl_bboxes, batch_mlvl_scores, batch_mlvl_centerness batch_mlvl_scores = batch_mlvl_scores * batch_mlvl_centerness - post_params = get_mmdet_params(deploy_cfg) + post_params = get_post_processing_params(deploy_cfg) max_output_boxes_per_class = post_params.max_output_boxes_per_class iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) score_threshold = cfg.get('score_thr', post_params.score_threshold) @@ -142,7 +141,7 @@ def fcos_head__get_bboxes(ctx, @FUNCTION_REWRITER.register_rewriter( - func_name='mmdet.models.FCOSHead.get_bboxes', backend='ncnn') + func_name='mmdet.models.dense_heads.FCOSHead.get_bboxes', backend='ncnn') def fcos_head__get_bboxes__ncnn(ctx, self, cls_scores, @@ -152,7 +151,7 @@ def fcos_head__get_bboxes__ncnn(ctx, with_nms=True, cfg=None, **kwargs): - """Rewrite `get_bboxes` of FCOSHead for ncnn backend. + """Rewrite `get_bboxes` of `FCOSHead` for ncnn backend. 1. Shape node and batch inference is not supported by ncnn. This function transform dynamic shape to constant shape and remove batch inference. @@ -257,7 +256,7 @@ def fcos_head__get_bboxes__ncnn(ctx, batch_mlvl_scores = (_batch_mlvl_scores * _batch_mlvl_centerness). \ reshape(batch_mlvl_scores.shape) batch_mlvl_bboxes = batch_mlvl_bboxes.reshape(batch_size, -1, 4) - post_params = get_mmdet_params(deploy_cfg) + post_params = get_post_processing_params(deploy_cfg) max_output_boxes_per_class = post_params.max_output_boxes_per_class iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) score_threshold = cfg.get('score_thr', post_params.score_threshold) diff --git a/mmdeploy/mmdet/models/dense_heads/fovea_head.py b/mmdeploy/codebase/mmdet/models/dense_heads/fovea_head.py similarity index 92% rename from mmdeploy/mmdet/models/dense_heads/fovea_head.py rename to mmdeploy/codebase/mmdet/models/dense_heads/fovea_head.py index ea1d8c16ce..81f4add4c6 100644 --- a/mmdeploy/mmdet/models/dense_heads/fovea_head.py +++ b/mmdeploy/codebase/mmdet/models/dense_heads/fovea_head.py @@ -1,11 +1,11 @@ import torch +from mmdeploy.codebase.mmdet import get_post_processing_params, multiclass_nms from mmdeploy.core import FUNCTION_REWRITER -from mmdeploy.mmdet.core import multiclass_nms -from mmdeploy.utils import get_mmdet_params -@FUNCTION_REWRITER.register_rewriter('mmdet.models.FoveaHead.get_bboxes') +@FUNCTION_REWRITER.register_rewriter( + 'mmdet.models.dense_heads.FoveaHead.get_bboxes') def fovea_head__get_bboxes(ctx, self, cls_scores, @@ -13,7 +13,7 @@ def fovea_head__get_bboxes(ctx, img_metas, cfg=None, rescale=None): - """Rewrite `get_bboxes` from FoveaHead for default backend. + """Rewrite `get_bboxes` of `FoveaHead` for default backend. Rewrite this function to deploy model, transform network output for a batch into bbox predictions. @@ -78,7 +78,7 @@ def fovea_head__get_bboxes(ctx, det_scores = torch.cat(det_scores, dim=1) deploy_cfg = ctx.cfg - post_params = get_mmdet_params(deploy_cfg) + post_params = get_post_processing_params(deploy_cfg) max_output_boxes_per_class = post_params.max_output_boxes_per_class iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) score_threshold = cfg.get('score_thr', post_params.score_threshold) diff --git a/mmdeploy/mmdet/models/dense_heads/rpn_head.py b/mmdeploy/codebase/mmdet/models/dense_heads/rpn_head.py similarity index 94% rename from mmdeploy/mmdet/models/dense_heads/rpn_head.py rename to mmdeploy/codebase/mmdet/models/dense_heads/rpn_head.py index 5604e83b90..1921d9a5fd 100644 --- a/mmdeploy/mmdet/models/dense_heads/rpn_head.py +++ b/mmdeploy/codebase/mmdet/models/dense_heads/rpn_head.py @@ -1,13 +1,13 @@ import torch +from mmdeploy.codebase.mmdet import (get_post_processing_params, + multiclass_nms, pad_with_value) from mmdeploy.core import FUNCTION_REWRITER -from mmdeploy.mmdet.core import multiclass_nms -from mmdeploy.mmdet.export import pad_with_value -from mmdeploy.utils import (Backend, get_backend, get_mmdet_params, - is_dynamic_shape) +from mmdeploy.utils import Backend, get_backend, is_dynamic_shape -@FUNCTION_REWRITER.register_rewriter('mmdet.models.RPNHead.get_bboxes') +@FUNCTION_REWRITER.register_rewriter( + 'mmdet.models.dense_heads.RPNHead.get_bboxes') def rpn_head__get_bboxes(ctx, self, cls_scores, @@ -16,7 +16,7 @@ def rpn_head__get_bboxes(ctx, with_nms=True, cfg=None, **kwargs): - """Rewrite `get_bboxes` of RPNHead for default backend. + """Rewrite `get_bboxes` of `RPNHead` for default backend. Rewrite this function to deploy model, transform network output for a batch into bbox predictions. @@ -120,7 +120,7 @@ def rpn_head__get_bboxes(ctx, if not with_nms: return batch_mlvl_bboxes, batch_mlvl_scores - post_params = get_mmdet_params(deploy_cfg) + post_params = get_post_processing_params(deploy_cfg) iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) score_threshold = cfg.get('score_thr', post_params.score_threshold) pre_top_k = post_params.pre_top_k @@ -138,7 +138,7 @@ def rpn_head__get_bboxes(ctx, @FUNCTION_REWRITER.register_rewriter( - 'mmdet.models.RPNHead.get_bboxes', backend='ncnn') + 'mmdet.models.dense_heads.RPNHead.get_bboxes', backend='ncnn') def rpn_head__get_bboxes__ncnn(ctx, self, cls_scores, @@ -147,7 +147,7 @@ def rpn_head__get_bboxes__ncnn(ctx, with_nms=True, cfg=None, **kwargs): - """Rewrite `get_bboxes` of RPNHead for ncnn backend. + """Rewrite `get_bboxes` of `RPNHead` for NCNN backend. Shape node and batch inference is not supported by ncnn. This function transform dynamic shape to constant shape and remove batch inference. @@ -242,7 +242,7 @@ def rpn_head__get_bboxes__ncnn(ctx, if not with_nms: return batch_mlvl_bboxes, batch_mlvl_scores - post_params = get_mmdet_params(deploy_cfg) + post_params = get_post_processing_params(deploy_cfg) iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) score_threshold = cfg.get('score_thr', post_params.score_threshold) pre_top_k = post_params.pre_top_k diff --git a/mmdeploy/mmdet/models/dense_heads/vfnet_head.py b/mmdeploy/codebase/mmdet/models/dense_heads/vfnet_head.py similarity index 93% rename from mmdeploy/mmdet/models/dense_heads/vfnet_head.py rename to mmdeploy/codebase/mmdet/models/dense_heads/vfnet_head.py index 8816231d62..2378727608 100644 --- a/mmdeploy/mmdet/models/dense_heads/vfnet_head.py +++ b/mmdeploy/codebase/mmdet/models/dense_heads/vfnet_head.py @@ -1,12 +1,12 @@ import torch +from mmdeploy.codebase.mmdet import (distance2bbox, get_post_processing_params, + multiclass_nms) from mmdeploy.core import FUNCTION_REWRITER -from mmdeploy.mmdet.core import distance2bbox, multiclass_nms -from mmdeploy.utils import get_mmdet_params @FUNCTION_REWRITER.register_rewriter( - func_name='mmdet.models.VFNetHead.get_bboxes') + 'mmdet.models.dense_heads.VFNetHead.get_bboxes') def vfnet_head__get_bboxes(ctx, self, cls_scores, @@ -16,7 +16,7 @@ def vfnet_head__get_bboxes(ctx, cfg=None, rescale=None, with_nms=True): - """Rewrite `get_bboxes` of VFNetHead for default backend. + """Rewrite `get_bboxes` of `VFNetHead` for default backend. Rewrite this function to deploy model, transform network output for a batch into bbox predictions. @@ -102,7 +102,7 @@ def vfnet_head__get_bboxes(ctx, return batch_mlvl_bboxes, batch_mlvl_scores deploy_cfg = ctx.cfg - post_params = get_mmdet_params(deploy_cfg) + post_params = get_post_processing_params(deploy_cfg) max_output_boxes_per_class = post_params.max_output_boxes_per_class iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) score_threshold = cfg.get('score_thr', post_params.score_threshold) diff --git a/mmdeploy/mmdet/models/dense_heads/yolo_head.py b/mmdeploy/codebase/mmdet/models/dense_heads/yolo_head.py similarity index 96% rename from mmdeploy/mmdet/models/dense_heads/yolo_head.py rename to mmdeploy/codebase/mmdet/models/dense_heads/yolo_head.py index 3983fa5828..24869be2b1 100644 --- a/mmdeploy/mmdet/models/dense_heads/yolo_head.py +++ b/mmdeploy/codebase/mmdet/models/dense_heads/yolo_head.py @@ -1,21 +1,20 @@ import torch +from mmdeploy.codebase.mmdet import (get_post_processing_params, + multiclass_nms, pad_with_value) from mmdeploy.core import FUNCTION_REWRITER -from mmdeploy.mmdet.core import multiclass_nms -from mmdeploy.mmdet.export import pad_with_value -from mmdeploy.utils import (Backend, get_backend, get_mmdet_params, - is_dynamic_shape) +from mmdeploy.utils import Backend, get_backend, is_dynamic_shape @FUNCTION_REWRITER.register_rewriter( - func_name='mmdet.models.YOLOV3Head.get_bboxes') + func_name='mmdet.models.dense_heads.YOLOV3Head.get_bboxes') def yolov3_head__get_bboxes(ctx, self, pred_maps, with_nms=True, cfg=None, **kwargs): - """Rewrite `get_bboxes` of YOLOV3Head for default backend. + """Rewrite `get_bboxes` of `YOLOV3Head` for default backend. Rewrite this function to deploy model, transform network output for a batch into bbox predictions. @@ -83,7 +82,6 @@ def yolov3_head__get_bboxes(ctx, conf_pred = torch.sigmoid(pred_map[..., 4]) cls_pred = torch.sigmoid(pred_map[..., 5:]).view( batch_size, -1, self.num_classes) # Cls pred one-hot. - backend = get_backend(ctx.cfg) # topk in tensorrt does not support shape BaseTask: + task = get_task_type(deploy_cfg) + return registry.module_dict[task.value](model_cfg, deploy_cfg, device) + + +MMEDIT_TASK = Registry('mmedit_tasks', build_func=__build_mmedit_task) + + +@CODEBASE.register_module(Codebase.MMEDIT.value) +class MMEditing(MMCodebase): + """mmediting codebase class.""" + + task_registry = MMEDIT_TASK + + def __init__(self): + super().__init__() + + @staticmethod + def build_task_processor(model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str) -> BaseTask: + """The interface to build the task processors of mmedit. + + Args: + model_cfg (mmcv.Config): Model config file. + deploy_cfg (mmcv.Config): Deployment config file. + device (str): A string specifying device type. + + Returns: + BaseTask: A task processor. + """ + return MMEDIT_TASK.build(model_cfg, deploy_cfg, device) + + @staticmethod + def build_dataset(dataset_cfg: Union[str, mmcv.Config], *args, + **kwargs) -> Dataset: + """Build dataset for processor. + + Args: + dataset_cfg (str | mmcv.Config): The input dataset config. + + Returns: + Dataset: A PyTorch dataset. + """ + from mmedit.datasets import build_dataset as build_dataset_mmedit + from mmdeploy.utils import load_config + dataset_cfg = load_config(dataset_cfg)[0] + data = dataset_cfg.data + + dataset = build_dataset_mmedit(data.test) + return dataset + + @staticmethod + def build_dataloader(dataset: Dataset, + samples_per_gpu: int, + workers_per_gpu: int, + num_gpus: int = 1, + dist: bool = False, + shuffle: bool = False, + seed: Optional[int] = None, + drop_last: bool = False, + pin_memory: bool = True, + persistent_workers: bool = True, + **kwargs) -> DataLoader: + """Build PyTorch DataLoader. + + In distributed training, each GPU/process has a dataloader. + In non-distributed training, there is only one dataloader for all GPUs. + + Args: + dataset (:obj:`Dataset`): A PyTorch dataset. + samples_per_gpu (int): Number of samples on each GPU, i.e., + batch size of each GPU. + workers_per_gpu (int): How many subprocesses to use for data + loading for each GPU. + num_gpus (int): Number of GPUs. Only used in non-distributed + training. Default: 1. + dist (bool): Distributed training/test or not. Default: True. + shuffle (bool): Whether to shuffle the data at every epoch. + Default: True. + seed (int | None): Seed to be used. Default: None. + drop_last (bool): Whether to drop the last incomplete batch + in epoch.　Default: False. + pin_memory (bool): Whether to use pin_memory in DataLoader. + Default: True. + persistent_workers (bool): If True, the data loader will not + shutdown　the worker processes after a dataset has been + consumed once. + This allows to maintain the workers Dataset instances alive. + The argument also has effect in PyTorch>=1.7.0. + Default: True. + kwargs (dict, optional): Any keyword argument to be used to + initialize　DataLoader. + + Returns: + DataLoader: A PyTorch dataloader. + """ + from mmedit.datasets import build_dataloader as build_dataloader_mmedit + return build_dataloader_mmedit(dataset, samples_per_gpu, + workers_per_gpu, num_gpus, dist, + shuffle, seed, drop_last, pin_memory, + persistent_workers, **kwargs) + + @staticmethod + def single_gpu_test(model: torch.nn.Module, + data_loader: DataLoader, + save_image: bool = False, + save_path: Optional[str] = None, + iteration: int = None) -> list: + """Run test with single gpu. + + Args: + model (torch.nn.Module): Input model from nn.Module. + data_loader (DataLoader): PyTorch data loader. + save_image (bool): Whether save image. Default: False. + save_path (str): The path to save image. Default: None. + iteration (int): Iteration number. It is used for the save + image name.　Default: None. + + Returns: + list: The prediction results. + """ + from mmedit.apis import single_gpu_test + outputs = single_gpu_test(model, data_loader, save_image, save_path, + iteration) + return outputs diff --git a/mmdeploy/codebase/mmedit/deploy/super_resolution.py b/mmdeploy/codebase/mmedit/deploy/super_resolution.py new file mode 100644 index 0000000000..2a7cfe1d70 --- /dev/null +++ b/mmdeploy/codebase/mmedit/deploy/super_resolution.py @@ -0,0 +1,280 @@ +import logging +import warnings +from typing import Any, Dict, Optional, Sequence, Tuple, Union + +import mmcv +import numpy as np +import torch +from mmcv.parallel import collate, scatter +from torch.utils.data import Dataset + +from mmdeploy.codebase.base import BaseTask +from mmdeploy.codebase.mmedit.deploy.mmediting import MMEDIT_TASK +from mmdeploy.utils import Task, load_config + + +@MMEDIT_TASK.register_module(Task.SUPER_RESOLUTION.value) +class SuperResolution(BaseTask): + """BaseTask class of super resolution task. + + Args: + model_cfg　(mmcv.Config):　Model config file. + deploy_cfg　(mmcv.Config):　Deployment config file. + device　(str):　A string specifying device type. + """ + + def __init__(self, model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str): + super().__init__(model_cfg, deploy_cfg, device) + + def init_backend_model(self, + model_files: Sequence[str] = None, + **kwargs) -> torch.nn.Module: + """Initialize backend model. + + Args: + model_files (Sequence[str]): Input model files. Default is None. + + Returns: + nn.Module: An initialized backend model. + """ + from .super_resolution_model import build_super_resolution_model + model = build_super_resolution_model( + model_files, self.model_cfg, self.deploy_cfg, device=self.device) + return model + + def init_pytorch_model(self, + model_checkpoint: Optional[str] = None, + **kwargs) -> torch.nn.Module: + """Initialize torch model. + + Args: + model_checkpoint (str): The checkpoint file of torch model, + defaults to `None`. + + Returns: + nn.Module: An initialized torch model generated by other OpenMMLab + codebases. + """ + from mmedit.apis import init_model + model = init_model(self.model_cfg, model_checkpoint, self.device) + model.forward = model.forward_dummy + return model.eval() + + def create_input(self, + imgs: Union[str, np.ndarray], + input_shape: Optional[Sequence[int]] = None, + **kwargs) -> Tuple[Dict, torch.Tensor]: + """Create input for editing processor. + + Args: + imgs (str | np.ndarray): Input image(s). + input_shape (Sequence[int] | None): A list of two integer in + (width, height) format specifying input shape. Defaults to `None`. + + Returns: + tuple: (data, img), meta information for the input image and input. + """ + from mmedit.datasets.pipelines import Compose + + if isinstance(imgs, (list, tuple)): + if not isinstance(imgs[0], (np.ndarray, str)): + raise AssertionError('imgs must be strings or numpy arrays') + elif isinstance(imgs, (np.ndarray, str)): + imgs = [imgs] + else: + raise AssertionError('imgs must be strings or numpy arrays') + + cfg = load_config(self.model_cfg)[0].copy() + + self._preprocess_cfg( + cfg, + load_from_file=isinstance(imgs[0], str), + is_static_cfg=input_shape is not None, + input_shape=input_shape) + + test_pipeline = Compose(cfg.test_pipeline) + + data_arr = [] + for img in imgs: + if isinstance(img, np.ndarray): + data = dict(lq=img) + else: + data = dict(lq_path=img) + + data = test_pipeline(data) + data_arr.append(data) + + data = collate(data_arr, samples_per_gpu=len(imgs)) + + data['img'] = data['lq'] + + if self.device != 'cpu': + data = scatter(data, [self.device])[0] + + return data, data['img'] + + def visualize(self, + model: torch.nn.Module, + image: Union[str, np.ndarray], + result: Union[list, np.ndarray], + output_file: str, + window_name: str = '', + show_result: bool = False, + **kwargs) -> np.ndarray: + """Visualize result of a model. + + Args: + model (nn.Module): Input model. + image (str | np.ndarray): Input image to draw predictions on. + result (list | np.ndarray): A list of result. + output_file (str): Output file to save drawn image. + window_name (str): The name of visualization window. Defaults to + an empty string. + show_result (bool): Whether to show result in windows, defaults + to `False`. + """ + if len(result.shape) == 4: + result = result[0] + + with torch.no_grad(): + result = result.transpose(1, 2, 0) + result = np.clip(result, 0, 1)[:, :, ::-1] + result = (result * 255.0).round() + + output_file = None if show_result else output_file + + if show_result: + int_result = result.astype(np.uint8) + mmcv.imshow(int_result, window_name, 0) + if output_file is not None: + mmcv.imwrite(result, output_file) + + if not (show_result or output_file): + warnings.warn( + 'show_result==False and output_file is not specified, only ' + 'result image will be returned') + return result + + @staticmethod + def run_inference(model: torch.nn.Module, + model_inputs: Dict[str, torch.Tensor]) -> list: + """Run inference once for a super resolution model of mmedit. + + Args: + model (nn.Module): Input model. + model_inputs (dict): A dict containing model inputs tensor and + meta info. + + Returns: + list: The predictions of model inference. + """ + result = model(model_inputs['lq']) + if not isinstance(result[0], np.ndarray): + result = [result[0].detach().cpu().numpy()] + return result + + @staticmethod + def get_partition_cfg(partition_type: str, **kwargs) -> Dict: + """Get a certain partition config for mmedit. + + Args: + partition_type (str): A string specifying partition type. + + Returns: + dict: A dictionary of partition config. + """ + raise NotImplementedError + + @staticmethod + def get_tensor_from_input(input_data: Dict[str, Any]) -> torch.Tensor: + """Get input tensor from input data. + + Args: + input_data (dict): Input data containing meta info + and image tensor. + Returns: + torch.Tensor: An image in `Tensor`. + """ + return input_data['lq'] + + @staticmethod + def evaluate_outputs(model_cfg, + outputs: list, + dataset: Dataset, + metrics: Optional[str] = None, + out: Optional[str] = None, + metric_options: Optional[dict] = None, + format_only: bool = False, + **kwargs) -> None: + """Evaluation function implemented in mmedit. + + Args: + model_cfg (mmcv.Config): The model config. + outputs (list): A list of result of model inference. + dataset (Dataset): Input dataset to run test. + metrics (str): Evaluation metrics, which depends on + the codebase and the dataset, e.g., "PSNR", "SSIM" in mmedit. + out (str): Output result file in pickle format, defaults to `None`. + metric_options (dict): Custom options for evaluation, will be + kwargs for dataset.evaluate() function. Defaults to `None`. + format_only (bool): Format the output results without perform + evaluation. It is useful when you want to format the result + to a specific format and submit it to the test server. Defaults + to `False`. + """ + if out: + logging.info(f'\nwriting results to {out}') + mmcv.dump(outputs, out) + # The Dataset doesn't need metrics + print('\n') + # print metrics + stats = dataset.evaluate(outputs) + for stat in stats: + print('Eval-{}: {}'.format(stat, stats[stat])) + + def _preprocess_cfg(self, config: mmcv.Config, load_from_file: bool, + is_static_cfg: bool, + input_shape: Sequence[int]) -> None: + """Remove unnecessary information in config. + + Args: + config (mmcv.Config): The input model config. + load_from_file (bool): Whether the input is a filename of a numpy + matrix. If this variable is True, extra preprocessing is + required. + is_static_cfg (bool): Whether the config specifys a static export. + If this variable if True, the input image will be resize to a + fix resolution. + input_shape (Sequence[int]): A list of two integer in + (width, height) format specifying input shape. + Defaults to `None`. + """ + keys_to_remove = ['gt', 'gt_path'] + # MMEdit doesn't support LoadImageFromWebcam. + # Remove "LoadImageFromFile" and related metakeys. + if not load_from_file: + config.test_pipeline.pop(0) + keys_to_remove.append('lq_path') + + # Fix the input shape by 'Resize' + if is_static_cfg: + resize = { + 'type': 'Resize', + 'scale': (input_shape[0], input_shape[1]), + 'keys': ['lq'] + } + config.test_pipeline.insert(1, resize) + + for key in keys_to_remove: + for pipeline in list(config.test_pipeline): + if 'key' in pipeline and key == pipeline['key']: + config.test_pipeline.remove(pipeline) + if 'keys' in pipeline: + while key in pipeline['keys']: + pipeline['keys'].remove(key) + if len(pipeline['keys']) == 0: + config.test_pipeline.remove(pipeline) + if 'meta_keys' in pipeline: + while key in pipeline['meta_keys']: + pipeline['meta_keys'].remove(key) diff --git a/mmdeploy/codebase/mmedit/deploy/super_resolution_model.py b/mmdeploy/codebase/mmedit/deploy/super_resolution_model.py new file mode 100644 index 0000000000..d8db5e143f --- /dev/null +++ b/mmdeploy/codebase/mmedit/deploy/super_resolution_model.py @@ -0,0 +1,181 @@ +from typing import List, Optional, Sequence, Union + +import mmcv +import numpy as np +import torch +from mmedit.core import psnr, ssim, tensor2img + +from mmdeploy.codebase.base import BaseBackendModel +from mmdeploy.utils import Backend, get_backend, get_onnx_config, load_config + + +class End2EndModel(BaseBackendModel): + """End to end model for inference of super resolution. + + Args: + backend (Backend): The backend enum, specifying backend type. + backend_files (Sequence[str]): Paths to all required backend files(e.g. + '.onnx' for ONNX Runtime, '.param' and '.bin' for ncnn). + device (str): A string represents device type. + model_cfg(mmcv.Config): Input model config object. + deploy_cfg(str | mmcv.Config):Deployment config file or loaded Config + object. + """ + + def __init__(self, + backend: Backend, + backend_files: Sequence[str], + device: str, + model_cfg: mmcv.Config, + deploy_cfg: Union[str, mmcv.Config] = None): + super().__init__() + self.deploy_cfg = deploy_cfg + self.test_cfg = model_cfg.test_cfg + self.allowed_metrics = {'PSNR': psnr, 'SSIM': ssim} + self._init_wrapper( + backend=backend, backend_files=backend_files, device=device) + + def _init_wrapper(self, backend: Backend, backend_files: Sequence[str], + device: str): + onnx_config = get_onnx_config(self.deploy_cfg) + output_names = onnx_config['output_names'] + self.wrapper = BaseBackendModel._build_wrapper( + backend=backend, + backend_files=backend_files, + device=device, + output_names=output_names) + + def forward(self, + lq: torch.Tensor, + test_mode: bool = False, + *args, + **kwargs) -> Union[list, dict]: + """Run test inference for restorer. + + We want forward() to output an image or a evaluation result. + When test_mode is set, the output is evaluation result. Otherwise + it is an image. + + Args: + lq (torch.Tensor): The input low-quality image of the model. + test_mode (bool): When test_mode is set, the output is evaluation + result. Otherwise it is an image. Default to `False`. + *args: Other arguments. + **kwargs: Other key-pair arguments. + + Returns: + list | dict: High resolution image or a evaluation results. + """ + + if test_mode: + return self.forward_test(lq, *args, **kwargs) + else: + return self.forward_dummy(lq, *args, **kwargs) + + def forward_test(self, + lq: torch.Tensor, + gt: Optional[torch.Tensor] = None, + *args, + **kwargs): + """Run inference for restorer to generate evaluation result. + + Args: + lq (torch.Tensor): The input low-quality image of the model. + gt (torch.Tensor): The ground truth of input image. Defaults to + `None`. + *args: Other arguments. + **kwargs: Other key-pair arguments. + + Returns: + dict: Evaluation results. + """ + outputs = self.forward_dummy(lq) + result = self.test_post_process(outputs, lq, gt) + return result + + def forward_dummy(self, lq: torch.Tensor, *args, **kwargs): + """Run test inference for restorer with backend wrapper. + + Args: + lq (torch.Tensor): The input low-quality image of the model. + + Returns: + list[np.ndarray] : High resolution image. + """ + outputs = self.wrapper({'input': lq}) + outputs = self.wrapper.output_to_list(outputs) + outputs = [out.detach().cpu().numpy() for out in outputs] + return outputs + + def evaluate(self, output: Union[torch.Tensor, np.ndarray], + gt: torch.Tensor): + """Evaluation function implemented in mmedit. + + Args: + output (torch.Tensor | np.ndarray): Model output with + shape (n, c, h, w). + gt (torch.Tensor): GT Tensor with shape (n, c, h, w). + + Returns: + dict: Evaluation results. + """ + crop_border = self.test_cfg.crop_border + + if isinstance(output, np.ndarray): + output = torch.from_numpy(output) + output = tensor2img(output) + gt = tensor2img(gt) + + eval_result = dict() + for metric in self.test_cfg.metrics: + eval_result[metric] = self.allowed_metrics[metric](output, gt, + crop_border) + return eval_result + + def test_post_process(self, + outputs: List[np.ndarray], + lq: torch.Tensor, + gt: Optional[torch.Tensor] = None): + """Get evaluation results by post-processing model outputs. + + Args: + output (list[np.ndarray]) : The output high resolution image. + lq (torch.Tensor): The input low-quality image of the model. + gt (torch.Tensor): The ground truth of input image, default is + `None`. + + Returns: + dict: Evaluation results. + """ + if self.test_cfg is not None and self.test_cfg.get('metrics', None): + assert gt is not None, ( + 'evaluation with metrics must have gt images.') + results = dict(eval_result=self.evaluate(outputs[0], gt)) + else: + results = dict(lq=lq.cpu(), output=outputs) + if gt is not None: + results['gt'] = gt.cpu() + + return results + + def show_result(self, *args, **kwargs): + raise NotImplementedError + + +def build_super_resolution_model(model_files: Sequence[str], + model_cfg: Union[str, mmcv.Config], + deploy_cfg: Union[str, + mmcv.Config], device: str): + model_cfg = load_config(model_cfg)[0] + deploy_cfg = load_config(deploy_cfg)[0] + + backend = get_backend(deploy_cfg) + + backend_model = End2EndModel( + backend=backend, + backend_files=model_files, + device=device, + model_cfg=model_cfg, + deploy_cfg=deploy_cfg) + + return backend_model diff --git a/mmdeploy/mmedit/models/__init__.py b/mmdeploy/codebase/mmedit/models/__init__.py similarity index 100% rename from mmdeploy/mmedit/models/__init__.py rename to mmdeploy/codebase/mmedit/models/__init__.py diff --git a/mmdeploy/mmedit/models/backbones/sr_backbones/__init__.py b/mmdeploy/codebase/mmedit/models/backbones/__init__.py similarity index 100% rename from mmdeploy/mmedit/models/backbones/sr_backbones/__init__.py rename to mmdeploy/codebase/mmedit/models/backbones/__init__.py diff --git a/mmdeploy/mmedit/models/backbones/sr_backbones/srcnn.py b/mmdeploy/codebase/mmedit/models/backbones/srcnn.py similarity index 100% rename from mmdeploy/mmedit/models/backbones/sr_backbones/srcnn.py rename to mmdeploy/codebase/mmedit/models/backbones/srcnn.py diff --git a/mmdeploy/mmedit/__init__.py b/mmdeploy/codebase/mmocr/__init__.py similarity index 50% rename from mmdeploy/mmedit/__init__.py rename to mmdeploy/codebase/mmocr/__init__.py index d2b62e1cb6..33b69c74df 100644 --- a/mmdeploy/mmedit/__init__.py +++ b/mmdeploy/codebase/mmocr/__init__.py @@ -1,2 +1,2 @@ -from .export import * # noqa: F401,F403 +from .deploy import * # noqa: F401,F403 from .models import * # noqa: F401,F403 diff --git a/mmdeploy/codebase/mmocr/deploy/__init__.py b/mmdeploy/codebase/mmocr/deploy/__init__.py new file mode 100644 index 0000000000..43b39c5ed0 --- /dev/null +++ b/mmdeploy/codebase/mmocr/deploy/__init__.py @@ -0,0 +1,5 @@ +from .mmocr import MMOCR +from .text_detection import TextDetection +from .text_recognition import TextRecognition + +__all__ = ['MMOCR', 'TextDetection', 'TextRecognition'] diff --git a/mmdeploy/codebase/mmocr/deploy/mmocr.py b/mmdeploy/codebase/mmocr/deploy/mmocr.py new file mode 100644 index 0000000000..71613217c6 --- /dev/null +++ b/mmdeploy/codebase/mmocr/deploy/mmocr.py @@ -0,0 +1,141 @@ +from typing import Optional, Union + +import mmcv +import torch +from mmcv.utils import Registry +from torch.utils.data import DataLoader, Dataset + +from mmdeploy.codebase.base import CODEBASE, BaseTask, MMCodebase +from mmdeploy.utils import Codebase, get_task_type +from mmdeploy.utils.config_utils import load_config + + +def __build_mmocr_task(model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str, registry: Registry) -> BaseTask: + task = get_task_type(deploy_cfg) + return registry.module_dict[task.value](model_cfg, deploy_cfg, device) + + +MMOCR_TASK = Registry('mmocr_tasks', build_func=__build_mmocr_task) + + +@CODEBASE.register_module(Codebase.MMOCR.value) +class MMOCR(MMCodebase): + """mmocr codebase class.""" + + task_registry = MMOCR_TASK + + def __init__(self): + super(MMOCR, self).__init__() + + @staticmethod + def build_task_processor(model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str): + """The interface to build the task processors of mmocr. + + Args: + model_cfg (str | mmcv.Config): Model config file or loaded Config + object. + deploy_cfg (str | mmcv.Config): Deployment config file or loaded + Config object. + device (str): A string specifying device type. + + Returns: + BaseTask: A task processor. + """ + return MMOCR_TASK.build(model_cfg, deploy_cfg, device) + + @staticmethod + def build_dataset(dataset_cfg: Union[str, mmcv.Config], + dataset_type: str = 'val', + **kwargs) -> Dataset: + """Build dataset for mmocr. + + Args: + dataset_cfg (str | mmcv.Config): The input dataset config. + dataset_type (str): A string represents dataset type, e.g.: 'train' + , 'test', 'val'. Defaults to 'val'. + + Returns: + Dataset: A PyTorch dataset. + """ + from mmocr.datasets import build_dataset as build_dataset_mmocr + + dataset_cfg = load_config(dataset_cfg)[0] + assert dataset_type in dataset_cfg.data + data_cfg = dataset_cfg.data[dataset_type] + dataset = build_dataset_mmocr(data_cfg) + return dataset + + @staticmethod + def build_dataloader(dataset: Dataset, + samples_per_gpu: int, + workers_per_gpu: int, + num_gpus: int = 1, + dist: bool = False, + shuffle: bool = False, + seed: Optional[int] = None, + drop_last: bool = False, + persistent_workers: bool = True, + **kwargs) -> DataLoader: + """Build dataloader for mmocr. + + Args: + dataset (Dataset): Input dataset. + samples_per_gpu (int): Number of training samples on each GPU, i.e. + ,batch size of each GPU. + workers_per_gpu (int): How many subprocesses to use for data + loading for each GPU. + num_gpus (int): Number of GPUs. Only used in non-distributed + training. + dist (bool): Distributed training/test or not. Defaults to `False`. + shuffle (bool): Whether to shuffle the data at every epoch. + Defaults to `False`. + seed (int): An integer set to be seed. Default is `None`. + drop_last (bool): Whether to drop the last incomplete batch in + epoch. Default to `False`. + persistent_workers (bool): If `True`, the data loader will not + shutdown the worker processes after a dataset has been + consumed once. This allows to maintain the workers Dataset + instances alive. The argument also has effect in + PyTorch>=1.7.0. Default is `True`. + kwargs: Any other keyword argument to be used to initialize + DataLoader. + + Returns: + DataLoader: A PyTorch dataloader. + """ + from mmocr.datasets import build_dataloader as build_dataloader_mmocr + return build_dataloader_mmocr( + dataset, + samples_per_gpu, + workers_per_gpu, + num_gpus=num_gpus, + dist=dist, + shuffle=shuffle, + seed=seed, + drop_last=drop_last, + persistent_workers=persistent_workers, + **kwargs) + + @staticmethod + def single_gpu_test(model: torch.nn.Module, + data_loader: DataLoader, + show: bool = False, + out_dir: Optional[str] = None, + **kwargs): + """Run test with single gpu. + + Args: + model (torch.nn.Module): Input model from nn.Module. + data_loader (DataLoader): PyTorch data loader. + show (bool): Specifying whether to show plotted results. Defaults + to `False`. + out_dir (str): A directory to save results, defaults to `None`. + + Returns: + list: The prediction results. + """ + from mmdet.apis import single_gpu_test + outputs = single_gpu_test(model, data_loader, show, out_dir, **kwargs) + return outputs diff --git a/mmdeploy/codebase/mmocr/deploy/text_detection.py b/mmdeploy/codebase/mmocr/deploy/text_detection.py new file mode 100644 index 0000000000..540426f569 --- /dev/null +++ b/mmdeploy/codebase/mmocr/deploy/text_detection.py @@ -0,0 +1,264 @@ +import logging +from typing import Any, Dict, Optional, Sequence, Tuple, Union + +import mmcv +import numpy as np +import torch +from mmcv.parallel import DataContainer, collate, scatter +from mmdet.datasets import replace_ImageToTensor +from torch import nn +from torch.utils.data import Dataset + +from mmdeploy.codebase.base import BaseTask +from mmdeploy.utils import Task +from .mmocr import MMOCR_TASK + + +@MMOCR_TASK.register_module(Task.TEXT_DETECTION.value) +class TextDetection(BaseTask): + """Text detection task class. + + Args: + model_cfg (mmcv.Config): Loaded model Config object.. + deploy_cfg (mmcv.Config): Loaded deployment Config object. + device (str): A string represents device type. + """ + + def __init__(self, model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str): + super(TextDetection, self).__init__(model_cfg, deploy_cfg, device) + + def init_backend_model(self, + model_files: Optional[str] = None, + **kwargs) -> torch.nn.Module: + """Initialize backend model. + + Args: + model_files (Sequence[str]): Input model files. + + Returns: + nn.Module: An initialized backend model. + """ + from .text_detection_model import build_text_detection_model + model = build_text_detection_model( + model_files, self.model_cfg, self.deploy_cfg, device=self.device) + return model.eval() + + def init_pytorch_model(self, + model_checkpoint: Optional[str] = None, + cfg_options: Optional[Dict] = None, + **kwargs) -> torch.nn.Module: + """Initialize torch model. + + Args: + model_checkpoint (str): The checkpoint file of torch model, + defaults to `None`. + cfg_options (dict): Optional config key-pair parameters. + + Returns: + nn.Module: An initialized torch model generated by OpenMMLab + codebases. + """ + from mmocr.apis import init_detector + model = init_detector(self.model_cfg, model_checkpoint, self.device, + cfg_options) + + return model.eval() + + def create_input(self, + imgs: Union[str, np.ndarray], + input_shape: Sequence[int] = None) \ + -> Tuple[Dict, torch.Tensor]: + """Create input for segmentor. + + Args: + imgs (str | np.ndarray): Input image(s), accepted data type are + `str`, `np.ndarray`. + input_shape (list[int]): A list of two integer in (width, height) + format specifying input shape. Defaults to `None`. + + Returns: + tuple: (data, img), meta information for the input image and input. + """ + if isinstance(imgs, (list, tuple)): + if not isinstance(imgs[0], (np.ndarray, str)): + raise AssertionError('imgs must be strings or numpy arrays') + + elif isinstance(imgs, (np.ndarray, str)): + imgs = [imgs] + else: + raise AssertionError('imgs must be strings or numpy arrays') + + if self.model_cfg.data.test['type'] == 'ConcatDataset': + self.model_cfg.data.test.pipeline = \ + self.model_cfg.data.test['datasets'][0].pipeline + + is_ndarray = isinstance(imgs[0], np.ndarray) + + if is_ndarray: + self.model_cfg.data.test.pipeline[0].type = 'LoadImageFromNdarray' + + test_pipeline = self.model_cfg.data.test.pipeline + test_pipeline = replace_ImageToTensor(test_pipeline) + # for static exporting + if input_shape is not None: + test_pipeline[1].img_scale = tuple(input_shape) + test_pipeline[1].transforms[0].keep_ratio = False + test_pipeline[1].transforms[0].img_scale = tuple(input_shape) + + from mmdet.datasets.pipelines import Compose + from mmocr.datasets import build_dataset # noqa: F401 + test_pipeline = Compose(test_pipeline) + + data_list = [] + for img in imgs: + # prepare data + if is_ndarray: + # directly add img + data = dict(img=img) + else: + # add information into dict + data = dict(img_info=dict(filename=img), img_prefix=None) + + # build the data pipeline + data = test_pipeline(data) + # get tensor from list to stack for batch mode (text detection) + data_list.append(data) + + if isinstance(data_list[0]['img'], list) and len(data_list) > 1: + raise Exception('aug test does not support ' + f'inference with batch size ' + f'{len(data_list)}') + + data = collate(data_list, samples_per_gpu=len(imgs)) + + # process img_metas + if isinstance(data['img_metas'], list): + data['img_metas'] = [ + img_metas.data[0] for img_metas in data['img_metas'] + ] + else: + data['img_metas'] = data['img_metas'].data + + if isinstance(data['img'], list): + data['img'] = [img.data for img in data['img']] + if isinstance(data['img'][0], list): + data['img'] = [img[0] for img in data['img']] + else: + data['img'] = data['img'].data + + if self.device != 'cpu': + data = scatter(data, [self.device])[0] + + return data, data['img'] + + def visualize(self, + model: nn.Module, + image: Union[str, np.ndarray], + result: list, + output_file: str, + window_name: str = '', + show_result: bool = False): + """Visualize predictions of a model. + + Args: + model (nn.Module): Input model. + image (str | np.ndarray): Input image to draw predictions on. + result (list): A list of predictions. + output_file (str): Output file to save drawn image. + window_name (str): The name of visualization window. Defaults to + an empty string. + show_result (bool): Whether to show result in windows, defaults + to `False`. + """ + show_img = mmcv.imread(image) if isinstance(image, str) else image + output_file = None if show_result else output_file + model.show_result( + show_img, + result, + out_file=output_file, + win_name=window_name, + show=show_result) + + @staticmethod + def run_inference(model: nn.Module, + model_inputs: Dict[str, torch.Tensor]) -> list: + """Run inference once for a segmentation model of mmseg. + + Args: + model (nn.Module): Input model. + model_inputs (dict): A dict containing model inputs tensor and + meta info. + + Returns: + list: The predictions of model inference. + """ + return model(**model_inputs, return_loss=False, rescale=True) + + @staticmethod + def get_partition_cfg(partition_type: str) -> Dict: + """Get a certain partition config. + + Args: + partition_type (str): A string specifying partition type. + + Returns: + dict: A dictionary of partition config. + """ + raise NotImplementedError('Not supported yet.') + + @staticmethod + def get_tensor_from_input(input_data: Dict[str, Any]) -> torch.Tensor: + """Get input tensor from input data. + + Args: + input_data (dict): Input data containing meta info and image + tensor. + Returns: + torch.Tensor: An image in `Tensor`. + """ + if isinstance(input_data['img'], DataContainer): + return input_data['img'].data[0] + return input_data['img'][0] + + @staticmethod + def evaluate_outputs(model_cfg, + outputs: Sequence, + dataset: Dataset, + metrics: Optional[str] = None, + out: Optional[str] = None, + metric_options: Optional[dict] = None, + format_only: bool = False): + """Perform post-processing to predictions of model. + + Args: + outputs (Sequence): A list of predictions of model inference. + dataset (Dataset): Input dataset to run test. + model_cfg (mmcv.Config): The model config. + metrics (str): Evaluation metrics, which depends on + the codebase and the dataset, e.g., e.g., "acc" for text + recognition, and "hmean-iou" for text detection. + out (str): Output result file in pickle format, defaults to `None`. + metric_options (dict): Custom options for evaluation, will be + kwargs for dataset.evaluate() function. Defaults to `None`. + format_only (bool): Format the output results without perform + evaluation. It is useful when you want to format the result + to a specific format and submit it to the test server. Defaults + to `False`. + """ + if out: + logging.info(f'\nwriting results to {out}') + mmcv.dump(outputs, out) + kwargs = {} if metric_options is None else metric_options + if format_only: + dataset.format_results(outputs, **kwargs) + if metrics: + eval_kwargs = model_cfg.get('evaluation', {}).copy() + # hard-code way to remove EvalHook args + for key in [ + 'interval', 'tmpdir', 'start', 'gpu_collect', 'save_best', + 'rule' + ]: + eval_kwargs.pop(key, None) + eval_kwargs.update(dict(metric=metrics, **kwargs)) + print(dataset.evaluate(outputs, **eval_kwargs)) diff --git a/mmdeploy/codebase/mmocr/deploy/text_detection_model.py b/mmdeploy/codebase/mmocr/deploy/text_detection_model.py new file mode 100644 index 0000000000..812f00f3ee --- /dev/null +++ b/mmdeploy/codebase/mmocr/deploy/text_detection_model.py @@ -0,0 +1,165 @@ +from typing import List, Sequence, Union + +import mmcv +import numpy as np +import torch +from mmocr.models.builder import build_head +from mmocr.models.textdet import TextDetectorMixin + +from mmdeploy.codebase.base import BaseBackendModel +from mmdeploy.utils import Backend, get_backend, get_onnx_config, load_config + + +class End2EndModel(BaseBackendModel): + """End to end model for inference of text detection. + + Args: + backend (Backend): The backend enum, specifying backend type. + backend_files (Sequence[str]): Paths to all required backend files(e.g. + '.onnx' for ONNX Runtime, '.param' and '.bin' for ncnn). + device (str): A string represents device type. + deploy_cfg (str | mmcv.Config): Deployment config file or loaded Config + object. + model_cfg (str | mmcv.Config): Model config file or loaded Config + object. + """ + + def __init__( + self, + backend: Backend, + backend_files: Sequence[str], + device: str, + deploy_cfg: Union[str, mmcv.Config] = None, + model_cfg: Union[str, mmcv.Config] = None, + ): + super(End2EndModel, self).__init__() + model_cfg, deploy_cfg = load_config(model_cfg, deploy_cfg) + self.deploy_cfg = deploy_cfg + self.show_score = False + self.bbox_head = build_head(model_cfg.model.bbox_head) + self._init_wrapper( + backend=backend, backend_files=backend_files, device=device) + + def _init_wrapper(self, backend: Backend, backend_files: Sequence[str], + device: str): + """Initialize the wrapper of backends. + + Args: + backend (Backend): The backend enum, specifying backend type. + backend_files (Sequence[str]): Paths to all required backend files + (e.g. .onnx' for ONNX Runtime, '.param' and '.bin' for ncnn). + device (str): A string represents device type. + """ + onnx_config = get_onnx_config(self.deploy_cfg) + output_names = onnx_config['output_names'] + self.wrapper = BaseBackendModel._build_wrapper( + backend=backend, + backend_files=backend_files, + device=device, + output_names=output_names) + + def forward(self, img: Sequence[torch.Tensor], + img_metas: Sequence[Sequence[dict]], *args, **kwargs) -> list: + """Run forward inference. + + Args: + img (Sequence[torch.Tensor]): A list contains input image(s) + in [N x C x H x W] format. + img_metas (Sequence[Sequence[dict]]): A list of meta info for + image(s). + + Returns: + list: A list contains predictions. + """ + input_img = img[0].contiguous() + img_metas = img_metas[0] + outputs = self.forward_test(input_img, img_metas, *args, **kwargs) + rescale = kwargs.get('rescale', False) + if len(img_metas) > 1: + boundaries = [ + self.bbox_head.get_boundary( + *(outputs[i].unsqueeze(0)), [img_metas[i]], + rescale=rescale) for i in range(len(img_metas)) + ] + + else: + boundaries = [ + self.bbox_head.get_boundary( + *outputs, img_metas, rescale=rescale) + ] + return boundaries + + def forward_test(self, imgs: torch.Tensor, *args, **kwargs) -> \ + List[np.ndarray]: + """The interface for forward test. + + Args: + imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. + + Returns: + List[np.ndarray]: A list of predictions of input images. + """ + outputs = self.wrapper({'input': imgs}) + outputs = self.wrapper.output_to_list(outputs) + return outputs + + def show_result(self, + img: np.ndarray, + result: list, + win_name: str, + show: bool = True, + score_thr: float = 0.3, + out_file: str = None): + """Show predictions of segmentation. + Args: + img: (np.ndarray): Input image to draw predictions. + result (list): A list of predictions. + win_name (str): The name of visualization window. + show (bool): Whether to show plotted image in windows. Defaults to + `True`. + score_thr: (float): The thresh of score. Defaults to `0.3`. + out_file (str): Output image file to save drawn predictions. + + Returns: + np.ndarray: Drawn image, only if not `show` or `out_file`. + """ + return TextDetectorMixin.show_result( + self, + img, + result, + score_thr=score_thr, + show=show, + win_name=win_name, + out_file=out_file) + + +def build_text_detection_model(model_files: Sequence[str], + model_cfg: Union[str, mmcv.Config], + deploy_cfg: Union[str, mmcv.Config], + device: str, **kwargs): + """Build text detection model for different backends. + + Args: + model_files (Sequence[str]): Input model file(s). + model_cfg (str | mmcv.Config): Input model config file or Config + object. + deploy_cfg (str | mmcv.Config): Input deployment config file or + Config object. + device (str): Device to input model. + + Returns: + BaseBackendModel: Text detector for a configured backend. + """ + # load cfg if necessary + deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg) + + backend = get_backend(deploy_cfg) + backend_text_detector = End2EndModel( + backend, + model_files, + device, + deploy_cfg=deploy_cfg, + model_cfg=model_cfg, + **kwargs) + + return backend_text_detector diff --git a/mmdeploy/codebase/mmocr/deploy/text_recognition.py b/mmdeploy/codebase/mmocr/deploy/text_recognition.py new file mode 100644 index 0000000000..c77a3d540a --- /dev/null +++ b/mmdeploy/codebase/mmocr/deploy/text_recognition.py @@ -0,0 +1,264 @@ +import logging +from typing import Any, Dict, Optional, Sequence, Tuple, Union + +import mmcv +import numpy as np +import torch +from mmcv.parallel import DataContainer, collate, scatter +from mmdet.datasets import replace_ImageToTensor +from torch import nn +from torch.utils.data import Dataset + +from mmdeploy.codebase.base import BaseTask +from mmdeploy.utils import Task +from .mmocr import MMOCR_TASK + + +@MMOCR_TASK.register_module(Task.TEXT_RECOGNITION.value) +class TextRecognition(BaseTask): + """Text detection task class. + + Args: + model_cfg (mmcv.Config): Original PyTorch model config file. + deploy_cfg (mmcv.Config): Loaded deployment config object. + device (str): A string represents device type. + """ + + def __init__(self, model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str): + super(TextRecognition, self).__init__(model_cfg, deploy_cfg, device) + + def init_backend_model(self, + model_files: Optional[str] = None, + **kwargs) -> torch.nn.Module: + """Initialize backend model. + + Args: + model_files (Sequence[str]): Input model files. + + Returns: + nn.Module: An initialized backend model. + """ + from .text_recognition_model import build_text_recognition_model + model = build_text_recognition_model( + model_files, self.model_cfg, self.deploy_cfg, device=self.device) + return model.eval() + + def init_pytorch_model(self, + model_checkpoint: Optional[str] = None, + cfg_options: Optional[Dict] = None, + **kwargs) -> torch.nn.Module: + """Initialize torch model. + + Args: + model_checkpoint (str): The checkpoint file of torch model, + defaults to `None`. + cfg_options (dict): Optional config key-pair parameters. + + Returns: + nn.Module: An initialized torch model generated by OpenMMLab + codebases. + """ + from mmocr.apis import init_detector + model = init_detector(self.model_cfg, model_checkpoint, self.device, + cfg_options) + + return model.eval() + + def create_input(self, + imgs: Union[str, np.ndarray], + input_shape: Sequence[int] = None) \ + -> Tuple[Dict, torch.Tensor]: + """Create input for segmentor. + + Args: + imgs (str | np.ndarray): Input image(s), accepted data type are + `str`, `np.ndarray`. + input_shape (list[int]): A list of two integer in (width, height) + format specifying input shape. Defaults to `None`. + + Returns: + tuple: (data, img), meta information for the input image and input. + """ + if isinstance(imgs, (list, tuple)): + if not isinstance(imgs[0], (np.ndarray, str)): + raise AssertionError('imgs must be strings or numpy arrays') + + elif isinstance(imgs, (np.ndarray, str)): + imgs = [imgs] + else: + raise AssertionError('imgs must be strings or numpy arrays') + + if self.model_cfg.data.test['type'] == 'ConcatDataset': + self.model_cfg.data.test.pipeline = \ + self.model_cfg.data.test['datasets'][0].pipeline + + is_ndarray = isinstance(imgs[0], np.ndarray) + + if is_ndarray: + self.model_cfg.data.test.pipeline[0].type = 'LoadImageFromNdarray' + + test_pipeline = self.model_cfg.data.test.pipeline + test_pipeline = replace_ImageToTensor(test_pipeline) + # for static exporting + if input_shape is not None: + test_pipeline[1].img_scale = tuple(input_shape) + test_pipeline[1].transforms[0].keep_ratio = False + test_pipeline[1].transforms[0].img_scale = tuple(input_shape) + + from mmdet.datasets.pipelines import Compose + from mmocr.datasets import build_dataset # noqa: F401 + test_pipeline = Compose(test_pipeline) + + data_list = [] + for img in imgs: + # prepare data + if is_ndarray: + # directly add img + data = dict(img=img) + else: + # add information into dict + data = dict(img_info=dict(filename=img), img_prefix=None) + + # build the data pipeline + data = test_pipeline(data) + # get tensor from list to stack for batch mode (text detection) + data_list.append(data) + + if isinstance(data_list[0]['img'], list) and len(data_list) > 1: + raise Exception('aug test does not support ' + f'inference with batch size ' + f'{len(data_list)}') + + data = collate(data_list, samples_per_gpu=len(imgs)) + + # process img_metas + if isinstance(data['img_metas'], list): + data['img_metas'] = [ + img_metas.data[0] for img_metas in data['img_metas'] + ] + else: + data['img_metas'] = data['img_metas'].data + + if isinstance(data['img'], list): + data['img'] = [img.data for img in data['img']] + if isinstance(data['img'][0], list): + data['img'] = [img[0] for img in data['img']] + else: + data['img'] = data['img'].data + + if self.device != 'cpu': + data = scatter(data, [self.device])[0] + + return data, data['img'] + + def visualize(self, + model: nn.Module, + image: Union[str, np.ndarray], + result: list, + output_file: str, + window_name: str = '', + show_result: bool = False): + """Visualize predictions of a model. + + Args: + model (nn.Module): Input model. + image (str | np.ndarray): Input image to draw predictions on. + result (list): A list of predictions. + output_file (str): Output file to save drawn image. + window_name (str): The name of visualization window. Defaults to + an empty string. + show_result (bool): Whether to show result in windows, defaults + to `False`. + """ + show_img = mmcv.imread(image) if isinstance(image, str) else image + output_file = None if show_result else output_file + model.show_result( + show_img, + result, + out_file=output_file, + win_name=window_name, + show=show_result) + + @staticmethod + def run_inference(model: nn.Module, + model_inputs: Dict[str, torch.Tensor]) -> list: + """Run inference once for a segmentation model of mmseg. + + Args: + model (nn.Module): Input model. + model_inputs (dict): A dict containing model inputs tensor and + meta info. + + Returns: + list: The predictions of model inference. + """ + return model(**model_inputs, return_loss=False, rescale=True) + + @staticmethod + def get_partition_cfg(partition_type: str) -> Dict: + """Get a certain partition config. + + Args: + partition_type (str): A string specifying partition type. + + Returns: + dict: A dictionary of partition config. + """ + raise NotImplementedError('Not supported yet.') + + @staticmethod + def get_tensor_from_input(input_data: Dict[str, Any]) -> torch.Tensor: + """Get input tensor from input data. + + Args: + input_data (dict): Input data containing meta info and image + tensor. + Returns: + torch.Tensor: An image in `Tensor`. + """ + if isinstance(input_data['img'], DataContainer): + return input_data['img'].data[0] + return input_data['img'][0] + + @staticmethod + def evaluate_outputs(model_cfg: mmcv.Config, + outputs: Sequence, + dataset: Dataset, + metrics: Optional[str] = None, + out: Optional[str] = None, + metric_options: Optional[dict] = None, + format_only: bool = False): + """Perform post-processing to predictions of model. + + Args: + model_cfg (mmcv.Config): The model config. + outputs (list): A list of predictions of model inference. + dataset (Dataset): Input dataset to run test. + metrics (str): Evaluation metrics, which depends on + the codebase and the dataset, e.g., e.g., "acc" for text + recognition, and "hmean-iou" for text detection. + out (str): Output result file in pickle format, defaults to `None`. + metric_options (dict): Custom options for evaluation, will be + kwargs for dataset.evaluate() function. Defaults to `None`. + format_only (bool): Format the output results without perform + evaluation. It is useful when you want to format the result + to a specific format and submit it to the test server. Defaults + to `False`. + """ + if out: + logging.info(f'\nwriting results to {out}') + mmcv.dump(outputs, out) + kwargs = {} if metric_options is None else metric_options + if format_only: + dataset.format_results(outputs, **kwargs) + if metrics: + eval_kwargs = model_cfg.get('evaluation', {}).copy() + # hard-code way to remove EvalHook args + for key in [ + 'interval', 'tmpdir', 'start', 'gpu_collect', 'save_best', + 'rule' + ]: + eval_kwargs.pop(key, None) + eval_kwargs.update(dict(metric=metrics, **kwargs)) + print(dataset.evaluate(outputs, **eval_kwargs)) diff --git a/mmdeploy/codebase/mmocr/deploy/text_recognition_model.py b/mmdeploy/codebase/mmocr/deploy/text_recognition_model.py new file mode 100644 index 0000000000..9f6b74017a --- /dev/null +++ b/mmdeploy/codebase/mmocr/deploy/text_recognition_model.py @@ -0,0 +1,172 @@ +from typing import List, Sequence, Union + +import mmcv +import numpy as np +import torch +from mmocr.models.builder import build_convertor +from mmocr.models.textrecog import BaseRecognizer + +from mmdeploy.codebase.base import BaseBackendModel +from mmdeploy.utils import Backend, get_backend, get_onnx_config, load_config + + +class End2EndModel(BaseBackendModel): + """End to end model for inference of text detection. + + Args: + backend (Backend): The backend enum, specifying backend type. + backend_files (Sequence[str]): Paths to all required backend files(e.g. + '.onnx' for ONNX Runtime, '.param' and '.bin' for ncnn). + device (str): A string represents device type. + deploy_cfg (str | mmcv.Config): Deployment config file or loaded Config + object. + model_cfg (str | mmcv.Config): Model config file or loaded Config + object. + """ + + def __init__( + self, + backend: Backend, + backend_files: Sequence[str], + device: str, + deploy_cfg: Union[str, mmcv.Config] = None, + model_cfg: Union[str, mmcv.Config] = None, + ): + super(End2EndModel, self).__init__() + model_cfg, deploy_cfg = load_config(model_cfg, deploy_cfg) + self.deploy_cfg = deploy_cfg + self.show_score = False + label_convertor = model_cfg.model.label_convertor + assert label_convertor is not None, 'model_cfg contains no label ' + 'convertor' + max_seq_len = 40 # default value in EncodeDecodeRecognizer of mmocr + label_convertor.update(max_seq_len=max_seq_len) + self.label_convertor = build_convertor(label_convertor) + self._init_wrapper( + backend=backend, backend_files=backend_files, device=device) + + def _init_wrapper(self, backend: Backend, backend_files: Sequence[str], + device: str): + """Initialize the wrapper of backends. + + Args: + backend (Backend): The backend enum, specifying backend type. + backend_files (Sequence[str]): Paths to all required backend files + (e.g. .onnx' for ONNX Runtime, '.param' and '.bin' for ncnn). + device (str): A string represents device type. + """ + onnx_config = get_onnx_config(self.deploy_cfg) + output_names = onnx_config['output_names'] + self.wrapper = BaseBackendModel._build_wrapper( + backend=backend, + backend_files=backend_files, + device=device, + output_names=output_names) + + def forward(self, img: Sequence[torch.Tensor], + img_metas: Sequence[Sequence[dict]], *args, **kwargs): + """Run forward inference. + + Args: + imgs (torch.Tensor | Sequence[torch.Tensor]): Image input tensor. + img_metas (Sequence[dict]): List of image information. + + Returns: + list[str]: Text label result of each image. + """ + if isinstance(img, list): + for idx, each_img in enumerate(img): + if each_img.dim() == 3: + img[idx] = each_img.unsqueeze(0) + img = img[0] # avoid aug_test + img_metas = img_metas[0] + else: + if len(img_metas) == 1 and isinstance(img_metas[0], list): + img_metas = img_metas[0] + + return self.forward_test(img, img_metas, **kwargs) + + def forward_test(self, imgs: torch.Tensor, + img_metas: Sequence[Sequence[dict]], *args, **kwargs) -> \ + List[np.ndarray]: + """The interface for forward test. + + Args: + imgs (torch.Tensor): Image input tensor. + img_metas (Sequence[dict]): List of image information. + + Returns: + list[str]: Text label result of each image. + """ + pred = self.wrapper({'input': imgs})['output'] + label_indexes, label_scores = self.label_convertor.tensor2idx( + pred, img_metas) + label_strings = self.label_convertor.idx2str(label_indexes) + + # flatten batch results + results = [] + for string, score in zip(label_strings, label_scores): + results.append(dict(text=string, score=score)) + + return results + + def show_result(self, + img: np.ndarray, + result: list, + win_name: str, + show: bool = True, + score_thr: float = 0.3, + out_file: str = None): + """Show predictions of segmentation. + Args: + img: (np.ndarray): Input image to draw predictions. + result (list): A list of predictions. + win_name (str): The name of visualization window. + show (bool): Whether to show plotted image in windows. Defaults to + `True`. + score_thr: (float): The thresh of score. Defaults to `0.3`. + out_file (str): Output image file to save drawn predictions. + + Returns: + np.ndarray: Drawn image, only if not `show` or `out_file`. + """ + return BaseRecognizer.show_result( + self, + img, + result, + score_thr=score_thr, + show=show, + win_name=win_name, + out_file=out_file) + + +def build_text_recognition_model(model_files: Sequence[str], + model_cfg: Union[str, mmcv.Config], + deploy_cfg: Union[str, mmcv.Config], + device: str, **kwargs): + """Build text recognition model for different backends. + + Args: + model_files (Sequence[str]): Input model file(s). + model_cfg (str | mmcv.Config): Input model config file or Config + object. + deploy_cfg (str | mmcv.Config): Input deployment config file or + Config object. + device (str): Device to input model. + + Returns: + BaseBackendModel: Text recognizer for a configured backend. + """ + # load cfg if necessary + deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg) + + backend = get_backend(deploy_cfg) + backend_text_recognizer = End2EndModel( + backend, + model_files, + device, + deploy_cfg=deploy_cfg, + model_cfg=model_cfg, + **kwargs) + + return backend_text_recognizer diff --git a/mmdeploy/codebase/mmocr/models/__init__.py b/mmdeploy/codebase/mmocr/models/__init__.py new file mode 100644 index 0000000000..dd24cc8e19 --- /dev/null +++ b/mmdeploy/codebase/mmocr/models/__init__.py @@ -0,0 +1,2 @@ +from .text_detection import * # noqa: F401,F403 +from .text_recognition import * # noqa: F401,F403 diff --git a/mmdeploy/codebase/mmocr/models/text_detection/__init__.py b/mmdeploy/codebase/mmocr/models/text_detection/__init__.py new file mode 100644 index 0000000000..fb8a50e832 --- /dev/null +++ b/mmdeploy/codebase/mmocr/models/text_detection/__init__.py @@ -0,0 +1,6 @@ +from .fpn_cat import fpnc__forward__tensorrt +from .single_stage_text_detector import single_stage_text_detector__simple_test + +__all__ = [ + 'fpnc__forward__tensorrt', 'single_stage_text_detector__simple_test' +] diff --git a/mmdeploy/mmocr/models/textdet/necks/fpn_cat.py b/mmdeploy/codebase/mmocr/models/text_detection/fpn_cat.py similarity index 100% rename from mmdeploy/mmocr/models/textdet/necks/fpn_cat.py rename to mmdeploy/codebase/mmocr/models/text_detection/fpn_cat.py diff --git a/mmdeploy/mmocr/models/textdet/detectors/single_stage_text_detector.py b/mmdeploy/codebase/mmocr/models/text_detection/single_stage_text_detector.py similarity index 100% rename from mmdeploy/mmocr/models/textdet/detectors/single_stage_text_detector.py rename to mmdeploy/codebase/mmocr/models/text_detection/single_stage_text_detector.py diff --git a/mmdeploy/codebase/mmocr/models/text_recognition/__init__.py b/mmdeploy/codebase/mmocr/models/text_recognition/__init__.py new file mode 100644 index 0000000000..dd0eb24297 --- /dev/null +++ b/mmdeploy/codebase/mmocr/models/text_recognition/__init__.py @@ -0,0 +1,13 @@ +from .base import base_recognizer__forward +from .crnn_decoder import crnndecoder__forward_train__ncnn +from .encode_decode_recognizer import encode_decode_recognizer__simple_test +from .lstm_layer import bidirectionallstm__forward__ncnn +from .sar import SARNet +from .sar_decoder import * # noqa: F401,F403 +from .sar_encoder import sar_encoder__forward + +__all__ = [ + 'base_recognizer__forward', 'crnndecoder__forward_train__ncnn', + 'encode_decode_recognizer__simple_test', + 'bidirectionallstm__forward__ncnn', 'sar_encoder__forward', 'SARNet' +] diff --git a/mmdeploy/mmocr/models/textrecog/recognizer/base.py b/mmdeploy/codebase/mmocr/models/text_recognition/base.py similarity index 100% rename from mmdeploy/mmocr/models/textrecog/recognizer/base.py rename to mmdeploy/codebase/mmocr/models/text_recognition/base.py diff --git a/mmdeploy/mmocr/models/textrecog/decoders/crnn_decoder.py b/mmdeploy/codebase/mmocr/models/text_recognition/crnn_decoder.py similarity index 100% rename from mmdeploy/mmocr/models/textrecog/decoders/crnn_decoder.py rename to mmdeploy/codebase/mmocr/models/text_recognition/crnn_decoder.py diff --git a/mmdeploy/mmocr/models/textrecog/recognizer/encode_decode_recognizer.py b/mmdeploy/codebase/mmocr/models/text_recognition/encode_decode_recognizer.py similarity index 100% rename from mmdeploy/mmocr/models/textrecog/recognizer/encode_decode_recognizer.py rename to mmdeploy/codebase/mmocr/models/text_recognition/encode_decode_recognizer.py diff --git a/mmdeploy/mmocr/models/textrecog/layers/lstm_layer.py b/mmdeploy/codebase/mmocr/models/text_recognition/lstm_layer.py similarity index 100% rename from mmdeploy/mmocr/models/textrecog/layers/lstm_layer.py rename to mmdeploy/codebase/mmocr/models/text_recognition/lstm_layer.py diff --git a/mmdeploy/mmocr/models/textrecog/recognizer/sar.py b/mmdeploy/codebase/mmocr/models/text_recognition/sar.py similarity index 97% rename from mmdeploy/mmocr/models/textrecog/recognizer/sar.py rename to mmdeploy/codebase/mmocr/models/text_recognition/sar.py index 9317e3d758..186c8a6c5c 100644 --- a/mmdeploy/mmocr/models/textrecog/recognizer/sar.py +++ b/mmdeploy/codebase/mmocr/models/text_recognition/sar.py @@ -4,8 +4,8 @@ import torch.nn as nn from mmdeploy.core import MODULE_REWRITER -from mmdeploy.mmocr.utils.cfg_utils import get_resize_ocr from mmdeploy.utils import is_dynamic_shape +from ..utils import get_resize_ocr @MODULE_REWRITER.register_rewrite_module( diff --git a/mmdeploy/mmocr/models/textrecog/decoders/sar_decoder.py b/mmdeploy/codebase/mmocr/models/text_recognition/sar_decoder.py similarity index 100% rename from mmdeploy/mmocr/models/textrecog/decoders/sar_decoder.py rename to mmdeploy/codebase/mmocr/models/text_recognition/sar_decoder.py diff --git a/mmdeploy/mmocr/models/textrecog/encoders/sar_encoder.py b/mmdeploy/codebase/mmocr/models/text_recognition/sar_encoder.py similarity index 100% rename from mmdeploy/mmocr/models/textrecog/encoders/sar_encoder.py rename to mmdeploy/codebase/mmocr/models/text_recognition/sar_encoder.py diff --git a/mmdeploy/mmocr/utils/cfg_utils.py b/mmdeploy/codebase/mmocr/models/utils.py similarity index 100% rename from mmdeploy/mmocr/utils/cfg_utils.py rename to mmdeploy/codebase/mmocr/models/utils.py diff --git a/mmdeploy/mmocr/__init__.py b/mmdeploy/codebase/mmseg/__init__.py similarity index 50% rename from mmdeploy/mmocr/__init__.py rename to mmdeploy/codebase/mmseg/__init__.py index d2b62e1cb6..33b69c74df 100644 --- a/mmdeploy/mmocr/__init__.py +++ b/mmdeploy/codebase/mmseg/__init__.py @@ -1,2 +1,2 @@ -from .export import * # noqa: F401,F403 +from .deploy import * # noqa: F401,F403 from .models import * # noqa: F401,F403 diff --git a/mmdeploy/codebase/mmseg/deploy/__init__.py b/mmdeploy/codebase/mmseg/deploy/__init__.py new file mode 100644 index 0000000000..3bf727b702 --- /dev/null +++ b/mmdeploy/codebase/mmseg/deploy/__init__.py @@ -0,0 +1,5 @@ +from .mmsegmentation import MMSegmentation +from .segmentation import Segmentation +from .utils import convert_syncbatchnorm + +__all__ = ['convert_syncbatchnorm', 'MMSegmentation', 'Segmentation'] diff --git a/mmdeploy/codebase/mmseg/deploy/mmsegmentation.py b/mmdeploy/codebase/mmseg/deploy/mmsegmentation.py new file mode 100644 index 0000000000..a1ee9f9a64 --- /dev/null +++ b/mmdeploy/codebase/mmseg/deploy/mmsegmentation.py @@ -0,0 +1,141 @@ +from typing import Optional, Union + +import mmcv +import torch +from mmcv.utils import Registry +from torch.utils.data import DataLoader, Dataset + +from mmdeploy.codebase.base import CODEBASE, BaseTask, MMCodebase +from mmdeploy.utils import Codebase, get_task_type + + +def __build_mmseg_task(model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str, registry: Registry) -> BaseTask: + task = get_task_type(deploy_cfg) + return registry.module_dict[task.value](model_cfg, deploy_cfg, device) + + +MMSEG_TASK = Registry('mmseg_tasks', build_func=__build_mmseg_task) + + +@CODEBASE.register_module(Codebase.MMSEG.value) +class MMSegmentation(MMCodebase): + """mmsegmentation codebase class.""" + + task_registry = MMSEG_TASK + + def __init__(self): + super(MMSegmentation, self).__init__() + + @staticmethod + def build_task_processor(model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str): + """The interface to build the task processors of mmseg. + + Args: + model_cfg (str | mmcv.Config): Model config file. + deploy_cfg (str | mmcv.Config): Deployment config file. + device (str): A string specifying device type. + + Returns: + BaseTask: A task processor. + """ + return MMSEG_TASK.build(model_cfg, deploy_cfg, device) + + @staticmethod + def build_dataset(dataset_cfg: Union[str, mmcv.Config], + dataset_type: str = 'val', + **kwargs) -> Dataset: + """Build dataset for segmentation. + + Args: + dataset_cfg (str | mmcv.Config): The input dataset config. + dataset_type (str): A string represents dataset type, e.g.: 'train' + , 'test', 'val'. Defaults to 'val'. + + Returns: + Dataset: A PyTorch dataset. + """ + from mmseg.datasets import build_dataset as build_dataset_mmseg + + assert dataset_type in dataset_cfg.data + data_cfg = dataset_cfg.data[dataset_type] + dataset = build_dataset_mmseg(data_cfg) + return dataset + + @staticmethod + def build_dataloader(dataset: Dataset, + samples_per_gpu: int, + workers_per_gpu: int, + num_gpus: int = 1, + dist: bool = False, + shuffle: bool = False, + seed: Optional[int] = None, + drop_last: bool = False, + pin_memory: bool = True, + persistent_workers: bool = True, + **kwargs) -> DataLoader: + """Build dataloader for segmentation. + + Args: + dataset (Dataset): Input dataset. + samples_per_gpu (int): Number of training samples on each GPU, i.e. + ,batch size of each GPU. + workers_per_gpu (int): How many subprocesses to use for data + loading for each GPU. + num_gpus (int): Number of GPUs. Only used in non-distributed + training. dist (bool): Distributed training/test or not. + Defaults to `False`. + shuffle (bool): Whether to shuffle the data at every epoch. + Defaults to `False`. + seed (int): An integer set to be seed. Default is `None`. + drop_last (bool): Whether to drop the last incomplete batch in + epoch. Default to `False`. + pin_memory (bool): Whether to use pin_memory in DataLoader. + Default is `True`. + persistent_workers (bool): If `True`, the data loader will not + shutdown the worker processes after a dataset has been + consumed once. This allows to maintain the workers Dataset + instances alive. The argument also has effect in + PyTorch>=1.7.0. Default is `True`. + kwargs: Any other keyword argument to be used to initialize + DataLoader. + + Returns: + DataLoader: A PyTorch dataloader. + """ + from mmseg.datasets import build_dataloader as build_dataloader_mmseg + return build_dataloader_mmseg( + dataset, + samples_per_gpu, + workers_per_gpu, + num_gpus=num_gpus, + dist=dist, + shuffle=shuffle, + seed=seed, + drop_last=drop_last, + pin_memory=pin_memory, + persistent_workers=persistent_workers, + **kwargs) + + @staticmethod + def single_gpu_test(model: torch.nn.Module, + data_loader: DataLoader, + show: bool = False, + out_dir: Optional[str] = None, + **kwargs): + """Run test with single gpu. + + Args: + model (torch.nn.Module): Input model from nn.Module. + data_loader (DataLoader): PyTorch data loader. + show (bool): Specifying whether to show plotted results. Defaults + to `False`. + out_dir (str): A directory to save results, defaults to `None`. + + Returns: + list: The prediction results. + """ + from mmseg.apis import single_gpu_test + outputs = single_gpu_test(model, data_loader, show, out_dir, **kwargs) + return outputs diff --git a/mmdeploy/codebase/mmseg/deploy/segmentation.py b/mmdeploy/codebase/mmseg/deploy/segmentation.py new file mode 100644 index 0000000000..8545c83555 --- /dev/null +++ b/mmdeploy/codebase/mmseg/deploy/segmentation.py @@ -0,0 +1,216 @@ +import logging +from typing import Any, Dict, Optional, Sequence, Tuple, Union + +import mmcv +import numpy as np +import torch +from torch.utils.data import Dataset + +from mmdeploy.codebase.base import BaseTask +from mmdeploy.utils import Task +from .mmsegmentation import MMSEG_TASK + + +@MMSEG_TASK.register_module(Task.SEGMENTATION.value) +class Segmentation(BaseTask): + """Segmentation task class. + + Args: + model_cfg (mmcv.Config): Original PyTorch model config file. + deploy_cfg (mmcv.Config): Deployment config file or loaded Config + object. + device (str): A string represents device type. + """ + + def __init__(self, model_cfg: mmcv.Config, deploy_cfg: mmcv.Config, + device: str): + super(Segmentation, self).__init__(model_cfg, deploy_cfg, device) + + def init_backend_model(self, + model_files: Optional[str] = None, + **kwargs) -> torch.nn.Module: + """Initialize backend model. + + Args: + model_files (Sequence[str]): Input model files. + + Returns: + nn.Module: An initialized backend model. + """ + from .segmentation_model import build_segmentation_model + model = build_segmentation_model( + model_files, self.model_cfg, self.deploy_cfg, device=self.device) + return model.eval() + + def init_pytorch_model(self, + model_checkpoint: Optional[str] = None, + cfg_options: Optional[Dict] = None, + **kwargs) -> torch.nn.Module: + """Initialize torch model. + + Args: + model_checkpoint (str): The checkpoint file of torch model, + defaults to `None`. + cfg_options (dict): Optional config key-pair parameters. + + Returns: + nn.Module: An initialized torch model generated by OpenMMLab + codebases. + """ + from mmseg.apis import init_segmentor + from mmdeploy.codebase.mmseg.deploy import convert_syncbatchnorm + model = init_segmentor(self.model_cfg, model_checkpoint, self.device) + model = convert_syncbatchnorm(model) + + return model.eval() + + def create_input(self, + imgs: Union[str, np.ndarray], + input_shape: Sequence[int] = None) \ + -> Tuple[Dict, torch.Tensor]: + """Create input for segmentor. + + Args: + imgs (Any): Input image(s), accepted data type are `str`, + `np.ndarray`, `torch.Tensor`. + input_shape (list[int]): A list of two integer in (width, height) + format specifying input shape. Defaults to `None`. + + Returns: + tuple: (data, img), meta information for the input image and input. + """ + from mmseg.apis.inference import LoadImage + from mmseg.datasets.pipelines import Compose + from mmcv.parallel import collate, scatter + + cfg = self.model_cfg.copy() + if not isinstance(imgs, (list, tuple)): + imgs = [imgs] + + if isinstance(imgs[0], np.ndarray): + cfg = cfg.copy() + # set loading pipeline type + cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam' + # for static exporting + if input_shape is not None: + cfg.data.test.pipeline[1]['img_scale'] = tuple(input_shape) + cfg.data.test.pipeline[1]['transforms'][0]['keep_ratio'] = False + cfg.data.test.pipeline = [LoadImage()] + cfg.data.test.pipeline[1:] + + test_pipeline = Compose(cfg.data.test.pipeline) + data_list = [] + for img in imgs: + # prepare data + data = dict(img=img) + # build the data pipeline + data = test_pipeline(data) + data_list.append(data) + + data = collate(data_list, samples_per_gpu=len(imgs)) + + data['img_metas'] = [ + img_metas.data[0] for img_metas in data['img_metas'] + ] + data['img'] = [img.data[0][None, :] for img in data['img']] + if self.device != 'cpu': + data = scatter(data, [self.device])[0] + + return data, data['img'] + + def visualize(self, + model, + image: Union[str, np.ndarray], + result: list, + output_file: str, + window_name: str = '', + show_result: bool = False, + opacity: float = 0.5): + """Visualize predictions of a model. + + Args: + model (nn.Module): Input model. + image (str | np.ndarray): Input image to draw predictions on. + result (list): A list of predictions. + output_file (str): Output file to save drawn image. + window_name (str): The name of visualization window. Defaults to + an empty string. + show_result (bool): Whether to show result in windows, defaults + to `False`. + opacity: (float): Opacity of painted segmentation map. + Defaults to `0.5`. + """ + show_img = mmcv.imread(image) if isinstance(image, str) else image + output_file = None if show_result else output_file + # Need to wrapper the result with list for mmseg + result = [result] + model.show_result( + show_img, + result, + out_file=output_file, + win_name=window_name, + show=show_result, + opacity=opacity) + + @staticmethod + def run_inference(model, model_inputs: Dict[str, torch.Tensor]): + """Run inference once for a segmentation model of mmseg. + + Args: + model (nn.Module): Input model. + model_inputs (dict): A dict containing model inputs tensor and + meta info. + + Returns: + list: The predictions of model inference. + """ + return model(**model_inputs, return_loss=False, rescale=True) + + @staticmethod + def get_partition_cfg(partition_type: str) -> Dict: + raise NotImplementedError('Not supported yet.') + + @staticmethod + def get_tensor_from_input(input_data: Dict[str, Any]) -> torch.Tensor: + """Get input tensor from input data. + + Args: + input_data (dict): Input data containing meta info and image + tensor. + Returns: + torch.Tensor: An image in `Tensor`. + """ + return input_data['img'][0] + + @staticmethod + def evaluate_outputs(model_cfg, + outputs: Sequence, + dataset: Dataset, + metrics: Optional[str] = None, + out: Optional[str] = None, + metric_options: Optional[dict] = None, + format_only: bool = False): + """Perform post-processing to predictions of model. + + Args: + outputs (list): A list of predictions of model inference. + dataset (Dataset): Input dataset to run test. + model_cfg (mmcv.Config): The model config. + metrics (str): Evaluation metrics, which depends on + the codebase and the dataset, e.g., e.g., "mIoU" for generic + datasets, and "cityscapes" for Cityscapes in mmseg. + out (str): Output result file in pickle format, defaults to `None`. + metric_options (dict): Custom options for evaluation, will be + kwargs for dataset.evaluate() function. Defaults to `None`. + format_only (bool): Format the output results without perform + evaluation. It is useful when you want to format the result + to a specific format and submit it to the test server. Defaults + to `False`. + """ + if out: + logging.info(f'\nwriting results to {out}') + mmcv.dump(outputs, out) + kwargs = {} if metric_options is None else metric_options + if format_only: + dataset.format_results(outputs, **kwargs) + if metrics: + dataset.evaluate(outputs, metrics, **kwargs) diff --git a/mmdeploy/codebase/mmseg/deploy/segmentation_model.py b/mmdeploy/codebase/mmseg/deploy/segmentation_model.py new file mode 100644 index 0000000000..b1a30bccaa --- /dev/null +++ b/mmdeploy/codebase/mmseg/deploy/segmentation_model.py @@ -0,0 +1,190 @@ +from typing import List, Sequence, Union + +import mmcv +import numpy as np +import torch +from mmseg.datasets import DATASETS +from mmseg.models.segmentors.base import BaseSegmentor +from mmseg.ops import resize + +from mmdeploy.codebase.base import BaseBackendModel +from mmdeploy.utils import Backend, get_backend, get_onnx_config, load_config + + +class End2EndModel(BaseBackendModel): + """End to end model for inference of segmentation. + + Args: + backend (Backend): The backend enum, specifying backend type. + backend_files (Sequence[str]): Paths to all required backend files(e.g. + '.onnx' for ONNX Runtime, '.param' and '.bin' for ncnn). + device (str): A string represents device type. + class_names (Sequence[str]): A list of string specifying class names. + palette (np.ndarray): The palette of segmentation map. + deploy_cfg (str | mmcv.Config): Deployment config file or loaded Config + object. + """ + + def __init__( + self, + backend: Backend, + backend_files: Sequence[str], + device: str, + class_names: Sequence[str], + palette: np.ndarray, + deploy_cfg: Union[str, mmcv.Config] = None, + ): + super(End2EndModel, self).__init__() + self.CLASSES = class_names + self.PALETTE = palette + self.deploy_cfg = deploy_cfg + self._init_wrapper( + backend=backend, backend_files=backend_files, device=device) + + def _init_wrapper(self, backend, backend_files, device): + onnx_config = get_onnx_config(self.deploy_cfg) + output_names = onnx_config['output_names'] + self.wrapper = BaseBackendModel._build_wrapper( + backend=backend, + backend_files=backend_files, + device=device, + output_names=output_names) + + def forward(self, img: Sequence[torch.Tensor], + img_metas: Sequence[Sequence[dict]], *args, **kwargs): + """Run forward inference. + + Args: + img (Sequence[torch.Tensor]): A list contains input image(s) + in [N x C x H x W] format. + img_metas (Sequence[Sequence[dict]]): A list of meta info for + image(s). + *args: Other arguments. + **kwargs: Other key-pair arguments. + + Returns: + list: A list contains predictions. + """ + input_img = img[0].contiguous() + outputs = self.forward_test(input_img, img_metas, *args, **kwargs) + seg_pred = outputs[0] + # whole mode supports dynamic shape + ori_shape = img_metas[0][0]['ori_shape'] + if not (ori_shape[0] == seg_pred.shape[-2] + and ori_shape[1] == seg_pred.shape[-1]): + seg_pred = torch.from_numpy(seg_pred).float() + seg_pred = resize( + seg_pred, size=tuple(ori_shape[:2]), mode='nearest') + seg_pred = seg_pred.long().detach().cpu().numpy() + # remove unnecessary dim + seg_pred = seg_pred.squeeze(1) + seg_pred = list(seg_pred) + return seg_pred + + def forward_test(self, imgs: torch.Tensor, *args, **kwargs) -> \ + List[np.ndarray]: + """The interface for forward test. + + Args: + imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. + + Returns: + List[np.ndarray]: A list of segmentation map. + """ + outputs = self.wrapper({'input': imgs}) + outputs = self.wrapper.output_to_list(outputs) + outputs = [out.detach().cpu().numpy() for out in outputs] + return outputs + + def show_result(self, + img: np.ndarray, + result: list, + win_name: str, + show: bool = True, + opacity: float = 0.5, + out_file: str = None): + """Show predictions of segmentation. + Args: + img: (np.ndarray): Input image to draw predictions. + result (list): A list of predictions. + win_name (str): The name of visualization window. + show (bool): Whether to show plotted image in windows. Defaults to + `True`. + opacity: (float): Opacity of painted segmentation map. + Defaults to `0.5`. + out_file (str): Output image file to save drawn predictions. + + Returns: + np.ndarray: Drawn image, only if not `show` or `out_file`. + """ + return BaseSegmentor.show_result( + self, + img, + result, + palette=self.PALETTE, + opacity=opacity, + show=show, + win_name=win_name, + out_file=out_file) + + +def get_classes_palette_from_config(model_cfg: Union[str, mmcv.Config]): + """Get class name and palette from config. + + Args: + model_cfg (str | mmcv.Config): Input model config file or + Config object. + Returns: + tuple(Sequence[str], np.ndarray): A list of string specifying names of + different class and the palette of segmentation map. + """ + # load cfg if necessary + model_cfg = load_config(model_cfg)[0] + + module_dict = DATASETS.module_dict + data_cfg = model_cfg.data + + if 'train' in data_cfg: + module = module_dict[data_cfg.train.type] + elif 'val' in data_cfg: + module = module_dict[data_cfg.val.type] + elif 'test' in data_cfg: + module = module_dict[data_cfg.test.type] + else: + raise RuntimeError(f'No dataset config found in: {model_cfg}') + + return module.CLASSES, module.PALETTE + + +def build_segmentation_model(model_files: Sequence[str], + model_cfg: Union[str, mmcv.Config], + deploy_cfg: Union[str, mmcv.Config], device: str, + **kwargs): + """Build object segmentation model for different backends. + + Args: + model_files (Sequence[str]): Input model file(s). + model_cfg (str | mmcv.Config): Input model config file or Config + object. + deploy_cfg (str | mmcv.Config): Input deployment config file or + Config object. + device (str): Device to input model. + + Returns: + BaseBackendModel: Segmentor for a configured backend. + """ + # load cfg if necessary + deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg) + + backend = get_backend(deploy_cfg) + class_names, palette = get_classes_palette_from_config(model_cfg) + backend_segmentor = End2EndModel( + backend, + model_files, + device, + class_names, + palette, + deploy_cfg=deploy_cfg, + **kwargs) + + return backend_segmentor diff --git a/mmdeploy/mmseg/export/onnx_utils.py b/mmdeploy/codebase/mmseg/deploy/utils.py similarity index 100% rename from mmdeploy/mmseg/export/onnx_utils.py rename to mmdeploy/codebase/mmseg/deploy/utils.py diff --git a/mmdeploy/mmseg/models/__init__.py b/mmdeploy/codebase/mmseg/models/__init__.py similarity index 100% rename from mmdeploy/mmseg/models/__init__.py rename to mmdeploy/codebase/mmseg/models/__init__.py diff --git a/mmdeploy/mmseg/models/decode_heads/__init__.py b/mmdeploy/codebase/mmseg/models/decode_heads/__init__.py similarity index 100% rename from mmdeploy/mmseg/models/decode_heads/__init__.py rename to mmdeploy/codebase/mmseg/models/decode_heads/__init__.py diff --git a/mmdeploy/mmseg/models/decode_heads/aspp_head.py b/mmdeploy/codebase/mmseg/models/decode_heads/aspp_head.py similarity index 100% rename from mmdeploy/mmseg/models/decode_heads/aspp_head.py rename to mmdeploy/codebase/mmseg/models/decode_heads/aspp_head.py diff --git a/mmdeploy/mmseg/models/decode_heads/psp_head.py b/mmdeploy/codebase/mmseg/models/decode_heads/psp_head.py similarity index 100% rename from mmdeploy/mmseg/models/decode_heads/psp_head.py rename to mmdeploy/codebase/mmseg/models/decode_heads/psp_head.py diff --git a/mmdeploy/mmseg/models/segmentors/__init__.py b/mmdeploy/codebase/mmseg/models/segmentors/__init__.py similarity index 100% rename from mmdeploy/mmseg/models/segmentors/__init__.py rename to mmdeploy/codebase/mmseg/models/segmentors/__init__.py diff --git a/mmdeploy/mmseg/models/segmentors/base.py b/mmdeploy/codebase/mmseg/models/segmentors/base.py similarity index 100% rename from mmdeploy/mmseg/models/segmentors/base.py rename to mmdeploy/codebase/mmseg/models/segmentors/base.py diff --git a/mmdeploy/mmseg/models/segmentors/encoder_decoder.py b/mmdeploy/codebase/mmseg/models/segmentors/encoder_decoder.py similarity index 100% rename from mmdeploy/mmseg/models/segmentors/encoder_decoder.py rename to mmdeploy/codebase/mmseg/models/segmentors/encoder_decoder.py diff --git a/mmdeploy/core/optimizers/extractor.py b/mmdeploy/core/optimizers/extractor.py index f15d1e4a2c..10f9bc5aea 100644 --- a/mmdeploy/core/optimizers/extractor.py +++ b/mmdeploy/core/optimizers/extractor.py @@ -4,7 +4,7 @@ from packaging import version -def parse_extractor_io_string(io_str): +def parse_extractor_io_string(io_str) -> tuple: """Parse IO string for extractor.""" name, io_type = io_str.split(':') assert io_type in ['input', 'output'] @@ -44,14 +44,14 @@ def impl(node_output_name, graph_input_nodes, reachable_nodes): impl(node_output_name, graph_input_nodes, reachable_nodes) -def create_extractor(model: onnx.ModelProto): +def create_extractor(model: onnx.ModelProto) -> onnx.utils.Extractor: """Create Extractor for ONNX. Args: model (onnx.ModelProto): An input onnx model. Returns: - Extractor: Extractor for the onnx. + onnx.utils.Extractor: Extractor for the onnx. """ assert version.parse(onnx.__version__) >= version.parse('1.8.0') # patch extractor diff --git a/mmdeploy/core/optimizers/function_marker.py b/mmdeploy/core/optimizers/function_marker.py index 500e5f2458..0a1c6f79da 100644 --- a/mmdeploy/core/optimizers/function_marker.py +++ b/mmdeploy/core/optimizers/function_marker.py @@ -1,10 +1,10 @@ import inspect -from typing import Any, Dict, Optional, Sequence +from typing import Any, Callable, Dict, Optional, Sequence import torch from mmdeploy.core.rewriters import FUNCTION_REWRITER -from mmdeploy.utils import cfg_apply_marks, get_codebase, get_partition_config +from mmdeploy.utils import cfg_apply_marks, get_partition_config MARK_FUNCTION_COUNT = dict() @@ -54,7 +54,7 @@ def symbolic(g, x, dtype, shape, func, func_id, type, name, id, attrs): return n @staticmethod - def forward(ctx, x, *args): + def forward(ctx, x, *args) -> torch.Tensor: """Run forward.""" return x @@ -71,19 +71,20 @@ def mark_symbolic(rewriter, g, x, *args): @FUNCTION_REWRITER.register_rewriter( 'mmdeploy.core.optimizers.function_marker.Mark.forward') def forward_of_mark(rewriter, ctx, x, dtype, shape, func, func_id, type, name, - id, attrs): + id, attrs) -> torch.Tensor: """Rewrite forward of mark op.""" deploy_cfg = rewriter.cfg # save calib data apply_marks = cfg_apply_marks(deploy_cfg) create_calib = getattr(rewriter, 'create_calib', False) if apply_marks and create_calib: - codebase = get_codebase(deploy_cfg) partition_params = get_partition_config(deploy_cfg) assert partition_params is not None, 'No partition config.' partition_type = partition_params['type'] - from mmdeploy.apis.utils import get_partition_cfg - partition_cfgs = get_partition_cfg(codebase, partition_type) + + from mmdeploy.apis import get_predefined_partition_cfg + partition_cfgs = get_predefined_partition_cfg(deploy_cfg, + partition_type) assert hasattr(rewriter, 'calib_file') for partition_id, partition_cfg in enumerate(partition_cfgs): @@ -123,7 +124,7 @@ def forward_of_mark(rewriter, ctx, x, dtype, shape, func, func_id, type, name, def mark_tensors(xs: Any, func: str, func_id: int, io_type: str, ctx: Any, - attrs: Dict, is_inspecting: bool, level: int): + attrs: Dict, is_inspecting: bool, level: int) -> tuple: """Add mark node recursively. Args: @@ -181,7 +182,7 @@ def impl(ys, prefix, level): def mark(func_name: Optional[str] = None, inputs: Optional[Sequence[str]] = None, outputs: Optional[Sequence[str]] = None, - **attrs): + **attrs) -> Callable: """The decorator used to add mark node. Mark node can be used to support model partition. diff --git a/mmdeploy/core/optimizers/optimize.py b/mmdeploy/core/optimizers/optimize.py index ed72282afb..57ce8cbe66 100644 --- a/mmdeploy/core/optimizers/optimize.py +++ b/mmdeploy/core/optimizers/optimize.py @@ -5,7 +5,7 @@ from onnx.helper import get_attribute_value -def attribute_to_dict(attr: onnx.AttributeProto): +def attribute_to_dict(attr: onnx.AttributeProto) -> Dict: """Convert onnx op attribute to dict. Args: @@ -23,7 +23,8 @@ def attribute_to_dict(attr: onnx.AttributeProto): return ret -def remove_nodes(model: onnx.ModelProto, predicate: Callable): +def remove_nodes(model: onnx.ModelProto, + predicate: Callable) -> onnx.ModelProto: """Remove nodes from ONNX model. Args: @@ -54,14 +55,14 @@ def remove_nodes(model: onnx.ModelProto, predicate: Callable): return model -def is_unused_mark(marks: Iterable[onnx.NodeProto]): +def is_unused_mark(marks: Iterable[onnx.NodeProto]) -> Callable: """Check whether a mark is unused. Args: marks (Iterable[onnx.NodeProto]): A list of onnx NodeProto. Returns: - bool: `True` if a mark node is in `marks`. + Callable: The function to check if a mark node is in `marks`. """ def f(node): @@ -75,14 +76,14 @@ def f(node): return f -def is_identity(node: onnx.NodeProto): +def is_identity(node: onnx.NodeProto) -> bool: """Check if an op is identity.""" return node.op_type == 'Identity' def get_new_name(attrs: Dict[str, str], mark_name: str = '', - name_map: Optional[Dict[str, str]] = None): + name_map: Optional[Dict[str, str]] = None) -> str: """Get new name for a node. Args: diff --git a/mmdeploy/core/rewriters/function_rewriter.py b/mmdeploy/core/rewriters/function_rewriter.py index a8d49b857a..175372b061 100644 --- a/mmdeploy/core/rewriters/function_rewriter.py +++ b/mmdeploy/core/rewriters/function_rewriter.py @@ -44,7 +44,7 @@ def __init__(self): self._registry = RewriterRegistry() def add_backend(self, backend: str): - """Add a beckend by calling the _registry.add_backend.""" + """Add a backend by calling the _registry.add_backend.""" self._registry.add_backend(backend) def register_rewriter(self, diff --git a/mmdeploy/core/rewriters/module_rewriter.py b/mmdeploy/core/rewriters/module_rewriter.py index adf20fbc21..7b45d4ae17 100644 --- a/mmdeploy/core/rewriters/module_rewriter.py +++ b/mmdeploy/core/rewriters/module_rewriter.py @@ -27,7 +27,7 @@ def __init__(self): self._registry = RewriterRegistry() def add_backend(self, backend: str): - """Add a beckend by calling the _registry.add_backend.""" + """Add a backend by calling the _registry.add_backend.""" self._registry.add_backend(backend) def register_rewrite_module(self, diff --git a/mmdeploy/core/rewriters/rewriter_manager.py b/mmdeploy/core/rewriters/rewriter_manager.py index 23eb424eff..4b0891b87a 100644 --- a/mmdeploy/core/rewriters/rewriter_manager.py +++ b/mmdeploy/core/rewriters/rewriter_manager.py @@ -28,12 +28,8 @@ def add_backend(self, backend: str): REWRITER_MANAGER = RewriterManager() -REWRITER_MANAGER.add_backend(Backend.ONNXRUNTIME.value) -REWRITER_MANAGER.add_backend(Backend.TENSORRT.value) -REWRITER_MANAGER.add_backend(Backend.NCNN.value) -REWRITER_MANAGER.add_backend(Backend.PPL.value) -REWRITER_MANAGER.add_backend(Backend.PYTORCH.value) -REWRITER_MANAGER.add_backend(Backend.OPENVINO.value) +for backend in Backend: + REWRITER_MANAGER.add_backend(backend.value) MODULE_REWRITER = REWRITER_MANAGER.module_rewriter FUNCTION_REWRITER = REWRITER_MANAGER.function_rewrite diff --git a/mmdeploy/core/rewriters/rewriter_utils.py b/mmdeploy/core/rewriters/rewriter_utils.py index f0eb8b4a87..e706b3b8a1 100644 --- a/mmdeploy/core/rewriters/rewriter_utils.py +++ b/mmdeploy/core/rewriters/rewriter_utils.py @@ -80,7 +80,7 @@ def decorator(object): return decorator -class ContextCaller(): +class ContextCaller: """A callable object used in RewriteContext. This class saves context variables as member variables. When a rewritten diff --git a/mmdeploy/core/rewriters/symbolic_rewriter.py b/mmdeploy/core/rewriters/symbolic_rewriter.py index 93fd640701..d5eb6f80ea 100644 --- a/mmdeploy/core/rewriters/symbolic_rewriter.py +++ b/mmdeploy/core/rewriters/symbolic_rewriter.py @@ -1,5 +1,5 @@ import logging -from typing import Dict, Optional, Sequence +from typing import Callable, Dict, Optional, Sequence from torch.autograd import Function from torch.onnx.symbolic_helper import parse_args @@ -36,7 +36,7 @@ def __init__(self) -> None: self._registry = RewriterRegistry() def add_backend(self, backend: str): - """Add a beckend by calling the _registry.add_backend.""" + """Add a backend by calling the _registry.add_backend.""" self._registry.add_backend(backend) def register_symbolic(self, @@ -44,7 +44,7 @@ def register_symbolic(self, backend: str = Backend.DEFAULT.value, is_pytorch: bool = False, arg_descriptors: Optional[Sequence[str]] = None, - **kwargs): + **kwargs) -> Callable: """The decorator of the custom symbolic. Args: diff --git a/mmdeploy/mmcls/apis/__init__.py b/mmdeploy/mmcls/apis/__init__.py deleted file mode 100644 index 1cfce5d45d..0000000000 --- a/mmdeploy/mmcls/apis/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .inference import build_classifier -from .visualize import show_result - -__all__ = ['build_classifier', 'show_result'] diff --git a/mmdeploy/mmcls/apis/inference.py b/mmdeploy/mmcls/apis/inference.py deleted file mode 100644 index b5fe05996b..0000000000 --- a/mmdeploy/mmcls/apis/inference.py +++ /dev/null @@ -1,239 +0,0 @@ -from typing import Sequence, Union - -import mmcv -import torch -from mmcls.datasets import DATASETS -from mmcls.models import BaseClassifier - -from mmdeploy.utils.config_utils import Backend, get_backend, load_config - - -class DeployBaseClassifier(BaseClassifier): - """Base Class of Wrapper for classifier's inference. - - Args: - class_names (Sequence[str]): A list of string specifying class names. - device_id (int): An integer represents device index. - """ - - def __init__(self, class_names: Sequence[str], device_id: int): - super(DeployBaseClassifier, self).__init__() - self.CLASSES = class_names - self.device_id = device_id - - def simple_test(self, img, *args, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def extract_feat(self, imgs): - raise NotImplementedError('This method is not implemented.') - - def forward_train(self, imgs, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def forward_test(self, imgs, *args, **kwargs): - raise NotImplementedError('This method is not implemented.') - - -class ONNXRuntimeClassifier(DeployBaseClassifier): - """Wrapper for classifier's inference with ONNXRuntime. - - Args: - model_file (str): The path of input model file. - class_names (Sequence[str]): A list of string specifying class names. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: str, class_names: Sequence[str], - device_id: int): - super(ONNXRuntimeClassifier, self).__init__(class_names, device_id) - from mmdeploy.apis.onnxruntime import ORTWrapper - self.model = ORTWrapper(model_file, device_id) - - def forward_test(self, imgs: torch.Tensor, *args, **kwargs): - """Run test inference. - - Args: - imgs (torch.Tensor): Input tensor of the model. - - Returns: - list[np.ndarray]: Predictions of a classifier. - """ - input_data = imgs - results = self.model({'input': input_data})[0] - return list(results) - - -class TensorRTClassifier(DeployBaseClassifier): - """Wrapper for classifier's inference with TensorRT. - - Args: - model_file (str): The path of input model file. - class_names (Sequence[str]): A list of string specifying class names. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: str, class_names: Sequence[str], - device_id: int): - super(TensorRTClassifier, self).__init__(class_names, device_id) - from mmdeploy.apis.tensorrt import TRTWrapper - model = TRTWrapper(model_file) - - self.model = model - - def forward_test(self, imgs: torch.Tensor, *args, **kwargs): - """Run test inference. - - Args: - imgs (torch.Tensor): Input tensor of the model. - - Returns: - list[np.ndarray]: Predictions of a classifier. - """ - input_data = imgs - with torch.cuda.device(self.device_id), torch.no_grad(): - results = self.model({'input': input_data})['output'] - results = results.detach().cpu().numpy() - - return list(results) - - -class NCNNClassifier(DeployBaseClassifier): - """Wrapper for classifier's inference with NCNN. - - Args: - param_file (str): Path of parameter file. - bin_file (str): Path of bin file. - class_names (Sequence[str]): A list of string specifying class names. - device_id (int): An integer represents device index. - """ - - def __init__(self, param_file: str, bin_file: str, - class_names: Sequence[str], device_id: int): - super(NCNNClassifier, self).__init__(class_names, device_id) - from mmdeploy.apis.ncnn import NCNNWrapper - self.model = NCNNWrapper(param_file, bin_file, output_names=['output']) - - def forward_test(self, imgs: torch.Tensor, *args, **kwargs): - """Run test inference. - - Args: - imgs (torch.Tensor): Input tensor of the model. - - Returns: - list[np.ndarray]: Predictions of a classifier. - """ - results = self.model({'input': imgs})['output'] - results = results.detach().cpu().numpy() - results_list = list(results) - return results_list - - -class PPLClassifier(DeployBaseClassifier): - """Wrapper for classifier's inference with PPL. - - Args: - onnx_file (str): Path of input ONNX model file. - algo_file (str): Path of PPL algorithm file. - class_names (Sequence[str]): A list of string specifying class names. - device_id (int): An integer represents device index. - """ - - def __init__(self, onnx_file, algo_file, class_names, device_id): - super(PPLClassifier, self).__init__(class_names, device_id) - from mmdeploy.apis.ppl import PPLWrapper - model = PPLWrapper( - onnx_file=onnx_file, algo_file=algo_file, device_id=device_id) - self.model = model - self.CLASSES = class_names - - def forward_test(self, imgs: torch.Tensor, *args, **kwargs): - """Run test inference. - - Args: - imgs (torch.Tensor): Input tensor of the model. - - Returns: - list[np.ndarray]: Predictions of a classifier. - """ - input_data = imgs - results = self.model({'input': input_data})[0] - - return list(results) - - -ONNXRUNTIME_CLASSIFIER_MAP = dict(end2end=ONNXRuntimeClassifier) - -TENSORRT_CLASSIFIER_MAP = dict(end2end=TensorRTClassifier) - -PPL_CLASSIFIER_MAP = dict(end2end=PPLClassifier) - -NCNN_CLASSIFIER_MAP = dict(end2end=NCNNClassifier) - -BACKEND_CLASSIFIER_MAP = { - Backend.ONNXRUNTIME: ONNXRUNTIME_CLASSIFIER_MAP, - Backend.TENSORRT: TENSORRT_CLASSIFIER_MAP, - Backend.PPL: PPL_CLASSIFIER_MAP, - Backend.NCNN: NCNN_CLASSIFIER_MAP -} - - -def get_classes_from_config(model_cfg: Union[str, mmcv.Config]): - """Get class name from config. - - Args: - model_cfg (str | mmcv.Config): Input model config file or - Config object. - - Returns: - list[str]: A list of string specifying names of different class. - """ - model_cfg = load_config(model_cfg)[0] - module_dict = DATASETS.module_dict - data_cfg = model_cfg.data - - if 'train' in data_cfg: - module = module_dict[data_cfg.train.type] - elif 'val' in data_cfg: - module = module_dict[data_cfg.val.type] - elif 'test' in data_cfg: - module = module_dict[data_cfg.test.type] - else: - raise RuntimeError(f'No dataset config found in: {model_cfg}') - - return module.CLASSES - - -def build_classifier(model_files: Sequence[str], model_cfg: Union[str, - mmcv.Config], - deploy_cfg: Union[str, - mmcv.Config], device_id: int, **kwargs): - """Build classifier for different backend. - - Args: - model_files (list[str]): Input model file(s). - model_cfg (str | mmcv.Config): Input model config file or Config - object. - deploy_cfg (str | mmcv.Config): Input deployment config file or - Config object. - device_id (int): An integer represents device index. - - Returns: - DeployBaseClassifier: Classifier for a configured backend. - """ - model_cfg, deploy_cfg = load_config(model_cfg, deploy_cfg) - - backend = get_backend(deploy_cfg) - class_names = get_classes_from_config(model_cfg) - - assert backend in BACKEND_CLASSIFIER_MAP, \ - f'Unsupported backend type: {backend.value}' - model_map = BACKEND_CLASSIFIER_MAP[backend] - - model_type = 'end2end' - assert model_type in model_map, f'Unsupported model type: {model_type}' - backend_classifier_class = model_map[model_type] - - backend_detector = backend_classifier_class( - *model_files, class_names=class_names, device_id=device_id) - - return backend_detector diff --git a/mmdeploy/mmcls/apis/visualize.py b/mmdeploy/mmcls/apis/visualize.py deleted file mode 100644 index f03d1a100d..0000000000 --- a/mmdeploy/mmcls/apis/visualize.py +++ /dev/null @@ -1,32 +0,0 @@ -import numpy as np -import torch - -from mmdeploy.utils import Backend - - -def show_result(model: torch.nn.Module, - image: np.ndarray, - result: list, - output_file: str, - backend: Backend, - show: bool = True): - """Show predictions of mmcls. - - Args: - model (nn.Module): Input model which has `show_result` method. - image: (np.ndarray): Input image to draw predictions. - result (list): A list of predictions. - output_file (str): Output image file to save drawn predictions. - backend (Backend): Specifying backend type. - show (bool): Whether to show plotted image in windows. Defaults to - `True`. - - Returns: - np.ndarray: Drawn image, only if not `show` or `out_file`. - """ - pred_score = np.max(result, axis=0) - pred_label = np.argmax(result, axis=0) - result = {'pred_label': pred_label, 'pred_score': float(pred_score)} - result['pred_class'] = model.CLASSES[result['pred_label']] - return model.show_result( - image, result, show=show, win_name=backend.value, out_file=output_file) diff --git a/mmdeploy/mmcls/export/__init__.py b/mmdeploy/mmcls/export/__init__.py deleted file mode 100644 index 2aa9aa73d0..0000000000 --- a/mmdeploy/mmcls/export/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from .prepare_input import (build_dataloader, build_dataset, create_input, - get_tensor_from_input) - -__all__ = [ - 'build_dataloader', 'build_dataset', 'create_input', - 'get_tensor_from_input' -] diff --git a/mmdeploy/mmcls/export/prepare_input.py b/mmdeploy/mmcls/export/prepare_input.py deleted file mode 100644 index f8f4378013..0000000000 --- a/mmdeploy/mmcls/export/prepare_input.py +++ /dev/null @@ -1,132 +0,0 @@ -import logging -from typing import Any, Optional, Sequence, Union - -import mmcv -from mmcls.datasets import build_dataloader as build_dataloader_mmcls -from mmcls.datasets import build_dataset as build_dataset_mmcls -from mmcls.datasets.pipelines import Compose -from mmcv.parallel import collate, scatter -from torch.utils.data import Dataset - -from mmdeploy.utils import Task, load_config - - -def create_input(task: Task, - model_cfg: Union[str, mmcv.Config], - imgs: Any, - input_shape: Optional[Sequence[int]] = None, - device: str = 'cuda:0'): - """Create input for classifier. - - Args: - task (Task): Specifying task type. - model_cfg (str | mmcv.Config): The input model config. - imgs (Any): Input image(s), accpeted data type are `str`, - `np.ndarray`, `torch.Tensor`. - input_shape (list[int]): A list of two integer in (width, height) - format specifying input shape. Defaults to `None`. - device (str): A string represents device type. Default is 'cuda:0'. - - Returns: - tuple: (data, img), meta information for the input image and input. - """ - assert task == Task.CLASSIFICATION - cfg = load_config(model_cfg)[0].copy() - if isinstance(imgs, str): - if cfg.data.test.pipeline[0]['type'] != 'LoadImageFromFile': - cfg.data.test.pipeline.insert(0, dict(type='LoadImageFromFile')) - data = dict(img_info=dict(filename=imgs), img_prefix=None) - else: - if cfg.data.test.pipeline[0]['type'] == 'LoadImageFromFile': - cfg.data.test.pipeline.pop(0) - data = dict(img=imgs) - # check whether input_shape is valid - if input_shape is not None: - if 'crop_size' in cfg.data.test.pipeline[2]: - crop_size = cfg.data.test.pipeline[2]['crop_size'] - if tuple(input_shape) != (crop_size, crop_size): - logging.warning( - f'`input shape` should be equal to `crop_size`: {crop_size},\ - but given: {input_shape}') - test_pipeline = Compose(cfg.data.test.pipeline) - data = test_pipeline(data) - data = collate([data], samples_per_gpu=1) - if device != 'cpu': - data = scatter(data, [device])[0] - return data, data['img'] - - -def build_dataset(dataset_cfg: Union[str, mmcv.Config], - dataset_type: str = 'val', - **kwargs): - """Build dataset for classifier. - - Args: - dataset_cfg (str | mmcv.Config): The input dataset config. - dataset_type (str): A string represents dataset type, e.g.: 'train', - 'test', 'val'. Defaults to 'val'. - - Returns: - Dataset: A PyTorch dataset. - """ - dataset_cfg = load_config(dataset_cfg)[0] - data = dataset_cfg.data - assert dataset_type in data - - dataset = build_dataset_mmcls(data[dataset_type]) - - return dataset - - -def build_dataloader(dataset: Dataset, - samples_per_gpu: int, - workers_per_gpu: int, - num_gpus: int = 1, - dist: bool = False, - shuffle: bool = False, - round_up: bool = True, - seed: Optional[int] = None, - pin_memory: bool = True, - persistent_workers: bool = True, - **kwargs): - """Build dataloader for classifier. - - Args: - dataset (Dataset): Input dataset. - samples_per_gpu (int): Number of training samples on each GPU, i.e., - batch size of each GPU. - workers_per_gpu (int): How many subprocesses to use for data loading - for each GPU. - num_gpus (int): Number of GPUs. Only used in non-distributed training. - dist (bool): Distributed training/test or not. Defaults to `False`. - shuffle (bool): Whether to shuffle the data at every epoch. - Defaults to `False`. - round_up (bool): Whether to round up the length of dataset by adding - extra samples to make it evenly divisible. Default is `True`. - seed (int): An integer set to be seed. Default is `None`. - pin_memory (bool): Whether to use pin_memory in DataLoader. - Default is `True`. - persistent_workers (bool): If `True`, the data loader will not shutdown - the worker processes after a dataset has been consumed once. - This allows to maintain the workers Dataset instances alive. - The argument also has effect in PyTorch>=1.7.0. - Default is `True`. - kwargs: Any other keyword argument to be used to initialize DataLoader. - - Returns: - DataLoader: A PyTorch dataloader. - """ - return build_dataloader_mmcls(dataset, samples_per_gpu, workers_per_gpu, - num_gpus, dist, shuffle, round_up, seed, - pin_memory, persistent_workers, **kwargs) - - -def get_tensor_from_input(input_data: tuple): - """Get input tensor from input data. - - Args: - input_data (tuple): Input data containing meta info and image tensor. - Returns: - torch.Tensor: An image in `Tensor`. - """ - return input_data['img'] diff --git a/mmdeploy/mmdet/__init__.py b/mmdeploy/mmdet/__init__.py deleted file mode 100644 index 72fcc2b119..0000000000 --- a/mmdeploy/mmdet/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .core import * # noqa: F401,F403 -from .export import * # noqa: F401,F403 -from .models import * # noqa: F401,F403 diff --git a/mmdeploy/mmdet/apis/__init__.py b/mmdeploy/mmdet/apis/__init__.py deleted file mode 100644 index 7c62082749..0000000000 --- a/mmdeploy/mmdet/apis/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .inference import build_detector -from .visualize import show_result - -__all__ = ['build_detector', 'show_result'] diff --git a/mmdeploy/mmdet/apis/inference.py b/mmdeploy/mmdet/apis/inference.py deleted file mode 100644 index 6d83c0ef11..0000000000 --- a/mmdeploy/mmdet/apis/inference.py +++ /dev/null @@ -1,977 +0,0 @@ -from functools import partial -from typing import List, Sequence, Tuple, Union - -import mmcv -import numpy as np -import torch -import torch.nn.functional as F -from mmdet.core import bbox2result -from mmdet.datasets import DATASETS -from mmdet.models import BaseDetector - -from mmdeploy.mmdet.core.post_processing import multiclass_nms -from mmdeploy.utils import (Backend, get_backend, get_mmdet_params, - get_partition_config, load_config) - - -class DeployBaseDetector(BaseDetector): - """Base Class of Wrapper for inference of detection. - - Args: - class_names (Sequence[str]): A list of string specifying class names. - device_id (int): An integer represents device index. - """ - - def __init__(self, class_names, device_id, deploy_cfg=None, **kwargs): - super(DeployBaseDetector, self).__init__() - self.CLASSES = class_names - self.device_id = device_id - self.deploy_cfg = deploy_cfg - - def simple_test(self, img, img_metas, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def aug_test(self, imgs, img_metas, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def extract_feat(self, imgs): - raise NotImplementedError('This method is not implemented.') - - def forward_train(self, imgs, img_metas, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def val_step(self, data, optimizer): - raise NotImplementedError('This method is not implemented.') - - def train_step(self, data, optimizer): - raise NotImplementedError('This method is not implemented.') - - def aforward_test(self, *, img, img_metas, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def async_simple_test(self, img, img_metas, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def __clear_outputs( - self, test_outputs: List[Union[torch.Tensor, np.ndarray]] - ) -> List[Union[List[torch.Tensor], List[np.ndarray]]]: - """Removes additional outputs and detections with zero score. - - Args: - test_outputs (List[Union[torch.Tensor, np.ndarray]]): - outputs of forward_test. - - Returns: - List[Union[List[torch.Tensor], List[np.ndarray]]]: - outputs with without zero score object. - """ - batch_size = len(test_outputs[0]) - - num_outputs = len(test_outputs) - outputs = [[None for _ in range(batch_size)] - for _ in range(num_outputs)] - - for i in range(batch_size): - inds = test_outputs[0][i, :, 4] > 0.0 - for output_id in range(num_outputs): - outputs[output_id][i] = test_outputs[output_id][i, inds, ...] - return outputs - - def __postprocessing_masks(self, - det_bboxes: np.ndarray, - det_masks: np.ndarray, - img_w: int, - img_h: int, - mask_thr_binary: float = 0.5) -> np.ndarray: - """Additional processing of masks. Resizes masks from [num_det, 28, 28] - to [num_det, img_w, img_h]. Analog of the 'mmdeploy.mmdet.models.roi_he - ads.mask_heads.fcn_mask_head._do_paste_mask' function. - - Args: - det_bboxes (np.ndarray): Bbox of shape [num_det, 5] - det_masks (np.ndarray): Masks of shape [num_det, 28, 28]. - img_w (int): Width of the original image. - img_h (int): Height of the original image. - mask_thr_binary (float): The threshold for the mask. - - Returns: - np.ndarray: masks of shape [N, num_det, img_w, img_h]. - """ - masks = det_masks - bboxes = det_bboxes - - num_det = bboxes.shape[0] - if num_det == 0: - return np.zeros((0, img_w, img_h)) - - if isinstance(masks, np.ndarray): - masks = torch.tensor(masks) - bboxes = torch.tensor(bboxes) - - result_masks = [] - for bbox, mask in zip(bboxes, masks): - - x0_int, y0_int = 0, 0 - x1_int, y1_int = img_w, img_h - - img_y = torch.arange(y0_int, y1_int, dtype=torch.float32) + 0.5 - img_x = torch.arange(x0_int, x1_int, dtype=torch.float32) + 0.5 - x0, y0, x1, y1 = bbox - - img_y = (img_y - y0) / (y1 - y0) * 2 - 1 - img_x = (img_x - x0) / (x1 - x0) * 2 - 1 - if torch.isinf(img_x).any(): - inds = torch.where(torch.isinf(img_x)) - img_x[inds] = 0 - if torch.isinf(img_y).any(): - inds = torch.where(torch.isinf(img_y)) - img_y[inds] = 0 - - gx = img_x[None, :].expand(img_y.size(0), img_x.size(0)) - gy = img_y[:, None].expand(img_y.size(0), img_x.size(0)) - grid = torch.stack([gx, gy], dim=2) - - img_masks = F.grid_sample( - mask.to(dtype=torch.float32)[None, None, :, :], - grid[None, :, :, :], - align_corners=False) - - mask = img_masks - mask = (mask >= mask_thr_binary).to(dtype=torch.bool) - result_masks.append(mask.numpy()) - result_masks = np.concatenate(result_masks, axis=1) - return result_masks.squeeze(0) - - def forward(self, img: Sequence[torch.Tensor], img_metas: Sequence[dict], - *args, **kwargs): - """Run forward inference. - - Args: - img (Sequence[torch.Tensor]): A list contains input image(s) - in [N x C x H x W] format. - img_metas (Sequence[dict]): A list of meta info for image(s). - *args: Other arguments. - **kwargs: Other key-pair arguments. - - Returns: - list: A list contains predictions. - """ - input_img = img[0].contiguous() - outputs = self.forward_test(input_img, img_metas, *args, **kwargs) - outputs = self.__clear_outputs(outputs) - batch_dets, batch_labels = outputs[:2] - batch_masks = outputs[2] if len(outputs) == 3 else None - batch_size = input_img.shape[0] - img_metas = img_metas[0] - results = [] - rescale = kwargs.get('rescale', True) - for i in range(batch_size): - dets, labels = batch_dets[i], batch_labels[i] - if rescale: - scale_factor = img_metas[i]['scale_factor'] - - if isinstance(scale_factor, (list, tuple, np.ndarray)): - assert len(scale_factor) == 4 - scale_factor = np.array(scale_factor)[None, :] # [1,4] - dets[:, :4] /= scale_factor - - if 'border' in img_metas[i]: - # offset pixel of the top-left corners between original image - # and padded/enlarged image, 'border' is used when exporting - # CornerNet and CentripetalNet to onnx - x_off = img_metas[i]['border'][2] - y_off = img_metas[i]['border'][0] - dets[:, [0, 2]] -= x_off - dets[:, [1, 3]] -= y_off - dets[:, :4] *= (dets[:, :4] > 0).astype(dets.dtype) - - dets_results = bbox2result(dets, labels, len(self.CLASSES)) - - if batch_masks is not None: - masks = batch_masks[i] - img_h, img_w = img_metas[i]['img_shape'][:2] - ori_h, ori_w = img_metas[i]['ori_shape'][:2] - export_postprocess_mask = True - if self.deploy_cfg is not None: - mmdet_deploy_cfg = get_mmdet_params(self.deploy_cfg) - # this flag enable postprocess when export. - export_postprocess_mask = mmdet_deploy_cfg.get( - 'export_postprocess_mask', True) - if not export_postprocess_mask: - masks = self.__postprocessing_masks( - dets[:, :4], masks, ori_w, ori_h) - else: - masks = masks[:, :img_h, :img_w] - # avoid to resize masks with zero dim - if rescale and masks.shape[0] != 0: - masks = masks.astype(np.float32) - masks = torch.from_numpy(masks) - masks = torch.nn.functional.interpolate( - masks.unsqueeze(0), size=(ori_h, ori_w)) - masks = masks.squeeze(0).detach().numpy() - if masks.dtype != np.bool: - masks = masks >= 0.5 - segms_results = [[] for _ in range(len(self.CLASSES))] - for j in range(len(dets)): - segms_results[labels[j]].append(masks[j]) - results.append((dets_results, segms_results)) - else: - results.append(dets_results) - return results - - -class ONNXRuntimeDetector(DeployBaseDetector): - """Wrapper for detection's inference with ONNXRuntime. - - Args: - model_file (str): The path of input model file. - class_names (Sequence[str]): A list of string specifying class names. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: str, class_names: Sequence[str], - device_id: int, **kwargs): - super(ONNXRuntimeDetector, self).__init__(class_names, device_id, - **kwargs) - from mmdeploy.apis.onnxruntime import ORTWrapper - self.model = ORTWrapper(model_file, device_id) - - def forward_test(self, imgs: torch.Tensor, *args, **kwargs): - """Implement forward test. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - - Returns: - tuple[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] - and class labels of shape [N, num_det]. - """ - ort_outputs = self.model({'input': imgs}) - return ort_outputs - - -class TensorRTDetector(DeployBaseDetector): - """Wrapper for detection's inference with TensorRT. - - Args: - model_file (str): The path of input model file. - class_names (Sequence[str]): A list of string specifying class names. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: str, class_names: Sequence[str], - device_id: int, **kwargs): - super(TensorRTDetector, self).__init__(class_names, device_id, - **kwargs) - from mmdeploy.apis.tensorrt import TRTWrapper - - self.model = TRTWrapper(model_file) - self.output_names = ['dets', 'labels'] - if len(self.model.output_names) == 3: - self.output_names.append('masks') - - def forward_test(self, imgs: torch.Tensor, *args, **kwargs): - """Implement forward test. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - - Returns: - tuple[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] and - class labels of shape [N, num_det]. - """ - with torch.cuda.device(self.device_id), torch.no_grad(): - outputs = self.model({'input': imgs}) - outputs = [outputs[name] for name in self.output_names] - outputs = [out.detach().cpu().numpy() for out in outputs] - # filtered out invalid output filled with -1 - batch_labels = outputs[1] - batch_size = batch_labels.shape[0] - inds = batch_labels.reshape(-1) != -1 - for i in range(len(outputs)): - ori_shape = outputs[i].shape - outputs[i] = outputs[i].reshape(-1, - *ori_shape[2:])[inds, ...].reshape( - batch_size, -1, *ori_shape[2:]) - return outputs - - -class PPLDetector(DeployBaseDetector): - """Wrapper for detection's inference with PPL. - - Args: - model_file (str): The path of input model file. - class_names (Sequence[str]): A list of string specifying class names. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file, class_names, device_id, **kwargs): - super(PPLDetector, self).__init__(class_names, device_id) - from mmdeploy.apis.ppl import PPLWrapper - self.model = PPLWrapper(*model_file, device_id=device_id) - - def forward_test(self, imgs: torch.Tensor, *args, **kwargs): - """Implement forward test. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - - Returns: - tuple[Tensor, Tensor]: dets of shape [N, num_det, 5] and class - labels of shape [N, num_det]. - """ - ppl_outputs = self.model({'input': imgs}) - return ppl_outputs - - -class OpenVINODetector(DeployBaseDetector): - """Wrapper for detector's inference with OpenVINO. - - Args: - model_file (str): The path of input model file (.xml). - class_names (Sequence[str]): A list of string specifying class names. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: str, class_names: Sequence[str], - device_id: int, **kwargs): - super(OpenVINODetector, self).__init__(class_names, device_id, - **kwargs) - from mmdeploy.apis.openvino import OpenVINOWrapper - self.model = OpenVINOWrapper(model_file) - - def forward_test(self, imgs: torch.Tensor, *args, **kwargs) -> Tuple: - """Implement forward test. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - - Returns: - If there are no masks in the output: - tuple[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] - and class labels of shape [N, num_det]. - If the output contains masks: - tuple[np.ndarray, np.ndarray, np.ndarray]: - dets of shape [N, num_det, 5], - class labels of shape [N, num_det] and - masks of shape [N, num_det, H, W]. - """ - openvino_outputs = self.model({'input': imgs}) - output_keys = ['dets', 'labels'] - if 'masks' in openvino_outputs: - output_keys += ['masks'] - openvino_outputs = [openvino_outputs[key] for key in output_keys] - return openvino_outputs - - -class PartitionSingleStageDetector(DeployBaseDetector): - """Base wrapper for partitioned single stage detector. - - Args: - model_file (str): The path of input model file. - class_names (Sequence[str]): A list of string specifying class names. - model_cfg: (str | mmcv.Config): Input model config. - deploy_cfg: (str | mmcv.Config): Input deployment config. - device_id (int): An integer represents device index. - """ - - def __init__(self, class_names: Sequence[str], - model_cfg: Union[str, mmcv.Config], - deploy_cfg: Union[str, - mmcv.Config], device_id: int, **kwargs): - super(PartitionSingleStageDetector, - self).__init__(class_names, device_id, **kwargs) - # load cfg if necessary - deploy_cfg = load_config(deploy_cfg)[0] - model_cfg = load_config(model_cfg)[0] - - self.model_cfg = model_cfg - self.deploy_cfg = deploy_cfg - - def partition0_postprocess(self, scores: torch.Tensor, - bboxes: torch.Tensor): - """Perform post-processing for partition 0. - - Args: - scores (Tensor): The detection scores of shape - [N, num_boxes, num_classes]. - bboxes (Tensor): The bounding boxes of shape [N, num_boxes, 4]. - - Returns: - tuple[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] and - class labels of shape [N, num_det]. - """ - cfg = self.model_cfg.model.test_cfg - deploy_cfg = self.deploy_cfg - - post_params = get_mmdet_params(deploy_cfg) - max_output_boxes_per_class = post_params.max_output_boxes_per_class - iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) - score_threshold = cfg.get('score_thr', post_params.score_threshold) - pre_top_k = -1 if post_params.pre_top_k >= bboxes.shape[1] \ - else post_params.pre_top_k - keep_top_k = cfg.get('max_per_img', post_params.keep_top_k) - ret = multiclass_nms( - bboxes, - scores, - max_output_boxes_per_class, - iou_threshold=iou_threshold, - score_threshold=score_threshold, - pre_top_k=pre_top_k, - keep_top_k=keep_top_k) - ret = [r.cpu() for r in ret] - return ret - - -class ONNXRuntimePSSDetector(PartitionSingleStageDetector): - """Wrapper for partitioned single stage detector with ONNX Runtime. - - Args: - model_file (str): The path of input model file. - class_names (Sequence[str]): A list of string specifying class names. - model_cfg: (str | mmcv.Config): Input model config. - deploy_cfg: (str | mmcv.Config): Input deployment config. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: str, class_names: Sequence[str], - model_cfg: Union[str, mmcv.Config], - deploy_cfg: Union[str, - mmcv.Config], device_id: int, **kwargs): - super(ONNXRuntimePSSDetector, - self).__init__(class_names, model_cfg, deploy_cfg, device_id, - **kwargs) - from mmdeploy.apis.onnxruntime import ORTWrapper - self.model = ORTWrapper( - model_file, device_id, output_names=['scores', 'boxes']) - - def forward_test(self, imgs: torch.Tensor, *args, **kwargs): - """Implement forward test. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - - Returns: - tuple[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] and - class labels of shape [N, num_det]. - """ - ort_outputs = self.model({'input': imgs}) - scores, bboxes = ort_outputs[:2] - scores = torch.from_numpy(scores).to(imgs.device) - bboxes = torch.from_numpy(bboxes).to(imgs.device) - return self.partition0_postprocess(scores, bboxes) - - -class TensorRTPSSDetector(PartitionSingleStageDetector): - """TensorRT Wrapper for partition single stage detector. - - Args: - model_file (str): Path of the engine file. - class_names (list[str] | tuple[str]): Class names of the detector. - model_cfg (str | mmcv.Config): Model config file or Config object. - deploy_cfg (str | mmcv.Config): Deployment config file or Config - object. - device_id (int): Device index, should be same as the engine. - """ - - def __init__(self, model_file: str, class_names: Sequence[str], - model_cfg: Union[str, mmcv.Config], - deploy_cfg: Union[str, - mmcv.Config], device_id: int, **kwargs): - super(TensorRTPSSDetector, - self).__init__(class_names, model_cfg, deploy_cfg, device_id, - **kwargs) - from mmdeploy.apis.tensorrt import TRTWrapper - - self.model = TRTWrapper(model_file) - self.output_names = ['scores', 'boxes'] - - def forward_test(self, imgs: torch.Tensor, *args, - **kwargs) -> Tuple[torch.Tensor, torch.Tensor]: - """Run forward test. - - Args: - imgs (torch.Tensor): The input image(s). - - Return: - tuple[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] and - class labels of shape [N, num_det]. - """ - with torch.cuda.device(self.device_id), torch.no_grad(): - outputs = self.model({'input': imgs}) - outputs = [outputs[name] for name in self.output_names] - scores, bboxes = outputs[:2] - return self.partition0_postprocess(scores, bboxes) - - -class NCNNPSSDetector(PartitionSingleStageDetector): - """Wrapper for partitioned single stage detector with NCNN. - - Args: - model_file (str): The path of input model file. - class_names (Sequence[str]): A list of string specifying class names. - model_cfg: (str | mmcv.Config): Input model config. - deploy_cfg: (str | mmcv.Config): Input deployment config. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: str, class_names: Sequence[str], - model_cfg: Union[str, mmcv.Config], - deploy_cfg: Union[str, - mmcv.Config], device_id: int, **kwargs): - super(NCNNPSSDetector, self).__init__(class_names, model_cfg, - deploy_cfg, device_id, **kwargs) - from mmdeploy.apis.ncnn import NCNNWrapper - assert len(model_file) == 2 - ncnn_param_file = model_file[0] - ncnn_bin_file = model_file[1] - self.model = NCNNWrapper( - ncnn_param_file, ncnn_bin_file, output_names=['boxes', 'scores']) - - def forward_test(self, imgs: torch.Tensor, *args, **kwargs): - """Run forward test. - - Args: - imgs (torch.Tensor): The input image(s). - - Return: - tuple[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] and - class labels of shape [N, num_det]. - """ - outputs = self.model({'input': imgs}) - boxes = outputs['boxes'] - scores = outputs['scores'] - return self.partition0_postprocess(scores, boxes) - - -class PartitionTwoStageDetector(DeployBaseDetector): - """Base wrapper for partitioned two stage detector. - - Args: - class_names (Sequence[str]): A list of string specifying class names. - model_cfg: (str | mmcv.Config): Input model config. - deploy_cfg: (str | mmcv.Config): Input deployment config. - device_id (int): An integer represents device index. - """ - - def __init__(self, class_names: Sequence[str], - model_cfg: Union[str, mmcv.Config], - deploy_cfg: Union[str, - mmcv.Config], device_id: int, **kwargs): - super(PartitionTwoStageDetector, - self).__init__(class_names, device_id, **kwargs) - from mmdet.models.builder import build_head, build_roi_extractor - - from mmdeploy.mmdet.models.roi_heads.bbox_heads import \ - bbox_head__get_bboxes - - # load cfg if necessary - deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg) - - self.model_cfg = model_cfg - self.deploy_cfg = deploy_cfg - - self.bbox_roi_extractor = build_roi_extractor( - model_cfg.model.roi_head.bbox_roi_extractor) - self.bbox_head = build_head(model_cfg.model.roi_head.bbox_head) - - class Context: - pass - - ctx = Context() - ctx.cfg = self.deploy_cfg - self.bbox_head__get_bboxes = partial(bbox_head__get_bboxes, ctx) - - def partition0_postprocess(self, x: Sequence[torch.Tensor], - scores: torch.Tensor, bboxes: torch.Tensor): - """Perform post-processing for partition 0. - - Args: - x (tuple[Tensor]): Feature maps of all scale levels. - scores (Tensor): The detection scores of shape - [N, num_boxes, num_classes]. - bboxes (Tensor): The bounding boxes of shape [N, num_boxes, 4]. - - Returns: - tuple(Tensor, Tensor): rois and bbox_feats. - """ - # rpn-nms + roi-extractor - cfg = self.model_cfg.model.test_cfg.rpn - deploy_cfg = self.deploy_cfg - - post_params = get_mmdet_params(deploy_cfg) - iou_threshold = cfg.nms.get('iou_threshold', post_params.iou_threshold) - score_threshold = cfg.get('score_thr', post_params.score_threshold) - pre_top_k = -1 if post_params.pre_top_k >= bboxes.shape[1] \ - else post_params.pre_top_k - keep_top_k = cfg.get('max_per_img', post_params.keep_top_k) - # only one class in rpn - max_output_boxes_per_class = keep_top_k - proposals, _ = multiclass_nms( - bboxes, - scores, - max_output_boxes_per_class, - iou_threshold=iou_threshold, - score_threshold=score_threshold, - pre_top_k=pre_top_k, - keep_top_k=keep_top_k) - - rois = proposals - batch_index = torch.arange( - rois.shape[0], device=rois.device).float().view(-1, 1, 1).expand( - rois.size(0), rois.size(1), 1) - rois = torch.cat([batch_index, rois[..., :4]], dim=-1) - batch_size = rois.shape[0] - num_proposals_per_img = rois.shape[1] - - # Eliminate the batch dimension - rois = rois.view(-1, 5) - bbox_feats = self.bbox_roi_extractor( - x[:self.bbox_roi_extractor.num_inputs], rois) - - rois = rois.reshape(batch_size, num_proposals_per_img, rois.size(-1)) - return rois, bbox_feats - - def partition1_postprocess(self, rois: torch.Tensor, - cls_score: torch.Tensor, - bbox_pred: torch.Tensor, - img_metas: Sequence[dict]): - """Perform post-processing for partition 1. - Args: - rois (torch.Tensor): Input tensor of roi. - cls_score (torch.Tensor): Scores of all classes. - bbox_pred (torch.Tensor): Bounding box proposals. - img_metas (Sequence[dict]): A list of image(s) meta information. - - Returns: - tuple[Tensor, Tensor]: dets of shape [N, num_det, 5] and class - labels of shape [N, num_det]. - """ - batch_size = rois.shape[0] - num_proposals_per_img = rois.shape[1] - - cls_score = cls_score.reshape(batch_size, num_proposals_per_img, - cls_score.size(-1)) - - bbox_pred = bbox_pred.reshape(batch_size, num_proposals_per_img, - bbox_pred.size(-1)) - - rcnn_test_cfg = self.model_cfg.model.test_cfg.rcnn - return self.bbox_head__get_bboxes(self.bbox_head, rois, cls_score, - bbox_pred, - img_metas[0][0]['img_shape'], - rcnn_test_cfg) - - -class ONNXRuntimePTSDetector(PartitionTwoStageDetector): - """Wrapper for partitioned two stage detector with ONNX Runtime. - - Args: - model_file (Sequence[str]): A list of paths of input model files. - class_names (Sequence[str]): A list of string specifying class names. - model_cfg: (str | mmcv.Config): Input model config. - deploy_cfg: (str | mmcv.Config): Input deployment config. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: Sequence[str], class_names: Sequence[str], - model_cfg: Union[str, mmcv.Config], - deploy_cfg: Union[str, - mmcv.Config], device_id: int, **kwargs): - super(ONNXRuntimePTSDetector, - self).__init__(class_names, model_cfg, deploy_cfg, device_id, - **kwargs) - from mmdeploy.apis.onnxruntime import ORTWrapper - self.model_list = [ - ORTWrapper(file, device_id=device_id) for file in model_file - ] - num_partition0_outputs = len(self.model_list[0].output_names) - num_feat = num_partition0_outputs - 2 - self.model_list[0].output_names = [ - 'feat/{}'.format(i) for i in range(num_feat) - ] + ['scores', 'boxes'] - self.model_list[1].output_names = ['cls_score', 'bbox_pred'] - - def forward_test(self, imgs: torch.Tensor, img_metas: Sequence[dict], - *args, **kwargs): - """Implement forward test. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - img_metas (Sequence[dict]): A list of image(s) meta information. - - Returns: - tuple[Tensor, Tensor]: dets of shape [N, num_det, 5] and class - labels of shape [N, num_det]. - """ - ort_outputs = self.model_list[0]({'input': imgs}) - feats = ort_outputs[:-2] - scores, bboxes = ort_outputs[-2:] - feats = [torch.from_numpy(feat).to(imgs.device) for feat in feats] - scores = torch.from_numpy(scores).to(imgs.device) - bboxes = torch.from_numpy(bboxes).to(imgs.device) - - # partition0_postprocess - rois, bbox_feats = self.partition0_postprocess(feats, scores, bboxes) - - # partition1 - ort_outputs = self.model_list[1]({'bbox_feats': bbox_feats}) - cls_score, bbox_pred = ort_outputs[:2] - cls_score = torch.from_numpy(cls_score).to(imgs.device) - bbox_pred = torch.from_numpy(bbox_pred).to(imgs.device) - - # partition1_postprocess - return self.partition1_postprocess(rois, cls_score, bbox_pred, - img_metas) - - -class TensorRTPTSDetector(PartitionTwoStageDetector): - """Wrapper for partitioned two stage detector with TensorRT. - - Args: - model_file (Sequence[str]): A list of paths of input model files. - class_names (Sequence[str]): A list of string specifying class names. - model_cfg: (str | mmcv.Config): Input model config. - deploy_cfg: (str | mmcv.Config): Input deployment config. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: Sequence[str], class_names: Sequence[str], - model_cfg: Union[str, mmcv.Config], - deploy_cfg: Union[str, - mmcv.Config], device_id: int, **kwargs): - super(TensorRTPTSDetector, - self).__init__(class_names, model_cfg, deploy_cfg, device_id, - **kwargs) - - from mmdeploy.apis.tensorrt import TRTWrapper - - model_list = [] - for m_file in model_file: - model = TRTWrapper(m_file) - model_list.append(model) - - self.model_list = model_list - - output_names_list = [] - num_partition0_outputs = len(model_list[0].output_names) - num_feat = num_partition0_outputs - 2 - output_names_list.append( - ['feat/{}'.format(i) - for i in range(num_feat)] + ['scores', 'boxes']) # partition0 - output_names_list.append(['cls_score', 'bbox_pred']) # partition1 - self.output_names_list = output_names_list - - def forward_test(self, imgs: torch.Tensor, img_metas: Sequence[dict], - *args, **kwargs): - """Implement forward test. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - img_metas (Sequence[dict]): A list of image(s) meta information. - - Returns: - tuple[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] and - class labels of shape [N, num_det]. - """ - with torch.cuda.device(self.device_id), torch.no_grad(): - outputs = self.model_list[0]({'input': imgs}) - outputs = [outputs[name] for name in self.output_names_list[0]] - feats = outputs[:-2] - scores, bboxes = outputs[-2:] - - # partition0_postprocess - rois, bbox_feats = self.partition0_postprocess(feats, scores, bboxes) - - # partition1 forward - bbox_feats = bbox_feats.contiguous() - with torch.cuda.device(self.device_id), torch.no_grad(): - outputs = self.model_list[1]({'bbox_feats': bbox_feats}) - outputs = [outputs[name] for name in self.output_names_list[1]] - cls_score, bbox_pred = outputs[:2] - - # partition1_postprocess - outputs = self.partition1_postprocess(rois, cls_score, bbox_pred, - img_metas) - outputs = [out.detach().cpu() for out in outputs] - return outputs - - -class NCNNPTSDetector(PartitionTwoStageDetector): - """Wrapper for partitioned two stage detector with NCNN. - - Args: - model_file (Sequence[str]): A list of paths of input model files. - class_names (Sequence[str]): A list of string specifying class names. - model_cfg: (str | mmcv.Config): Input model config. - deploy_cfg: (str | mmcv.Config): Input deployment config. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: Sequence[str], class_names: Sequence[str], - model_cfg: Union[str, mmcv.Config], - deploy_cfg: Union[str, - mmcv.Config], device_id: int, **kwargs): - super(NCNNPTSDetector, self).__init__(class_names, model_cfg, - deploy_cfg, device_id, **kwargs) - from mmdeploy.apis.ncnn import NCNNWrapper - assert self.device_id == -1 - assert len(model_file) == 4 - - model_list = [] - for ncnn_param_file, ncnn_bin_file in zip(model_file[::2], - model_file[1::2]): - model = NCNNWrapper(ncnn_param_file, ncnn_bin_file) - model_list.append(model) - - model_cfg = load_config(model_cfg)[0] - num_output_stage1 = model_cfg['model']['neck']['num_outs'] - - output_names_list = [] - output_names_list.append( - ['feat/{}'.format(i) - for i in range(num_output_stage1)] + ['scores', 'boxes']) - output_names_list.append(['cls_score', 'bbox_pred']) - - model_list[0].set_output_names(output_names_list[0]) - model_list[1].set_output_names(output_names_list[1]) - - self.model_list = model_list - self.output_names_list = output_names_list - - def forward_test(self, imgs: torch.Tensor, img_metas: Sequence[dict], - *args, **kwargs): - """Implement forward test. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - img_metas (Sequence[dict]): A list of image(s) meta information. - - Returns: - tuple[np.ndarray, np.ndarray]: dets of shape [N, num_det, 5] and - class labels of shape [N, num_det]. - """ - # stage0 forward - out_stage0 = self.model_list[0]({'input': imgs}) - - outputs = [] - for name in self.output_names_list[0]: - out = out_stage0[name] - outputs.append(out) - feats = outputs[:-2] - scores, bboxes = outputs[-2:] - - # stage0_postprocess - rois, bbox_feats = self.partition0_postprocess(feats, scores, bboxes) - - # stage1 forward - out_stage1 = self.model_list[1]({'bbox_feats': bbox_feats}) - cls_score = out_stage1['cls_score'] - bbox_pred = out_stage1['bbox_pred'] - - # stage1_postprocess - outputs = self.partition1_postprocess(rois, cls_score, bbox_pred, - img_metas) - outputs = [out.detach().cpu() for out in outputs] - return outputs - - -def get_classes_from_config(model_cfg: Union[str, mmcv.Config], **kwargs): - """Get class name from config. - - Args: - model_cfg (str | mmcv.Config): Input model config file or - Config object. - - Returns: - list[str]: A list of string specifying names of different class. - """ - # load cfg if necessary - model_cfg = load_config(model_cfg)[0] - module_dict = DATASETS.module_dict - data_cfg = model_cfg.data - - if 'test' in data_cfg: - module = module_dict[data_cfg.test.type] - elif 'val' in data_cfg: - module = module_dict[data_cfg.val.type] - elif 'train' in data_cfg: - module = module_dict[data_cfg.train.type] - else: - raise RuntimeError(f'No dataset config found in: {model_cfg}') - - return module.CLASSES - - -ONNXRUNTIME_DETECTOR_MAP = dict( - end2end=ONNXRuntimeDetector, - single_stage=ONNXRuntimePSSDetector, - two_stage=ONNXRuntimePTSDetector) - -TENSORRT_DETECTOR_MAP = dict( - end2end=TensorRTDetector, - single_stage=TensorRTPSSDetector, - two_stage=TensorRTPTSDetector) - -PPL_DETECTOR_MAP = dict(end2end=PPLDetector) - -NCNN_DETECTOR_MAP = dict( - single_stage=NCNNPSSDetector, two_stage=NCNNPTSDetector) - -OPENVINO_MAP = dict(end2end=OpenVINODetector) - -BACKEND_DETECTOR_MAP = { - Backend.ONNXRUNTIME: ONNXRUNTIME_DETECTOR_MAP, - Backend.TENSORRT: TENSORRT_DETECTOR_MAP, - Backend.PPL: PPL_DETECTOR_MAP, - Backend.NCNN: NCNN_DETECTOR_MAP, - Backend.OPENVINO: OPENVINO_MAP -} - - -def build_detector(model_files: Sequence[str], model_cfg: Union[str, - mmcv.Config], - deploy_cfg: Union[str, - mmcv.Config], device_id: int, **kwargs): - """Build detector for different backend. - - Args: - model_files (list[str]): Input model file(s). - model_cfg (str | mmcv.Config): Input model config file or Config - object. - deploy_cfg (str | mmcv.Config): Input deployment config file or - Config object. - device_id (int): An integer represents device index. - - Returns: - DeployBaseDetector: Detector for a configured backend. - """ - # load cfg if necessary - deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg) - - backend = get_backend(deploy_cfg) - class_names = get_classes_from_config(model_cfg) - - assert backend in BACKEND_DETECTOR_MAP, \ - f'Unsupported backend type: {backend.value}' - detector_map = BACKEND_DETECTOR_MAP[backend] - - partition_type = 'end2end' - partition_config = get_partition_config(deploy_cfg) - if partition_config is not None: - partition_type = partition_config.get('type', None) - - assert partition_type in detector_map,\ - f'Unsupported partition type: {partition_type}' - backend_detector_class = detector_map[partition_type] - - model_files = model_files[0] if len(model_files) == 1 else model_files - backend_detector = backend_detector_class( - model_file=model_files, - class_names=class_names, - device_id=device_id, - model_cfg=model_cfg, - deploy_cfg=deploy_cfg, - **kwargs) - - return backend_detector diff --git a/mmdeploy/mmdet/apis/visualize.py b/mmdeploy/mmdet/apis/visualize.py deleted file mode 100644 index b42f4dc750..0000000000 --- a/mmdeploy/mmdet/apis/visualize.py +++ /dev/null @@ -1,34 +0,0 @@ -import numpy as np - -from mmdeploy.utils import Backend - - -def show_result(model, - image: np.ndarray, - result: list, - output_file: str, - backend: Backend, - show: bool = True, - score_thr: float = 0.3): - """Show predictions of detection. - - Args: - model (nn.Module): Input model which has `show_result` method. - image: (np.ndarray): Input image to draw predictions. - result (list): A list of predictions. - output_file (str): Output image file to save drawn predictions. - backend (Backend): Specifying backend type. - show (bool): Whether to show plotted image in windows. Defaults to - `True`. - score_thr (float): Score threshold for detection, defaults to `0.3`. - - Returns: - np.ndarray: Drawn image, only if not `show` or `out_file`. - """ - return model.show_result( - image, - result, - score_thr=score_thr, - show=show, - win_name=backend.value, - out_file=output_file) diff --git a/mmdeploy/mmdet/core/bbox/__init__.py b/mmdeploy/mmdet/core/bbox/__init__.py deleted file mode 100644 index 7a66740087..0000000000 --- a/mmdeploy/mmdet/core/bbox/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .coder import * # noqa: F401,F403 -from .transforms import distance2bbox # noqa: F401,F403 diff --git a/mmdeploy/mmdet/core/bbox/coder/__init__.py b/mmdeploy/mmdet/core/bbox/coder/__init__.py deleted file mode 100644 index fb83388204..0000000000 --- a/mmdeploy/mmdet/core/bbox/coder/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .delta_xywh_bbox_coder import * # noqa: F401,F403 -from .tblr_bbox_coder import * # noqa: F401, F403 diff --git a/mmdeploy/mmdet/export/__init__.py b/mmdeploy/mmdet/export/__init__.py deleted file mode 100644 index 1092eaeaac..0000000000 --- a/mmdeploy/mmdet/export/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -from .model_partition import get_partition_cfg -from .onnx_utils import clip_bboxes -from .prepare_input import (build_dataloader, build_dataset, create_input, - get_tensor_from_input) -from .tensorrt_helper import pad_with_value - -__all__ = [ - 'get_partition_cfg', 'clip_bboxes', 'create_input', 'build_dataloader', - 'build_dataset', 'get_tensor_from_input', 'pad_with_value' -] diff --git a/mmdeploy/mmdet/export/onnx_utils.py b/mmdeploy/mmdet/export/onnx_utils.py deleted file mode 100644 index f08bd39f12..0000000000 --- a/mmdeploy/mmdet/export/onnx_utils.py +++ /dev/null @@ -1,47 +0,0 @@ -from typing import Sequence, Union - -import torch -from torch import Tensor - - -def clip_bboxes(x1: Tensor, y1: Tensor, x2: Tensor, y2: Tensor, - max_shape: Union[Tensor, Sequence[int]]): - """Clip bboxes for onnx. - - Since torch.clamp cannot have dynamic `min` and `max`, we scale the - boxes by 1/max_shape and clamp in the range [0, 1] if necessary. - - Args: - x1 (Tensor): The x1 for bounding boxes. - y1 (Tensor): The y1 for bounding boxes. - x2 (Tensor): The x2 for bounding boxes. - y2 (Tensor): The y2 for bounding boxes. - max_shape (Tensor | Sequence[int]): The (H,W) of original image. - Returns: - tuple(Tensor): The clipped x1, y1, x2, y2. - """ - assert len(max_shape) == 2, '`max_shape` should be [h, w]' - if isinstance(max_shape, torch.Tensor): - # scale by 1/max_shape - x1 = x1 / max_shape[1] - y1 = y1 / max_shape[0] - x2 = x2 / max_shape[1] - y2 = y2 / max_shape[0] - - # clamp [0, 1] - x1 = torch.clamp(x1, 0, 1) - y1 = torch.clamp(y1, 0, 1) - x2 = torch.clamp(x2, 0, 1) - y2 = torch.clamp(y2, 0, 1) - - # scale back - x1 = x1 * max_shape[1] - y1 = y1 * max_shape[0] - x2 = x2 * max_shape[1] - y2 = y2 * max_shape[0] - else: - x1 = torch.clamp(x1, 0, max_shape[1]) - y1 = torch.clamp(y1, 0, max_shape[0]) - x2 = torch.clamp(x2, 0, max_shape[1]) - y2 = torch.clamp(y2, 0, max_shape[0]) - return x1, y1, x2, y2 diff --git a/mmdeploy/mmdet/export/prepare_input.py b/mmdeploy/mmdet/export/prepare_input.py deleted file mode 100644 index 2d2f473c16..0000000000 --- a/mmdeploy/mmdet/export/prepare_input.py +++ /dev/null @@ -1,162 +0,0 @@ -from typing import Any, Dict, Optional, Sequence, Union - -import mmcv -import numpy as np -from mmcv.parallel import collate, scatter -from mmdet.datasets import build_dataloader as build_dataloader_mmdet -from mmdet.datasets import build_dataset as build_dataset_mmdet -from mmdet.datasets import replace_ImageToTensor -from mmdet.datasets.pipelines import Compose -from torch.utils.data import Dataset - -from mmdeploy.utils import Task, load_config - - -def create_input(task: Task, - model_cfg: Union[str, mmcv.Config], - imgs: Any, - input_shape: Sequence[int] = None, - device: str = 'cuda:0'): - """Create input for detector. - - Args: - task (Task): Specifying task type. - model_cfg (str | mmcv.Config): The input model config. - imgs (Any): Input image(s), accpeted data type are `str`, - `np.ndarray`, `torch.Tensor`. - input_shape (list[int]): A list of two integer in (width, height) - format specifying input shape. Defaults to `None`. - device (str): A string represents device type. Default is 'cuda:0'. - - Returns: - tuple: (data, img), meta information for the input image and input. - """ - assert task == Task.OBJECT_DETECTION - cfg = load_config(model_cfg)[0].copy() - - if not isinstance(imgs, (list, tuple)): - imgs = [imgs] - - if isinstance(imgs[0], np.ndarray): - cfg = cfg.copy() - # set loading pipeline type - cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam' - # for static exporting - if input_shape is not None: - cfg.data.test.pipeline[1]['img_scale'] = tuple(input_shape) - transforms = cfg.data.test.pipeline[1]['transforms'] - for trans in transforms: - trans_type = trans['type'] - if trans_type == 'Resize': - trans['keep_ratio'] = False - elif trans_type == 'Pad': - trans['size_divisor'] = 1 - - cfg.data.test.pipeline = replace_ImageToTensor(cfg.data.test.pipeline) - test_pipeline = Compose(cfg.data.test.pipeline) - data_list = [] - for img in imgs: - # prepare data - if isinstance(img, np.ndarray): - # directly add img - data = dict(img=img) - else: - # add information into dict - data = dict(img_info=dict(filename=img), img_prefix=None) - # build the data pipeline - data = test_pipeline(data) - data_list.append(data) - - data = collate(data_list, samples_per_gpu=len(imgs)) - - data['img_metas'] = [img_metas.data[0] for img_metas in data['img_metas']] - data['img'] = [img.data[0] for img in data['img']] - if device != 'cpu': - data = scatter(data, [device])[0] - - return data, data['img'] - - -def build_dataset(dataset_cfg: Union[str, mmcv.Config], - dataset_type: str = 'val', - **kwargs): - """Build dataset for detection. - - Args: - dataset_cfg (str | mmcv.Config): The input dataset config. - dataset_type (str): A string represents dataset type, e.g.: 'train', - 'test', 'val'. Defaults to 'val'. - - Returns: - Dataset: A PyTorch dataset. - """ - dataset_cfg = load_config(dataset_cfg)[0].copy() - - assert dataset_type in dataset_cfg.data - data_cfg = dataset_cfg.data[dataset_type] - # in case the dataset is concatenated - if isinstance(data_cfg, dict): - data_cfg.test_mode = True - samples_per_gpu = data_cfg.get('samples_per_gpu', 1) - if samples_per_gpu > 1: - # Replace 'ImageToTensor' to 'DefaultFormatBundle' - data_cfg.pipeline = replace_ImageToTensor(data_cfg.pipeline) - elif isinstance(data_cfg, list): - for ds_cfg in data_cfg: - ds_cfg.test_mode = True - samples_per_gpu = max( - [ds_cfg.get('samples_per_gpu', 1) for ds_cfg in data_cfg]) - if samples_per_gpu > 1: - for ds_cfg in data_cfg: - ds_cfg.pipeline = replace_ImageToTensor(ds_cfg.pipeline) - dataset = build_dataset_mmdet(data_cfg) - - return dataset - - -def build_dataloader(dataset: Dataset, - samples_per_gpu: int, - workers_per_gpu: int, - num_gpus: int = 1, - dist: bool = False, - shuffle: bool = False, - seed: Optional[int] = None, - **kwargs): - """Build dataloader for detection. - - Args: - dataset (Dataset): Input dataset. - samples_per_gpu (int): Number of training samples on each GPU, i.e., - batch size of each GPU. - workers_per_gpu (int): How many subprocesses to use for data loading - for each GPU. - num_gpus (int): Number of GPUs. Only used in non-distributed training. - dist (bool): Distributed training/test or not. Defaults to `False`. - shuffle (bool): Whether to shuffle the data at every epoch. - Defaults to `False`. - seed (int): An integer set to be seed. Default is `None`. - kwargs: Any other keyword argument to be used to initialize DataLoader. - - Returns: - DataLoader: A PyTorch dataloader. - """ - return build_dataloader_mmdet( - dataset, - samples_per_gpu, - workers_per_gpu, - num_gpus=num_gpus, - dist=dist, - shuffle=shuffle, - seed=seed, - **kwargs) - - -def get_tensor_from_input(input_data: Dict[str, Any]): - """Get input tensor from input data. - - Args: - input_data (dict): Input data containing meta info and image tensor. - Returns: - torch.Tensor: An image in `Tensor`. - """ - return input_data['img'][0] diff --git a/mmdeploy/mmdet/export/tensorrt_helper.py b/mmdeploy/mmdet/export/tensorrt_helper.py deleted file mode 100644 index 0b7dc0871c..0000000000 --- a/mmdeploy/mmdet/export/tensorrt_helper.py +++ /dev/null @@ -1,35 +0,0 @@ -from typing import Any, Optional - -import torch -from torch import Tensor - - -def pad_with_value(x: Tensor, - pad_dim: int, - pad_size: int, - pad_value: Optional[Any] = None): - """Pad a tensor with a value along some dim. - - Args: - x (Tensor): Input tensor. - pad_dim (int): Along which dim to pad. - pad_size (int): To which size to pad. - pad_value (Any): Filled value for padding. Defaults to `None`. - - Returns: - Tensor: Padded tensor. - """ - num_dims = len(x.shape) - pad_slice = (slice(None, None, None), ) * num_dims - pad_slice = pad_slice[:pad_dim] + (slice(0, 1, - 1), ) + pad_slice[pad_dim + 1:] - repeat_size = [1] * num_dims - repeat_size[pad_dim] = pad_size - - x_pad = x.__getitem__(pad_slice) - if pad_value is not None: - x_pad = x_pad * 0 + pad_value - - x_pad = x_pad.repeat(*repeat_size) - x = torch.cat([x, x_pad], dim=pad_dim) - return x diff --git a/mmdeploy/mmdet/models/dense_heads/__init__.py b/mmdeploy/mmdet/models/dense_heads/__init__.py deleted file mode 100644 index 467b387b28..0000000000 --- a/mmdeploy/mmdet/models/dense_heads/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -from .anchor_head import anchor_head__get_bboxes -from .atss_head import atss_head__get_bboxes -from .fcos_head import fcos_head__get_bboxes -from .fovea_head import fovea_head__get_bboxes -from .rpn_head import rpn_head__get_bboxes -from .vfnet_head import vfnet_head__get_bboxes -from .yolo_head import yolov3_head__get_bboxes, yolov3_head__get_bboxes__ncnn -from .yolox_head import yolox_head__get_bboxes - -__all__ = [ - 'anchor_head__get_bboxes', 'atss_head__get_bboxes', - 'fcos_head__get_bboxes', 'fovea_head__get_bboxes', 'rpn_head__get_bboxes', - 'vfnet_head__get_bboxes', 'yolov3_head__get_bboxes', - 'yolov3_head__get_bboxes__ncnn', 'yolox_head__get_bboxes' -] diff --git a/mmdeploy/mmdet/models/detectors/__init__.py b/mmdeploy/mmdet/models/detectors/__init__.py deleted file mode 100644 index d206d63afe..0000000000 --- a/mmdeploy/mmdet/models/detectors/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -from .base import base_detector__forward -from .rpn import rpn__simple_test -from .single_stage import single_stage__simple_test -from .two_stage import two_stage__extract_feat - -__all__ = [ - 'single_stage__simple_test', 'two_stage__extract_feat', - 'base_detector__forward', 'rpn__simple_test' -] diff --git a/mmdeploy/mmdet/models/roi_heads/__init__.py b/mmdeploy/mmdet/models/roi_heads/__init__.py deleted file mode 100644 index b5ac1221bf..0000000000 --- a/mmdeploy/mmdet/models/roi_heads/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from .bbox_heads import * # noqa: F401, F403 -from .cascade_roi_head import * # noqa: F401, F403 -from .mask_heads import * # noqa: F401, F403 -from .roi_extractors import * # noqa: F401, F403 -from .standard_roi_head import * # noqa: F401, F403 -from .test_mixins import * # noqa: F401, F403 diff --git a/mmdeploy/mmdet/models/roi_heads/bbox_heads/__init__.py b/mmdeploy/mmdet/models/roi_heads/bbox_heads/__init__.py deleted file mode 100644 index 7b7bdf48c9..0000000000 --- a/mmdeploy/mmdet/models/roi_heads/bbox_heads/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .bbox_head import bbox_head__get_bboxes - -__all__ = ['bbox_head__get_bboxes'] diff --git a/mmdeploy/mmdet/models/roi_heads/mask_heads/__init__.py b/mmdeploy/mmdet/models/roi_heads/mask_heads/__init__.py deleted file mode 100644 index 7bafcfc008..0000000000 --- a/mmdeploy/mmdet/models/roi_heads/mask_heads/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .fcn_mask_head import fcn_mask_head__get_seg_masks - -__all__ = ['fcn_mask_head__get_seg_masks'] diff --git a/mmdeploy/mmdet/models/roi_heads/roi_extractors/__init__.py b/mmdeploy/mmdet/models/roi_heads/roi_extractors/__init__.py deleted file mode 100644 index 105e406d1d..0000000000 --- a/mmdeploy/mmdet/models/roi_heads/roi_extractors/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from .single_level_roi_extractor import ( - single_roi_extractor__forward, single_roi_extractor__forward__openvino, - single_roi_extractor__forward__tensorrt) - -__all__ = [ - 'single_roi_extractor__forward', 'single_roi_extractor__forward__openvino', - 'single_roi_extractor__forward__tensorrt' -] diff --git a/mmdeploy/mmedit/apis/__init__.py b/mmdeploy/mmedit/apis/__init__.py deleted file mode 100644 index 61a020cfa7..0000000000 --- a/mmdeploy/mmedit/apis/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .inference import build_editing_processor -from .visualize import show_result - -__all__ = ['build_editing_processor', 'show_result'] diff --git a/mmdeploy/mmedit/apis/inference.py b/mmdeploy/mmedit/apis/inference.py deleted file mode 100644 index c675574f1b..0000000000 --- a/mmdeploy/mmedit/apis/inference.py +++ /dev/null @@ -1,314 +0,0 @@ -import warnings -from typing import Optional, Sequence, Union - -import mmcv -import numpy as np -import torch -from mmedit.core import psnr, ssim, tensor2img -from mmedit.models import BaseModel - -from mmdeploy.utils.config_utils import Backend, get_backend, load_config - - -class DeployBaseRestorer(BaseModel): - """Base Class of Wrapper for restorer's inference. - - Args: - device_id (int): An integer represents device index. - test_cfg (mmcv.Config) : The test config in model config, which is used - in evaluation. Defaults to `None`. - """ - - allowed_metrics = {'PSNR': psnr, 'SSIM': ssim} - - def __init__(self, - device_id: int, - test_cfg: Optional[mmcv.Config] = None, - **kwargs): - super(DeployBaseRestorer, self).__init__(**kwargs) - self.test_cfg = test_cfg - self.device_id = device_id - - def init_weights(self): - raise NotImplementedError('This method is not implemented.') - - def forward(self, lq: torch.Tensor, test_mode: bool = False, **kwargs): - """Run test inference for restorer. - - We want forward() to output an image or a evaluation result. - When test_mode is set, the output is evaluation result. Otherwise - it is an image. - - Args: - lq (torch.Tensor): The input low-quality image of the model. - test_mode (bool): When test_mode is set, the output is evaluation - result. Otherwise it is an image. Default to `False`. - - Returns: - torch.Tensor | dict: High resolution image or a evaluation results. - """ - - if test_mode: - return self.forward_test(lq, **kwargs) - else: - return self.forward_dummy(lq, **kwargs) - - def forward_train(self, imgs, labels): - raise NotImplementedError('This method is not implemented.') - - def forward_test(self, - lq: torch.Tensor, - gt: Optional[torch.Tensor] = None, - **kwargs): - """Run inference for restorer to generate evaluation result. - - Args: - lq (torch.Tensor): The input low-quality image of the model. - gt (torch.Tensor): The ground truth of input image. Defaults to - `None`. - - Returns: - dict: Evaluation results. - """ - outputs = self.forward_dummy(lq) - result = self._test_post_process(outputs, lq, gt) - return result - - def train_step(self, data_batch, optimizer): - raise NotImplementedError('This method is not implemented.') - - def evaluate(self, output: torch.Tensor, gt: torch.Tensor): - """Evaluation function implemented in mmedit. - - Args: - output (torch.Tensor): Model output with shape (n, c, h, w). - gt (torch.Tensor): GT Tensor with shape (n, c, h, w). - - Returns: - dict: Evaluation results. - """ - crop_border = self.test_cfg.crop_border - - if isinstance(output, np.ndarray): - output = torch.from_numpy(output) - output = tensor2img(output) - gt = tensor2img(gt) - - eval_result = dict() - for metric in self.test_cfg.metrics: - eval_result[metric] = self.allowed_metrics[metric](output, gt, - crop_border) - return eval_result - - def _test_post_process(self, - outputs: torch.Tensor, - lq: torch.Tensor, - gt: Optional[torch.Tensor] = None): - """Get evaluation results by post-processing model outputs. - - Args: - output (torch.Tensor) : The output high resolution image. - lq (torch.Tensor): The input low-quality image of the model. - gt (torch.Tensor): The ground truth of input image, default is - `None`. - - Returns: - dict: Evaluation results. - """ - if self.test_cfg is not None and self.test_cfg.get('metrics', None): - assert gt is not None, ( - 'evaluation with metrics must have gt images.') - results = dict(eval_result=self.evaluate(outputs, gt)) - else: - results = dict(lq=lq.cpu(), output=outputs) - if gt is not None: - results['gt'] = gt.cpu() - - return results - - -class ONNXRuntimeRestorer(DeployBaseRestorer): - """Wrapper for restorer's inference with ONNXRuntime. - - Args: - model_file (str): The path of an input model file. - device_id (int): An integer represents device index. - test_cfg (mmcv.Config) : The test config in model config, which is - used in evaluation. Defaults to `None`. - """ - - def __init__(self, - model_file: str, - device_id: int, - test_cfg: Optional[mmcv.Config] = None, - **kwargs): - super(ONNXRuntimeRestorer, self).__init__( - device_id, test_cfg=test_cfg, **kwargs) - - from mmdeploy.apis.onnxruntime import ORTWrapper - self.model = ORTWrapper(model_file, device_id) - - def forward_dummy(self, lq: torch.Tensor, *args, **kwargs): - """Run test inference for restorer with ONNXRuntime. - - Args: - lq (torch.Tensor): The input low-quality image of the model. - - Returns: - list[np.ndarray] : High resolution image. - """ - ort_outputs = self.model({'input': lq}) - # only concern pred_alpha value - if isinstance(ort_outputs, (tuple, list)): - ort_outputs = ort_outputs[0] - return ort_outputs - - -class TensorRTRestorer(DeployBaseRestorer): - """Wrapper for restorer's inference with TensorRT. - - Args: - trt_file (str): The path of an input model file. - device_id (int): An integer represents device index. - test_cfg (mmcv.Config) : The test config in model config, which is - used in evaluation. - """ - - def __init__(self, - trt_file: str, - device_id: int, - test_cfg: Optional[mmcv.Config] = None, - **kwargs): - super(TensorRTRestorer, self).__init__( - device_id, test_cfg=test_cfg, **kwargs) - - from mmdeploy.apis.tensorrt import TRTWrapper, load_tensorrt_plugin - try: - load_tensorrt_plugin() - except (ImportError, ModuleNotFoundError): - warnings.warn('If input model has custom plugins, \ - you may have to build backend ops with TensorRT') - model = TRTWrapper(trt_file) - self.model = model - - def forward_dummy(self, lq: torch.Tensor, *args, **kwargs): - """Run test inference for restorer with TensorRT. - - Args: - lq (torch.Tensor): The input low-quality image of the model. - - Returns: - list[np.ndarray]: High resolution image. - """ - input_data = lq.contiguous() - with torch.cuda.device(self.device_id), torch.no_grad(): - pred = self.model({'input': input_data})['output'] - pred = pred.detach().cpu().numpy() - return pred - - -class PPLRestorer(DeployBaseRestorer): - """Wrapper for restorer's inference with ppl. - - Args: - onnx_file (str): Path of input ONNX model file. - algo_file (str): Path of PPL algorithm file. - device_id (int): An integer represents device index. - test_cfg (mmcv.Config): The test config in model config, which is - used in evaluation. - """ - - def __init__(self, - onnx_file: str, - algo_file: str, - device_id: int, - test_cfg: Optional[mmcv.Config] = None, - **kwargs): - super(PPLRestorer, self).__init__( - device_id, test_cfg=test_cfg, **kwargs) - - from mmdeploy.apis.ppl import PPLWrapper - self.model = PPLWrapper(onnx_file, algo_file, device_id) - - def forward_dummy(self, lq: torch.Tensor, *args, **kwargs): - """Run test inference for restorer with PPL. - - Args: - lq (torch.Tensor): Input low-quality image of the model. - - Returns: - list[np.ndarray]: High resolution image. - """ - ppl_outputs = self.model({'input': lq}) - # only concern pred_alpha value - if isinstance(ppl_outputs, (tuple, list)): - ppl_outputs = ppl_outputs[0] - return ppl_outputs - - -ONNXRUNTIME_RESTORER_MAP = dict(end2end=ONNXRuntimeRestorer) - -TENSORRT_RESTORER_MAP = dict(end2end=TensorRTRestorer) - -PPL_RESTORER_MAP = dict(end2end=PPLRestorer) - -BACKEND_RESTORER_MAP = { - Backend.ONNXRUNTIME: ONNXRUNTIME_RESTORER_MAP, - Backend.TENSORRT: TENSORRT_RESTORER_MAP, - Backend.PPL: PPL_RESTORER_MAP, -} - - -def build_restorer(model_files: Sequence[str], backend: Backend, - model_cfg: Union[str, mmcv.Config], device_id: int): - """Build restorer for different backend. - - Args: - model_files (Sequence[str]): Input model file(s). - backend (Backend): Target backend. - model_cfg (str | mmcv.Config): Input model config file or config - object. - device_id (int): An integer represents device index. - - Returns: - DeployBaseRestorer: Restorer for a configured backend. - """ - model_map = BACKEND_RESTORER_MAP[backend] - - model_type = 'end2end' - assert model_type in model_map, f'Unsupported model type: {model_type}' - backend_model_class = model_map[model_type] - - backend_model = backend_model_class( - *model_files, device_id=device_id, test_cfg=model_cfg.test_cfg) - - return backend_model - - -def build_editing_processor(model_files: Sequence[str], - model_cfg: Union[str, mmcv.Config], - deploy_cfg: Union[str, - mmcv.Config], device_id: int): - """Build editing processor for different backend. - - Args: - model_files (Sequence[str]): Input model file(s). - model_cfg (str | mmcv.Config): Input model config file or Config - object. - deploy_cfg (str | mmcv.Config): Input deployment config file or - Config object. - device_id (int): An integer represents device index. - - Returns: - BaseModel: Editing processor for a configured backend. - """ - model_cfg = load_config(model_cfg)[0] - deploy_cfg = load_config(deploy_cfg)[0] - - backend = get_backend(deploy_cfg) - - assert backend in BACKEND_RESTORER_MAP, \ - f'Unsupported backend type: {backend.value}' - - # TODO: Add other tasks - return build_restorer(model_files, backend, model_cfg, device_id) diff --git a/mmdeploy/mmedit/apis/visualize.py b/mmdeploy/mmedit/apis/visualize.py deleted file mode 100644 index dc2995afb3..0000000000 --- a/mmdeploy/mmedit/apis/visualize.py +++ /dev/null @@ -1,46 +0,0 @@ -import warnings - -import mmcv -import numpy as np -import torch - -from mmdeploy.utils import Backend - - -# BaseModel in mmedit doesn't implement show_result -# TODO: add show_result to different tasks -def show_result(result: np.ndarray, - output_file: str, - backend: Backend, - show: bool = True): - """Show high resolution image of mmedit. - - Args: - result: (np.ndarray): Input high resolution image. - output_file (str): Output image file to save image. - backend (Backend): Specifying backend type. - show (bool): Whether to show plotted image in windows. Defaults to - `True`. - - Returns: - np.ndarray: Drawn image, only if not `show` or `out_file`. - """ - win_name = backend.value - with torch.no_grad(): - result = result.transpose(1, 2, 0) - result = np.clip(result, 0, 1)[:, :, ::-1] - result = (result * 255.0).round() - - if output_file is not None: - show = False - - if show: - int_result = result.astype(np.uint8) - mmcv.imshow(int_result, win_name, 0) - if output_file is not None: - mmcv.imwrite(result, output_file) - - if not (show or output_file): - warnings.warn('show==False and output_file is not specified, only ' - 'result image will be returned') - return result diff --git a/mmdeploy/mmedit/export/__init__.py b/mmdeploy/mmedit/export/__init__.py deleted file mode 100644 index 0ed7292c77..0000000000 --- a/mmdeploy/mmedit/export/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from .prepare_input import (build_dataloader, build_dataset, create_input, - get_tensor_from_input) - -__all__ = [ - 'create_input', 'build_dataset', 'build_dataloader', - 'get_tensor_from_input' -] diff --git a/mmdeploy/mmedit/export/prepare_input.py b/mmdeploy/mmedit/export/prepare_input.py deleted file mode 100644 index 63efc94910..0000000000 --- a/mmdeploy/mmedit/export/prepare_input.py +++ /dev/null @@ -1,201 +0,0 @@ -from typing import Any, Dict, Optional, Sequence, Union - -import mmcv -import numpy as np -from mmcv.parallel import collate, scatter -from mmedit.datasets import build_dataloader as build_dataloader_mmedit -from mmedit.datasets import build_dataset as build_dataset_mmedit -from mmedit.datasets.pipelines import Compose -from torch.utils.data.dataset import Dataset - -from mmdeploy.utils import Task, load_config - - -def _preprocess_cfg(config: Union[str, mmcv.Config], task: Task, - load_from_file: bool, is_static_cfg: bool, - input_shape: Sequence[int]): - """Remove unnecessary information in config. - - Args: - model_cfg (str | mmcv.Config): The input model config. - task (Task): Specifying editing task type. - load_from_file (bool): Whether the input is a filename of a numpy - matrix. If this variable is True, extra preprocessing is required. - is_static_cfg (bool): Whether the config specifys a static export. - If this variable if True, the input image will be resize to a fix - resolution. - input_shape (Sequence[int]): A list of two integer in (width, height) - format specifying input shape. Defaults to `None`. - """ - - # TODO: Differentiate the editing tasks (e.g. restorers and mattors - # preprocess the data in differenet ways) - - if task == Task.SUPER_RESOLUTION: - keys_to_remove = ['gt', 'gt_path'] - else: - raise NotImplementedError(f'Unknown task type: {task.value}') - - # MMEdit doesn't support LoadImageFromWebcam. - # Remove "LoadImageFromFile" and related metakeys. - if not load_from_file: - config.test_pipeline.pop(0) - if task == Task.SUPER_RESOLUTION: - keys_to_remove.append('lq_path') - - # Fix the input shape by 'Resize' - if is_static_cfg: - if task == Task.SUPER_RESOLUTION: - resize = { - 'type': 'Resize', - 'scale': (input_shape[0], input_shape[1]), - 'keys': ['lq'] - } - config.test_pipeline.insert(1, resize) - - for key in keys_to_remove: - for pipeline in list(config.test_pipeline): - if 'key' in pipeline and key == pipeline['key']: - config.test_pipeline.remove(pipeline) - if 'keys' in pipeline: - while key in pipeline['keys']: - pipeline['keys'].remove(key) - if len(pipeline['keys']) == 0: - config.test_pipeline.remove(pipeline) - if 'meta_keys' in pipeline: - while key in pipeline['meta_keys']: - pipeline['meta_keys'].remove(key) - - -def create_input(task: Task, - model_cfg: Union[str, mmcv.Config], - imgs: Union[str, np.ndarray], - input_shape: Optional[Sequence[int]] = None, - device: Optional[str] = 'cuda:0'): - """Create input for editing processor. - - Args: - task (Task): Specifying editing task type. - model_cfg (str | mmcv.Config): The input model config. - imgs (str | np.ndarray): Input image(s). - input_shape (Sequence[int]): A list of two integer in (width, height) - format specifying input shape. Defaults to `None`. - device (str): A string represents device type. Default is 'cuda:0'. - - Returns: - tuple: (data, img), meta information for the input image and input. - """ - if isinstance(imgs, (list, tuple)): - if not isinstance(imgs[0], (np.ndarray, str)): - raise AssertionError('imgs must be strings or numpy arrays') - elif isinstance(imgs, (np.ndarray, str)): - imgs = [imgs] - else: - raise AssertionError('imgs must be strings or numpy arrays') - - cfg = load_config(model_cfg)[0].copy() - - _preprocess_cfg( - cfg, - task=task, - load_from_file=isinstance(imgs[0], str), - is_static_cfg=input_shape is not None, - input_shape=input_shape) - - test_pipeline = Compose(cfg.test_pipeline) - - data_arr = [] - for img in imgs: - # TODO: This is only for restore. Add condiction statement. - if isinstance(img, np.ndarray): - data = dict(lq=img) - else: - data = dict(lq_path=img) - - data = test_pipeline(data) - data_arr.append(data) - - data = collate(data_arr, samples_per_gpu=len(imgs)) - - # TODO: This is only for restore. Add condiction statement. - data['img'] = data['lq'] - - if device != 'cpu': - data = scatter(data, [device])[0] - - return data, data['img'] - - -def build_dataset(dataset_cfg: Union[str, mmcv.Config], **kwargs): - """Build dataset for processor. - - Args: - dataset_cfg (str | mmcv.Config): The input dataset config. - - Returns: - Dataset: A PyTorch dataset. - """ - dataset_cfg = load_config(dataset_cfg)[0] - data = dataset_cfg.data - - dataset = build_dataset_mmedit(data.test) - return dataset - - -def build_dataloader(dataset: Dataset, - samples_per_gpu: int, - workers_per_gpu: int, - num_gpus: int = 1, - dist: bool = False, - shuffle: bool = False, - seed: Optional[int] = None, - drop_last: bool = False, - pin_memory: bool = True, - persistent_workers: bool = True, - **kwargs): - """Build PyTorch DataLoader. - - In distributed training, each GPU/process has a dataloader. - In non-distributed training, there is only one dataloader for all GPUs. - - Args: - dataset (:obj:`Dataset`): A PyTorch dataset. - samples_per_gpu (int): Number of samples on each GPU, i.e., - batch size of each GPU. - workers_per_gpu (int): How many subprocesses to use for data - loading for each GPU. - num_gpus (int): Number of GPUs. Only used in non-distributed - training. Default: 1. - dist (bool): Distributed training/test or not. Default: True. - shuffle (bool): Whether to shuffle the data at every epoch. - Default: True. - seed (int | None): Seed to be used. Default: None. - drop_last (bool): Whether to drop the last incomplete batch in epoch. - Default: False - pin_memory (bool): Whether to use pin_memory in DataLoader. - Default: True - persistent_workers (bool): If True, the data loader will not shutdown - the worker processes after a dataset has been consumed once. - This allows to maintain the workers Dataset instances alive. - The argument also has effect in PyTorch>=1.7.0. - Default: True - kwargs (dict, optional): Any keyword argument to be used to initialize - DataLoader. - - Returns: - DataLoader: A PyTorch dataloader. - """ - return build_dataloader_mmedit(dataset, samples_per_gpu, workers_per_gpu, - num_gpus, dist, shuffle, seed, drop_last, - pin_memory, persistent_workers, **kwargs) - - -def get_tensor_from_input(input_data: Dict[str, Any]): - """Get input tensor from input data. - - Args: - input_data (dict): Input data containing meta info and image tensor. - Returns: - torch.Tensor: An image in `Tensor`. - """ - return input_data['lq'] diff --git a/mmdeploy/mmedit/models/backbones/__init__.py b/mmdeploy/mmedit/models/backbones/__init__.py deleted file mode 100644 index aeaf91486a..0000000000 --- a/mmdeploy/mmedit/models/backbones/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .sr_backbones import * # noqa: F401,F403 diff --git a/mmdeploy/mmocr/apis/__init__.py b/mmdeploy/mmocr/apis/__init__.py deleted file mode 100644 index 5cfb27301c..0000000000 --- a/mmdeploy/mmocr/apis/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .inference import build_ocr_processor -from .visualize import show_result - -__all__ = ['build_ocr_processor', 'show_result'] diff --git a/mmdeploy/mmocr/apis/inference.py b/mmdeploy/mmocr/apis/inference.py deleted file mode 100644 index ab8e8997f9..0000000000 --- a/mmdeploy/mmocr/apis/inference.py +++ /dev/null @@ -1,538 +0,0 @@ -from typing import Iterable, Sequence, Union - -import mmcv -import torch -from mmdet.models.builder import DETECTORS -from mmocr.datasets import DATASETS -from mmocr.models.textdet.detectors import (SingleStageTextDetector, - TextDetectorMixin) -from mmocr.models.textrecog.recognizer import EncodeDecodeRecognizer - -from mmdeploy.utils.config_utils import (Backend, Task, get_backend, - get_task_type, load_config) - - -@DETECTORS.register_module() -class DeployBaseTextDetector(TextDetectorMixin, SingleStageTextDetector): - """Base Class of Wrapper for TextDetector. - - Args: - cfg (str | mmcv.ConfigDict): Input model config. - device_id (int): An integer represents device index. - show_score (bool): Whether to show scores. Defaults to `False`. - """ - - def __init__(self, - cfg: Union[mmcv.Config, mmcv.ConfigDict], - device_id: int, - show_score: bool = False, - *args, - **kwargs): - SingleStageTextDetector.__init__(self, cfg.model.backbone, - cfg.model.neck, cfg.model.bbox_head) - TextDetectorMixin.__init__(self, show_score) - self.device_id = device_id - self.show_score = show_score - self.cfg = cfg - - def forward_train(self, img, img_metas, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def aug_test(self, imgs, img_metas, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def extract_feat(self, imgs): - raise NotImplementedError('This method is not implemented.') - - def simple_test(self, - img: torch.Tensor, - img_metas: Sequence[dict], - rescale: bool = False, - *args, - **kwargs): - """Run forward test. - - Args: - img (torch.Tensor): Input image tensor. - img_metas (Sequence[dict]): A list of meta info for image(s). - - Returns: - list: A list of predictions. - """ - pred = self.forward_of_backend(img, img_metas, *args, **kwargs) - if len(img_metas) > 1: - boundaries = [ - self.bbox_head.get_boundary( - *(pred[i].unsqueeze(0)), [img_metas[i]], rescale=rescale) - for i in range(len(img_metas)) - ] - - else: - boundaries = [ - self.bbox_head.get_boundary(*pred, img_metas, rescale=rescale) - ] - return boundaries - - -@DETECTORS.register_module() -class DeployBaseRecognizer(EncodeDecodeRecognizer): - """Base Class of Wrapper for TextRecognizer. - - Args: - cfg (str | mmcv.ConfigDict): Input model config. - device_id (int): An integer represents device index. - show_score (bool): Whether to show scores. Defaults to `False`. - """ - - def __init__(self, - cfg: Union[mmcv.Config, mmcv.ConfigDict], - device_id: int, - show_score: bool = False, - *args, - **kwargs): - super(DeployBaseRecognizer, - self).__init__(None, cfg.model.backbone, cfg.model.encoder, - cfg.model.decoder, cfg.model.loss, - cfg.model.label_convertor, None, None, 40, None) - self.device_id = device_id - self.show_score = show_score - self.cfg = cfg - - def forward_train(self, img, img_metas, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def aug_test(self, imgs, img_metas, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def extract_feat(self, imgs): - raise NotImplementedError('This method is not implemented.') - - def forward(self, img: Union[torch.Tensor, Sequence[torch.Tensor]], - img_metas: Sequence[dict], *args, **kwargs): - """Run forward. - - Args: - imgs (torch.Tensor | Sequence[torch.Tensor]): Image input tensor. - img_metas (Sequence[dict]): List of image information. - - Returns: - list[str]: Text label result of each image. - """ - - if isinstance(img, list): - for idx, each_img in enumerate(img): - if each_img.dim() == 3: - img[idx] = each_img.unsqueeze(0) - img = img[0] # avoid aug_test - img_metas = img_metas[0] - else: - if len(img_metas) == 1 and isinstance(img_metas[0], list): - img_metas = img_metas[0] - - return self.simple_test(img, img_metas, **kwargs) - - def simple_test(self, img: torch.Tensor, img_metas: Sequence[dict], *args, - **kwargs): - """Run forward test. - - Args: - imgs (torch.Tensor): Image input tensor. - img_metas (Sequence[dict]): List of image information. - - Returns: - list[str]: Text label result of each image. - """ - pred = self.forward_of_backend(img, img_metas, *args, **kwargs) - label_indexes, label_scores = self.label_convertor.tensor2idx( - pred, img_metas) - label_strings = self.label_convertor.idx2str(label_indexes) - - # flatten batch results - results = [] - for string, score in zip(label_strings, label_scores): - results.append(dict(text=string, score=score)) - - return results - - -class ONNXRuntimeDetector(DeployBaseTextDetector): - """Wrapper for TextDetector with ONNX Runtime. - - Args: - model_file (str): The path of input model file. - cfg (str | mmcv.ConfigDict): Input model config. - device_id (int): An integer represents device index. - show_score (bool): Whether to show scores. Defaults to `False`. - """ - - def __init__(self, - model_file: str, - cfg: Union[mmcv.Config, mmcv.ConfigDict], - device_id: int, - show_score: bool = False, - *args, - **kwargs): - super(ONNXRuntimeDetector, self).__init__(cfg, device_id, show_score) - from mmdeploy.apis.onnxruntime import ORTWrapper - self.model = ORTWrapper(model_file, device_id) - - def forward_of_backend(self, img: torch.Tensor, img_metas: Iterable, *args, - **kwargs): - """Implement forward test with a backend. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - img_metas (Sequence[dict]]): List of image information. - Returns: - np.ndarray: Prediction of input model. - """ - onnx_pred = self.model({'input': img}) - onnx_pred = torch.from_numpy(onnx_pred[0]) - return onnx_pred - - -class ONNXRuntimeRecognizer(DeployBaseRecognizer): - """Wrapper for TextRecognizer with ONNX Runtime. - - Args: - model_file (str): The path of input model file. - cfg (str | mmcv.ConfigDict): Input model config. - device_id (int): An integer represents device index. - show_score (bool): Whether to show scores. Defaults to `False`. - """ - - def __init__(self, - model_file: str, - cfg: Union[mmcv.Config, mmcv.ConfigDict], - device_id: int, - show_score: bool = False, - *args, - **kwargs): - super(ONNXRuntimeRecognizer, self).__init__(cfg, device_id, show_score) - from mmdeploy.apis.onnxruntime import ORTWrapper - self.model = ORTWrapper(model_file, device_id) - - def forward_of_backend(self, img: torch.Tensor, img_metas: Sequence[dict], - *args, **kwargs): - """Implement forward test with a backend. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - img_metas (Sequence[dict]]): List of image information. - Returns: - np.ndarray: Prediction of input model. - """ - onnx_pred = self.model({'input': img}) - onnx_pred = torch.from_numpy(onnx_pred[0]) - return onnx_pred - - -class TensorRTDetector(DeployBaseTextDetector): - """Wrapper for TextDetector with TensorRT. - - Args: - model_file (str): The path of input model file. - cfg (str | mmcv.ConfigDict): Input model config. - device_id (int): An integer represents device index. - show_score (bool): Whether to show scores. Defaults to `False`. - """ - - def __init__(self, - model_file: str, - cfg: Union[mmcv.Config, mmcv.ConfigDict], - device_id: int, - show_score: bool = False, - *args, - **kwargs): - super(TensorRTDetector, self).__init__(cfg, device_id, show_score) - from mmdeploy.apis.tensorrt import TRTWrapper - model = TRTWrapper(model_file) - self.model = model - - def forward_of_backend(self, img: torch.Tensor, img_metas: Sequence[dict], - *args, **kwargs): - """Implement forward test with a backend. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - img_metas (Sequence[dict]]): List of image information. - Returns: - np.ndarray: Prediction of input model. - """ - with torch.cuda.device(self.device_id), torch.no_grad(): - trt_pred = self.model({'input': img})['output'] - return trt_pred - - -class TensorRTRecognizer(DeployBaseRecognizer): - """Wrapper for TextRecognizer with TensorRT. - - Args: - model_file (str): The path of input model file. - cfg (str | mmcv.ConfigDict): Input model config. - device_id (int): An integer represents device index. - show_score (bool): Whether to show scores. Defaults to `False`. - """ - - def __init__(self, - model_file: str, - cfg: Union[mmcv.Config, mmcv.ConfigDict], - device_id: int, - show_score: bool = False, - *args, - **kwargs): - super(TensorRTRecognizer, self).__init__(cfg, device_id, show_score) - from mmdeploy.apis.tensorrt import TRTWrapper - model = TRTWrapper(model_file) - self.model = model - - def forward_of_backend(self, img: torch.Tensor, img_metas: Sequence[dict], - *args, **kwargs): - """Implement forward test with a backend. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - img_metas (Sequence[dict]]): List of image information. - Returns: - torch.Tensor: Prediction of input model. - """ - with torch.cuda.device(self.device_id), torch.no_grad(): - trt_pred = self.model({'input': img})['output'] - return trt_pred - - -class NCNNDetector(DeployBaseTextDetector): - """Wrapper for TextDetector with NCNN. - - Args: - model_file (Sequence[str]): Paths of input model files. - cfg (str | mmcv.ConfigDict): Input model config. - device_id (int): An integer represents device index. - show_score (bool): Whether to show scores. Defaults to `False`. - """ - - def __init__(self, - model_file: Sequence[str], - cfg: Union[mmcv.Config, mmcv.ConfigDict], - device_id: int, - show_score: bool = False): - super(NCNNDetector, self).__init__(cfg, device_id, show_score) - from mmdeploy.apis.ncnn import NCNNWrapper - self.model = NCNNWrapper( - model_file[0], model_file[1], output_names=['output']) - - def forward_of_backend(self, img: torch.Tensor, img_metas: Sequence[dict], - *args, **kwargs): - """Implement forward test with a backend. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - img_metas (Sequence[dict]]): List of image information. - Returns: - torch.Tensor: Prediction of input model. - """ - pred = self.model({'input': img})['output'] - return pred - - -class NCNNRecognizer(DeployBaseRecognizer): - """Wrapper for TextRecognizer with NCNN. - - Args: - model_file (Sequence[str]): Paths of input model files. - cfg (str | mmcv.ConfigDict): Input model config. - device_id (int): An integer represents device index. - show_score (bool): Whether to show scores. Defaults to `False`. - """ - - def __init__(self, - model_file: Sequence[str], - cfg: Union[mmcv.Config, mmcv.ConfigDict], - device_id: int, - show_score: bool = False): - super(NCNNRecognizer, self).__init__(cfg, device_id, show_score) - from mmdeploy.apis.ncnn import NCNNWrapper - self.model = NCNNWrapper( - model_file[0], model_file[1], output_names=['output']) - - def forward_of_backend(self, img: torch.Tensor, img_metas: Sequence[dict], - *args, **kwargs): - """Implement forward test with a backend. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - img_metas (Sequence[dict]]): List of image information. - Returns: - torch.Tensor: Prediction of input model. - """ - pred = self.model({'input': img})['output'] - return pred - - -class PPLDetector(DeployBaseTextDetector): - """Wrapper for TextDetector with PPL. - - Args: - model_file (Sequence[str]): Paths of input model files. - cfg (str | mmcv.ConfigDict): Input model config. - device_id (int): An integer represents device index. - show_score (bool): Whether to show scores. Defaults to `False`. - """ - - def __init__(self, - model_file: str, - cfg: Union[mmcv.Config, mmcv.ConfigDict], - device_id: int, - show_score: bool = False, - *args, - **kwargs): - super(PPLDetector, self).__init__(cfg, device_id, show_score) - from mmdeploy.apis.ppl import PPLWrapper - model = PPLWrapper(model_file[0], model_file[1], device_id) - self.model = model - - def forward_of_backend(self, img: torch.Tensor, img_metas: Sequence[dict], - *args, **kwargs): - """Implement forward test with a backend. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - img_metas (Sequence[dict]]): List of image information. - Returns: - torch.Tensor: Prediction of input model. - """ - with torch.cuda.device(self.device_id), torch.no_grad(): - ppl_pred = self.model({'input': img}) - ppl_pred = torch.from_numpy(ppl_pred[0]) - return ppl_pred - - -class PPLRecognizer(DeployBaseRecognizer): - """Wrapper for TextRecognizer with PPL. - - Args: - onnx_file (str): Path of input ONNX model file. - algo_file (str): Path of PPL algorithm file. - cfg (str | mmcv.ConfigDict): Input model config. - device_id (int): An integer represents device index. - show_score (bool): Whether to show scores. Defaults to `False`. - """ - - def __init__(self, - model_file: str, - algo_file: str, - cfg: Union[mmcv.Config, mmcv.ConfigDict], - device_id: int, - show_score: bool = False, - *args, - **kwargs): - super(PPLRecognizer, self).__init__(cfg, device_id, show_score) - from mmdeploy.apis.ppl import PPLWrapper - model = PPLWrapper(model_file, algo_file, device_id) - self.model = model - - def forward_of_backend(self, img: torch.Tensor, img_metas: Sequence[dict], - *args, **kwargs): - """Implement forward test with a backend. - - Args: - imgs (torch.Tensor): Input image(s) in [N x C x H x W] format. - img_metas (Sequence[dict]]): List of image information. - Returns: - torch.Tensor: Prediction of input model. - """ - with torch.cuda.device(self.device_id), torch.no_grad(): - ppl_pred = self.model({'input': img})[0] - ppl_pred = torch.from_numpy(ppl_pred[0]) - return ppl_pred - - -def get_classes_from_config(model_cfg: Union[str, mmcv.Config], **kwargs): - """Get class name from config. - - Args: - model_cfg (str | mmcv.Config): Input model config file or - Config object. - - Returns: - list[str]: A list of string specifying names of different class. - """ - # load cfg if necessary - model_cfg = load_config(model_cfg)[0] - module_dict = DATASETS.module_dict - data_cfg = model_cfg.data - - if 'train' in data_cfg: - module = module_dict[data_cfg.train.type] - elif 'val' in data_cfg: - module = module_dict[data_cfg.val.type] - elif 'test' in data_cfg: - module = module_dict[data_cfg.test.type] - else: - raise RuntimeError(f'No dataset config found in: {model_cfg}') - - return module.CLASSES - - -TASK_ONNXRUNTIME_MAP = { - Task.TEXT_DETECTION: ONNXRuntimeDetector, - Task.TEXT_RECOGNITION: ONNXRuntimeRecognizer -} - -TASK_TENSORRT_MAP = { - Task.TEXT_DETECTION: TensorRTDetector, - Task.TEXT_RECOGNITION: TensorRTRecognizer -} - -TASK_PPL_MAP = { - Task.TEXT_DETECTION: PPLDetector, - Task.TEXT_RECOGNITION: PPLRecognizer -} - -TASK_NCNN_MAP = { - Task.TEXT_DETECTION: NCNNDetector, - Task.TEXT_RECOGNITION: NCNNRecognizer -} - -BACKEND_TASK_MAP = { - Backend.ONNXRUNTIME: TASK_ONNXRUNTIME_MAP, - Backend.TENSORRT: TASK_TENSORRT_MAP, - Backend.PPL: TASK_PPL_MAP, - Backend.NCNN: TASK_NCNN_MAP -} - - -def build_ocr_processor(model_files: Sequence[str], - model_cfg: Union[str, mmcv.Config], - deploy_cfg: Union[str, mmcv.Config], device_id: int, - **kwargs): - """Build text detector or recognizer for a backend. - - Args: - model_files (Sequence[str]): Input model file(s). - model_cfg (str | mmcv.Config): Input model config file or Config - object. - deploy_cfg (str | mmcv.Config): Input deployment config file or - Config object. - device_id (int): An integer represents device index. - - Returns: - nn.Module: An instance of text detector or recognizer. - """ - # load cfg if necessary - deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg) - - backend = get_backend(deploy_cfg) - task = get_task_type(deploy_cfg) - - assert backend in BACKEND_TASK_MAP, \ - f'Unsupported backend type: {backend.value}' - assert task in BACKEND_TASK_MAP[backend], \ - f'Unsupported task type: {task.value}' - backend_task_class = BACKEND_TASK_MAP[backend][task] - - model_files = model_files[0] if len(model_files) == 1 else model_files - backend_detector = backend_task_class( - model_file=model_files, cfg=model_cfg, device_id=device_id, **kwargs) - - return backend_detector diff --git a/mmdeploy/mmocr/apis/visualize.py b/mmdeploy/mmocr/apis/visualize.py deleted file mode 100644 index 7b584a8fcb..0000000000 --- a/mmdeploy/mmocr/apis/visualize.py +++ /dev/null @@ -1,35 +0,0 @@ -import numpy as np -import torch - -from mmdeploy.utils import Backend - - -def show_result(model: torch.nn.Module, - image: np.ndarray, - result: list, - output_file: str, - backend: Backend, - show: bool = True, - score_thr: float = 0.3): - """Show predictions of text detector or recognizer. - - Args: - model (nn.Module): Input model which has `show_result` method. - image: (np.ndarray): Input image to draw predictions. - result (list): A list of predictions. - output_file (str): Output image file to save drawn predictions. - backend (Backend): Specifying backend type. - show (bool): Whether to show plotted image in windows. Defaults to - `True`. - score_thr (float): Score threshold for result, defaults to `0.3`. - - Returns: - np.ndarray: Drawn image, only if not `show` or `out_file`. - """ - return model.show_result( - image, - result, - score_thr=score_thr, - show=show, - win_name=backend.value, - out_file=output_file) diff --git a/mmdeploy/mmocr/export/__init__.py b/mmdeploy/mmocr/export/__init__.py deleted file mode 100644 index 0ed7292c77..0000000000 --- a/mmdeploy/mmocr/export/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from .prepare_input import (build_dataloader, build_dataset, create_input, - get_tensor_from_input) - -__all__ = [ - 'create_input', 'build_dataset', 'build_dataloader', - 'get_tensor_from_input' -] diff --git a/mmdeploy/mmocr/export/prepare_input.py b/mmdeploy/mmocr/export/prepare_input.py deleted file mode 100644 index f562366950..0000000000 --- a/mmdeploy/mmocr/export/prepare_input.py +++ /dev/null @@ -1,194 +0,0 @@ -from typing import Any, Optional, Sequence, Union - -import mmcv -import numpy as np -from mmcv.parallel import DataContainer, collate, scatter -from mmdet.datasets import replace_ImageToTensor -from mmocr.datasets import build_dataloader as build_dataloader_mmocr -from mmocr.datasets import build_dataset as build_dataset_mmocr -from torch.utils.data import Dataset - -from mmdeploy.utils import Task, load_config - - -def create_input(task: Task, - model_cfg: Union[str, mmcv.Config], - imgs: Any, - input_shape: Sequence[int] = None, - device: str = 'cuda:0'): - """Create input for text detector/recognizer. - - Args: - task (Task): Specifying task type. - model_cfg (str | mmcv.Config): The input model config. - imgs (Any): Input image(s), accpeted data type are `str`, - `np.ndarray`, `torch.Tensor`. - input_shape (list[int]): A list of two integer in (width, height) - format specifying input shape. Defaults to `None`. - device (str): A string represents device type. Default is 'cuda:0'. - - Returns: - tuple: (data, img), meta information for the input image and input. - """ - if isinstance(imgs, (list, tuple)): - if not isinstance(imgs[0], (np.ndarray, str)): - raise AssertionError('imgs must be strings or numpy arrays') - - elif isinstance(imgs, (np.ndarray, str)): - imgs = [imgs] - else: - raise AssertionError('imgs must be strings or numpy arrays') - - if model_cfg.data.test['type'] == 'ConcatDataset': - model_cfg.data.test.pipeline = \ - model_cfg.data.test['datasets'][0].pipeline - - is_ndarray = isinstance(imgs[0], np.ndarray) - - if is_ndarray: - model_cfg = model_cfg.copy() - # set loading pipeline type - model_cfg.data.test.pipeline[0].type = 'LoadImageFromNdarray' - - test_pipeline = model_cfg.data.test.pipeline - test_pipeline = replace_ImageToTensor(test_pipeline) - # for static exporting - if input_shape is not None: - if task == Task.TEXT_DETECTION: - test_pipeline[1].img_scale = tuple(input_shape) - test_pipeline[1].transforms[0].keep_ratio = False - test_pipeline[1].transforms[0].img_scale = tuple(input_shape) - elif task == Task.TEXT_RECOGNITION: - resize = { - 'height': input_shape[1], - 'min_width': input_shape[0], - 'max_width': input_shape[0], - 'keep_aspect_ratio': False - } - if 'transforms' in test_pipeline[1]: - if test_pipeline[1].transforms[0].type == 'ResizeOCR': - test_pipeline[1].transforms[0].height = input_shape[1] - test_pipeline[1].transforms[0].max_width = input_shape[0] - else: - raise ValueError( - f'Transforms[0] should be ResizeOCR, but got\ - {test_pipeline[1].transforms[0].type}') - else: - test_pipeline[1].update(resize) - from mmdet.datasets.pipelines import Compose - from mmocr.datasets import build_dataset # noqa: F401 - test_pipeline = Compose(test_pipeline) - - data_list = [] - for img in imgs: - # prepare data - if is_ndarray: - # directly add img - data = dict(img=img) - else: - # add information into dict - data = dict(img_info=dict(filename=img), img_prefix=None) - - # build the data pipeline - data = test_pipeline(data) - # get tensor from list to stack for batch mode (text detection) - data_list.append(data) - - if isinstance(data_list[0]['img'], list) and len(data_list) > 1: - raise Exception('aug test does not support ' - f'inference with batch size ' - f'{len(data_list)}') - - data = collate(data_list, samples_per_gpu=len(imgs)) - - # process img_metas - if isinstance(data['img_metas'], list): - data['img_metas'] = [ - img_metas.data[0] for img_metas in data['img_metas'] - ] - else: - data['img_metas'] = data['img_metas'].data - - if isinstance(data['img'], list): - data['img'] = [img.data for img in data['img']] - if isinstance(data['img'][0], list): - data['img'] = [img[0] for img in data['img']] - else: - data['img'] = data['img'].data - - if device != 'cpu': - data = scatter(data, [device])[0] - - return data, data['img'] - - -def build_dataset(dataset_cfg: Union[str, mmcv.Config], - dataset_type: str = 'val', - **kwargs): - """Build dataset for detector/recognizer. - - Args: - dataset_cfg (str | mmcv.Config): The input dataset config. - dataset_type (str): A string represents dataset type, e.g.: 'train', - 'test', 'val'. Defaults to 'val'. - - Returns: - Dataset: A PyTorch dataset. - """ - dataset_cfg = load_config(dataset_cfg)[0].copy() - - data = dataset_cfg.data - assert dataset_type in data - dataset = build_dataset_mmocr(data[dataset_type]) - - return dataset - - -def build_dataloader(dataset: Dataset, - samples_per_gpu: int, - workers_per_gpu: int, - num_gpus: int = 1, - dist: bool = False, - shuffle: bool = False, - seed: Optional[int] = None, - **kwargs): - """Build dataloader for detector/recognizer. - - Args: - dataset (Dataset): Input dataset. - samples_per_gpu (int): Number of training samples on each GPU, i.e., - batch size of each GPU. - workers_per_gpu (int): How many subprocesses to use for data loading - for each GPU. - num_gpus (int): Number of GPUs. Only used in non-distributed training. - dist (bool): Distributed training/test or not. Defaults to `False`. - shuffle (bool): Whether to shuffle the data at every epoch. - Defaults to `False`. - seed (int): An integer set to be seed. Default is `None`. - kwargs: Any other keyword argument to be used to initialize DataLoader. - - Returns: - DataLoader: A PyTorch dataloader. - """ - return build_dataloader_mmocr( - dataset, - samples_per_gpu, - workers_per_gpu, - num_gpus=num_gpus, - dist=dist, - shuffle=shuffle, - seed=seed, - **kwargs) - - -def get_tensor_from_input(input_data: tuple): - """Get input tensor from input data. - - Args: - input_data (tuple): Input data containing meta info and image tensor. - Returns: - torch.Tensor: An image in `Tensor`. - """ - if isinstance(input_data['img'], DataContainer): - return input_data['img'].data[0] - return input_data['img'][0] diff --git a/mmdeploy/mmocr/models/__init__.py b/mmdeploy/mmocr/models/__init__.py deleted file mode 100644 index 296cfe2dfd..0000000000 --- a/mmdeploy/mmocr/models/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .textdet import * # noqa: F401,F403 -from .textrecog import * # noqa: F401,F403 diff --git a/mmdeploy/mmocr/models/textdet/__init__.py b/mmdeploy/mmocr/models/textdet/__init__.py deleted file mode 100644 index b64298155e..0000000000 --- a/mmdeploy/mmocr/models/textdet/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .detectors.single_stage_text_detector import * # noqa: F401,F403 -from .necks.fpn_cat import * # noqa: F401,F403 diff --git a/mmdeploy/mmocr/models/textrecog/__init__.py b/mmdeploy/mmocr/models/textrecog/__init__.py deleted file mode 100644 index 075e9cc92c..0000000000 --- a/mmdeploy/mmocr/models/textrecog/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -from .decoders import * # noqa: F401, F403 -from .encoders import sar_encoder__forward -from .layers import * # noqa: F401, F403 -from .recognizer.base import base_recognizer__forward -from .recognizer.encode_decode_recognizer import \ - encode_decode_recognizer__simple_test -from .recognizer.sar import SARNet - -__all__ = [ - 'encode_decode_recognizer__simple_test', 'base_recognizer__forward', - 'sar_encoder__forward', 'SARNet' -] diff --git a/mmdeploy/mmocr/models/textrecog/decoders/__init__.py b/mmdeploy/mmocr/models/textrecog/decoders/__init__.py deleted file mode 100644 index 032de25fb6..0000000000 --- a/mmdeploy/mmocr/models/textrecog/decoders/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .crnn_decoder import crnndecoder__forward_train__ncnn -from .sar_decoder import * # noqa: F401, F403 - -__all__ = ['crnndecoder__forward_train__ncnn'] diff --git a/mmdeploy/mmocr/models/textrecog/encoders/__init__.py b/mmdeploy/mmocr/models/textrecog/encoders/__init__.py deleted file mode 100644 index 41462dcab8..0000000000 --- a/mmdeploy/mmocr/models/textrecog/encoders/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .sar_encoder import sar_encoder__forward - -__all__ = ['sar_encoder__forward'] diff --git a/mmdeploy/mmocr/models/textrecog/layers/__init__.py b/mmdeploy/mmocr/models/textrecog/layers/__init__.py deleted file mode 100644 index 3dd6cef15f..0000000000 --- a/mmdeploy/mmocr/models/textrecog/layers/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .lstm_layer import bidirectionallstm__forward__ncnn - -__all__ = ['bidirectionallstm__forward__ncnn'] diff --git a/mmdeploy/mmseg/__init__.py b/mmdeploy/mmseg/__init__.py deleted file mode 100644 index d2b62e1cb6..0000000000 --- a/mmdeploy/mmseg/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .export import * # noqa: F401,F403 -from .models import * # noqa: F401,F403 diff --git a/mmdeploy/mmseg/apis/__init__.py b/mmdeploy/mmseg/apis/__init__.py deleted file mode 100644 index 9d5a605822..0000000000 --- a/mmdeploy/mmseg/apis/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .inference import build_segmentor -from .visualize import show_result - -__all__ = ['build_segmentor', 'show_result'] diff --git a/mmdeploy/mmseg/apis/inference.py b/mmdeploy/mmseg/apis/inference.py deleted file mode 100644 index 1b9b53c01b..0000000000 --- a/mmdeploy/mmseg/apis/inference.py +++ /dev/null @@ -1,289 +0,0 @@ -from typing import Sequence, Union - -import mmcv -import numpy as np -import torch -from mmseg.datasets import DATASETS -from mmseg.models.segmentors.base import BaseSegmentor -from mmseg.ops import resize - -from mmdeploy.utils.config_utils import Backend, get_backend, load_config - - -class DeployBaseSegmentor(BaseSegmentor): - """Base Class of wrapper for segmentation's inference. - - Args: - class_names (Sequence[str]): A list of string specifying class names. - palette (np.ndarray): The palette of segmentation map. - device_id (int): An integer represents device index. - """ - - def __init__(self, class_names: Sequence[str], palette: np.ndarray, - device_id: int): - super(DeployBaseSegmentor, self).__init__(init_cfg=None) - self.CLASSES = class_names - self.device_id = device_id - self.PALETTE = palette - - def extract_feat(self, imgs): - raise NotImplementedError('This method is not implemented.') - - def encode_decode(self, img, img_metas): - raise NotImplementedError('This method is not implemented.') - - def forward_train(self, imgs, img_metas, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def simple_test(self, img, img_meta, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def aug_test(self, imgs, img_metas, **kwargs): - raise NotImplementedError('This method is not implemented.') - - def forward(self, img: Sequence[torch.Tensor], img_metas: Sequence[dict], - **kwargs): - """Run forward test. - - Args: - img (Sequence[torch.Tensor]]): A list of input image tensor(s). - img_metas (Sequence[dict]]): A list of dict containing image(s) - meta information. - - Returns: - list[np.ndarray]: A list of segmentation result. - """ - if isinstance(img, (list, tuple)): - img = img[0] - img = img.contiguous() - seg_pred = self.forward_test(img, img_metas, **kwargs) - # whole mode supports dynamic shape - ori_shape = img_metas[0][0]['ori_shape'] - if not (ori_shape[0] == seg_pred.shape[-2] - and ori_shape[1] == seg_pred.shape[-1]): - seg_pred = torch.from_numpy(seg_pred).float() - seg_pred = resize( - seg_pred, size=tuple(ori_shape[:2]), mode='nearest') - seg_pred = seg_pred.long().detach().cpu().numpy() - # remove unnecessary dim - seg_pred = seg_pred.squeeze(1) - seg_pred = list(seg_pred) - return seg_pred - - -class ONNXRuntimeSegmentor(DeployBaseSegmentor): - """Wrapper for segmentation's inference with ONNX Runtime. - - Args: - model_file (str): The path of input model file. - class_names (Sequence[str]): A list of string specifying class names. - palette (np.ndarray): The palette of segmentation map. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: str, class_names: Sequence[str], - palette: np.ndarray, device_id: int): - super(ONNXRuntimeSegmentor, self).__init__(class_names, palette, - device_id) - from mmdeploy.apis.onnxruntime import ORTWrapper - self.model = ORTWrapper(model_file, device_id) - - def forward_test(self, imgs: torch.Tensor, img_metas: Sequence[dict], - **kwargs): - """Run forward test to get predictions. - - Args: - imgs (torch.Tensor): Input tensor of the model. - img_metas (Sequence[dict]]): A list of dict containing image(s) - meta information. - Returns: - torch.Tensor: Segmentation result. - """ - seg_pred = self.model({'input': imgs})[0] - return seg_pred - - -class TensorRTSegmentor(DeployBaseSegmentor): - """Wrapper for segmentation's inference with TensorRT. - - Args: - model_file (str): The path of input model file. - class_names (Sequence[str]): A list of string specifying class names. - palette (np.ndarray): The palette of segmentation map. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: str, class_names: Sequence[str], - palette: np.ndarray, device_id: int): - super(TensorRTSegmentor, self).__init__(class_names, palette, - device_id) - from mmdeploy.apis.tensorrt import TRTWrapper - - model = TRTWrapper(model_file) - self.model = model - self.output_name = self.model.output_names[0] - - def forward_test(self, imgs: torch.Tensor, img_metas: Sequence[dict], - **kwargs): - """Run forward test to get predictions. - - Args: - imgs (torch.Tensor): Input tensor of the model. - img_metas (Sequence[dict]]): A list of dict containing image(s) - meta information. - Returns: - np.ndarray: Segmentation result. - """ - with torch.cuda.device(self.device_id), torch.no_grad(): - seg_pred = self.model({'input': imgs})[self.output_name] - seg_pred = seg_pred.detach().cpu().numpy() - return seg_pred - - -class PPLSegmentor(DeployBaseSegmentor): - """Wrapper for segmentation's inference with PPL. - - Args: - model_file (Sequence[str]): Paths of input params and bin files. - class_names (Sequence[str]): A list of string specifying class names. - palette (np.ndarray): The palette of segmentation map. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: str, class_names: Sequence[str], - palette: np.ndarray, device_id: int): - super(PPLSegmentor, self).__init__(class_names, palette, device_id) - from mmdeploy.apis.ppl import PPLWrapper - self.model = PPLWrapper(model_file[0], model_file[1], device_id) - - def forward_test(self, imgs: torch.Tensor, img_metas: Sequence[dict], - **kwargs): - """Run forward test to get predictions. - - Args: - imgs (torch.Tensor): Input tensor of the model. - img_metas (Sequence[dict]]): A list of dict containing image(s) - meta information. - Returns: - np.ndarray: Segmentation result. - """ - seg_pred = self.model({'input': imgs})[0] - return seg_pred - - -class NCNNSegmentor(DeployBaseSegmentor): - """Wrapper for segmentation's inference with NCNN. - - Args: - model_file (Sequence[str]): Paths of input params and bin files. - class_names (Sequence[str]): A list of string specifying class names. - palette (np.ndarray): The palette of segmentation map. - device_id (int): An integer represents device index. - """ - - def __init__(self, model_file: Sequence[str], class_names: Sequence[str], - palette: np.ndarray, device_id: int): - super(NCNNSegmentor, self).__init__(class_names, palette, device_id) - from mmdeploy.apis.ncnn import NCNNWrapper - assert len(model_file) == 2, f'`model_file` should be [param_file, \ - bin_file], but given {model_file}' - - ncnn_param_file = model_file[0] - ncnn_bin_file = model_file[1] - self.model = NCNNWrapper( - ncnn_param_file, ncnn_bin_file, output_names=['output']) - - def forward_test(self, imgs: torch.Tensor, img_metas: Sequence[dict], - **kwargs): - """Run forward test to get predictions. - - Args: - imgs (torch.Tensor): Input tensor of the model. - img_metas (Sequence[dict]]): A list of dict containing image(s) - meta information. - Returns: - np.ndarray: Segmentation result. - """ - results = self.model({'input': imgs})['output'] - results = results.detach().cpu().numpy() - return results - - -ONNXRUNTIME_SEGMENTOR_MAP = dict(end2end=ONNXRuntimeSegmentor) - -TENSORRT_SEGMENTOR_MAP = dict(end2end=TensorRTSegmentor) - -PPL_SEGMENTOR_MAP = dict(end2end=PPLSegmentor) -NCNN_SEGMENTOR_MAP = dict(end2end=NCNNSegmentor) - -BACKEND_SEGMENTOR_MAP = { - Backend.ONNXRUNTIME: ONNXRUNTIME_SEGMENTOR_MAP, - Backend.TENSORRT: TENSORRT_SEGMENTOR_MAP, - Backend.PPL: PPL_SEGMENTOR_MAP, - Backend.NCNN: NCNN_SEGMENTOR_MAP -} - - -def get_classes_palette_from_config(model_cfg: Union[str, mmcv.Config]): - """Get class name and palette from config. - - Args: - model_cfg (str | mmcv.Config): Input model config file or - Config object. - - Returns: - tuple(Sequence[str], np.ndarray): A list of string specifying names of - different class and the palette of segmentation map. - """ - # load cfg if necessary - model_cfg = load_config(model_cfg)[0] - - module_dict = DATASETS.module_dict - data_cfg = model_cfg.data - - if 'train' in data_cfg: - module = module_dict[data_cfg.train.type] - elif 'val' in data_cfg: - module = module_dict[data_cfg.val.type] - elif 'test' in data_cfg: - module = module_dict[data_cfg.test.type] - else: - raise RuntimeError(f'No dataset config found in: {model_cfg}') - - return module.CLASSES, module.PALETTE - - -def build_segmentor(model_files, model_cfg, deploy_cfg, device_id): - """Build segmentor for different backend. - - Args: - model_files (list[str]): Input model file(s). - model_cfg (str | mmcv.Config): Input model config file or Config - object. - deploy_cfg (str | mmcv.Config): Input deployment config file or - Config object. - device_id (int): An integer represents device index. - - Returns: - DeployBaseSegmentor: Segmentor for a configured backend. - """ - # load cfg if necessary - model_cfg, deploy_cfg = load_config(model_cfg, deploy_cfg) - - backend = get_backend(deploy_cfg) - class_names, palette = get_classes_palette_from_config(model_cfg) - assert backend in BACKEND_SEGMENTOR_MAP, \ - f'Unsupported backend type: {backend.value}' - segmentor_map = BACKEND_SEGMENTOR_MAP[backend] - - model_type = 'end2end' - assert model_type in segmentor_map, f'Unsupported model type: {model_type}' - backend_segmentor_class = segmentor_map[model_type] - model_files = model_files[0] if len(model_files) == 1 else model_files - backend_segmentor = backend_segmentor_class( - model_files, - class_names=class_names, - device_id=device_id, - palette=palette) - - return backend_segmentor diff --git a/mmdeploy/mmseg/apis/visualize.py b/mmdeploy/mmseg/apis/visualize.py deleted file mode 100644 index c1257369c1..0000000000 --- a/mmdeploy/mmseg/apis/visualize.py +++ /dev/null @@ -1,36 +0,0 @@ -import numpy as np -import torch - -from mmdeploy.utils import Backend - - -def show_result(model: torch.nn.Module, - image: np.ndarray, - result: list, - output_file: str, - backend: Backend, - show: bool = True, - opacity: float = 0.5): - """Show predictions of segmentation. - - Args: - model (nn.Module): Input model which has `show_result` method. - image: (np.ndarray): Input image to draw predictions. - result (list): A list of predictions. - output_file (str): Output image file to save drawn predictions. - backend (Backend): Specifying backend type. - show (bool): Whether to show plotted image in windows. Defaults to - `True`. - opacity: (float): Opacity of painted segmentation map. - Defaults to `0.5`. - - Returns: - np.ndarray: Drawn image, only if not `show` or `out_file`. - """ - return model.show_result( - image, - result, - opacity=opacity, - show=show, - win_name=backend.value, - out_file=output_file) diff --git a/mmdeploy/mmseg/export/__init__.py b/mmdeploy/mmseg/export/__init__.py deleted file mode 100644 index 82f32aa792..0000000000 --- a/mmdeploy/mmseg/export/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from .onnx_utils import convert_syncbatchnorm -from .prepare_input import (build_dataloader, build_dataset, create_input, - get_tensor_from_input) - -__all__ = [ - 'create_input', 'convert_syncbatchnorm', 'build_dataloader', - 'build_dataset', 'get_tensor_from_input' -] diff --git a/mmdeploy/mmseg/export/prepare_input.py b/mmdeploy/mmseg/export/prepare_input.py deleted file mode 100644 index aa32d2dc56..0000000000 --- a/mmdeploy/mmseg/export/prepare_input.py +++ /dev/null @@ -1,141 +0,0 @@ -from typing import Any, Optional, Sequence, Union - -import mmcv -import numpy as np -from mmcv.parallel import collate, scatter -from mmseg.apis.inference import LoadImage -from mmseg.datasets import build_dataloader as build_dataloader_mmseg -from mmseg.datasets import build_dataset as build_dataset_mmseg -from mmseg.datasets.pipelines import Compose -from torch.utils.data import Dataset - -from mmdeploy.utils import Task, load_config - - -def create_input(task: Task, - model_cfg: Union[str, mmcv.Config], - imgs: Any, - input_shape: Optional[Sequence[int]] = None, - device: str = 'cuda:0'): - """Create input for segmentation. - - Args: - task (Task): Specifying task type. - model_cfg (str | mmcv.Config): The input model config. - imgs (Any): Input image(s), accpeted data type are `str`, - `np.ndarray`, `torch.Tensor`. - input_shape (list[int]): A list of two integer in (width, height) - format specifying input shape. Defaults to `None`. - device (str): A string represents device type. Default is 'cuda:0'. - - Returns: - tuple: (data, img), meta information for the input image and input. - """ - assert task == Task.SEGMENTATION - cfg = load_config(model_cfg)[0].copy() - if not isinstance(imgs, (list, tuple)): - imgs = [imgs] - - if isinstance(imgs[0], np.ndarray): - cfg = cfg.copy() - # set loading pipeline type - cfg.data.test.pipeline[0].type = 'LoadImageFromWebcam' - # for static exporting - if input_shape is not None: - cfg.data.test.pipeline[1]['img_scale'] = tuple(input_shape) - cfg.data.test.pipeline[1]['transforms'][0]['keep_ratio'] = False - cfg.data.test.pipeline = [LoadImage()] + cfg.data.test.pipeline[1:] - - test_pipeline = Compose(cfg.data.test.pipeline) - data_list = [] - for img in imgs: - # prepare data - data = dict(img=img) - # build the data pipeline - data = test_pipeline(data) - data_list.append(data) - - data = collate(data_list, samples_per_gpu=len(imgs)) - - data['img_metas'] = [img_metas.data[0] for img_metas in data['img_metas']] - data['img'] = [img.data[0][None, :] for img in data['img']] - if device != 'cpu': - data = scatter(data, [device])[0] - - return data, data['img'] - - -def build_dataset(dataset_cfg: Union[str, mmcv.Config], - dataset_type: str = 'val', - **kwargs): - """Build dataset for segmentation. - - Args: - dataset_cfg (str | mmcv.Config): The input dataset config. - dataset_type (str): A string represents dataset type, e.g.: 'train', - 'test', 'val'. Defaults to 'val'. - - Returns: - Dataset: A PyTorch dataset. - """ - dataset_cfg = load_config(dataset_cfg)[0] - data = dataset_cfg.data - assert dataset_type in data - - dataset = build_dataset_mmseg(data[dataset_type]) - - return dataset - - -def build_dataloader(dataset: Dataset, - samples_per_gpu: int, - workers_per_gpu: int, - num_gpus: int = 1, - dist: bool = False, - shuffle: bool = False, - seed: Optional[int] = None, - drop_last: bool = False, - pin_memory: bool = True, - persistent_workers: bool = True, - **kwargs): - """Build dataloader for segmentation. - - Args: - dataset (Dataset): Input dataset. - samples_per_gpu (int): Number of training samples on each GPU, i.e., - batch size of each GPU. - workers_per_gpu (int): How many subprocesses to use for data loading - for each GPU. - num_gpus (int): Number of GPUs. Only used in non-distributed training. - dist (bool): Distributed training/test or not. Defaults to `False`. - shuffle (bool): Whether to shuffle the data at every epoch. - Defaults to `False`. - seed (int): An integer set to be seed. Default is `None`. - drop_last (bool): Whether to drop the last incomplete batch in epoch. - Default to `False`. - pin_memory (bool): Whether to use pin_memory in DataLoader. - Default is `True`. - persistent_workers (bool): If `True`, the data loader will not shutdown - the worker processes after a dataset has been consumed once. - This allows to maintain the workers Dataset instances alive. - The argument also has effect in PyTorch>=1.7.0. - Default is `True`. - kwargs: Any other keyword argument to be used to initialize DataLoader. - - Returns: - DataLoader: A PyTorch dataloader. - """ - return build_dataloader_mmseg(dataset, samples_per_gpu, workers_per_gpu, - num_gpus, dist, shuffle, seed, drop_last, - pin_memory, persistent_workers, **kwargs) - - -def get_tensor_from_input(input_data: tuple): - """Get input tensor from input data. - - Args: - input_data (tuple): Input data containing meta info and image tensor. - Returns: - torch.Tensor: An image in `Tensor`. - """ - return input_data['img'][0] diff --git a/mmdeploy/utils/__init__.py b/mmdeploy/utils/__init__.py index ec36ead2f8..b114bc31f7 100644 --- a/mmdeploy/utils/__init__.py +++ b/mmdeploy/utils/__init__.py @@ -1,15 +1,16 @@ from .config_utils import (cfg_apply_marks, get_backend, get_calib_config, get_calib_filename, get_codebase, get_common_config, - get_input_shape, get_mmdet_params, get_model_inputs, - get_onnx_config, get_partition_config, - get_task_type, is_dynamic_batch, is_dynamic_shape, - load_config) + get_input_shape, get_model_inputs, get_onnx_config, + get_partition_config, get_task_type, + is_dynamic_batch, is_dynamic_shape, load_config) from .constants import Backend, Codebase, Task +from .device import parse_cuda_device_id, parse_device_id __all__ = [ 'is_dynamic_batch', 'is_dynamic_shape', 'get_task_type', 'get_codebase', 'get_backend', 'load_config', 'Backend', 'Codebase', 'Task', 'get_onnx_config', 'get_partition_config', 'get_calib_config', 'get_calib_filename', 'get_common_config', 'get_model_inputs', - 'cfg_apply_marks', 'get_mmdet_params', 'get_input_shape' + 'cfg_apply_marks', 'get_input_shape', 'parse_device_id', + 'parse_cuda_device_id' ] diff --git a/mmdeploy/utils/config_utils.py b/mmdeploy/utils/config_utils.py index 82cc8ce392..59c8ab2665 100644 --- a/mmdeploy/utils/config_utils.py +++ b/mmdeploy/utils/config_utils.py @@ -1,11 +1,11 @@ -from typing import Optional, Union +from typing import Dict, List, Optional, Union import mmcv from .constants import Backend, Codebase, Task -def load_config(*args): +def load_config(*args) -> List[mmcv.Config]: """Load the configuration and check the validity. Args: @@ -93,7 +93,7 @@ def get_backend(deploy_cfg: Union[str, mmcv.Config], default=None) -> Backend: return backend -def get_onnx_config(deploy_cfg: Union[str, mmcv.Config]) -> str: +def get_onnx_config(deploy_cfg: Union[str, mmcv.Config]) -> Dict: """Get the onnx parameters in export() from config. Args: @@ -167,14 +167,15 @@ def is_dynamic_shape(deploy_cfg: Union[str, mmcv.Config], return False -def get_input_shape(deploy_cfg: Union[str, mmcv.Config]): +def get_input_shape(deploy_cfg: Union[str, mmcv.Config]) -> List[int]: """Get the input shape for static exporting. Args: deploy_cfg (str | mmcv.Config): The path or content of config. Returns: - List: The input shape for backend model (axis 2 and 3), e.g [512, 512]. + List[int]: The input shape for backend model (axis 2 and 3), + e.g [512, 512]. """ input_shape = get_onnx_config(deploy_cfg)['input_shape'] if input_shape is not None: @@ -182,7 +183,7 @@ def get_input_shape(deploy_cfg: Union[str, mmcv.Config]): return input_shape -def cfg_apply_marks(deploy_cfg: Union[str, mmcv.Config]) -> bool: +def cfg_apply_marks(deploy_cfg: Union[str, mmcv.Config]) -> Union[bool, None]: """Check if the model needs to be partitioned by checking if the config contains 'apply_marks'. @@ -190,7 +191,7 @@ def cfg_apply_marks(deploy_cfg: Union[str, mmcv.Config]) -> bool: deploy_cfg (str | mmcv.Config): The path or content of config. Returns: - bool: Whether config contains 'apply_marks'. + bool or None: Whether config contains 'apply_marks'. """ partition_config = deploy_cfg.get('partition_config', None) if partition_config is None: @@ -200,7 +201,7 @@ def cfg_apply_marks(deploy_cfg: Union[str, mmcv.Config]) -> bool: return apply_marks -def get_partition_config(deploy_cfg: Union[str, mmcv.Config]): +def get_partition_config(deploy_cfg: Union[str, mmcv.Config]) -> Dict: """Check if the model needs to be partitioned and get the config of partition. @@ -208,7 +209,7 @@ def get_partition_config(deploy_cfg: Union[str, mmcv.Config]): deploy_cfg (str | mmcv.Config): The path or content of config. Returns: - dict: The config of partition + dict: The config of partition. """ partition_config = deploy_cfg.get('partition_config', None) if partition_config is None: @@ -221,29 +222,28 @@ def get_partition_config(deploy_cfg: Union[str, mmcv.Config]): return partition_config -def get_calib_config(deploy_cfg: Union[str, mmcv.Config]): +def get_calib_config(deploy_cfg: Union[str, mmcv.Config]) -> Dict: """Check if the model has calibration configs. Args: deploy_cfg (str | mmcv.Config): The path or content of config. Returns: - dict: The config of calibration + dict: The config of calibration. """ calib_config = deploy_cfg.get('calib_config', None) return calib_config -def get_calib_filename(deploy_cfg: Union[str, mmcv.Config]): - """Check if the model needs to create calib and get output filename of - calib. +def get_calib_filename(deploy_cfg: Union[str, mmcv.Config]) -> str: + """Check if the model needs to create calib and get filename of calib. Args: deploy_cfg (str | mmcv.Config): The path or content of config. Returns: - str: The filename of output calib file + str: The filename of output calib file. """ calib_config = get_calib_config(deploy_cfg) @@ -257,7 +257,7 @@ def get_calib_filename(deploy_cfg: Union[str, mmcv.Config]): return None -def get_common_config(deploy_cfg: Union[str, mmcv.Config]): +def get_common_config(deploy_cfg: Union[str, mmcv.Config]) -> Dict: """Get common parameters from config. Args: @@ -271,7 +271,7 @@ def get_common_config(deploy_cfg: Union[str, mmcv.Config]): return model_params -def get_model_inputs(deploy_cfg: Union[str, mmcv.Config]): +def get_model_inputs(deploy_cfg: Union[str, mmcv.Config]) -> List[Dict]: """Get model input parameters from config. Args: @@ -283,21 +283,3 @@ def get_model_inputs(deploy_cfg: Union[str, mmcv.Config]): backend_config = deploy_cfg['backend_config'] model_params = backend_config.get('model_inputs', []) return model_params - - -def get_mmdet_params(deploy_cfg: Union[str, mmcv.Config]): - """Get mmdet post-processing parameters from config. - - Args: - deploy_cfg (str | mmcv.Config): The path or content of config. - - Returns: - dict: A dict of parameters for mmdet. - """ - deploy_cfg = load_config(deploy_cfg)[0] - codebase_key = 'codebase_config' - assert codebase_key in deploy_cfg - codebase_config = deploy_cfg[codebase_key] - post_params = codebase_config.get('post_processing', None) - assert post_params is not None, 'Failed to get `post_processing`.' - return post_params diff --git a/mmdeploy/utils/device.py b/mmdeploy/utils/device.py new file mode 100644 index 0000000000..3e851ad728 --- /dev/null +++ b/mmdeploy/utils/device.py @@ -0,0 +1,37 @@ +import torch + + +def parse_device_id(device: str) -> int: + """Parse cuda device index from a string. + + Args: + device (str): The typical style of string specifying cuda device, + e.g.: 'cuda:0'. + + Returns: + int: The parsed device id, defaults to `0`. + """ + if device == 'cpu': + return -1 + device_id = 0 + if len(device) >= 6: + device_id = torch.device(device).index + return device_id + + +def parse_cuda_device_id(device: str) -> int: + """Parse cuda device index from a string. + + Args: + device (str): The typical style of string specifying cuda device, + e.g.: 'cuda:0'. + + Returns: + int: The parsed device id, defaults to `0`. + """ + device = torch.device(device) + assert device.type == 'cuda', 'Not cuda device.' + + device_id = 0 if device.index is None else device.index + + return device_id diff --git a/mmdeploy/utils/test.py b/mmdeploy/utils/test.py index ccd812f9fa..6e41240f4b 100644 --- a/mmdeploy/utils/test.py +++ b/mmdeploy/utils/test.py @@ -33,7 +33,7 @@ def __init__(self, wrapped_function: Callable, **kwargs): self.wrapped_function = wrapped_function self.kwargs = kwargs - def forward(self, *args, **kwargs): + def forward(self, *args, **kwargs) -> Any: """Call the wrapped function.""" kwargs.update(self.kwargs) return self.wrapped_function(*args, **kwargs) @@ -73,11 +73,31 @@ def forward(self, *args, **kwargs): return func(*args, **kwargs) +class DummyModel(torch.nn.Module): + """A dummy model for unit tests. + + Args: + outputs (Any): Predefined output variables. + """ + + def __init__(self, outputs=None, *args, **kwargs): + torch.nn.Module.__init__(self) + self.outputs = outputs + + def forward(self, *args, **kwargs): + """Run forward.""" + return self.outputs + + def __call__(self, *args, **kwds): + """Call the forward method.""" + return self.forward(*args, **kwds) + + class SwitchBackendWrapper: """A switcher for backend wrapper for unit tests. Examples: >>> from mmdeploy.utils.test import SwitchBackendWrapper - >>> from mmdeploy.apis.onnxruntime.onnxruntime_utils import ORTWrapper + >>> from mmdeploy.backend.onnxruntime import ORTWrapper >>> with SwitchBackendWrapper(ORTWrapper) as wrapper: >>> wrapper.set(ORTWrapper, outputs=outputs) >>> ... @@ -89,10 +109,20 @@ class SwitchBackendWrapper: call = None class BackendWrapper(torch.nn.Module): - """A dummy wrapper for unit tests.""" + """A dummy backend wrapper for unit tests. - def __init__(self, *args, **kwargs): - self.output_names = ['dets', 'labels'] + To enable BaseWrapper.output_to_list(), the wrapper needs member + variable `_output_names` that is set in constructor. Therefore, + the dummy BackendWrapper needs a constructor that receives + output_names. + + Args: + output_names (Any): `output_name` of BaseWrapper + """ + + def __init__(self, output_names=['dets', 'labels'], *args, **kwargs): + torch.nn.Module.__init__(self) + self._output_names = output_names def forward(self, *args, **kwargs): """Run forward.""" @@ -165,7 +195,8 @@ def assert_allclose(expected: List[Union[torch.Tensor, np.ndarray]], raise -def get_model_outputs(model: nn.Module, func_name: str, model_inputs: dict): +def get_model_outputs(model: nn.Module, func_name: str, + model_inputs: dict) -> Any: """To get outputs of pytorch model. Args: @@ -268,6 +299,7 @@ def get_backend_outputs(onnx_file_path: str, flatten_model_inputs = get_flatten_inputs(model_inputs) input_names = [k for k, v in flatten_model_inputs.items() if k != 'ctx'] output_names = get_onnx_config(deploy_cfg).get('output_names', None) + backend_files = [onnx_file_path] # prepare backend model and input features if backend == Backend.TENSORRT: # convert to engine @@ -281,16 +313,16 @@ def get_backend_outputs(onnx_file_path: str, 0, deploy_cfg=deploy_cfg, onnx_model=onnx_file_path) - backend_model = trt_apis.TRTWrapper(trt_file_path) + backend_files = [trt_file_path] for k, v in model_inputs.items(): model_inputs[k] = model_inputs[k].cuda() backend_feats = model_inputs + device = 'cuda:0' elif backend == Backend.ONNXRUNTIME: import mmdeploy.apis.onnxruntime as ort_apis if not ort_apis.is_available(): return None - backend_model = ort_apis.ORTWrapper(onnx_file_path, 0, None) feature_list = [] backend_feats = {} for k, item in model_inputs.items(): @@ -313,6 +345,7 @@ def get_backend_outputs(onnx_file_path: str, backend_feats[input_names[i]] = feature_list[i] else: backend_feats[str(i)] = feature_list[i] + device = 'cpu' elif backend == Backend.NCNN: return None elif backend == Backend.OPENVINO: @@ -328,17 +361,21 @@ def get_backend_outputs(onnx_file_path: str, } openvino_apis.onnx2openvino(input_info, output_names, onnx_file_path, openvino_work_dir) - backend_model = openvino_apis.OpenVINOWrapper(openvino_file_path) - + backend_files = [openvino_file_path] backend_feats = flatten_model_inputs + device = 'cpu' elif backend == Backend.DEFAULT: return None else: raise NotImplementedError( f'Unimplemented backend type: {backend.value}') + from mmdeploy.codebase.base import BaseBackendModel + backend_model = BaseBackendModel._build_wrapper(backend, backend_files, + device, output_names) with torch.no_grad(): - backend_outputs = backend_model.forward(backend_feats) + backend_outputs = backend_model(backend_feats) + backend_outputs = backend_model.output_to_list(backend_outputs) return backend_outputs @@ -354,7 +391,7 @@ def get_rewrite_outputs(wrapped_model: nn.Module, deploy_cfg (mmcv.Config): Deployment config. Returns: - Any: The outputs of model, decided by the backend wrapper. + List[torch.Tensor]: The outputs of model. bool: A flag indicate the type of outputs. If the flag is True, then the outputs are backend output, otherwise they are outputs of wrapped pytorch model. diff --git a/mmdeploy/utils/timer.py b/mmdeploy/utils/timer.py index b24d002f28..900a35cf73 100644 --- a/mmdeploy/utils/timer.py +++ b/mmdeploy/utils/timer.py @@ -3,6 +3,7 @@ import time import warnings from contextlib import contextmanager +from typing import Union import torch @@ -92,7 +93,7 @@ def activate(cls, warmup: int = 1, log_interval: int = 1, with_sync: bool = False, - file: io.TextIOWrapper = sys.stdout): + file: Union[str, io.TextIOWrapper] = sys.stdout): """Activate the time counter. Args: @@ -102,8 +103,8 @@ def activate(cls, log_interval (int): Interval between each log, default 1. with_sync (bool): Whether use cuda synchronize for time counting, default False. - file (io.TextIOWrapper): A file or file-like object to save output - messages. The default is `sys.stdout`. + file (str | io.TextIOWrapper): A file or file-like object to save + output messages. The default is `sys.stdout`. """ assert warmup >= 1 if file != sys.stdout: @@ -112,7 +113,7 @@ def activate(cls, if func_name is not None: warnings.warn('func_name must be globally unique if you call ' 'activate multiple times') - assert func_name in cls.names, '{} must be registried before '\ + assert func_name in cls.names, '{} must be registered before '\ 'setting params'.format(func_name) cls.names[func_name]['warmup'] = warmup cls.names[func_name]['log_interval'] = log_interval diff --git a/mmdeploy/version.py b/mmdeploy/version.py index f74e40ba30..8db80794c7 100644 --- a/mmdeploy/version.py +++ b/mmdeploy/version.py @@ -1,10 +1,11 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import Tuple __version__ = '0.1.0' short_version = __version__ -def parse_version_info(version_str: str): +def parse_version_info(version_str: str) -> Tuple: """Parse version from a string. Args: diff --git a/tests/test_mmdet/data/coco_sample.json b/tests/test_codebase/test_mmdet/data/coco_sample.json similarity index 100% rename from tests/test_mmdet/data/coco_sample.json rename to tests/test_codebase/test_mmdet/data/coco_sample.json diff --git a/tests/test_mmdet/data/mask_model.json b/tests/test_codebase/test_mmdet/data/mask_model.json similarity index 100% rename from tests/test_mmdet/data/mask_model.json rename to tests/test_codebase/test_mmdet/data/mask_model.json diff --git a/tests/test_codebase/test_mmdet/data/model.py b/tests/test_codebase/test_mmdet/data/model.py new file mode 100644 index 0000000000..95811df86d --- /dev/null +++ b/tests/test_codebase/test_mmdet/data/model.py @@ -0,0 +1,108 @@ +model = dict( + type='YOLOV3', + backbone=dict( + type='MobileNetV2', + out_indices=(2, 4, 6), + act_cfg=dict(type='LeakyReLU', negative_slope=0.1), + init_cfg=dict( + type='Pretrained', checkpoint='open-mmlab://mmdet/mobilenet_v2')), + neck=dict( + type='YOLOV3Neck', + num_scales=3, + in_channels=[320, 96, 32], + out_channels=[96, 96, 96]), + bbox_head=dict( + type='YOLOV3Head', + num_classes=80, + in_channels=[96, 96, 96], + out_channels=[96, 96, 96], + anchor_generator=dict( + type='YOLOAnchorGenerator', + base_sizes=[[(116, 90), (156, 198), (373, 326)], + [(30, 61), (62, 45), (59, 119)], + [(10, 13), (16, 30), (33, 23)]], + strides=[32, 16, 8]), + bbox_coder=dict(type='YOLOBBoxCoder'), + featmap_strides=[32, 16, 8], + loss_cls=dict( + type='CrossEntropyLoss', + use_sigmoid=True, + loss_weight=1.0, + reduction='sum'), + loss_conf=dict( + type='CrossEntropyLoss', + use_sigmoid=True, + loss_weight=1.0, + reduction='sum'), + loss_xy=dict( + type='CrossEntropyLoss', + use_sigmoid=True, + loss_weight=2.0, + reduction='sum'), + loss_wh=dict(type='MSELoss', loss_weight=2.0, reduction='sum')), + # training and testing settings + train_cfg=dict( + assigner=dict( + type='GridAssigner', + pos_iou_thr=0.5, + neg_iou_thr=0.5, + min_pos_iou=0)), + test_cfg=dict( + nms_pre=1000, + min_bbox_size=0, + score_thr=0.05, + conf_thr=0.005, + nms=dict(type='nms', iou_threshold=0.45), + max_per_img=100)) +# dataset settings +dataset_type = 'CocoDataset' +data_root = '.' +img_norm_cfg = dict( + mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True) +train_pipeline = [ + dict(type='LoadImageFromFile', to_float32=True), + dict(type='LoadAnnotations', with_bbox=True), + dict(type='PhotoMetricDistortion'), + dict( + type='Expand', + mean=img_norm_cfg['mean'], + to_rgb=img_norm_cfg['to_rgb'], + ratio_range=(1, 2)), + dict( + type='MinIoURandomCrop', + min_ious=(0.4, 0.5, 0.6, 0.7, 0.8, 0.9), + min_crop_size=0.3), + dict( + type='Resize', + img_scale=[(320, 320), (416, 416)], + multiscale_mode='range', + keep_ratio=True), + dict(type='RandomFlip', flip_ratio=0.5), + dict(type='Normalize', **img_norm_cfg), + dict(type='Pad', size_divisor=32), + dict(type='DefaultFormatBundle'), + dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']) +] +test_pipeline = [ + dict(type='LoadImageFromFile'), + dict( + type='MultiScaleFlipAug', + img_scale=(416, 416), + flip=False, + transforms=[ + dict(type='Resize', keep_ratio=True), + dict(type='RandomFlip'), + dict(type='Normalize', **img_norm_cfg), + dict(type='Pad', size_divisor=32), + dict(type='DefaultFormatBundle'), + dict(type='Collect', keys=['img']) + ]) +] +data = dict( + samples_per_gpu=24, + workers_per_gpu=4, + test=dict( + type=dataset_type, + ann_file='tests/test_codebase/test_mmdet/data/coco_sample.json', + img_prefix=data_root, + pipeline=test_pipeline)) diff --git a/tests/test_mmdet/data/single_stage_model.json b/tests/test_codebase/test_mmdet/data/single_stage_model.json similarity index 100% rename from tests/test_mmdet/data/single_stage_model.json rename to tests/test_codebase/test_mmdet/data/single_stage_model.json diff --git a/tests/test_mmdet/test_mmdet_core.py b/tests/test_codebase/test_mmdet/test_mmdet_core.py similarity index 89% rename from tests/test_mmdet/test_mmdet_core.py rename to tests/test_codebase/test_mmdet/test_mmdet_core.py index 8dab90fb4d..1db30d0aee 100644 --- a/tests/test_mmdet/test_mmdet_core.py +++ b/tests/test_codebase/test_mmdet/test_mmdet_core.py @@ -15,7 +15,7 @@ def test_multiclass_nms_static(): import tensorrt as trt - from mmdeploy.mmdet.core import multiclass_nms + from mmdeploy.codebase.mmdet.core import multiclass_nms deploy_cfg = mmcv.Config( dict( onnx_config=dict(output_names=None, input_shape=None), @@ -70,7 +70,8 @@ def test_multiclass_nms_static(): @pytest.mark.parametrize('backend_type', ['onnxruntime', 'ncnn']) -def test_delta2bbox(backend_type): +@pytest.mark.parametrize('add_ctr_clamp', [True, False]) +def test_delta2bbox(backend_type, add_ctr_clamp): pytest.importorskip(backend_type, reason=f'requires {backend_type}') deploy_cfg = mmcv.Config( dict( @@ -86,10 +87,10 @@ def delta2bbox(*args, **kwargs): rois = torch.rand(1, 5, 4) deltas = torch.rand(1, 5, 4) - original_outputs = delta2bbox(rois, deltas) + original_outputs = delta2bbox(rois, deltas, add_ctr_clamp=add_ctr_clamp) # wrap function to nn.Module, enable torch.onnx.export - wrapped_func = WrapFunction(delta2bbox) + wrapped_func = WrapFunction(delta2bbox, add_ctr_clamp=add_ctr_clamp) rewrite_outputs, is_backend_output = get_rewrite_outputs( wrapped_func, model_inputs={ @@ -146,7 +147,7 @@ def tblr2bboxes(*args, **kwargs): def test_distance2bbox(): - from mmdeploy.mmdet.core import distance2bbox + from mmdeploy.codebase.mmdet.core import distance2bbox points = torch.rand(3, 2) distance = torch.rand(3, 4) bbox = distance2bbox(points, distance) @@ -155,10 +156,11 @@ def test_distance2bbox(): @pytest.mark.skipif( not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') -def test_multiclass_nms_with_keep_top_k(): +@pytest.mark.parametrize('pre_top_k', [-1, 1000]) +def test_multiclass_nms_with_keep_top_k(pre_top_k): backend_type = 'onnxruntime' - from mmdeploy.mmdet.core import multiclass_nms + from mmdeploy.codebase.mmdet.core import multiclass_nms max_output_boxes_per_class = 20 keep_top_k = 15 deploy_cfg = mmcv.Config( @@ -186,7 +188,7 @@ def test_multiclass_nms_with_keep_top_k(): score_threshold=0.05, iou_threshold=0.5, max_output_boxes_per_class=max_output_boxes_per_class, - pre_top_k=-1, + pre_top_k=pre_top_k, keep_top_k=keep_top_k, background_label_id=-1, )))) @@ -211,9 +213,11 @@ def test_multiclass_nms_with_keep_top_k(): test_scores = torch.ones(batch_size, num_boxes, num_classes) model_inputs = {'boxes': test_boxes, 'scores': test_scores} - import mmdeploy.apis.onnxruntime as ort_apis - backend_model = ort_apis.ORTWrapper(onnx_model_path, 0, None) - dets, _ = backend_model.forward(model_inputs) + import mmdeploy.backend.onnxruntime as ort_apis + backend_model = ort_apis.ORTWrapper(onnx_model_path, 'cuda:0', None) + output = backend_model.forward(model_inputs) + output = backend_model.output_to_list(output) + dets = output[0] assert dets.shape[1] < keep_top_k, \ 'multiclass_nms returned more values than "keep_top_k"\n' \ diff --git a/tests/test_mmdet/test_mmdet_models.py b/tests/test_codebase/test_mmdet/test_mmdet_models.py similarity index 93% rename from tests/test_mmdet/test_mmdet_models.py rename to tests/test_codebase/test_mmdet/test_mmdet_models.py index f76803642d..931b913cec 100644 --- a/tests/test_mmdet/test_mmdet_models.py +++ b/tests/test_codebase/test_mmdet/test_mmdet_models.py @@ -2,7 +2,6 @@ import importlib import os import random -import tempfile from typing import Dict, List import mmcv @@ -10,7 +9,6 @@ import pytest import torch -from mmdeploy.utils.constants import Backend, Codebase from mmdeploy.utils.test import (WrapModel, get_model_outputs, get_rewrite_outputs) @@ -79,10 +77,10 @@ def get_rpn_head_model(): test_cfg = mmcv.Config( dict( deploy_nms_pre=0, - min_bbox_size=0, - score_thr=0.05, - nms=dict(type='nms', iou_threshold=0.5), - max_per_img=100)) + nms_pre=0, + max_per_img=100, + nms=dict(type='nms', iou_threshold=0.7), + min_bbox_size=0)) from mmdet.models import RPNHead model = RPNHead(in_channels=1, test_cfg=test_cfg) @@ -126,7 +124,7 @@ def test_anchor_head_get_bboxes(backend_type): score_threshold=0.05, iou_threshold=0.5, max_output_boxes_per_class=200, - pre_top_k=-1, + pre_top_k=5000, keep_top_k=100, background_label_id=-1, )))) @@ -205,7 +203,7 @@ def test_get_bboxes_of_fcos_head(backend_type): score_threshold=0.05, iou_threshold=0.5, max_output_boxes_per_class=200, - pre_top_k=-1, + pre_top_k=5000, keep_top_k=100, background_label_id=-1, )))) @@ -270,6 +268,61 @@ def test_get_bboxes_of_fcos_head(backend_type): assert rewrite_outputs is not None +@pytest.mark.parametrize('backend_type', ['onnxruntime', 'ncnn']) +def test_get_bboxes_of_rpn_head(backend_type): + pytest.importorskip(backend_type, reason=f'requires {backend_type}') + head = get_rpn_head_model() + head.cpu().eval() + s = 4 + img_metas = [{ + 'scale_factor': np.ones(4), + 'pad_shape': (s, s, 3), + 'img_shape': (s, s, 3) + }] + + output_names = ['dets'] + deploy_cfg = mmcv.Config( + dict( + backend_config=dict(type=backend_type), + onnx_config=dict(output_names=output_names, input_shape=None), + codebase_config=dict( + type='mmdet', + task='ObjectDetection', + post_processing=dict( + score_threshold=0.05, + iou_threshold=0.5, + max_output_boxes_per_class=200, + pre_top_k=5000, + keep_top_k=100, + background_label_id=-1, + )))) + + # the cls_score's size: (1, 36, 32, 32), (1, 36, 16, 16), + # (1, 36, 8, 8), (1, 36, 4, 4), (1, 36, 2, 2). + # the bboxes's size: (1, 36, 32, 32), (1, 36, 16, 16), + # (1, 36, 8, 8), (1, 36, 4, 4), (1, 36, 2, 2) + seed_everything(1234) + cls_score = [ + torch.rand(1, 9, pow(2, i), pow(2, i)) for i in range(5, 0, -1) + ] + seed_everything(5678) + bboxes = [torch.rand(1, 36, pow(2, i), pow(2, i)) for i in range(5, 0, -1)] + + # to get outputs of onnx model after rewrite + img_metas[0]['img_shape'] = torch.Tensor([s, s]) + wrapped_model = WrapModel( + head, 'get_bboxes', img_metas=img_metas[0], with_nms=True) + rewrite_inputs = { + 'cls_scores': cls_score, + 'bbox_preds': bboxes, + } + rewrite_outputs, is_backend_output = get_rewrite_outputs( + wrapped_model=wrapped_model, + model_inputs=rewrite_inputs, + deploy_cfg=deploy_cfg) + assert rewrite_outputs is not None + + def _replace_r50_with_r18(model): """Replace ResNet50 with ResNet18 in config.""" model = copy.deepcopy(model) @@ -281,12 +334,12 @@ def _replace_r50_with_r18(model): @pytest.mark.parametrize('model_cfg_path', [ - 'tests/test_mmdet/data/single_stage_model.json', - 'tests/test_mmdet/data/mask_model.json' + 'tests/test_codebase/test_mmdet/data/single_stage_model.json', + 'tests/test_codebase/test_mmdet/data/mask_model.json' ]) @pytest.mark.skipif( not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') -def test_forward_of_base_detector_and_visualize(model_cfg_path): +def test_forward_of_base_detector(model_cfg_path): deploy_cfg = mmcv.Config( dict( backend_config=dict(type='onnxruntime'), @@ -316,22 +369,10 @@ def test_forward_of_base_detector_and_visualize(model_cfg_path): model_inputs=rewrite_inputs, deploy_cfg=deploy_cfg) - from mmdeploy.apis.utils import visualize - output_file = tempfile.NamedTemporaryFile(suffix='.jpg').name - model.CLASSES = [''] * 80 - visualize( - Codebase.MMDET, - img.squeeze().permute(1, 2, 0).numpy(), - result=[torch.rand(0, 5).numpy()] * 80, - model=model, - output_file=output_file, - backend=Backend.ONNXRUNTIME, - show_result=False) - assert rewrite_outputs is not None -@pytest.mark.parametrize('backend_type', ['openvino']) +@pytest.mark.parametrize('backend_type', ['onnxruntime', 'openvino']) def test_single_roi_extractor(backend_type): pytest.importorskip(backend_type, reason=f'requires {backend_type}') @@ -485,7 +526,8 @@ def test_cascade_roi_head(backend_type): model_outputs = get_model_outputs(cascade_roi_head, 'simple_test', model_inputs) processed_model_outputs = [] - for output in model_outputs[0]: + outputs = model_outputs[0] + for output in outputs: if output.shape == (0, 5): processed_model_outputs.append(np.zeros((1, 5))) else: diff --git a/tests/test_codebase/test_mmdet/test_mmdet_utils.py b/tests/test_codebase/test_mmdet/test_mmdet_utils.py new file mode 100644 index 0000000000..2c24a9a8f2 --- /dev/null +++ b/tests/test_codebase/test_mmdet/test_mmdet_utils.py @@ -0,0 +1,49 @@ +import mmcv +import numpy as np +import torch + +from mmdeploy.codebase.mmdet import (clip_bboxes, get_post_processing_params, + pad_with_value) + + +def test_clip_bboxes(): + x1 = torch.rand(3, 2) * 224 + y1 = torch.rand(3, 2) * 224 + x2 = x1 * 2 + y2 = y1 * 2 + outs = clip_bboxes(x1, y1, x2, y2, [224, 224]) + for out in outs: + assert int(out.max()) <= 224 + + +def test_pad_with_value(): + x = torch.rand(3, 2) + padded_x = pad_with_value(x, pad_dim=1, pad_size=4, pad_value=0) + assert np.allclose( + padded_x.shape, torch.Size([3, 6]), rtol=1e-03, atol=1e-05) + assert np.allclose(padded_x.sum(), x.sum(), rtol=1e-03, atol=1e-05) + + +config_with_mmdet_params = mmcv.Config( + dict( + codebase_config=dict( + type='mmdet', + task='ObjectDetection', + post_processing=dict( + score_threshold=0.05, + iou_threshold=0.5, + max_output_boxes_per_class=200, + pre_top_k=-1, + keep_top_k=100, + background_label_id=-1, + )))) + + +def test_get_mmdet_params(): + assert get_post_processing_params(config_with_mmdet_params) == dict( + score_threshold=0.05, + iou_threshold=0.5, + max_output_boxes_per_class=200, + pre_top_k=-1, + keep_top_k=100, + background_label_id=-1) diff --git a/tests/test_codebase/test_mmdet/test_object_detection.py b/tests/test_codebase/test_mmdet/test_object_detection.py new file mode 100644 index 0000000000..99b7364d9e --- /dev/null +++ b/tests/test_codebase/test_mmdet/test_object_detection.py @@ -0,0 +1,155 @@ +import os +from tempfile import NamedTemporaryFile, TemporaryDirectory + +import mmcv +import numpy as np +import pytest +import torch +from torch.utils.data import DataLoader +from torch.utils.data.dataset import Dataset + +import mmdeploy.backend.onnxruntime as ort_apis +from mmdeploy.apis import build_task_processor +from mmdeploy.utils import load_config +from mmdeploy.utils.test import DummyModel, SwitchBackendWrapper + +model_cfg_path = 'tests/test_codebase/test_mmdet/data/model.py' +model_cfg = load_config(model_cfg_path)[0] +deploy_cfg = mmcv.Config( + dict( + backend_config=dict(type='onnxruntime'), + codebase_config=dict( + type='mmdet', + 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, + )), + 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 = build_task_processor(model_cfg, deploy_cfg, 'cpu') +img_shape = (32, 32) +img = np.random.rand(*img_shape, 3) + + +def test_init_pytorch_model(): + from mmdet.models import BaseDetector + model = task_processor.init_pytorch_model(None) + assert isinstance(model, BaseDetector) + + +@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), + 'labels': torch.rand(1, 10) + }) + + yield task_processor.init_backend_model(['']) + + wrapper.recover() + + +def test_init_backend_model(backend_model): + from mmdeploy.codebase.mmdet.deploy.object_detection_model \ + import End2EndModel + assert isinstance(backend_model, End2EndModel) + + +@pytest.mark.parametrize('device', ['cpu', 'cuda:0']) +def test_create_input(device): + original_device = task_processor.device + task_processor.device = device + inputs = task_processor.create_input(img, input_shape=img_shape) + assert len(inputs) == 2 + task_processor.device = original_device + + +def test_run_inference(backend_model): + torch_model = task_processor.init_pytorch_model(None) + input_dict, _ = task_processor.create_input(img, input_shape=img_shape) + torch_results = task_processor.run_inference(torch_model, input_dict) + backend_results = task_processor.run_inference(backend_model, input_dict) + assert torch_results is not None + assert backend_results is not None + assert len(torch_results[0]) == len(backend_results[0]) + + +def test_visualize(backend_model): + input_dict, _ = task_processor.create_input(img, input_shape=img_shape) + results = task_processor.run_inference(backend_model, input_dict) + with TemporaryDirectory() as dir: + filename = dir + 'tmp.jpg' + task_processor.visualize(backend_model, img, results[0], filename, '') + assert os.path.exists(filename) + + +@pytest.mark.parametrize('partition_type', ['single_stage', 'two_stage']) +# Currently only mmdet implements get_partition_cfg +def test_get_partition_cfg(partition_type): + from mmdeploy.codebase.mmdet.deploy.model_partition_cfg import \ + MMDET_PARTITION_CFG + partition_cfg = task_processor.get_partition_cfg( + partition_type=partition_type) + assert partition_cfg == MMDET_PARTITION_CFG[partition_type] + + +def test_get_tensort_from_input(): + input_data = {'img': [torch.ones(3, 4, 5)]} + inputs = task_processor.get_tensor_from_input(input_data) + assert torch.equal(inputs, torch.ones(3, 4, 5)) + + +def test_build_dataset_and_dataloader(): + dataset = task_processor.build_dataset( + dataset_cfg=model_cfg, dataset_type='test') + assert isinstance(dataset, Dataset), 'Failed to build dataset' + dataloader = task_processor.build_dataloader(dataset, 1, 1) + assert isinstance(dataloader, DataLoader), 'Failed to build dataloader' + + +def test_single_gpu_test_and_evaluate(): + from mmcv.parallel import MMDataParallel + + class DummyDataset(Dataset): + + def __getitem__(self, index): + return 0 + + def __len__(self): + return 0 + + def evaluate(self, *args, **kwargs): + return 0 + + def format_results(self, *args, **kwargs): + return 0 + + dataset = DummyDataset() + # Prepare dataloader + dataloader = DataLoader(dataset) + + # Prepare dummy model + model = DummyModel(outputs=[torch.rand([1, 10, 5]), torch.rand([1, 10])]) + model = MMDataParallel(model, device_ids=[0]) + # Run test + outputs = task_processor.single_gpu_test(model, dataloader) + assert isinstance(outputs, list) + output_file = NamedTemporaryFile(suffix='.pkl').name + task_processor.evaluate_outputs( + model_cfg, outputs, dataset, 'bbox', out=output_file, format_only=True) diff --git a/tests/test_codebase/test_mmdet/test_object_detection_model.py b/tests/test_codebase/test_mmdet/test_object_detection_model.py new file mode 100644 index 0000000000..bb16f8970b --- /dev/null +++ b/tests/test_codebase/test_mmdet/test_object_detection_model.py @@ -0,0 +1,446 @@ +import importlib +import os.path as osp +from tempfile import NamedTemporaryFile + +import mmcv +import numpy as np +import pytest +import torch + +import mmdeploy.backend.onnxruntime as ort_apis +from mmdeploy.codebase.mmdet.deploy.object_detection_model import End2EndModel +from mmdeploy.utils import Backend +from mmdeploy.utils.test import SwitchBackendWrapper + + +def assert_det_results(results, module_name: str = 'model'): + assert results is not None, f'failed to get output using {module_name}' + assert isinstance(results, list) + assert len(results) == 2 + assert results[0].shape[0] == results[1].shape[0] + assert results[0].shape[1] == results[1].shape[1] + + +def assert_forward_results(results, module_name: str = 'model'): + assert results is not None, f'failed to get output using {module_name}' + assert isinstance(results, list) + assert len(results) == 1 + if isinstance(results[0], tuple): # mask + assert len(results[0][0]) == 80 + else: + assert len(results[0]) == 80 + + +@pytest.mark.skipif( + not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') +class TestEnd2EndModel: + + @classmethod + def setup_class(cls): + # force add backend wrapper regardless of plugins + # make sure ONNXRuntimeDetector can use ORTWrapper inside itself + from mmdeploy.backend.onnxruntime import ORTWrapper + ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + + # simplify backend inference + cls.wrapper = SwitchBackendWrapper(ORTWrapper) + cls.outputs = { + 'dets': torch.rand(1, 10, 5), + 'labels': torch.rand(1, 10) + } + cls.wrapper.set(outputs=cls.outputs) + deploy_cfg = mmcv.Config( + {'onnx_config': { + 'output_names': ['dets', 'labels'] + }}) + + from mmdeploy.codebase.mmdet.deploy.object_detection_model \ + import End2EndModel + cls.end2end_model = End2EndModel(Backend.ONNXRUNTIME, [''], 'cpu', + ['' for i in range(80)], deploy_cfg) + + @classmethod + def teardown_class(cls): + cls.wrapper.recover() + + def test_forward(self): + imgs = [torch.rand(1, 3, 64, 64)] + img_metas = [[{ + 'ori_shape': [64, 64, 3], + 'img_shape': [64, 64, 3], + 'scale_factor': [1, 1, 1, 1], + 'border': [0, 0, 0] + }]] + results = self.end2end_model.forward(imgs, img_metas) + assert_forward_results(results, 'End2EndModel') + + def test_show_result(self): + input_img = np.zeros([64, 64, 3]) + img_path = NamedTemporaryFile(suffix='.jpg').name + + result = (torch.rand(1, 10, 5), torch.rand(1, 10)) + self.end2end_model.show_result( + input_img, result, '', show=False, out_file=img_path) + assert osp.exists(img_path) + + +@pytest.mark.skipif( + not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') +class TestMaskEnd2EndModel: + + @classmethod + def setup_class(cls): + # force add backend wrapper regardless of plugins + # make sure ONNXRuntimeDetector can use ORTWrapper inside itself + from mmdeploy.backend.onnxruntime import ORTWrapper + ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + + # simplify backend inference + num_classes = 80 + num_dets = 10 + cls.wrapper = SwitchBackendWrapper(ORTWrapper) + cls.outputs = { + 'dets': torch.rand(1, num_dets, 5), + 'labels': torch.randint(num_classes, (1, num_dets)), + 'masks': torch.rand(1, num_dets, 28, 28) + } + cls.wrapper.set(outputs=cls.outputs) + deploy_cfg = mmcv.Config({ + 'onnx_config': { + 'output_names': ['dets', 'labels', 'masks'] + }, + 'codebase_config': { + 'post_processing': { + 'export_postprocess_mask': False + } + } + }) + + from mmdeploy.codebase.mmdet.deploy.object_detection_model \ + import End2EndModel + cls.end2end_model = End2EndModel(Backend.ONNXRUNTIME, [''], 'cpu', + ['' for i in range(80)], deploy_cfg) + + @classmethod + def teardown_class(cls): + cls.wrapper.recover() + + def test_forward(self): + imgs = [torch.rand(1, 3, 64, 64)] + img_metas = [[{ + 'ori_shape': [64, 64, 3], + 'img_shape': [64, 64, 3], + 'scale_factor': [1, 1, 1, 1], + }]] + results = self.end2end_model.forward(imgs, img_metas) + assert_forward_results(results, 'mask End2EndModel') + + +def get_test_cfg_and_post_processing(): + test_cfg = { + 'nms_pre': 100, + 'min_bbox_size': 0, + 'score_thr': 0.05, + 'nms': { + 'type': 'nms', + 'iou_threshold': 0.5 + }, + 'max_per_img': 10 + } + post_processing = { + 'score_threshold': 0.05, + 'iou_threshold': 0.5, + 'max_output_boxes_per_class': 20, + 'pre_top_k': -1, + 'keep_top_k': 10, + 'background_label_id': -1 + } + return test_cfg, post_processing + + +@pytest.mark.skipif( + not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') +class TestPartitionSingleStageModel: + + @classmethod + def setup_class(cls): + # force add backend wrapper regardless of plugins + # make sure ONNXRuntimeDetector can use ORTWrapper inside itself + from mmdeploy.backend.onnxruntime import ORTWrapper + ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + + # simplify backend inference + cls.wrapper = SwitchBackendWrapper(ORTWrapper) + cls.outputs = { + 'scores': torch.rand(1, 10, 80), + 'boxes': torch.rand(1, 10, 4) + } + cls.wrapper.set(outputs=cls.outputs) + + test_cfg, post_processing = get_test_cfg_and_post_processing() + model_cfg = mmcv.Config(dict(model=dict(test_cfg=test_cfg))) + deploy_cfg = mmcv.Config( + dict(codebase_config=dict(post_processing=post_processing))) + + from mmdeploy.codebase.mmdet.deploy.object_detection_model \ + import PartitionSingleStageModel + cls.model = PartitionSingleStageModel( + Backend.ONNXRUNTIME, [''], + 'cpu', ['' for i in range(80)], + model_cfg=model_cfg, + deploy_cfg=deploy_cfg) + + @classmethod + def teardown_class(cls): + cls.wrapper.recover() + + def test_forward_test(self): + imgs = [torch.rand(1, 3, 64, 64)] + img_metas = [[{ + 'ori_shape': [64, 64, 3], + 'img_shape': [64, 64, 3], + 'scale_factor': [1, 1, 1, 1], + }]] + results = self.model.forward_test(imgs, img_metas) + assert_det_results(results, 'PartitionSingleStageModel') + + def test_postprocess(self): + scores = torch.rand(1, 120, 80) + bboxes = torch.rand(1, 120, 4) + + results = self.model.partition0_postprocess( + scores=scores, bboxes=bboxes) + assert_det_results( + results, '.partition0_postprocess of' + 'PartitionSingleStageModel') + + +def prepare_model_deploy_cfgs(): + test_cfg, post_processing = get_test_cfg_and_post_processing() + bbox_roi_extractor = { + 'type': 'SingleRoIExtractor', + 'roi_layer': { + 'type': 'RoIAlign', + 'output_size': 7, + 'sampling_ratio': 0 + }, + 'out_channels': 8, + 'featmap_strides': [4] + } + bbox_head = { + 'type': 'Shared2FCBBoxHead', + 'in_channels': 8, + 'fc_out_channels': 1024, + 'roi_feat_size': 7, + 'num_classes': 80, + 'bbox_coder': { + 'type': 'DeltaXYWHBBoxCoder', + 'target_means': [0.0, 0.0, 0.0, 0.0], + 'target_stds': [0.1, 0.1, 0.2, 0.2] + }, + 'reg_class_agnostic': False, + 'loss_cls': { + 'type': 'CrossEntropyLoss', + 'use_sigmoid': False, + 'loss_weight': 1.0 + }, + 'loss_bbox': { + 'type': 'L1Loss', + 'loss_weight': 1.0 + } + } + roi_head = dict(bbox_roi_extractor=bbox_roi_extractor, bbox_head=bbox_head) + model_cfg = mmcv.Config( + dict( + model=dict( + neck=dict(num_outs=0), + test_cfg=dict(rpn=test_cfg, rcnn=test_cfg), + roi_head=roi_head))) + deploy_cfg = mmcv.Config( + dict(codebase_config=dict(post_processing=post_processing))) + return model_cfg, deploy_cfg + + +class DummyWrapper(torch.nn.Module): + + def __init__(self, outputs): + self.outputs = outputs + + def __call__(self, *arg, **kwargs): + return 0 + + def output_to_list(self, *arg, **kwargs): + return self.outputs + + +@pytest.mark.skipif( + not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') +class TestPartitionTwoStageModel: + + @classmethod + def setup_class(cls): + # force add backend wrapper regardless of plugins + # make sure ONNXRuntimeDetector can use ORTWrapper inside itself + from mmdeploy.backend.onnxruntime import ORTWrapper + ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + + # simplify backend inference + cls.wrapper = SwitchBackendWrapper(ORTWrapper) + outputs = [ + np.random.rand(1, 12, 80).astype(np.float32), + np.random.rand(1, 12, 4).astype(np.float32), + ] * 2 + + model_cfg, deploy_cfg = prepare_model_deploy_cfgs() + + cls.wrapper.set( + outputs=outputs, model_cfg=model_cfg, deploy_cfg=deploy_cfg) + + # replace original function in PartitionTwoStageModel + from mmdeploy.codebase.mmdet.deploy.object_detection_model \ + import PartitionTwoStageModel + + cls.model = PartitionTwoStageModel( + Backend.ONNXRUNTIME, ['', ''], + 'cpu', ['' for i in range(80)], + model_cfg=model_cfg, + deploy_cfg=deploy_cfg) + feats = [torch.randn(1, 8, 14, 14) for i in range(5)] + scores = torch.rand(1, 10, 1) + bboxes = torch.rand(1, 10, 4) + bboxes[..., 2:4] = 2 * bboxes[..., :2] + + cls_score = torch.rand(10, 81) + bbox_pred = torch.rand(10, 320) + + cls.model.device = 'cpu' + cls.model.CLASSES = ['' for i in range(80)] + cls.model.first_wrapper = DummyWrapper([*feats, scores, bboxes]) + cls.model.second_wrapper = DummyWrapper([cls_score, bbox_pred]) + + @classmethod + def teardown_class(cls): + cls.wrapper.recover() + + def test_postprocess(self): + feats = [torch.randn(1, 8, 14, 14) for i in range(5)] + scores = torch.rand(1, 50, 1) + bboxes = torch.rand(1, 50, 4) + bboxes[..., 2:4] = 2 * bboxes[..., :2] + + results = self.model.partition0_postprocess( + x=feats, scores=scores, bboxes=bboxes) + assert results is not None, 'failed to get output using '\ + 'partition0_postprocess of PartitionTwoStageDetector' + assert isinstance(results, tuple) + assert len(results) == 2 + + rois = torch.rand(1, 10, 5) + cls_score = torch.rand(10, 81) + bbox_pred = torch.rand(10, 320) + img_metas = [[{ + 'ori_shape': [32, 32, 3], + 'img_shape': [32, 32, 3], + 'scale_factor': [1, 1, 1, 1], + }]] + results = self.model.partition1_postprocess( + rois=rois, + cls_score=cls_score, + bbox_pred=bbox_pred, + img_metas=img_metas) + assert results is not None, 'failed to get output using '\ + 'partition1_postprocess of PartitionTwoStageDetector' + assert isinstance(results, tuple) + assert len(results) == 2 + + def test_forward(self): + + class DummyPTSDetector(torch.nn.Module): + """A dummy wrapper for unit tests.""" + + def __init__(self, *args, **kwargs): + self.output_names = ['dets', 'labels'] + + def partition0_postprocess(self, *args, **kwargs): + return self.outputs0 + + def partition1_postprocess(self, *args, **kwargs): + return self.outputs1 + + import types + self.model.partition0_postprocess = types.MethodType( + DummyPTSDetector.partition0_postprocess, self.model) + self.model.partition1_postprocess = types.MethodType( + DummyPTSDetector.partition1_postprocess, self.model) + self.model.outputs0 = [torch.rand(2, 3).cuda()] * 2 + self.model.outputs1 = [ + torch.rand(1, 9, 5).cuda(), + torch.rand(1, 9).cuda() + ] + + imgs = [torch.rand(1, 3, 32, 32)] + img_metas = [[{ + 'ori_shape': [32, 32, 3], + 'img_shape': [32, 32, 3], + 'scale_factor': [1, 1, 1, 1], + }]] + results = self.model.forward(imgs, img_metas) + assert_forward_results(results, 'PartitionTwoStageModel') + + +data_cfg1 = mmcv.Config( + dict( + data=dict( + test=dict(type='CocoDataset'), + val=dict(type='CityscapesDataset'), + train=dict(type='CityscapesDataset')))) +data_cfg2 = mmcv.Config( + dict( + data=dict( + val=dict(type='CocoDataset'), train=dict( + type='CityscapesDataset')))) +data_cfg3 = mmcv.Config(dict(data=dict(train=dict(type='CocoDataset')))) +data_cfg4 = mmcv.Config(dict(data=dict(error=dict(type='CocoDataset')))) + + +@pytest.mark.parametrize('cfg', [data_cfg1, data_cfg2, data_cfg3, data_cfg4]) +def test_get_classes_from_cfg(cfg): + from mmdet.datasets import DATASETS + from mmdeploy.codebase.mmdet.deploy.object_detection_model import \ + get_classes_from_config + + if 'error' in cfg.data: + with pytest.raises(RuntimeError): + get_classes_from_config(cfg) + else: + assert get_classes_from_config( + cfg) == DATASETS.module_dict['CocoDataset'].CLASSES + + +@pytest.mark.skipif( + not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') +@pytest.mark.parametrize('partition_type', [None, 'end2end']) +def test_build_object_detection_model(partition_type): + _, post_processing = get_test_cfg_and_post_processing() + model_cfg = mmcv.Config(dict(data=dict(test={'type': 'CocoDataset'}))) + deploy_cfg = mmcv.Config( + dict( + backend_config=dict(type='onnxruntime'), + onnx_config=dict(output_names=['dets', 'labels']), + codebase_config=dict( + type='mmdet', post_processing=post_processing))) + if partition_type: + deploy_cfg.partition_config = dict( + apply_marks=True, type=partition_type) + + from mmdeploy.backend.onnxruntime import ORTWrapper + ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + + # simplify backend inference + with SwitchBackendWrapper(ORTWrapper) as wrapper: + wrapper.set(model_cfg=model_cfg, deploy_cfg=deploy_cfg) + from mmdeploy.codebase.mmdet.deploy.object_detection_model import \ + build_object_detection_model + detector = build_object_detection_model([''], model_cfg, deploy_cfg, + 'cpu') + assert isinstance(detector, End2EndModel) diff --git a/tests/test_mmedit/data/imgs/blank.jpg b/tests/test_codebase/test_mmedit/data/imgs/blank.jpg similarity index 100% rename from tests/test_mmedit/data/imgs/blank.jpg rename to tests/test_codebase/test_mmedit/data/imgs/blank.jpg diff --git a/tests/test_mmedit/data/model.py b/tests/test_codebase/test_mmedit/data/model.py similarity index 96% rename from tests/test_mmedit/data/model.py rename to tests/test_codebase/test_mmedit/data/model.py index 289ece5728..eecc3f7fb3 100644 --- a/tests/test_mmedit/data/model.py +++ b/tests/test_codebase/test_mmedit/data/model.py @@ -73,8 +73,8 @@ test_dataloader=dict(samples_per_gpu=1), test=dict( type=val_dataset_type, - lq_folder='tests/test_mmedit/data/imgs', - gt_folder='tests/test_mmedit/data/imgs', + lq_folder='tests/test_codebase/test_mmedit/data/imgs', + gt_folder='tests/test_codebase/test_mmedit/data/imgs', pipeline=test_pipeline, scale=scale, filename_tmpl='{}')) diff --git a/tests/test_mmedit/test_mmedit_models.py b/tests/test_codebase/test_mmedit/test_mmedit_models.py similarity index 100% rename from tests/test_mmedit/test_mmedit_models.py rename to tests/test_codebase/test_mmedit/test_mmedit_models.py diff --git a/tests/test_codebase/test_mmedit/test_super_resolution.py b/tests/test_codebase/test_mmedit/test_super_resolution.py new file mode 100644 index 0000000000..1839dee317 --- /dev/null +++ b/tests/test_codebase/test_mmedit/test_super_resolution.py @@ -0,0 +1,117 @@ +import os +import tempfile +from tempfile import NamedTemporaryFile + +import mmcv +import numpy as np +import pytest +import torch + +import mmdeploy.apis.onnxruntime as ort_apis +from mmdeploy.apis import build_task_processor +from mmdeploy.utils import load_config +from mmdeploy.utils.test import SwitchBackendWrapper + +model_cfg = 'tests/test_codebase/test_mmedit/data/model.py' +model_cfg = load_config(model_cfg)[0] +deploy_cfg = mmcv.Config( + dict( + backend_config=dict(type='onnxruntime'), + codebase_config=dict(type='mmedit', task='SuperResolution'), + onnx_config=dict( + type='onnx', + export_params=True, + keep_initializers_as_inputs=False, + opset_version=11, + input_shape=None, + input_names=['input'], + output_names=['output']))) +input_img = np.random.rand(32, 32, 3) +img_shape = [32, 32] +input = {'lq': input_img} +onnx_file = NamedTemporaryFile(suffix='.onnx').name +task_processor = build_task_processor(model_cfg, deploy_cfg, 'cpu') + + +def test_init_pytorch_model(): + torch_model = task_processor.init_pytorch_model(None) + assert torch_model is not None + + +@pytest.fixture +def backend_model(): + from mmdeploy.backend.onnxruntime import ORTWrapper + ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + wrapper = SwitchBackendWrapper(ORTWrapper) + wrapper.set(outputs={ + 'output': torch.rand(3, 50, 50), + }) + + yield task_processor.init_backend_model(['']) + + wrapper.recover() + + +def test_init_backend_model(backend_model): + assert backend_model is not None + + +def test_create_input(): + inputs = task_processor.create_input(input_img, img_shape=img_shape) + assert inputs is not None + + +def test_visualize(backend_model): + result = task_processor.run_inference(backend_model, input) + with tempfile.TemporaryDirectory() as dir: + filename = dir + 'tmp.jpg' + task_processor.visualize(backend_model, input_img, result[0], filename, + 'onnxruntime') + assert os.path.exists(filename) + + +def test_run_inference(backend_model): + results = task_processor.run_inference(backend_model, input) + assert results is not None + + +def test_get_tensor_from_input(): + assert type(task_processor.get_tensor_from_input(input)) is not dict + + +def test_get_partition_cfg(): + with pytest.raises(NotImplementedError): + task_processor.get_partition_cfg(None) + + +def test_build_dataset(): + data = dict( + test={ + 'type': 'SRFolderDataset', + 'lq_folder': 'tests/test_codebase/test_mmedit/data/imgs', + 'gt_folder': 'tests/test_codebase/test_mmedit/data/imgs', + 'scale': 1, + 'filename_tmpl': '{}', + 'pipeline': [ + { + 'type': 'LoadImageFromFile' + }, + ] + }) + dataset_cfg = mmcv.Config(dict(data=data)) + dataset = task_processor.build_dataset( + dataset_cfg=dataset_cfg, dataset_type='test') + assert dataset is not None, 'Failed to build dataset' + dataloader = task_processor.build_dataloader(dataset, 1, 1) + assert dataloader is not None, 'Failed to build dataloader' + + +def test_single_gpu_test(backend_model): + from mmcv.parallel import MMDataParallel + dataset = task_processor.build_dataset(model_cfg, dataset_type='test') + assert dataset is not None, 'Failed to build dataset' + dataloader = task_processor.build_dataloader(dataset, 1, 1) + assert dataloader is not None, 'Failed to build dataloader' + backend_model = MMDataParallel(backend_model, device_ids=[0]) + outputs = task_processor.single_gpu_test(backend_model, dataloader) + assert outputs is not None, 'Failed to test model' diff --git a/tests/test_codebase/test_mmedit/test_super_resolution_model.py b/tests/test_codebase/test_mmedit/test_super_resolution_model.py new file mode 100644 index 0000000000..0c12260a3c --- /dev/null +++ b/tests/test_codebase/test_mmedit/test_super_resolution_model.py @@ -0,0 +1,53 @@ +import importlib + +import mmcv +import numpy as np +import pytest +import torch + +import mmdeploy.backend.onnxruntime as ort_apis +from mmdeploy.utils import Backend, load_config +from mmdeploy.utils.test import SwitchBackendWrapper + + +@pytest.mark.skipif( + not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') +class TestEnd2EndModel: + + @classmethod + def setup_class(cls): + # force add backend wrapper regardless of plugins + # make sure ONNXRuntimeEditor can use ORTWrapper inside itself + from mmdeploy.backend.onnxruntime import ORTWrapper + ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + + # simplify backend inference + cls.wrapper = SwitchBackendWrapper(ORTWrapper) + cls.outputs = { + 'outputs': torch.rand(3, 64, 64), + } + cls.wrapper.set(outputs=cls.outputs) + deploy_cfg = mmcv.Config( + {'onnx_config': { + 'output_names': ['outputs'] + }}) + model_cfg = 'tests/test_codebase/test_mmedit/data/model.py' + model_cfg = load_config(model_cfg)[0] + from mmdeploy.codebase.mmedit.deploy.super_resolution_model\ + import End2EndModel + cls.end2end_model = End2EndModel(Backend.ONNXRUNTIME, [''], 'cpu', + model_cfg, deploy_cfg) + + @classmethod + def teardown_class(cls): + cls.wrapper.recover() + + def test_forward(self): + input_img = np.random.rand(3, 32, 32) + + results = self.end2end_model.forward(input_img, test_mode=False) + assert results is not None + + results = self.end2end_model.forward( + input_img, test_mode=True, gt=torch.tensor(results[0])) + assert results is not None diff --git a/tests/test_mmseg/data/model.py b/tests/test_codebase/test_mmseg/data/model.py similarity index 50% rename from tests/test_mmseg/data/model.py rename to tests/test_codebase/test_mmseg/data/model.py index 925c36db5f..d839a24fa0 100644 --- a/tests/test_mmseg/data/model.py +++ b/tests/test_codebase/test_mmseg/data/model.py @@ -1,15 +1,15 @@ # dataset settings dataset_type = 'CityscapesDataset' -data_root = 'data/cityscapes/' +data_root = '.' img_norm_cfg = dict( mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True) -crop_size = (512, 1024) +crop_size = (128, 128) test_pipeline = [ dict(type='LoadImageFromFile'), dict( type='MultiScaleFlipAug', - img_scale=(2048, 1024), + img_scale=(128, 128), flip=False, transforms=[ dict(type='Resize', keep_ratio=True), @@ -25,58 +25,44 @@ val=dict( type=dataset_type, data_root=data_root, - img_dir='leftImg8bit/val', - ann_dir='gtFine/val', + img_dir='', + ann_dir='', pipeline=test_pipeline), test=dict( type=dataset_type, data_root=data_root, - img_dir='leftImg8bit/val', - ann_dir='gtFine/val', + img_dir='', + ann_dir='', pipeline=test_pipeline)) # model settings -norm_cfg = dict(type='SyncBN', requires_grad=True) +norm_cfg = dict(type='SyncBN', requires_grad=True, momentum=0.01) model = dict( type='EncoderDecoder', - pretrained='open-mmlab://resnet50_v1c', backbone=dict( - type='ResNetV1c', - depth=50, - num_stages=4, - out_indices=(0, 1, 2, 3), - dilations=(1, 1, 2, 4), - strides=(1, 2, 1, 1), + type='FastSCNN', + downsample_dw_channels=(32, 48), + global_in_channels=64, + global_block_channels=(64, 96, 128), + global_block_strides=(2, 2, 1), + global_out_channels=128, + higher_in_channels=64, + lower_in_channels=128, + fusion_out_channels=128, + out_indices=(0, 1, 2), norm_cfg=norm_cfg, - norm_eval=False, - style='pytorch', - contract_dilation=True), + align_corners=False), decode_head=dict( - type='FCNHead', - in_channels=2048, - in_index=3, - channels=512, - num_convs=2, - concat_input=True, - dropout_ratio=0.1, - num_classes=19, - norm_cfg=norm_cfg, - align_corners=False, - loss_decode=dict( - type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)), - auxiliary_head=dict( - type='FCNHead', - in_channels=1024, - in_index=2, - channels=256, - num_convs=1, + type='DepthwiseSeparableFCNHead', + in_channels=128, + channels=128, concat_input=False, - dropout_ratio=0.1, num_classes=19, + in_index=-1, norm_cfg=norm_cfg, align_corners=False, loss_decode=dict( - type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)), + type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1)), # model training and testing settings train_cfg=dict(), test_cfg=dict(mode='whole')) diff --git a/tests/test_mmseg/test_mmseg_models.py b/tests/test_codebase/test_mmseg/test_mmseg_models.py similarity index 88% rename from tests/test_mmseg/test_mmseg_models.py rename to tests/test_codebase/test_mmseg/test_mmseg_models.py index a6338aa723..de40e81041 100644 --- a/tests/test_mmseg/test_mmseg_models.py +++ b/tests/test_codebase/test_mmseg/test_mmseg_models.py @@ -1,5 +1,6 @@ import mmcv import numpy as np +import pytest import torch import torch.nn as nn from mmcv import ConfigDict @@ -87,13 +88,14 @@ def _demo_mm_inputs(input_shape=(1, 3, 8, 16), num_classes=10): return mm_inputs -def test_encoderdecoder_simple_test(): +@pytest.mark.parametrize('backend_type', ['onnxruntime', 'ncnn']) +def test_encoderdecoder_simple_test(backend_type): segmentor = get_model() segmentor.cpu().eval() deploy_cfg = mmcv.Config( dict( - backend_config=dict(type='onnxruntime'), + backend_config=dict(type=backend_type), onnx_config=dict(output_names=['result'], input_shape=None), codebase_config=dict(type='mmseg', task='Segmentation'))) @@ -119,22 +121,24 @@ def test_encoderdecoder_simple_test(): wrapped_model=wrapped_model, model_inputs=rewrite_inputs, deploy_cfg=deploy_cfg) + if is_backend_output: - rewrite_outputs = torch.tensor(rewrite_outputs) - model_outputs = torch.tensor(model_outputs) + rewrite_outputs = rewrite_outputs[0] + model_outputs = torch.tensor(model_outputs[0]) model_outputs = model_outputs.unsqueeze(0).unsqueeze(0) assert torch.allclose(rewrite_outputs, model_outputs) else: assert rewrite_outputs is not None -def test_basesegmentor_forward(): +@pytest.mark.parametrize('backend_type', ['onnxruntime', 'ncnn']) +def test_basesegmentor_forward(backend_type): segmentor = get_model() segmentor.cpu().eval() deploy_cfg = mmcv.Config( dict( - backend_config=dict(type='onnxruntime'), + backend_config=dict(type=backend_type), onnx_config=dict(output_names=['result'], input_shape=None), codebase_config=dict(type='mmseg', task='Segmentation'))) @@ -159,15 +163,16 @@ def test_basesegmentor_forward(): model_inputs=rewrite_inputs, deploy_cfg=deploy_cfg) if is_backend_output: - rewrite_outputs = torch.tensor(rewrite_outputs) - model_outputs = torch.tensor(model_outputs) + rewrite_outputs = torch.tensor(rewrite_outputs[0]) + model_outputs = torch.tensor(model_outputs[0]) model_outputs = model_outputs.unsqueeze(0).unsqueeze(0) assert torch.allclose(rewrite_outputs, model_outputs) else: assert rewrite_outputs is not None -def test_aspphead_forward(): +@pytest.mark.parametrize('backend_type', ['onnxruntime', 'ncnn']) +def test_aspphead_forward(backend_type): from mmseg.models.decode_heads import ASPPHead head = ASPPHead( in_channels=32, channels=16, num_classes=19, @@ -175,7 +180,7 @@ def test_aspphead_forward(): deploy_cfg = mmcv.Config( dict( - backend_config=dict(type='onnxruntime'), + backend_config=dict(type=backend_type), onnx_config=dict( output_names=['result'], input_shape=(1, 32, 45, 45)), codebase_config=dict(type='mmseg', task='Segmentation'))) @@ -197,13 +202,14 @@ def test_aspphead_forward(): assert rewrite_outputs is not None -def test_psphead_forward(): +@pytest.mark.parametrize('backend_type', ['onnxruntime', 'ncnn']) +def test_psphead_forward(backend_type): from mmseg.models.decode_heads import PSPHead head = PSPHead(in_channels=32, channels=16, num_classes=19).eval() deploy_cfg = mmcv.Config( dict( - backend_config=dict(type='onnxruntime'), + backend_config=dict(type=backend_type), onnx_config=dict(output_names=['result'], input_shape=None), codebase_config=dict(type='mmseg', task='Segmentation'))) inputs = [torch.randn(1, 32, 45, 45)] diff --git a/tests/test_codebase/test_mmseg/test_segmentation.py b/tests/test_codebase/test_mmseg/test_segmentation.py new file mode 100644 index 0000000000..b0c987ad99 --- /dev/null +++ b/tests/test_codebase/test_mmseg/test_segmentation.py @@ -0,0 +1,115 @@ +import os +from tempfile import NamedTemporaryFile, TemporaryDirectory + +import mmcv +import numpy as np +import pytest +import torch +from torch.utils.data import DataLoader + +import mmdeploy.backend.onnxruntime as ort_apis +from mmdeploy.apis import build_task_processor +from mmdeploy.utils import load_config +from mmdeploy.utils.test import DummyModel, SwitchBackendWrapper + +model_cfg_path = 'tests/test_codebase/test_mmseg/data/model.py' +model_cfg = load_config(model_cfg_path)[0] +deploy_cfg = mmcv.Config( + dict( + backend_config=dict(type='onnxruntime'), + codebase_config=dict(type='mmseg', task='Segmentation'), + onnx_config=dict( + type='onnx', + export_params=True, + keep_initializers_as_inputs=False, + opset_version=11, + input_shape=None, + input_names=['input'], + output_names=['output']))) + +onnx_file = NamedTemporaryFile(suffix='.onnx').name +task_processor = build_task_processor(model_cfg, deploy_cfg, 'cpu') +img_shape = (32, 32) +img = np.random.rand(*img_shape, 3) + + +def test_init_pytorch_model(): + from mmseg.models.segmentors.base import BaseSegmentor + model = task_processor.init_pytorch_model(None) + assert isinstance(model, BaseSegmentor) + + +@pytest.fixture +def backend_model(): + from mmdeploy.backend.onnxruntime import ORTWrapper + ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + wrapper = SwitchBackendWrapper(ORTWrapper) + wrapper.set(outputs={ + 'output': torch.rand(1, 1, *img_shape), + }) + + yield task_processor.init_backend_model(['']) + + wrapper.recover() + + +def test_init_backend_model(backend_model): + assert isinstance(backend_model, torch.nn.Module) + + +def test_create_input(): + inputs = task_processor.create_input(img, input_shape=img_shape) + assert isinstance(inputs, tuple) and len(inputs) == 2 + + +def test_run_inference(backend_model): + input_dict, _ = task_processor.create_input(img, input_shape=img_shape) + results = task_processor.run_inference(backend_model, input_dict) + assert results is not None + + +def test_visualize(backend_model): + input_dict, _ = task_processor.create_input(img, input_shape=img_shape) + results = task_processor.run_inference(backend_model, input_dict) + with TemporaryDirectory() as dir: + filename = dir + 'tmp.jpg' + task_processor.visualize(backend_model, img, results[0], filename, '') + assert os.path.exists(filename) + + +def test_get_tensort_from_input(): + input_data = {'img': [torch.ones(3, 4, 5)]} + inputs = task_processor.get_tensor_from_input(input_data) + assert torch.equal(inputs, torch.ones(3, 4, 5)) + + +def test_get_partition_cfg(): + try: + _ = task_processor.get_partition_cfg(partition_type='') + except NotImplementedError: + pass + + +def test_build_dataset_and_dataloader(): + from torch.utils.data import Dataset, DataLoader + dataset = task_processor.build_dataset( + dataset_cfg=model_cfg, dataset_type='test') + assert isinstance(dataset, Dataset), 'Failed to build dataset' + dataloader = task_processor.build_dataloader(dataset, 1, 1) + assert isinstance(dataloader, DataLoader), 'Failed to build dataloader' + + +def test_single_gpu_test_and_evaluate(): + from mmcv.parallel import MMDataParallel + + # Prepare dataloader + dataloader = DataLoader([]) + + # Prepare dummy model + model = DummyModel(outputs=[torch.rand([1, 1, *img_shape])]) + model = MMDataParallel(model, device_ids=[0]) + assert model is not None + # Run test + outputs = task_processor.single_gpu_test(model, dataloader) + assert outputs is not None + task_processor.evaluate_outputs(model_cfg, outputs, []) diff --git a/tests/test_codebase/test_mmseg/test_segmentation_model.py b/tests/test_codebase/test_mmseg/test_segmentation_model.py new file mode 100644 index 0000000000..b8a5c61bc8 --- /dev/null +++ b/tests/test_codebase/test_mmseg/test_segmentation_model.py @@ -0,0 +1,137 @@ +import importlib +import os.path as osp +from tempfile import NamedTemporaryFile + +import mmcv +import numpy as np +import pytest +import torch + +import mmdeploy.backend.onnxruntime as ort_apis +from mmdeploy.utils import Backend +from mmdeploy.utils.test import SwitchBackendWrapper + +NUM_CLASS = 19 +IMAGE_SIZE = 32 + + +@pytest.mark.skipif( + not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') +class TestEnd2EndModel: + + @classmethod + def setup_class(cls): + # force add backend wrapper regardless of plugins + from mmdeploy.backend.onnxruntime import ORTWrapper + ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + + # simplify backend inference + cls.wrapper = SwitchBackendWrapper(ORTWrapper) + cls.outputs = { + 'outputs': torch.rand(1, 1, IMAGE_SIZE, IMAGE_SIZE), + } + cls.wrapper.set(outputs=cls.outputs) + deploy_cfg = mmcv.Config( + {'onnx_config': { + 'output_names': ['outputs'] + }}) + + from mmdeploy.codebase.mmseg.deploy.segmentation_model \ + import End2EndModel + class_names = ['' for i in range(NUM_CLASS)] + palette = np.random.randint(0, 255, size=(NUM_CLASS, 3)) + cls.end2end_model = End2EndModel( + Backend.ONNXRUNTIME, [''], + device='cpu', + class_names=class_names, + palette=palette, + deploy_cfg=deploy_cfg) + + @classmethod + def teardown_class(cls): + cls.wrapper.recover() + + @pytest.mark.parametrize( + 'ori_shape', + [[IMAGE_SIZE, IMAGE_SIZE, 3], [2 * IMAGE_SIZE, 2 * IMAGE_SIZE, 3]]) + def test_forward(self, ori_shape): + imgs = [torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE)] + img_metas = [[{ + 'ori_shape': ori_shape, + 'img_shape': [IMAGE_SIZE, IMAGE_SIZE, 3], + 'scale_factor': [1., 1., 1., 1.], + }]] + results = self.end2end_model.forward(imgs, img_metas) + assert results is not None, 'failed to get output using '\ + 'End2EndModel' + + def test_forward_test(self): + imgs = torch.rand(2, 3, IMAGE_SIZE, IMAGE_SIZE) + results = self.end2end_model.forward_test(imgs) + assert isinstance(results[0], np.ndarray) + + def test_show_result(self): + input_img = np.zeros([IMAGE_SIZE, IMAGE_SIZE, 3]) + img_path = NamedTemporaryFile(suffix='.jpg').name + + result = [torch.rand(IMAGE_SIZE, IMAGE_SIZE)] + self.end2end_model.show_result( + input_img, result, '', show=False, out_file=img_path) + assert osp.exists(img_path), 'Fails to create drawn image.' + + +@pytest.mark.parametrize('from_file', [True, False]) +@pytest.mark.parametrize('data_type', ['train', 'val', 'test']) +def test_get_classes_palette_from_config(from_file, data_type): + from mmseg.datasets import DATASETS + from mmdeploy.codebase.mmseg.deploy.segmentation_model \ + import get_classes_palette_from_config + dataset_type = 'CityscapesDataset' + data_cfg = mmcv.Config({ + 'data': { + data_type: + dict( + type=dataset_type, + data_root='', + img_dir='', + ann_dir='', + pipeline=None) + } + }) + + if from_file: + config_path = NamedTemporaryFile(suffix='.py').name + with open(config_path, 'w') as file: + file.write(data_cfg.pretty_text) + data_cfg = config_path + + classes, palette = get_classes_palette_from_config(data_cfg) + module = DATASETS.module_dict[dataset_type] + assert classes == module.CLASSES, \ + f'fail to get CLASSES of dataset: {dataset_type}' + assert palette == module.PALETTE, \ + f'fail to get PALETTE of dataset: {dataset_type}' + + +@pytest.mark.skipif( + not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') +def test_build_segmentation_model(): + model_cfg = mmcv.Config( + dict(data=dict(test={'type': 'CityscapesDataset'}))) + deploy_cfg = mmcv.Config( + dict( + backend_config=dict(type='onnxruntime'), + onnx_config=dict(output_names=['outputs']), + codebase_config=dict(type='mmseg'))) + + from mmdeploy.backend.onnxruntime import ORTWrapper + ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) + + # simplify backend inference + with SwitchBackendWrapper(ORTWrapper) as wrapper: + wrapper.set(model_cfg=model_cfg, deploy_cfg=deploy_cfg) + from mmdeploy.codebase.mmseg.deploy.segmentation_model import \ + build_segmentation_model, End2EndModel + segmentor = build_segmentation_model([''], model_cfg, deploy_cfg, + 'cpu') + assert isinstance(segmentor, End2EndModel) diff --git a/tests/test_codebase/test_mmseg/test_utils.py b/tests/test_codebase/test_mmseg/test_utils.py new file mode 100644 index 0000000000..38e8e47e9b --- /dev/null +++ b/tests/test_codebase/test_mmseg/test_utils.py @@ -0,0 +1,25 @@ +import torch +import torch.nn as nn + +from mmdeploy.codebase.mmseg.deploy import convert_syncbatchnorm + + +def test_convert_syncbatchnorm(): + + class ExampleModel(nn.Module): + + def __init__(self): + super(ExampleModel, self).__init__() + self.model = nn.Sequential( + nn.Linear(2, 4), nn.SyncBatchNorm(4), nn.Sigmoid(), + nn.Linear(4, 6), nn.SyncBatchNorm(6), nn.Sigmoid()) + + def forward(self, x): + return self.model(x) + + model = ExampleModel() + out_model = convert_syncbatchnorm(model) + assert isinstance(out_model.model[1], + torch.nn.modules.batchnorm.BatchNorm2d) and isinstance( + out_model.model[4], + torch.nn.modules.batchnorm.BatchNorm2d) diff --git a/tests/test_mmdet/data/imgs/000000000139.jpg b/tests/test_mmdet/data/imgs/000000000139.jpg deleted file mode 100755 index 19023f7183..0000000000 Binary files a/tests/test_mmdet/data/imgs/000000000139.jpg and /dev/null differ diff --git a/tests/test_mmdet/test_mmdet_apis.py b/tests/test_mmdet/test_mmdet_apis.py deleted file mode 100644 index 4e29af5eb6..0000000000 --- a/tests/test_mmdet/test_mmdet_apis.py +++ /dev/null @@ -1,574 +0,0 @@ -import importlib - -import mmcv -import numpy as np -import pytest -import torch - -import mmdeploy.apis.ncnn as ncnn_apis -import mmdeploy.apis.onnxruntime as ort_apis -import mmdeploy.apis.openvino as openvino_apis -import mmdeploy.apis.ppl as ppl_apis -import mmdeploy.apis.tensorrt as trt_apis -from mmdeploy.utils.test import SwitchBackendWrapper - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda') -@pytest.mark.skipif( - not importlib.util.find_spec('tensorrt'), reason='requires tensorrt') -def test_TensorRTDetector(): - # force add backend wrapper regardless of plugins - # make sure TensorRTDetector can use TRTWrapper inside itself - from mmdeploy.apis.tensorrt.tensorrt_utils import TRTWrapper - trt_apis.__dict__.update({'TRTWrapper': TRTWrapper}) - - # simplify backend inference - outputs = { - 'dets': torch.rand(1, 100, 5).cuda(), - 'labels': torch.rand(1, 100).cuda() - } - with SwitchBackendWrapper(TRTWrapper) as wrapper: - wrapper.set(outputs=outputs) - - from mmdeploy.mmdet.apis.inference import TensorRTDetector - trt_detector = TensorRTDetector('', ['' for i in range(80)], 0) - imgs = [torch.rand(1, 3, 64, 64).cuda()] - img_metas = [[{ - 'ori_shape': [64, 64, 3], - 'img_shape': [64, 64, 3], - 'scale_factor': [2.09, 1.87, 2.09, 1.87], - }]] - - results = trt_detector.forward(imgs, img_metas) - assert results is not None, ('failed to get output using ' - 'TensorRTDetector') - - -@pytest.mark.skipif( - not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') -def test_ONNXRuntimeDetector(): - # force add backend wrapper regardless of plugins - # make sure ONNXRuntimeDetector can use ORTWrapper inside itself - from mmdeploy.apis.onnxruntime.onnxruntime_utils import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - - # simplify backend inference - outputs = (torch.rand(1, 100, 5), torch.rand(1, 100)) - with SwitchBackendWrapper(ORTWrapper) as wrapper: - wrapper.set(outputs=outputs) - - from mmdeploy.mmdet.apis.inference import ONNXRuntimeDetector - ort_detector = ONNXRuntimeDetector('', ['' for i in range(80)], 0) - imgs = [torch.rand(1, 3, 64, 64)] - img_metas = [[{ - 'ori_shape': [64, 64, 3], - 'img_shape': [64, 64, 3], - 'scale_factor': [2.09, 1.87, 2.09, 1.87], - }]] - - results = ort_detector.forward(imgs, img_metas) - assert results is not None, 'failed to get output using '\ - 'ONNXRuntimeDetector' - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda') -@pytest.mark.skipif( - not importlib.util.find_spec('pyppl'), reason='requires pyppl') -def test_PPLDetector(): - # force add backend wrapper regardless of plugins - # make sure PPLDetector can use PPLWrapper inside itself - from mmdeploy.apis.ppl.ppl_utils import PPLWrapper - ppl_apis.__dict__.update({'PPLWrapper': PPLWrapper}) - - # simplify backend inference - outputs = (torch.rand(1, 100, 5), torch.rand(1, 100)) - with SwitchBackendWrapper(PPLWrapper) as wrapper: - wrapper.set(outputs=outputs) - - from mmdeploy.mmdet.apis.inference import PPLDetector - ppl_detector = PPLDetector('', ['' for i in range(80)], 0) - imgs = [torch.rand(1, 3, 64, 64)] - img_metas = [[{ - 'ori_shape': [64, 64, 3], - 'img_shape': [64, 64, 3], - 'scale_factor': [2.09, 1.87, 2.09, 1.87], - }]] - - results = ppl_detector.forward(imgs, img_metas) - assert results is not None, 'failed to get output using PPLDetector' - - -@pytest.mark.skipif( - not importlib.util.find_spec('openvino'), reason='requires openvino') -def test_OpenVINODetector(): - # force add backend wrapper regardless of plugins - # make sure OpenVINODetector can use OpenVINOWrapper inside itself - from mmdeploy.apis.openvino.openvino_utils import OpenVINOWrapper - openvino_apis.__dict__.update({'OpenVINOWrapper': OpenVINOWrapper}) - - # simplify backend inference - num_classes = 80 - num_dets = 10 - outputs = { - 'dets': torch.rand(1, num_dets, 5), - 'labels': torch.randint(num_classes, (1, num_dets)), - 'masks': np.random.rand(1, num_dets, 28, 28) - } - deploy_cfg = mmcv.Config( - dict( - codebase_config=dict( - post_processing=dict(export_postprocess_mask=False)))) - with SwitchBackendWrapper(OpenVINOWrapper) as wrapper: - wrapper.set(outputs=outputs) - - from mmdeploy.mmdet.apis.inference import OpenVINODetector - openvino_detector = OpenVINODetector( - '', ['' for i in range(80)], 0, deploy_cfg=deploy_cfg) - imgs = [torch.rand(1, 3, 64, 64)] - img_metas = [[{ - 'ori_shape': [64, 64, 3], - 'img_shape': [64, 64, 3], - 'scale_factor': [2.09, 1.87, 2.09, 1.87], - }]] - - results = openvino_detector.forward(imgs, img_metas) - assert results is not None, 'failed to get output using '\ - 'OpenVINODetector' - - -def get_test_cfg_and_post_processing(): - test_cfg = { - 'nms_pre': 100, - 'min_bbox_size': 0, - 'score_thr': 0.05, - 'nms': { - 'type': 'nms', - 'iou_threshold': 0.5 - }, - 'max_per_img': 10 - } - post_processing = { - 'score_threshold': 0.05, - 'iou_threshold': 0.5, - 'max_output_boxes_per_class': 20, - 'pre_top_k': -1, - 'keep_top_k': 10, - 'background_label_id': -1 - } - return test_cfg, post_processing - - -def test_PartitionSingleStageDetector(): - test_cfg, post_processing = get_test_cfg_and_post_processing() - model_cfg = mmcv.Config(dict(model=dict(test_cfg=test_cfg))) - deploy_cfg = mmcv.Config( - dict(codebase_config=dict(post_processing=post_processing))) - - from mmdeploy.mmdet.apis.inference import PartitionSingleStageDetector - pss_detector = PartitionSingleStageDetector(['' for i in range(80)], - model_cfg=model_cfg, - deploy_cfg=deploy_cfg, - device_id=0) - scores = torch.rand(1, 120, 80) - bboxes = torch.rand(1, 120, 4) - - results = pss_detector.partition0_postprocess(scores=scores, bboxes=bboxes) - assert results is not None, 'failed to get output using '\ - 'partition0_postprocess of PartitionSingleStageDetector' - - -@pytest.mark.skipif( - not importlib.util.find_spec('ncnn'), reason='requires ncnn') -def test_NCNNPSSDetector(): - test_cfg, post_processing = get_test_cfg_and_post_processing() - model_cfg = mmcv.Config(dict(model=dict(test_cfg=test_cfg))) - deploy_cfg = mmcv.Config( - dict(codebase_config=dict(post_processing=post_processing))) - - # force add backend wrapper regardless of plugins - # make sure NCNNPSSDetector can use NCNNWrapper inside itself - from mmdeploy.apis.ncnn.ncnn_utils import NCNNWrapper - ncnn_apis.__dict__.update({'NCNNWrapper': NCNNWrapper}) - - # simplify backend inference - outputs = { - 'scores': torch.rand(1, 120, 80), - 'boxes': torch.rand(1, 120, 4) - } - with SwitchBackendWrapper(NCNNWrapper) as wrapper: - wrapper.set( - outputs=outputs, model_cfg=model_cfg, deploy_cfg=deploy_cfg) - - from mmdeploy.mmdet.apis.inference import NCNNPSSDetector - - ncnn_pss_detector = NCNNPSSDetector(['', ''], ['' for i in range(80)], - model_cfg=model_cfg, - deploy_cfg=deploy_cfg, - device_id=0) - imgs = [torch.rand(1, 3, 32, 32)] - img_metas = [[{ - 'ori_shape': [32, 32, 3], - 'img_shape': [32, 32, 3], - 'scale_factor': [2.09, 1.87, 2.09, 1.87], - }]] - - results = ncnn_pss_detector.forward(imgs, img_metas) - assert results is not None, ('failed to get output using ' - 'NCNNPSSDetector') - - -@pytest.mark.skipif( - not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') -def test_ONNXRuntimePSSDetector(): - test_cfg, post_processing = get_test_cfg_and_post_processing() - model_cfg = mmcv.Config(dict(model=dict(test_cfg=test_cfg))) - deploy_cfg = mmcv.Config( - dict(codebase_config=dict(post_processing=post_processing))) - - # force add backend wrapper regardless of plugins - # make sure ONNXRuntimePSSDetector can use ORTWrapper inside itself - from mmdeploy.apis.onnxruntime.onnxruntime_utils import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - - # simplify backend inference - outputs = [ - np.random.rand(1, 120, 80).astype(np.float32), - np.random.rand(1, 120, 4).astype(np.float32) - ] - with SwitchBackendWrapper(ORTWrapper) as wrapper: - wrapper.set( - outputs=outputs, model_cfg=model_cfg, deploy_cfg=deploy_cfg) - - from mmdeploy.mmdet.apis.inference import ONNXRuntimePSSDetector - - ort_pss_detector = ONNXRuntimePSSDetector( - '', ['' for i in range(80)], - model_cfg=model_cfg, - deploy_cfg=deploy_cfg, - device_id=0) - imgs = [torch.rand(1, 3, 32, 32)] - img_metas = [[{ - 'ori_shape': [32, 32, 3], - 'img_shape': [32, 32, 3], - 'scale_factor': [2.09, 1.87, 2.09, 1.87], - }]] - - results = ort_pss_detector.forward(imgs, img_metas) - assert results is not None, 'failed to get output using ' - 'ONNXRuntimePSSDetector' - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda') -@pytest.mark.skipif( - not importlib.util.find_spec('tensorrt'), reason='requires tensorrt') -def test_TensorRTPSSDetector(): - test_cfg, post_processing = get_test_cfg_and_post_processing() - model_cfg = mmcv.Config(dict(model=dict(test_cfg=test_cfg))) - deploy_cfg = mmcv.Config( - dict(codebase_config=dict(post_processing=post_processing))) - - # force add backend wrapper regardless of plugins - # make sure TensorRTPSSDetector can use TRTWrapper inside itself - from mmdeploy.apis.tensorrt.tensorrt_utils import TRTWrapper - trt_apis.__dict__.update({'TRTWrapper': TRTWrapper}) - - # simplify backend inference - outputs = { - 'scores': torch.rand(1, 120, 80).cuda(), - 'boxes': torch.rand(1, 120, 4).cuda() - } - with SwitchBackendWrapper(TRTWrapper) as wrapper: - wrapper.set( - outputs=outputs, model_cfg=model_cfg, deploy_cfg=deploy_cfg) - - from mmdeploy.mmdet.apis.inference import TensorRTPSSDetector - - trt_pss_detector = TensorRTPSSDetector( - '', ['' for i in range(80)], - model_cfg=model_cfg, - deploy_cfg=deploy_cfg, - device_id=0) - imgs = [torch.rand(1, 3, 32, 32).cuda()] - img_metas = [[{ - 'ori_shape': [32, 32, 3], - 'img_shape': [32, 32, 3], - 'scale_factor': [2.09, 1.87, 2.09, 1.87], - }]] - - results = trt_pss_detector.forward(imgs, img_metas) - assert results is not None, 'failed to get output using ' - 'TensorRTPSSDetector' - - -def prepare_model_deploy_cfgs(): - test_cfg, post_processing = get_test_cfg_and_post_processing() - bbox_roi_extractor = { - 'type': 'SingleRoIExtractor', - 'roi_layer': { - 'type': 'RoIAlign', - 'output_size': 7, - 'sampling_ratio': 0 - }, - 'out_channels': 8, - 'featmap_strides': [4] - } - bbox_head = { - 'type': 'Shared2FCBBoxHead', - 'in_channels': 8, - 'fc_out_channels': 1024, - 'roi_feat_size': 7, - 'num_classes': 80, - 'bbox_coder': { - 'type': 'DeltaXYWHBBoxCoder', - 'target_means': [0.0, 0.0, 0.0, 0.0], - 'target_stds': [0.1, 0.1, 0.2, 0.2] - }, - 'reg_class_agnostic': False, - 'loss_cls': { - 'type': 'CrossEntropyLoss', - 'use_sigmoid': False, - 'loss_weight': 1.0 - }, - 'loss_bbox': { - 'type': 'L1Loss', - 'loss_weight': 1.0 - } - } - roi_head = dict(bbox_roi_extractor=bbox_roi_extractor, bbox_head=bbox_head) - model_cfg = mmcv.Config( - dict( - model=dict( - test_cfg=dict(rpn=test_cfg, rcnn=test_cfg), - roi_head=roi_head))) - deploy_cfg = mmcv.Config( - dict(codebase_config=dict(post_processing=post_processing))) - return model_cfg, deploy_cfg - - -def test_PartitionTwoStageDetector(): - model_cfg, deploy_cfg = prepare_model_deploy_cfgs() - from mmdeploy.mmdet.apis.inference import PartitionTwoStageDetector - pts_detector = PartitionTwoStageDetector(['' for i in range(80)], - model_cfg=model_cfg, - deploy_cfg=deploy_cfg, - device_id=0) - feats = [torch.randn(1, 8, 14, 14) for i in range(5)] - scores = torch.rand(1, 50, 1) - bboxes = torch.rand(1, 50, 4) - bboxes[..., 2:4] = 2 * bboxes[..., :2] - results = pts_detector.partition0_postprocess( - x=feats, scores=scores, bboxes=bboxes) - assert results is not None, 'failed to get output using '\ - 'partition0_postprocess of PartitionTwoStageDetector' - - rois = torch.rand(1, 10, 5) - cls_score = torch.rand(10, 81) - bbox_pred = torch.rand(10, 320) - img_metas = [[{ - 'ori_shape': [32, 32, 3], - 'img_shape': [32, 32, 3], - 'scale_factor': [2.09, 1.87, 2.09, 1.87], - }]] - results = pts_detector.partition1_postprocess( - rois=rois, - cls_score=cls_score, - bbox_pred=bbox_pred, - img_metas=img_metas) - assert results is not None, 'failed to get output using '\ - 'partition1_postprocess of PartitionTwoStageDetector' - - -class DummyPTSDetector(torch.nn.Module): - """A dummy wrapper for unit tests.""" - - def __init__(self, *args, **kwargs): - self.output_names = ['dets', 'labels'] - - def partition0_postprocess(self, *args, **kwargs): - return self.outputs0 - - def partition1_postprocess(self, *args, **kwargs): - return self.outputs1 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda') -@pytest.mark.skipif( - not importlib.util.find_spec('tensorrt'), reason='requires tensorrt') -def test_TensorRTPTSDetector(): - model_cfg, deploy_cfg = prepare_model_deploy_cfgs() - - # force add backend wrapper regardless of plugins - # make sure TensorRTPTSDetector can use TRTWrapper inside itself - from mmdeploy.apis.tensorrt.tensorrt_utils import TRTWrapper - trt_apis.__dict__.update({'TRTWrapper': TRTWrapper}) - - # simplify backend inference - outputs = { - 'scores': torch.rand(1, 12, 80).cuda(), - 'boxes': torch.rand(1, 12, 4).cuda(), - 'cls_score': torch.rand(1, 12, 80).cuda(), - 'bbox_pred': torch.rand(1, 12, 4).cuda() - } - with SwitchBackendWrapper(TRTWrapper) as wrapper: - wrapper.set( - outputs=outputs, model_cfg=model_cfg, deploy_cfg=deploy_cfg) - - # replace original function in PartitionTwoStageDetector - from mmdeploy.mmdet.apis.inference import PartitionTwoStageDetector - PartitionTwoStageDetector.__init__ = DummyPTSDetector.__init__ - PartitionTwoStageDetector.partition0_postprocess = \ - DummyPTSDetector.partition0_postprocess - PartitionTwoStageDetector.partition1_postprocess = \ - DummyPTSDetector.partition1_postprocess - PartitionTwoStageDetector.outputs0 = [torch.rand(2, 3).cuda()] * 2 - PartitionTwoStageDetector.outputs1 = [ - torch.rand(1, 9, 5).cuda(), - torch.rand(1, 9).cuda() - ] - PartitionTwoStageDetector.device_id = 0 - PartitionTwoStageDetector.CLASSES = ['' for i in range(80)] - - from mmdeploy.mmdet.apis.inference import TensorRTPTSDetector - trt_pts_detector = TensorRTPTSDetector(['', ''], - ['' for i in range(80)], - model_cfg=model_cfg, - deploy_cfg=deploy_cfg, - device_id=0) - - imgs = [torch.rand(1, 3, 32, 32).cuda()] - img_metas = [[{ - 'ori_shape': [32, 32, 3], - 'img_shape': [32, 32, 3], - 'scale_factor': [2.09, 1.87, 2.09, 1.87], - }]] - results = trt_pts_detector.forward(imgs, img_metas) - assert results is not None, 'failed to get output using ' - 'TensorRTPTSDetector' - - -@pytest.mark.skipif( - not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') -def test_ONNXRuntimePTSDetector(): - model_cfg, deploy_cfg = prepare_model_deploy_cfgs() - - # force add backend wrapper regardless of plugins - # make sure ONNXRuntimePTSDetector can use TRTWrapper inside itself - from mmdeploy.apis.onnxruntime.onnxruntime_utils import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - - # simplify backend inference - outputs = [ - np.random.rand(1, 12, 80).astype(np.float32), - np.random.rand(1, 12, 4).astype(np.float32), - ] * 2 - with SwitchBackendWrapper(ORTWrapper) as wrapper: - wrapper.set( - outputs=outputs, model_cfg=model_cfg, deploy_cfg=deploy_cfg) - - # replace original function in PartitionTwoStageDetector - from mmdeploy.mmdet.apis.inference import PartitionTwoStageDetector - PartitionTwoStageDetector.__init__ = DummyPTSDetector.__init__ - PartitionTwoStageDetector.partition0_postprocess = \ - DummyPTSDetector.partition0_postprocess - PartitionTwoStageDetector.partition1_postprocess = \ - DummyPTSDetector.partition1_postprocess - PartitionTwoStageDetector.outputs0 = [torch.rand(2, 3)] * 2 - PartitionTwoStageDetector.outputs1 = [ - torch.rand(1, 9, 5), torch.rand(1, 9) - ] - PartitionTwoStageDetector.device_id = -1 - PartitionTwoStageDetector.CLASSES = ['' for i in range(80)] - - from mmdeploy.mmdet.apis.inference import ONNXRuntimePTSDetector - ort_pts_detector = ONNXRuntimePTSDetector(['', ''], - ['' for i in range(80)], - model_cfg=model_cfg, - deploy_cfg=deploy_cfg, - device_id=0) - - imgs = [torch.rand(1, 3, 32, 32)] - img_metas = [[{ - 'ori_shape': [32, 32, 3], - 'img_shape': [32, 32, 3], - 'scale_factor': [2.09, 1.87, 2.09, 1.87], - }]] - results = ort_pts_detector.forward(imgs, img_metas) - assert results is not None, 'failed to get output using ' - 'ONNXRuntimePTSDetector' - - -@pytest.mark.skipif( - not importlib.util.find_spec('ncnn'), reason='requires ncnn') -def test_NCNNPTSDetector(): - model_cfg, deploy_cfg = prepare_model_deploy_cfgs() - num_outs = dict(model=dict(neck=dict(num_outs=0))) - model_cfg.update(num_outs) - - # force add backend wrapper regardless of plugins - # make sure NCNNPTSDetector can use TRTWrapper inside itself - from mmdeploy.apis.ncnn.ncnn_utils import NCNNWrapper - ncnn_apis.__dict__.update({'NCNNWrapper': NCNNWrapper}) - - # simplify backend inference - outputs = { - 'scores': torch.rand(1, 12, 80), - 'boxes': torch.rand(1, 12, 4), - 'cls_score': torch.rand(1, 12, 80), - 'bbox_pred': torch.rand(1, 12, 4) - } - with SwitchBackendWrapper(NCNNWrapper) as wrapper: - wrapper.set( - outputs=outputs, model_cfg=model_cfg, deploy_cfg=deploy_cfg) - - # replace original function in PartitionTwoStageDetector - from mmdeploy.mmdet.apis.inference import PartitionTwoStageDetector - PartitionTwoStageDetector.__init__ = DummyPTSDetector.__init__ - PartitionTwoStageDetector.partition0_postprocess = \ - DummyPTSDetector.partition0_postprocess - PartitionTwoStageDetector.partition1_postprocess = \ - DummyPTSDetector.partition1_postprocess - PartitionTwoStageDetector.outputs0 = [torch.rand(2, 3)] * 2 - PartitionTwoStageDetector.outputs1 = [ - torch.rand(1, 9, 5), torch.rand(1, 9) - ] - PartitionTwoStageDetector.device_id = -1 - PartitionTwoStageDetector.CLASSES = ['' for i in range(80)] - - from mmdeploy.mmdet.apis.inference import NCNNPTSDetector - ncnn_pts_detector = NCNNPTSDetector( - [''] * 4, [''] * 80, - model_cfg=model_cfg, - deploy_cfg=deploy_cfg, - device_id=0) - - imgs = [torch.rand(1, 3, 32, 32)] - img_metas = [[{ - 'ori_shape': [32, 32, 3], - 'img_shape': [32, 32, 3], - 'scale_factor': [2.09, 1.87, 2.09, 1.87], - }]] - results = ncnn_pts_detector.forward(imgs, img_metas) - assert results is not None, 'failed to get output using ' - 'NCNNPTSDetector' - - -@pytest.mark.skipif( - not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime') -def test_build_detector(): - _, post_processing = get_test_cfg_and_post_processing() - model_cfg = mmcv.Config(dict(data=dict(test={'type': 'CocoDataset'}))) - deploy_cfg = mmcv.Config( - dict( - backend_config=dict(type='onnxruntime'), - codebase_config=dict( - type='mmdet', post_processing=post_processing))) - - from mmdeploy.apis.onnxruntime.onnxruntime_utils import ORTWrapper - ort_apis.__dict__.update({'ORTWrapper': ORTWrapper}) - - # simplify backend inference - with SwitchBackendWrapper(ORTWrapper) as wrapper: - wrapper.set(model_cfg=model_cfg, deploy_cfg=deploy_cfg) - from mmdeploy.apis.utils import init_backend_model - detector = init_backend_model([''], model_cfg, deploy_cfg, -1) - assert detector is not None diff --git a/tests/test_mmdet/test_mmdet_export.py b/tests/test_mmdet/test_mmdet_export.py deleted file mode 100644 index c7c00a484f..0000000000 --- a/tests/test_mmdet/test_mmdet_export.py +++ /dev/null @@ -1,104 +0,0 @@ -import mmcv -import numpy as np -import pytest -import torch - -from mmdeploy.apis.utils import (build_dataloader, build_dataset, create_input, - get_tensor_from_input) -from mmdeploy.utils.constants import Codebase, Task - - -def test_create_input(): - task = Task.OBJECT_DETECTION - test = dict(pipeline=[{ - 'type': 'LoadImageFromWebcam' - }, { - 'type': - 'MultiScaleFlipAug', - 'img_scale': [32, 32], - 'flip': - False, - 'transforms': [{ - 'type': 'Resize', - 'keep_ratio': True - }, { - 'type': 'RandomFlip' - }, { - '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': 'DefaultFormatBundle' - }, { - 'type': 'Collect', - 'keys': ['img'] - }] - }]) - data = dict(test=test) - model_cfg = mmcv.Config(dict(data=data)) - imgs = [np.random.rand(32, 32, 3)] - inputs = create_input( - Codebase.MMDET, - task, - model_cfg, - imgs, - input_shape=(32, 32), - device='cpu') - assert inputs is not None, 'Failed to create input' - - -@pytest.mark.parametrize('input_data', [{'img': [torch.ones(3, 4, 5)]}]) -def test_get_tensor_from_input(input_data): - inputs = get_tensor_from_input(Codebase.MMDET, input_data) - assert inputs is not None, 'Failed to get tensor from input' - - -def test_build_dataset(): - data = dict( - test={ - 'type': 'CocoDataset', - 'ann_file': 'tests/test_mmdet/data/coco_sample.json', - 'img_prefix': 'tests/test_mmdet/data/imgs/', - 'pipeline': [ - { - 'type': 'LoadImageFromFile' - }, - ] - }) - dataset_cfg = mmcv.Config(dict(data=data)) - dataset = build_dataset( - Codebase.MMDET, dataset_cfg=dataset_cfg, dataset_type='test') - assert dataset is not None, 'Failed to build dataset' - dataloader = build_dataloader(Codebase.MMDET, dataset, 1, 1) - assert dataloader is not None, 'Failed to build dataloader' - - -def test_clip_bboxes(): - from mmdeploy.mmdet.export import clip_bboxes - x1 = torch.rand(3, 2) * 224 - y1 = torch.rand(3, 2) * 224 - x2 = x1 * 2 - y2 = y1 * 2 - outs = clip_bboxes(x1, y1, x2, y2, [224, 224]) - for out in outs: - assert int(out.max()) <= 224 - - -def test_pad_with_value(): - from mmdeploy.mmdet.export import pad_with_value - x = torch.rand(3, 2) - padded_x = pad_with_value(x, pad_dim=1, pad_size=4, pad_value=0) - assert np.allclose( - padded_x.shape, torch.Size([3, 6]), rtol=1e-03, atol=1e-05) - assert np.allclose(padded_x.sum(), x.sum(), rtol=1e-03, atol=1e-05) - - -@pytest.mark.parametrize('partition_type', ['single_stage', 'two_stage']) -def test_get_partition_cfg(partition_type): - from mmdeploy.mmdet.export import get_partition_cfg - partition_cfg = get_partition_cfg(partition_type=partition_type) - assert partition_cfg is not None diff --git a/tests/test_mmedit/test_mmedit_export.py b/tests/test_mmedit/test_mmedit_export.py deleted file mode 100644 index 75861191c3..0000000000 --- a/tests/test_mmedit/test_mmedit_export.py +++ /dev/null @@ -1,97 +0,0 @@ -import mmcv -import numpy as np - -from mmdeploy.apis.utils import build_dataloader, build_dataset, create_input -from mmdeploy.utils.constants import Codebase, Task - - -class TestCreateInput: - task = Task.SUPER_RESOLUTION - img_test_pipeline = [ - dict( - type='LoadImageFromFile', - io_backend='disk', - key='lq', - flag='unchanged'), - dict( - type='LoadImageFromFile', - io_backend='disk', - key='gt', - flag='unchanged'), - dict(type='RescaleToZeroOne', keys=['lq', 'gt']), - dict( - type='Normalize', - keys=['lq', 'gt'], - mean=[0, 0, 0], - std=[1, 1, 1], - to_rgb=True), - dict( - type='Collect', - keys=['lq', 'gt'], - meta_keys=['lq_path', 'lq_path']), - dict(type='ImageToTensor', keys=['lq', 'gt']) - ] - - imgs = np.random.rand(32, 32, 3) - img_path = 'tests/test_mmedit/data/imgs/blank.jpg' - - def test_create_input_static(self): - data = dict(test=dict(pipeline=TestCreateInput.img_test_pipeline)) - model_cfg = mmcv.Config( - dict(data=data, test_pipeline=TestCreateInput.img_test_pipeline)) - inputs = create_input( - Codebase.MMEDIT, - TestCreateInput.task, - model_cfg, - TestCreateInput.imgs, - input_shape=(32, 32), - device='cpu') - assert inputs is not None, 'Failed to create input' - - def test_create_input_dynamic(self): - data = dict(test=dict(pipeline=TestCreateInput.img_test_pipeline)) - model_cfg = mmcv.Config( - dict(data=data, test_pipeline=TestCreateInput.img_test_pipeline)) - inputs = create_input( - Codebase.MMEDIT, - TestCreateInput.task, - model_cfg, - TestCreateInput.imgs, - input_shape=None, - device='cpu') - assert inputs is not None, 'Failed to create input' - - def test_create_input_from_file(self): - data = dict(test=dict(pipeline=TestCreateInput.img_test_pipeline)) - model_cfg = mmcv.Config( - dict(data=data, test_pipeline=TestCreateInput.img_test_pipeline)) - inputs = create_input( - Codebase.MMEDIT, - TestCreateInput.task, - model_cfg, - TestCreateInput.img_path, - input_shape=None, - device='cpu') - assert inputs is not None, 'Failed to create input' - - -def test_build_dataset(): - data = dict( - test={ - 'type': 'SRFolderDataset', - 'lq_folder': 'tests/test_mmedit/data/imgs', - 'gt_folder': 'tests/test_mmedit/data/imgs', - 'scale': 1, - 'filename_tmpl': '{}', - 'pipeline': [ - { - 'type': 'LoadImageFromFile' - }, - ] - }) - dataset_cfg = mmcv.Config(dict(data=data)) - dataset = build_dataset( - Codebase.MMEDIT, dataset_cfg=dataset_cfg, dataset_type='test') - assert dataset is not None, 'Failed to build dataset' - dataloader = build_dataloader(Codebase.MMEDIT, dataset, 1, 1) - assert dataloader is not None, 'Failed to build dataloader' diff --git a/tests/test_mmseg/data/imgs/blank.jpg b/tests/test_mmseg/data/imgs/blank.jpg deleted file mode 100644 index ac446f47d9..0000000000 Binary files a/tests/test_mmseg/data/imgs/blank.jpg and /dev/null differ diff --git a/tests/test_mmseg/test_mmseg_export.py b/tests/test_mmseg/test_mmseg_export.py deleted file mode 100644 index 80028fb783..0000000000 --- a/tests/test_mmseg/test_mmseg_export.py +++ /dev/null @@ -1,114 +0,0 @@ -import mmcv -import numpy as np -import torch -import torch.nn as nn - -from mmdeploy.apis.utils import build_dataloader, build_dataset, create_input -from mmdeploy.mmseg.export import convert_syncbatchnorm -from mmdeploy.utils.constants import Codebase, Task - - -def test_convert_syncbatchnorm(): - - class ExampleModel(nn.Module): - - def __init__(self): - super(ExampleModel, self).__init__() - self.model = nn.Sequential( - nn.Linear(2, 4), nn.SyncBatchNorm(4), nn.Sigmoid(), - nn.Linear(4, 6), nn.SyncBatchNorm(6), nn.Sigmoid()) - - def forward(self, x): - return self.model(x) - - model = ExampleModel() - out_model = convert_syncbatchnorm(model) - assert isinstance(out_model.model[1], - torch.nn.modules.batchnorm.BatchNorm2d) and isinstance( - out_model.model[4], - torch.nn.modules.batchnorm.BatchNorm2d) - - -class TestCreateInput: - task = Task.SEGMENTATION - img_norm_cfg = dict( - mean=[123.675, 116.28, 103.53], - std=[58.395, 57.12, 57.375], - to_rgb=True) - img_test_pipeline = [ - dict(type='LoadImageFromFile'), - dict( - type='MultiScaleFlipAug', - img_scale=(50, 50), - flip=False, - transforms=[ - dict(type='Resize', keep_ratio=True), - dict(type='RandomFlip'), - dict(type='Normalize', **img_norm_cfg), - dict(type='ImageToTensor', keys=['img']), - dict(type='Collect', keys=['img']), - ]) - ] - - imgs = np.random.rand(32, 32, 3) - img_path = 'tests/test_mmseg/data/imgs/blank.jpg' - - def test_create_input_static(self): - data = dict(test=dict(pipeline=TestCreateInput.img_test_pipeline)) - model_cfg = mmcv.Config( - dict(data=data, test_pipeline=TestCreateInput.img_test_pipeline)) - inputs = create_input( - Codebase.MMSEG, - TestCreateInput.task, - model_cfg, - TestCreateInput.imgs, - input_shape=(32, 32), - device='cpu') - assert inputs is not None, 'Failed to create input' - - def test_create_input_dynamic(self): - data = dict(test=dict(pipeline=TestCreateInput.img_test_pipeline)) - model_cfg = mmcv.Config( - dict(data=data, test_pipeline=TestCreateInput.img_test_pipeline)) - inputs = create_input( - Codebase.MMSEG, - TestCreateInput.task, - model_cfg, - TestCreateInput.imgs, - input_shape=None, - device='cpu') - assert inputs is not None, 'Failed to create input' - - def test_create_input_from_file(self): - data = dict(test=dict(pipeline=TestCreateInput.img_test_pipeline)) - model_cfg = mmcv.Config( - dict(data=data, test_pipeline=TestCreateInput.img_test_pipeline)) - inputs = create_input( - Codebase.MMSEG, - TestCreateInput.task, - model_cfg, - TestCreateInput.img_path, - input_shape=None, - device='cpu') - assert inputs is not None, 'Failed to create input' - - -def test_build_dataset(): - data = dict( - test={ - 'type': 'CityscapesDataset', - 'data_root': 'tests/data', - 'img_dir': '', - 'ann_dir': '', - 'pipeline': [ - { - 'type': 'LoadImageFromFile' - }, - ] - }) - dataset_cfg = mmcv.Config(dict(data=data)) - dataset = build_dataset( - Codebase.MMSEG, dataset_cfg=dataset_cfg, dataset_type='test') - assert dataset is not None, 'Failed to build dataset' - dataloader = build_dataloader(Codebase.MMSEG, dataset, 1, 1) - assert dataloader is not None, 'Failed to build dataloader' diff --git a/tests/test_pytorch/test_pytorch_ops.py b/tests/test_pytorch/test_pytorch_ops.py index 4c8352c34e..b31575b776 100644 --- a/tests/test_pytorch/test_pytorch_ops.py +++ b/tests/test_pytorch/test_pytorch_ops.py @@ -102,7 +102,6 @@ def test_squeeze_default(self): x = torch.rand(1, 1, 2, 2) model = OpModel(torch.squeeze) nodes = get_model_onnx_nodes(model, x) - print(nodes) assert nodes[0].attribute[0].ints == [0, 1] assert nodes[0].op_type == 'Squeeze' @@ -110,6 +109,5 @@ def test_squeeze(self): x = torch.rand(1, 1, 2, 2) model = OpModel(torch.squeeze, 0) nodes = get_model_onnx_nodes(model, x) - print(nodes) assert nodes[0].attribute[0].ints == [0] assert nodes[0].op_type == 'Squeeze' diff --git a/tools/deploy.py b/tools/deploy.py index ad56e7e179..1333280054 100644 --- a/tools/deploy.py +++ b/tools/deploy.py @@ -1,7 +1,6 @@ import argparse import logging import os.path as osp -import subprocess import sys import traceback from functools import partial @@ -10,11 +9,11 @@ import torch.multiprocessing as mp from torch.multiprocessing import Process, set_start_method -from mmdeploy.apis import (create_calib_table, extract_model, inference_model, - torch2onnx) -from mmdeploy.apis.utils import get_partition_cfg as parse_partition_cfg +from mmdeploy.apis import (create_calib_table, extract_model, + get_predefined_partition_cfg, torch2onnx, + visualize_model) from mmdeploy.utils import (Backend, get_backend, get_calib_filename, - get_codebase, get_model_inputs, get_onnx_config, + get_model_inputs, get_onnx_config, get_partition_config, load_config) from mmdeploy.utils.export_info import dump_info @@ -51,7 +50,7 @@ def parse_args(): def target_wrapper(target, log_level, ret_value, *args, **kwargs): logger = logging.getLogger() logging.basicConfig( - format='%(asctime)s,%(msecs)d %(levelname)-8s' + format='%(asctime)s,%(name)s %(levelname)-8s' ' [%(filename)s:%(lineno)d] %(message)s', datefmt='%Y-%m-%d:%H:%M:%S') logger.level @@ -90,7 +89,7 @@ def main(): args = parse_args() set_start_method('spawn') logging.basicConfig( - format='%(asctime)s,%(msecs)d %(levelname)-8s' + format='%(asctime)s,%(name)s %(levelname)-8s' ' [%(filename)s:%(lineno)d] %(message)s', datefmt='%Y-%m-%d:%H:%M:%S') logger = logging.getLogger() @@ -133,8 +132,8 @@ def main(): partition_cfgs = partition_cfgs.get('partition_cfg', None) else: assert 'type' in partition_cfgs - partition_cfgs = parse_partition_cfg( - get_codebase(deploy_cfg), partition_cfgs['type']) + partition_cfgs = get_predefined_partition_cfg( + deploy_cfg, partition_cfgs['type']) origin_onnx_file = onnx_files[0] onnx_files = [] @@ -201,27 +200,23 @@ def main(): backend_files.append(osp.join(args.work_dir, save_file)) elif backend == Backend.NCNN: - from mmdeploy.apis.ncnn import get_onnx2ncnn_path from mmdeploy.apis.ncnn import is_available as is_available_ncnn if not is_available_ncnn(): logging.error('ncnn support is not available.') exit(-1) - onnx2ncnn_path = get_onnx2ncnn_path() + from mmdeploy.apis.ncnn import onnx2ncnn, get_output_model_file backend_files = [] for onnx_path in onnx_files: - onnx_name = osp.splitext(osp.split(onnx_path)[1])[0] - save_param = onnx_name + '.param' - save_bin = onnx_name + '.bin' - - save_param = osp.join(args.work_dir, save_param) - save_bin = osp.join(args.work_dir, save_bin) - - subprocess.call([onnx2ncnn_path, onnx_path, save_param, save_bin]) - - backend_files += [save_param, save_bin] + create_process( + f'onnx2ncnn with {onnx_path}', + target=onnx2ncnn, + args=(onnx_path, args.work_dir), + kwargs=dict(), + ret_value=ret_value) + backend_files += get_output_model_file(onnx_path, args.work_dir) elif backend == Backend.OPENVINO: from mmdeploy.apis.openvino import \ @@ -278,7 +273,7 @@ def main(): # visualize model of the backend create_process( f'visualize {backend.value} model', - target=inference_model, + target=visualize_model, args=(model_cfg_path, deploy_cfg_path, backend_files, args.test_img, args.device), kwargs=dict( @@ -290,7 +285,7 @@ def main(): # visualize pytorch model create_process( 'visualize pytorch model', - target=inference_model, + target=visualize_model, args=(model_cfg_path, deploy_cfg_path, [checkpoint_path], args.test_img, args.device), kwargs=dict( diff --git a/tools/test.py b/tools/test.py index 675b058b91..c365f1a35a 100644 --- a/tools/test.py +++ b/tools/test.py @@ -4,9 +4,9 @@ from mmcv import DictAction from mmcv.parallel import MMDataParallel -from mmdeploy.apis import (build_dataloader, build_dataset, init_backend_model, - post_process_outputs, single_gpu_test) -from mmdeploy.utils.config_utils import get_codebase, load_config +from mmdeploy.apis import build_task_processor +from mmdeploy.utils.config_utils import load_config +from mmdeploy.utils.device import parse_device_id from mmdeploy.utils.timer import TimeCounter @@ -36,11 +36,6 @@ def parse_args(): parser.add_argument('--show', action='store_true', help='show results') parser.add_argument( '--show-dir', help='directory where painted images will be saved') - parser.add_argument( - '--show-score-thr', - type=float, - default=0.3, - help='score threshold (default: 0.3)') parser.add_argument( '--device', help='device used for conversion', default='cpu') parser.add_argument( @@ -95,23 +90,20 @@ def main(): if args.cfg_options is not None: model_cfg.merge_from_dict(args.cfg_options) + task_processor = build_task_processor(model_cfg, deploy_cfg, args.device) + # prepare the dataset loader - codebase = get_codebase(deploy_cfg) dataset_type = 'test' - dataset = build_dataset(codebase, model_cfg, dataset_type) - data_loader = build_dataloader( - codebase, + dataset = task_processor.build_dataset(model_cfg, dataset_type) + data_loader = task_processor.build_dataloader( dataset, samples_per_gpu=1, workers_per_gpu=model_cfg.data.workers_per_gpu) # load the model of the backend - device_id = -1 if args.device == 'cpu' else 0 - model = init_backend_model( - args.model, - model_cfg=args.model_cfg, - deploy_cfg=args.deploy_cfg, - device_id=device_id) + model = task_processor.init_backend_model(args.model) + + device_id = parse_device_id(args.device) model = MMDataParallel(model, device_ids=[0]) if args.speed_test: @@ -125,13 +117,14 @@ def main(): log_interval=args.log_interval, with_sync=with_sync, file=output_file): - outputs = single_gpu_test(codebase, model, data_loader, args.show, - args.show_dir, args.show_score_thr) + outputs = task_processor.single_gpu_test(model, data_loader, + args.show, args.show_dir) else: - outputs = single_gpu_test(codebase, model, data_loader, args.show, - args.show_dir, args.show_score_thr) - post_process_outputs(outputs, dataset, model_cfg, codebase, args.metrics, - args.out, args.metric_options, args.format_only) + outputs = task_processor.single_gpu_test(model, data_loader, args.show, + args.show_dir) + task_processor.evaluate_outputs(model_cfg, outputs, dataset, args.metrics, + args.out, args.metric_options, + args.format_only) if __name__ == '__main__':