mmsegmentation/demo/image_demo.py
谢昕辰 039ba5d4ca
[Feature] Support auto import modules from registry. (#2481)
## Motivation

The registry now supports auto-import modules from the given location.

register_all_modules before running is no longer needed. The modules
will be lazy-imported during building.

- [x] This PR can be merged after
https://github.com/open-mmlab/mmengine/pull/643. The MMEngine version
should be updated.

Ref: https://github.com/open-mmlab/mmdetection/pull/9143
2023-02-23 20:33:17 +08:00

46 lines
1.4 KiB
Python

# Copyright (c) OpenMMLab. All rights reserved.
from argparse import ArgumentParser
from mmengine.model import revert_sync_batchnorm
from mmseg.apis import inference_model, init_model, show_result_pyplot
def main():
parser = ArgumentParser()
parser.add_argument('img', help='Image file')
parser.add_argument('config', help='Config file')
parser.add_argument('checkpoint', help='Checkpoint file')
parser.add_argument('--out-file', default=None, help='Path to output file')
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.')
parser.add_argument(
'--title', default='result', help='The image identifier.')
args = parser.parse_args()
# build the model from a config file and a checkpoint file
model = init_model(args.config, args.checkpoint, device=args.device)
if args.device == 'cpu':
model = revert_sync_batchnorm(model)
# test a single image
result = inference_model(model, args.img)
# show the results
show_result_pyplot(
model,
args.img,
result,
title=args.title,
opacity=args.opacity,
draw_gt=False,
show=False if args.out_file is not None else True,
out_file=args.out_file)
if __name__ == '__main__':
main()