mmsegmentation/demo/image_demo_with_inferencer.py
谢昕辰 53fe1ccf39
[Feature] Support MMSegInferencer (#2413)
## Motivation

Support `MMSegInferencer` for providing an easy and clean interface for
single or multiple images inferencing.

Ref: https://github.com/open-mmlab/mmengine/pull/773
https://github.com/open-mmlab/mmocr/pull/1608

## Modification

- mmseg/apis/mmseg_inferencer.py
- mmseg/visualization/local_visualizer.py
- demo/image_demo_with_inferencer.py

## Use cases (Optional)

Based on https://github.com/open-mmlab/mmengine/tree/inference

Add a new image inference demo with `MMSegInferencer`

- demo/image_demo_with_inferencer.py

```shell
python demo/image_demo_with_inferencer.py demo/demo.png fcn_r50-d8_4xb2-40k_cityscapes-512x1024
```

---------

Co-authored-by: MeowZheng <meowzheng@outlook.com>
2023-02-23 21:16:19 +08:00

55 lines
1.6 KiB
Python

# Copyright (c) OpenMMLab. All rights reserved.
from argparse import ArgumentParser
from mmseg.apis import MMSegInferencer
def main():
parser = ArgumentParser()
parser.add_argument('img', help='Image file')
parser.add_argument('model', help='Config file')
parser.add_argument('--checkpoint', default=None, help='Checkpoint file')
parser.add_argument(
'--out-dir', default='', help='Path to save result file')
parser.add_argument(
'--show',
action='store_true',
default=False,
help='Whether to display the drawn image.')
parser.add_argument(
'--save-mask',
action='store_true',
default=False,
help='Enable save the mask file')
parser.add_argument(
'--dataset-name',
default='cityscapes',
help='Color palette used for segmentation map')
parser.add_argument(
'--device', default='cuda:0', help='Device used for inference')
parser.add_argument(
'--opacity',
type=float,
default=0.5,
help='Opacity of painted segmentation map. In (0, 1] range.')
args = parser.parse_args()
# build the model from a config file and a checkpoint file
mmseg_inferencer = MMSegInferencer(
args.model,
args.checkpoint,
dataset_name=args.dataset_name,
device=args.device)
# test a single image
mmseg_inferencer(
args.img,
show=args.show,
out_dir=args.out_dir,
save_mask=args.save_mask,
opacity=args.opacity)
if __name__ == '__main__':
main()