diff --git a/docs/en/02-how-to-run/useful_tools.md b/docs/en/02-how-to-run/useful_tools.md index c18a0e3048..bcf12e2fdc 100644 --- a/docs/en/02-how-to-run/useful_tools.md +++ b/docs/en/02-how-to-run/useful_tools.md @@ -203,6 +203,23 @@ And the output look like this: +--------+------------+---------+ ``` +## 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} +``` + ## generate_md_table This tool can be used to generate supported-backends markdown table. @@ -244,3 +261,32 @@ And the output look like this: | [SAR](https://github.com/open-mmlab/mmocr/tree/main/configs/textrecog/sar) | TextRecognition | Y | N | Y | N | N | N | | [SATRN](https://github.com/open-mmlab/mmocr/tree/main/configs/textrecog/satrn) | TextRecognition | Y | Y | Y | N | N | N | | [ABINet](https://github.com/open-mmlab/mmocr/tree/main/configs/textrecog/abinet) | TextRecognition | Y | Y | Y | N | N | N | + +## 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} \ + --batch {BATCH} + --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. +- `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 ae746f4fe0..d35844726e 100644 --- a/docs/zh_cn/02-how-to-run/useful_tools.md +++ b/docs/zh_cn/02-how-to-run/useful_tools.md @@ -244,3 +244,32 @@ python tools/generate_md_table.py tests/regression/mmocr.yml tests/regression/mm | [SAR](https://github.com/open-mmlab/mmocr/tree/main/configs/textrecog/sar) | TextRecognition | Y | N | Y | N | N | N | | [SATRN](https://github.com/open-mmlab/mmocr/tree/main/configs/textrecog/satrn) | TextRecognition | Y | Y | Y | N | N | N | | [ABINet](https://github.com/open-mmlab/mmocr/tree/main/configs/textrecog/abinet) | TextRecognition | Y | Y | Y | N | N | N | + +## visualize + +这个工具可以用来可视化不同推理引擎的推理结果。 + +### 用法 + +```bash +python tools/visualize.py \ + --deploy-cfg {DEPLOY_CFG} \ + --model-cfg {MODEL_CFG} \ + --deploy-path {DEPLOY_PATH} \ + --test-img {TEST_IMGS} \ + --checkpoint {CHECKPOINTS} \ + --batch {BATCH} \ + --save-dir {SAVE_DIR} \ + --device {DEVICE} +``` + +### 参数说明 + +- `deploy-cfg` : 输入的配置文件路径 +- `model-cfg` : 输入的模型配置文件路径 +- `deploy-path` : 测试的模型文件路径，如果部分模型含有多个文件，请多次使用该参数 +- `test-img` : 测试的图片路径，可以多次使用测试测试多张图片 +- `checkpoint` : PyTorch的权重文件，如果使用这个参数，推理的结果会被拼接到图像的右侧 +- `batch` : 可视化的batch大小，默认为1 +- `save-dir` : 保存可视化结果，如果没有指定则会被设为当前目录 +- `device` : 运行的设备类型，如果没有指定则会默认设置为`cpu` diff --git a/tools/visualize.py b/tools/visualize.py new file mode 100644 index 0000000000..eb73e87d9d --- /dev/null +++ b/tools/visualize.py @@ -0,0 +1,125 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import logging +import os.path as osp + +import mmcv +import mmengine +import numpy as np + +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( + '--model', + type=str, + nargs='+', + required=True, + help='deploy model path') + parser.add_argument( + '--checkpoint', default=None, help='model checkpoint path') + parser.add_argument( + '--device', help='device type for inference', default='cpu') + parser.add_argument( + '--test-img', + 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='workdir', + help='the dir to save inference results') + 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())) + 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) + + # 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__': + main()