diff --git a/demo/image_demo.py b/demo/image_demo.py index 183f23871b..2698e422eb 100644 --- a/demo/image_demo.py +++ b/demo/image_demo.py @@ -15,6 +15,11 @@ def main(): '--palette', default='cityscapes', help='Color palette used for segmentation map') + 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 @@ -22,7 +27,12 @@ def main(): # test a single image result = inference_segmentor(model, args.img) # show the results - show_result_pyplot(model, args.img, result, get_palette(args.palette)) + show_result_pyplot( + model, + args.img, + result, + get_palette(args.palette), + opacity=args.opacity) if __name__ == '__main__': diff --git a/docs/get_started.md b/docs/get_started.md index 3182c53451..5ef2442ac5 100644 --- a/docs/get_started.md +++ b/docs/get_started.md @@ -166,7 +166,8 @@ result = inference_segmentor(model, img) # visualize the results in a new window model.show_result(img, result, show=True) # or save the visualization results to image files -model.show_result(img, result, out_file='result.jpg') +# you can change the opacity of the painted segmentation map in (0, 1]. +model.show_result(img, result, out_file='result.jpg', opacity=0.5) # test a video and show the results video = mmcv.VideoReader('video.mp4') diff --git a/mmseg/apis/inference.py b/mmseg/apis/inference.py index 20c20dccda..9052cdd32a 100644 --- a/mmseg/apis/inference.py +++ b/mmseg/apis/inference.py @@ -98,7 +98,12 @@ def inference_segmentor(model, img): return result -def show_result_pyplot(model, img, result, palette=None, fig_size=(15, 10)): +def show_result_pyplot(model, + img, + result, + palette=None, + fig_size=(15, 10), + opacity=0.5): """Visualize the segmentation results on the image. Args: @@ -109,10 +114,14 @@ def show_result_pyplot(model, img, result, palette=None, fig_size=(15, 10)): map. If None is given, random palette will be generated. Default: None fig_size (tuple): Figure size of the pyplot figure. + opacity(float): Opacity of painted segmentation map. + Default 0.5. + Must be in (0, 1] range. """ if hasattr(model, 'module'): model = model.module - img = model.show_result(img, result, palette=palette, show=False) + img = model.show_result( + img, result, palette=palette, show=False, opacity=opacity) plt.figure(figsize=fig_size) plt.imshow(mmcv.bgr2rgb(img)) plt.show() diff --git a/mmseg/apis/test.py b/mmseg/apis/test.py index 148df7680e..cbf80e7a8e 100644 --- a/mmseg/apis/test.py +++ b/mmseg/apis/test.py @@ -35,7 +35,8 @@ def single_gpu_test(model, data_loader, show=False, out_dir=None, - efficient_test=False): + efficient_test=False, + opacity=0.5): """Test with single GPU. Args: @@ -46,7 +47,9 @@ def single_gpu_test(model, the directory to save output results. efficient_test (bool): Whether save the results as local numpy files to save CPU memory during evaluation. Default: False. - + opacity(float): Opacity of painted segmentation map. + Default 0.5. + Must be in (0, 1] range. Returns: list: The prediction results. """ @@ -82,7 +85,8 @@ def single_gpu_test(model, result, palette=dataset.PALETTE, show=show, - out_file=out_file) + out_file=out_file, + opacity=opacity) if isinstance(result, list): if efficient_test: diff --git a/mmseg/models/segmentors/base.py b/mmseg/models/segmentors/base.py index 6f59dbc72e..58c31887f3 100644 --- a/mmseg/models/segmentors/base.py +++ b/mmseg/models/segmentors/base.py @@ -212,7 +212,8 @@ def show_result(self, win_name='', show=False, wait_time=0, - out_file=None): + out_file=None, + opacity=0.5): """Draw `result` over `img`. Args: @@ -229,7 +230,9 @@ def show_result(self, Default: False. out_file (str or None): The filename to write the image. Default: None. - + opacity(float): Opacity of painted segmentation map. + Default 0.5. + Must be in (0, 1] range. Returns: img (Tensor): Only if not `show` or `out_file` """ @@ -246,13 +249,14 @@ def show_result(self, assert palette.shape[0] == len(self.CLASSES) assert palette.shape[1] == 3 assert len(palette.shape) == 2 + assert 0 < opacity <= 1.0 color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) for label, color in enumerate(palette): color_seg[seg == label, :] = color # convert to BGR color_seg = color_seg[..., ::-1] - img = img * 0.5 + color_seg * 0.5 + img = img * (1 - opacity) + color_seg * opacity img = img.astype(np.uint8) # if out_file specified, do not show image in window if out_file is not None: diff --git a/tools/test.py b/tools/test.py index a106f04626..c074fcc4bb 100644 --- a/tools/test.py +++ b/tools/test.py @@ -55,6 +55,11 @@ def parse_args(): choices=['none', 'pytorch', 'slurm', 'mpi'], default='none', help='job launcher') + parser.add_argument( + '--opacity', + type=float, + default=0.5, + help='Opacity of painted segmentation map. In (0, 1] range.') parser.add_argument('--local_rank', type=int, default=0) args = parser.parse_args() if 'LOCAL_RANK' not in os.environ: @@ -123,7 +128,7 @@ def main(): if not distributed: model = MMDataParallel(model, device_ids=[0]) outputs = single_gpu_test(model, data_loader, args.show, args.show_dir, - efficient_test) + efficient_test, args.opacity) else: model = MMDistributedDataParallel( model.cuda(),