Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Enhancement] Add a visualize script for inference result comparation. #2012

Draft
wants to merge 7 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 46 additions & 0 deletions docs/en/02-how-to-run/useful_tools.md
Original file line number Diff line number Diff line change
Expand Up @@ -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} \
RunningLeon marked this conversation as resolved.
Show resolved Hide resolved
--checkpoint {CHECKPOINTS} \
--save-dir {SAVE_DIR} \
--device {DEVICE}
```

## generate_md_table

This tool can be used to generate supported-backends markdown table.
Expand Down Expand Up @@ -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`.
29 changes: 29 additions & 0 deletions docs/zh_cn/02-how-to-run/useful_tools.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`
125 changes: 125 additions & 0 deletions tools/visualize.py
Original file line number Diff line number Diff line change
@@ -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()