From fbae550adf581ce57635ea5b7553f95106408484 Mon Sep 17 00:00:00 2001 From: Bo Date: Wed, 19 Apr 2023 22:10:46 +0800 Subject: [PATCH 1/6] Add visualize scripts. --- tools/visualize.py | 125 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 125 insertions(+) create mode 100644 tools/visualize.py diff --git a/tools/visualize.py b/tools/visualize.py new file mode 100644 index 0000000000..278049c4a0 --- /dev/null +++ b/tools/visualize.py @@ -0,0 +1,125 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import logging +import os +import os.path as osp +import time +from functools import partial + +from tqdm import tqdm +import numpy as np +import mmcv +import mmengine + +from mmdeploy.apis import visualize_model +from mmdeploy.utils import (Backend, get_backend, get_root_logger, + load_config) + + +def parse_args(): + parser = argparse.ArgumentParser(description='Model inference visualization.') + parser.add_argument('--deploy-cfg', help='deploy config path') + parser.add_argument('--model-cfg', help='model config path') + parser.add_argument('--deploy-path', + type=str, + nargs='+', + help='deploy model path') + parser.add_argument( + '--checkpoint', + default=None, + help='model checkpoint path') + parser.add_argument( + '--test-img', + default=None, + type=str, + nargs='+', + help='image used to test model') + parser.add_argument( + '--save-dir', + default=None, + help='the dir to save inference results') + parser.add_argument( + '--device', help='device to run model', default='cpu') + parser.add_argument( + '--log-level', + help='set log level', + default='INFO', + choices=list(logging._nameToLevel.keys())) + parser.add_argument( + '--uri', + default='192.168.1.1:60000', + help='Remote ipv4:port or ipv6:port for inference on edge device.') + args = parser.parse_args() + return args + + +def main(): + args = parse_args() + logger = get_root_logger() + log_level = logging.getLevelName(args.log_level) + logger.setLevel(log_level) + + deploy_cfg_path = args.deploy_cfg + model_cfg_path = args.model_cfg + checkpoint_path = args.checkpoint + deploy_model_path = args.deploy_path + if not isinstance(deploy_model_path, list): + deploy_model_path = [deploy_model_path] + + # load deploy_cfg + deploy_cfg = load_config(deploy_cfg_path)[0] + + # create save_dir or generate default save_dir + save_dir = args.save_dir + if save_dir: + # generate default dir + current_time = time.localtime() + save_dir = osp.join( + os.getcwd(), time.strftime("%Y_%m_%d_%H_%M_%S", current_time) + ) + mmengine.mkdir_or_exist(save_dir) + + # get backend info + backend = get_backend(deploy_cfg) + extra = dict() + if backend == Backend.SNPE: + extra['uri'] = args.uri + + # iterate single_img + for single_img in tqdm(args.test_img): + filename = osp.basename(single_img) + output_file=osp.join(save_dir, filename) + visualize_model( + model_cfg_path, + deploy_cfg_path, + deploy_model_path, + single_img, + args.device, + backend, + output_file, + False, + **extra) + + if checkpoint_path: + pytorch_output_file = osp.join(save_dir, 'pytorch_out.jpg') + visualize_model( + model_cfg_path, + deploy_cfg_path, + [checkpoint_path], + single_img, + args.device, + Backend.PYTORCH, + pytorch_output_file, + False) + + # concat pytorch result and backend result + backend_result = mmcv.imread(output_file) + pytorch_result = mmcv.imread(pytorch_output_file) + result = np.concatenate((backend_result, pytorch_result), axis=1) + mmcv.imwrite(result, output_file) + + # remove temp pytorch result + os.remove(osp.join(save_dir, pytorch_output_file)) + +if __name__ == '__main__': + main() \ No newline at end of file From 0a8e585004429cf504e2b1737fee4868ea420c43 Mon Sep 17 00:00:00 2001 From: Bo Date: Wed, 19 Apr 2023 22:17:09 +0800 Subject: [PATCH 2/6] Fix visualize script. --- tools/visualize.py | 75 +++++++++++++++++----------------------------- 1 file changed, 27 insertions(+), 48 deletions(-) diff --git a/tools/visualize.py b/tools/visualize.py index 278049c4a0..89e3930485 100644 --- a/tools/visualize.py +++ b/tools/visualize.py @@ -4,30 +4,25 @@ import os import os.path as osp import time -from functools import partial -from tqdm import tqdm -import numpy as np import mmcv import mmengine +import numpy as np +from tqdm import tqdm from mmdeploy.apis import visualize_model -from mmdeploy.utils import (Backend, get_backend, get_root_logger, - load_config) +from mmdeploy.utils import Backend, get_backend, get_root_logger, load_config def parse_args(): - parser = argparse.ArgumentParser(description='Model inference visualization.') + parser = argparse.ArgumentParser( + description='Model inference visualization.') parser.add_argument('--deploy-cfg', help='deploy config path') parser.add_argument('--model-cfg', help='model config path') - parser.add_argument('--deploy-path', - type=str, - nargs='+', - help='deploy model path') parser.add_argument( - '--checkpoint', - default=None, - help='model checkpoint path') + '--deploy-path', type=str, nargs='+', help='deploy model path') + parser.add_argument( + '--checkpoint', default=None, help='model checkpoint path') parser.add_argument( '--test-img', default=None, @@ -35,11 +30,8 @@ def parse_args(): nargs='+', help='image used to test model') parser.add_argument( - '--save-dir', - default=None, - help='the dir to save inference results') - parser.add_argument( - '--device', help='device to run model', default='cpu') + '--save-dir', default=None, help='the dir to save inference results') + parser.add_argument('--device', help='device to run model', default='cpu') parser.add_argument( '--log-level', help='set log level', @@ -65,18 +57,17 @@ def main(): deploy_model_path = args.deploy_path if not isinstance(deploy_model_path, list): deploy_model_path = [deploy_model_path] - + # load deploy_cfg deploy_cfg = load_config(deploy_cfg_path)[0] - + # create save_dir or generate default save_dir save_dir = args.save_dir if save_dir: # generate default dir current_time = time.localtime() - save_dir = osp.join( - os.getcwd(), time.strftime("%Y_%m_%d_%H_%M_%S", current_time) - ) + save_dir = osp.join(os.getcwd(), + time.strftime('%Y_%m_%d_%H_%M_%S', current_time)) mmengine.mkdir_or_exist(save_dir) # get backend info @@ -84,42 +75,30 @@ def main(): extra = dict() if backend == Backend.SNPE: extra['uri'] = args.uri - + # iterate single_img for single_img in tqdm(args.test_img): filename = osp.basename(single_img) - output_file=osp.join(save_dir, filename) - visualize_model( - model_cfg_path, - deploy_cfg_path, - deploy_model_path, - single_img, - args.device, - backend, - output_file, - False, - **extra) - + output_file = osp.join(save_dir, filename) + visualize_model(model_cfg_path, deploy_cfg_path, deploy_model_path, + single_img, args.device, backend, output_file, False, + **extra) + if checkpoint_path: pytorch_output_file = osp.join(save_dir, 'pytorch_out.jpg') - visualize_model( - model_cfg_path, - deploy_cfg_path, - [checkpoint_path], - single_img, - args.device, - Backend.PYTORCH, - pytorch_output_file, - False) - + visualize_model(model_cfg_path, deploy_cfg_path, [checkpoint_path], + single_img, args.device, Backend.PYTORCH, + pytorch_output_file, False) + # concat pytorch result and backend result backend_result = mmcv.imread(output_file) pytorch_result = mmcv.imread(pytorch_output_file) result = np.concatenate((backend_result, pytorch_result), axis=1) mmcv.imwrite(result, output_file) - + # remove temp pytorch result os.remove(osp.join(save_dir, pytorch_output_file)) + if __name__ == '__main__': - main() \ No newline at end of file + main() From bece673403cc32f536548b28e83f025ab1ab6e91 Mon Sep 17 00:00:00 2001 From: Bo Date: Fri, 21 Apr 2023 11:04:40 +0800 Subject: [PATCH 3/6] Fix default save_dir. --- tools/visualize.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/tools/visualize.py b/tools/visualize.py index 89e3930485..567a36172b 100644 --- a/tools/visualize.py +++ b/tools/visualize.py @@ -30,7 +30,9 @@ def parse_args(): nargs='+', help='image used to test model') parser.add_argument( - '--save-dir', default=None, help='the dir to save inference results') + '--save-dir', + default=os.getcwd(), + help='the dir to save inference results') parser.add_argument('--device', help='device to run model', default='cpu') parser.add_argument( '--log-level', @@ -62,12 +64,9 @@ def main(): deploy_cfg = load_config(deploy_cfg_path)[0] # create save_dir or generate default save_dir - save_dir = args.save_dir - if save_dir: - # generate default dir - current_time = time.localtime() - save_dir = osp.join(os.getcwd(), - time.strftime('%Y_%m_%d_%H_%M_%S', current_time)) + current_time = time.localtime() + save_dir = osp.join(os.getcwd(), + time.strftime('%Y_%m_%d_%H_%M_%S', current_time)) mmengine.mkdir_or_exist(save_dir) # get backend info From 9514b47c13e59ef96473ba891fa38c0378f144a3 Mon Sep 17 00:00:00 2001 From: Bo Date: Fri, 21 Apr 2023 11:37:09 +0800 Subject: [PATCH 4/6] Update visualize tool's doc. --- docs/en/02-how-to-run/useful_tools.md | 27 ++++++++++++++++++++++++ docs/zh_cn/02-how-to-run/useful_tools.md | 27 ++++++++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/docs/en/02-how-to-run/useful_tools.md b/docs/en/02-how-to-run/useful_tools.md index a402f2113d..b603df5f44 100644 --- a/docs/en/02-how-to-run/useful_tools.md +++ b/docs/en/02-how-to-run/useful_tools.md @@ -202,3 +202,30 @@ And the output look like this: | Max | 1.689 | 591.983 | +--------+------------+---------+ ``` + +## visualize + +This tool can be used to visualize model inference results in different backend. + +### Usage + +```bash +python tools/visualize.py \ + --deploy-cfg {DEPLOY_CFG} \ + --model-cfg {MODEL_CFG} \ + --deploy-path {DEPLOY_PATH} \ + --test-img {TEST_IMGS} \ + --checkpoint {CHECKPOINTS} \ + --save-dir {SAVE_DIR} \ + --device {DEVICE} +``` + +### Description of all arguments + +- `deploy-cfg` : The path of the deploy config file in MMDeploy codebase. +- `model-cfg` : The path of model config file in OpenMMLab codebase. +- `deploy-path` : The path of the model to be tested, if the backend contains multiple files, you can use it multiple times. +- `test-img` : The path of the images to be tested, you can use it multiple times. +- `checkpoint` : The path of the checkpoint to be tested, if it is used, the result will be cancated to right part. +- `save-dir` : The path to save the visualization results, if it not specified, it will be set to '.'. +- `device` : The device type. If not specified, it will be set to `cpu`. diff --git a/docs/zh_cn/02-how-to-run/useful_tools.md b/docs/zh_cn/02-how-to-run/useful_tools.md index 2f08cefa9f..5069196c06 100644 --- a/docs/zh_cn/02-how-to-run/useful_tools.md +++ b/docs/zh_cn/02-how-to-run/useful_tools.md @@ -202,3 +202,30 @@ python tools/profiler.py \ | Max | 1.689 | 591.983 | +--------+------------+---------+ ``` + +## visualize + +这个工具可以用来可视化不同推理引擎的推理结果。 + +### 用法 + +```bash +python tools/visualize.py \ + --deploy-cfg {DEPLOY_CFG} \ + --model-cfg {MODEL_CFG} \ + --deploy-path {DEPLOY_PATH} \ + --test-img {TEST_IMGS} \ + --checkpoint {CHECKPOINTS} \ + --save-dir {SAVE_DIR} \ + --device {DEVICE} +``` + +### 参数说明 + +- `deploy-cfg` : 输入的配置文件路径 +- `model-cfg` : 输入的模型配置文件路径 +- `deploy-path` : 测试的模型文件路径，如果部分模型含有多个文件，请多次使用该参数 +- `test-img` : 测试的图片路径，可以多次使用测试测试多张图片 +- `checkpoint` : PyTorch的权重文件，如果使用这个参数，推理的结果会被拼接到图像的右侧 +- `save-dir` : 保存可视化结果，如果没有指定则会被设为当前目录 +- `device` : 运行的设备类型，如果没有指定则会默认设置为`cpu` From f8e7fbea7b326cefa105f2ff0bccf987b94f5805 Mon Sep 17 00:00:00 2001 From: Bo Date: Fri, 26 May 2023 18:20:55 +0800 Subject: [PATCH 5/6] Update batch visualize. --- docs/en/02-how-to-run/useful_tools.md | 2 + docs/zh_cn/02-how-to-run/useful_tools.md | 2 + tools/visualize.py | 121 +++++++++++++++++++---- 3 files changed, 105 insertions(+), 20 deletions(-) diff --git a/docs/en/02-how-to-run/useful_tools.md b/docs/en/02-how-to-run/useful_tools.md index 7a2381fead..bcf12e2fdc 100644 --- a/docs/en/02-how-to-run/useful_tools.md +++ b/docs/en/02-how-to-run/useful_tools.md @@ -275,6 +275,7 @@ python tools/visualize.py \ --deploy-path {DEPLOY_PATH} \ --test-img {TEST_IMGS} \ --checkpoint {CHECKPOINTS} \ + --batch {BATCH} --save-dir {SAVE_DIR} \ --device {DEVICE} ``` @@ -286,5 +287,6 @@ python tools/visualize.py \ - `deploy-path` : The path of the model to be tested, if the backend contains multiple files, you can use it multiple times. - `test-img` : The path of the images to be tested, you can use it multiple times. - `checkpoint` : The path of the checkpoint to be tested, if it is used, the result will be cancated to right part. +- `batch` : The batch size of the visual result, if not specified, it will be set to 1. - `save-dir` : The path to save the visualization results, if it not specified, it will be set to '.'. - `device` : The device type. If not specified, it will be set to `cpu`. diff --git a/docs/zh_cn/02-how-to-run/useful_tools.md b/docs/zh_cn/02-how-to-run/useful_tools.md index f6f439737e..d35844726e 100644 --- a/docs/zh_cn/02-how-to-run/useful_tools.md +++ b/docs/zh_cn/02-how-to-run/useful_tools.md @@ -258,6 +258,7 @@ python tools/visualize.py \ --deploy-path {DEPLOY_PATH} \ --test-img {TEST_IMGS} \ --checkpoint {CHECKPOINTS} \ + --batch {BATCH} \ --save-dir {SAVE_DIR} \ --device {DEVICE} ``` @@ -269,5 +270,6 @@ python tools/visualize.py \ - `deploy-path` : 测试的模型文件路径，如果部分模型含有多个文件，请多次使用该参数 - `test-img` : 测试的图片路径，可以多次使用测试测试多张图片 - `checkpoint` : PyTorch的权重文件，如果使用这个参数，推理的结果会被拼接到图像的右侧 +- `batch` : 可视化的batch大小，默认为1 - `save-dir` : 保存可视化结果，如果没有指定则会被设为当前目录 - `device` : 运行的设备类型，如果没有指定则会默认设置为`cpu` diff --git a/tools/visualize.py b/tools/visualize.py index 567a36172b..9fed221f36 100644 --- a/tools/visualize.py +++ b/tools/visualize.py @@ -4,14 +4,16 @@ import os import os.path as osp import time +from typing import Optional, Sequence, Union import mmcv import mmengine import numpy as np +import torch from tqdm import tqdm -from mmdeploy.apis import visualize_model -from mmdeploy.utils import Backend, get_backend, get_root_logger, load_config +from mmdeploy.utils import (Backend, get_backend, get_input_shape, + get_root_logger, load_config) def parse_args(): @@ -23,6 +25,11 @@ def parse_args(): '--deploy-path', type=str, nargs='+', help='deploy model path') parser.add_argument( '--checkpoint', default=None, help='model checkpoint path') + parser.add_argument( + '--batch', + type=int, + choices=[1, 2], + help='batch size for inference, accepts only 1 or 2') parser.add_argument( '--test-img', default=None, @@ -47,6 +54,89 @@ def parse_args(): return args +def batch_visualize_model( + model_cfg: Union[str, mmengine.Config], + deploy_cfg: Union[str, mmengine.Config], + imgs: Union[str, np.ndarray, Sequence[str]], + device: str, + backend_model_path: Union[str, Sequence[str]], + checkpoint_path: Optional[Union[str, Sequence[str]]] = None, + backend: Optional[Backend] = None, + output_file: Optional[str] = None): + """Run inference with PyTorch or backend model and show results. + + Args: + model_cfg (str | mmengine.Config): Model config file or Config object. + deploy_cfg (str | mmengine.Config): Deployment config file or Config + object. + img (str | np.ndarray | Sequence[str]): Input image file or numpy array + for inference. + device (str): A string specifying device type. + backend_model_path (str | Sequence[str]): Input backend model or + file(s). + checkpoint_path (str | Sequence[str]): Input pytorch checkpoint + model or file(s), defaults to `None`. + 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`. + """ + 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(backend_model_path, str): + backend_model_path = [backend_model_path] + + assert len( + backend_model_path) > 0, 'Model should have at least one element.' + + # build model + if checkpoint_path is not None: + pytorch_model = task_processor.build_pytorch_model(checkpoint_path) + backend_model = task_processor.build_backend_model( + backend_model_path, + data_preprocessor_updater=task_processor.update_data_preprocessor) + + if isinstance(imgs, str) or not isinstance(imgs, Sequence): + imgs = [imgs] + + # batch inference + for batch_img in imgs: + model_inputs, _ = task_processor.create_input(batch_img, input_shape) + with torch.no_grad(): + if checkpoint_path is not None: + pytorch_result = pytorch_model.test_step(model_inputs)[0] + backend_result = backend_model.test_step(model_inputs)[0] + + task_processor.visualize( + image=batch_img, + model=backend_model, + result=backend_result, + output_file=output_file, + window_name=backend.value, + show_result=False) + backend_result_img = mmcv.imread(output_file) + + if checkpoint_path is not None: + task_processor.visualize( + image=batch_img, + model=pytorch_model, + result=pytorch_result, + output_file=output_file, + window_name=backend.value, + show_result=False) + pytorch_result_img = mmcv.imread(output_file) + + result = np.concatenate( + (backend_result_img, pytorch_result_img), axis=1) + mmcv.imwrite(result, output_file) + + def main(): args = parse_args() logger = get_root_logger() @@ -79,24 +169,15 @@ def main(): for single_img in tqdm(args.test_img): filename = osp.basename(single_img) output_file = osp.join(save_dir, filename) - visualize_model(model_cfg_path, deploy_cfg_path, deploy_model_path, - single_img, args.device, backend, output_file, False, - **extra) - - if checkpoint_path: - pytorch_output_file = osp.join(save_dir, 'pytorch_out.jpg') - visualize_model(model_cfg_path, deploy_cfg_path, [checkpoint_path], - single_img, args.device, Backend.PYTORCH, - pytorch_output_file, False) - - # concat pytorch result and backend result - backend_result = mmcv.imread(output_file) - pytorch_result = mmcv.imread(pytorch_output_file) - result = np.concatenate((backend_result, pytorch_result), axis=1) - mmcv.imwrite(result, output_file) - - # remove temp pytorch result - os.remove(osp.join(save_dir, pytorch_output_file)) + + if args.batch < 2: + batch_visualize_model(model_cfg_path, deploy_cfg_path, single_img, + args.device, deploy_model_path, None, + backend, output_file) + else: + batch_visualize_model(model_cfg_path, deploy_cfg_path, single_img, + args.device, deploy_model_path, + checkpoint_path, backend, output_file) if __name__ == '__main__': From 0e1591787439c952ea35a8b567a0e1b959d54df8 Mon Sep 17 00:00:00 2001 From: RunningLeon Date: Wed, 7 Jun 2023 11:31:27 +0800 Subject: [PATCH 6/6] reconstruct --- tools/visualize.py | 219 +++++++++++++++++---------------------------- 1 file changed, 80 insertions(+), 139 deletions(-) diff --git a/tools/visualize.py b/tools/visualize.py index 9fed221f36..eb73e87d9d 100644 --- a/tools/visualize.py +++ b/tools/visualize.py @@ -1,183 +1,124 @@ # Copyright (c) OpenMMLab. All rights reserved. import argparse import logging -import os import os.path as osp -import time -from typing import Optional, Sequence, Union import mmcv import mmengine import numpy as np -import torch -from tqdm import tqdm -from mmdeploy.utils import (Backend, get_backend, get_input_shape, - get_root_logger, load_config) +from mmdeploy.apis.utils import build_task_processor +from mmdeploy.utils import get_input_shape, get_root_logger, load_config def parse_args(): parser = argparse.ArgumentParser( description='Model inference visualization.') - parser.add_argument('--deploy-cfg', help='deploy config path') - parser.add_argument('--model-cfg', help='model config path') + parser.add_argument('deploy_cfg', help='deploy config path') + parser.add_argument('model_cfg', help='model config path') parser.add_argument( - '--deploy-path', type=str, nargs='+', help='deploy model path') + '--model', + type=str, + nargs='+', + required=True, + help='deploy model path') parser.add_argument( '--checkpoint', default=None, help='model checkpoint path') parser.add_argument( - '--batch', - type=int, - choices=[1, 2], - help='batch size for inference, accepts only 1 or 2') + '--device', help='device type for inference', default='cpu') parser.add_argument( '--test-img', - default=None, type=str, nargs='+', + required=True, help='image used to test model') + parser.add_argument( + '--batch', + type=int, + choices=[1, 2], + help='batch size for inference, accepts only 1 or 2') parser.add_argument( '--save-dir', - default=os.getcwd(), + default='workdir', help='the dir to save inference results') - parser.add_argument('--device', help='device to run model', default='cpu') + parser.add_argument( + '--show', action='store_true', help='Show detection outputs') parser.add_argument( '--log-level', help='set log level', default='INFO', choices=list(logging._nameToLevel.keys())) - parser.add_argument( - '--uri', - default='192.168.1.1:60000', - help='Remote ipv4:port or ipv6:port for inference on edge device.') args = parser.parse_args() return args -def batch_visualize_model( - model_cfg: Union[str, mmengine.Config], - deploy_cfg: Union[str, mmengine.Config], - imgs: Union[str, np.ndarray, Sequence[str]], - device: str, - backend_model_path: Union[str, Sequence[str]], - checkpoint_path: Optional[Union[str, Sequence[str]]] = None, - backend: Optional[Backend] = None, - output_file: Optional[str] = None): - """Run inference with PyTorch or backend model and show results. - - Args: - model_cfg (str | mmengine.Config): Model config file or Config object. - deploy_cfg (str | mmengine.Config): Deployment config file or Config - object. - img (str | np.ndarray | Sequence[str]): Input image file or numpy array - for inference. - device (str): A string specifying device type. - backend_model_path (str | Sequence[str]): Input backend model or - file(s). - checkpoint_path (str | Sequence[str]): Input pytorch checkpoint - model or file(s), defaults to `None`. - 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`. - """ - 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(backend_model_path, str): - backend_model_path = [backend_model_path] - - assert len( - backend_model_path) > 0, 'Model should have at least one element.' - - # build model - if checkpoint_path is not None: - pytorch_model = task_processor.build_pytorch_model(checkpoint_path) - backend_model = task_processor.build_backend_model( - backend_model_path, - data_preprocessor_updater=task_processor.update_data_preprocessor) - - if isinstance(imgs, str) or not isinstance(imgs, Sequence): - imgs = [imgs] - - # batch inference - for batch_img in imgs: - model_inputs, _ = task_processor.create_input(batch_img, input_shape) - with torch.no_grad(): - if checkpoint_path is not None: - pytorch_result = pytorch_model.test_step(model_inputs)[0] - backend_result = backend_model.test_step(model_inputs)[0] - - task_processor.visualize( - image=batch_img, - model=backend_model, - result=backend_result, - output_file=output_file, - window_name=backend.value, - show_result=False) - backend_result_img = mmcv.imread(output_file) - - if checkpoint_path is not None: - task_processor.visualize( - image=batch_img, - model=pytorch_model, - result=pytorch_result, - output_file=output_file, - window_name=backend.value, - show_result=False) - pytorch_result_img = mmcv.imread(output_file) - - result = np.concatenate( - (backend_result_img, pytorch_result_img), axis=1) - mmcv.imwrite(result, output_file) - - def main(): args = parse_args() logger = get_root_logger() log_level = logging.getLevelName(args.log_level) logger.setLevel(log_level) - deploy_cfg_path = args.deploy_cfg - model_cfg_path = args.model_cfg - checkpoint_path = args.checkpoint - deploy_model_path = args.deploy_path - if not isinstance(deploy_model_path, list): - deploy_model_path = [deploy_model_path] - - # load deploy_cfg - deploy_cfg = load_config(deploy_cfg_path)[0] - - # create save_dir or generate default save_dir - current_time = time.localtime() - save_dir = osp.join(os.getcwd(), - time.strftime('%Y_%m_%d_%H_%M_%S', current_time)) - mmengine.mkdir_or_exist(save_dir) - - # get backend info - backend = get_backend(deploy_cfg) - extra = dict() - if backend == Backend.SNPE: - extra['uri'] = args.uri - - # iterate single_img - for single_img in tqdm(args.test_img): - filename = osp.basename(single_img) - output_file = osp.join(save_dir, filename) - - if args.batch < 2: - batch_visualize_model(model_cfg_path, deploy_cfg_path, single_img, - args.device, deploy_model_path, None, - backend, output_file) - else: - batch_visualize_model(model_cfg_path, deploy_cfg_path, single_img, - args.device, deploy_model_path, - checkpoint_path, backend, output_file) + # load cfgs + deploy_cfg, model_cfg = load_config(args.deploy_cfg, args.model_cfg) + task_processor = build_task_processor(model_cfg, deploy_cfg, args.device) + input_shape = get_input_shape(deploy_cfg) + backend_model = task_processor.build_backend_model( + args.model, + data_preprocessor_updater=task_processor.update_data_preprocessor) + torch_model = None + if args.checkpoint is not None: + torch_model = task_processor.build_pytorch_model(args.checkpoint) + + mmengine.mkdir_or_exist(args.save_dir) + # get visualizer + visualizer = task_processor.get_visualizer('mmdeploy', args.save_dir) + + for i in range(0, len(args.test_img), args.batch): + imgs = args.test_img[i:(i + args.batch)] + model_inputs, _ = task_processor.create_input( + imgs, + input_shape, + data_preprocessor=getattr(backend_model, 'data_preprocessor', + None)) + backend_results = backend_model.test_step(model_inputs) + torch_results = [None] * len(imgs) + if torch_model is not None: + torch_results = torch_model.test_step(model_inputs) + + # get visualized results + for img_path, torch_res, backend_res in zip(imgs, torch_results, + backend_results): + _, filename = osp.split(img_path) + output_file = osp.join(args.save_dir, filename) + image = mmcv.imread(img_path, channel_order='rgb') + visualizer.add_datasample( + filename, + image, + data_sample=backend_res, + draw_gt=False, + show=False, + out_file=None) + drawn_img = visualizer.get_image() + if torch_res: + visualizer.add_datasample( + filename, + image, + data_sample=torch_res, + draw_gt=False, + show=False, + out_file=None) + drawn_img_torch = visualizer.get_image() + shape = drawn_img.shape + dummy_img = np.full((shape[0], 20, shape[2]), + 255, + dtype=np.uint8) + drawn_img = np.concatenate( + (drawn_img, dummy_img, drawn_img_torch), axis=1) + if args.show: + visualizer.show(drawn_img, win_name=filename, wait_time=0) + drawn_img = mmcv.image.rgb2bgr(drawn_img) + mmcv.imwrite(drawn_img, output_file) + logger.info(f'Saved to {output_file}') if __name__ == '__main__':