From f8089601139cb65fa01ac5b13ff0af5b4795ac82 Mon Sep 17 00:00:00 2001 From: MeowZheng Date: Mon, 19 Sep 2022 21:27:06 +0800 Subject: [PATCH 1/3] [Enhancement]Add `out_file` in add_datasample to for save vis image directly --- mmseg/apis/inference.py | 3 +-- mmseg/visualization/local_visualizer.py | 29 +++++++++++++++++-------- 2 files changed, 21 insertions(+), 11 deletions(-) diff --git a/mmseg/apis/inference.py b/mmseg/apis/inference.py index 1544c52516..8d718c4724 100644 --- a/mmseg/apis/inference.py +++ b/mmseg/apis/inference.py @@ -199,9 +199,8 @@ def show_result_pyplot(model: BaseSegmentor, draw_gt=draw_gt, draw_pred=draw_pred, wait_time=wait_time, + out_file=out_file, show=show) vis_img = visualizer.get_image() - if out_file is not None: - mmcv.imwrite(vis_img, out_file) return vis_img diff --git a/mmseg/visualization/local_visualizer.py b/mmseg/visualization/local_visualizer.py index 72b868ba3b..232d6c54f6 100644 --- a/mmseg/visualization/local_visualizer.py +++ b/mmseg/visualization/local_visualizer.py @@ -1,6 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. from typing import Dict, List, Optional, Tuple +import mmcv import numpy as np from mmengine.dist import master_only from mmengine.structures import PixelData @@ -99,15 +100,18 @@ def _draw_sem_seg(self, image: np.ndarray, sem_seg: PixelData, return self.get_image() @master_only - def add_datasample(self, - name: str, - image: np.ndarray, - data_sample: Optional[SegDataSample] = None, - draw_gt: bool = True, - draw_pred: bool = True, - show: bool = False, - wait_time: float = 0, - step: int = 0) -> None: + def add_datasample( + self, + name: str, + image: np.ndarray, + data_sample: Optional[SegDataSample] = None, + draw_gt: bool = True, + draw_pred: bool = True, + show: bool = False, + wait_time: float = 0, + # TODO: Supported in mmengine's Viusalizer. + out_file: Optional[str] = None, + step: int = 0) -> None: """Draw datasample and save to all backends. - If GT and prediction are plotted at the same time, they are @@ -115,6 +119,9 @@ def add_datasample(self, ground truth and the right image is the prediction. - If ``show`` is True, all storage backends are ignored, and the images will be displayed in a local window. + - If ``out_file`` is specified, the drawn image will be + saved to ``out_file``. t is usually used when the display + is not available. Args: name (str): The image identifier. @@ -128,6 +135,7 @@ def add_datasample(self, Defaults to True. show (bool): Whether to display the drawn image. Default to False. wait_time (float): The interval of show (s). Defaults to 0. + out_file (str): Path to output file. Defaults to None. step (int): Global step value to record. Defaults to 0. """ classes = self.dataset_meta.get('classes', None) @@ -166,5 +174,8 @@ def add_datasample(self, if show: self.show(drawn_img, win_name=name, wait_time=wait_time) + + if out_file is not None: + mmcv.imwrite(drawn_img, out_file) else: self.add_image(name, drawn_img, step) From 20714e45bb342a839f8da44f63eeac16cd021ab2 Mon Sep 17 00:00:00 2001 From: MeowZheng Date: Tue, 20 Sep 2022 13:38:18 +0800 Subject: [PATCH 2/3] comments --- mmseg/visualization/local_visualizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mmseg/visualization/local_visualizer.py b/mmseg/visualization/local_visualizer.py index 232d6c54f6..070b06b73b 100644 --- a/mmseg/visualization/local_visualizer.py +++ b/mmseg/visualization/local_visualizer.py @@ -120,7 +120,7 @@ def add_datasample( - If ``show`` is True, all storage backends are ignored, and the images will be displayed in a local window. - If ``out_file`` is specified, the drawn image will be - saved to ``out_file``. t is usually used when the display + saved to ``out_file``. it is usually used when the display is not available. Args: From 9fab2a16b050e7617d4c253e937cb448396e8ec1 Mon Sep 17 00:00:00 2001 From: MeowZheng Date: Tue, 20 Sep 2022 14:38:20 +0800 Subject: [PATCH 3/3] ut --- .../test_local_visualizer.py | 22 ++++++++----------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/tests/test_visualization/test_local_visualizer.py b/tests/test_visualization/test_local_visualizer.py index 7754c30ed6..b60a9b8750 100644 --- a/tests/test_visualization/test_local_visualizer.py +++ b/tests/test_visualization/test_local_visualizer.py @@ -118,19 +118,14 @@ def test_cityscapes_add_datasample_forward(gt_sem_seg): [255, 0, 0], [0, 0, 142], [0, 0, 70], [0, 60, 100], [0, 80, 100], [0, 0, 230], [119, 11, 32]]) - seg_local_visualizer.add_datasample(out_file, image, - data_sample) - # test out_file - seg_local_visualizer.add_datasample(out_file, image, - data_sample) - assert os.path.exists( - osp.join(tmp_dir, 'vis_data', 'vis_image', - out_file + '_0.png')) - drawn_img = cv2.imread( - osp.join(tmp_dir, 'vis_data', 'vis_image', - out_file + '_0.png')) - assert drawn_img.shape == (h, w, 3) + seg_local_visualizer.add_datasample( + out_file, + image, + data_sample, + out_file=osp.join(tmp_dir, 'test.png')) + self._assert_image_and_shape( + osp.join(tmp_dir, 'test.png'), (h, w, 3)) # test gt_instances and pred_instances pred_sem_seg_data = dict( @@ -139,12 +134,13 @@ def test_cityscapes_add_datasample_forward(gt_sem_seg): data_sample.pred_sem_seg = pred_sem_seg + # test draw prediction with gt seg_local_visualizer.add_datasample(out_file, image, data_sample) self._assert_image_and_shape( osp.join(tmp_dir, 'vis_data', 'vis_image', out_file + '_0.png'), (h, w * 2, 3)) - + # test draw prediction without gt seg_local_visualizer.add_datasample( out_file, image, data_sample, draw_gt=False) self._assert_image_and_shape(