diff --git a/mmseg/apis/inference.py b/mmseg/apis/inference.py index 1544c5251..8d718c472 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 72b868ba3..070b06b73 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 @@ class SegLocalVisualizer(Visualizer): 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 @@ class SegLocalVisualizer(Visualizer): 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``. it is usually used when the display + is not available. Args: name (str): The image identifier. @@ -128,6 +135,7 @@ class SegLocalVisualizer(Visualizer): 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 @@ class SegLocalVisualizer(Visualizer): 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) diff --git a/tests/test_visualization/test_local_visualizer.py b/tests/test_visualization/test_local_visualizer.py index 7754c30ed..b60a9b875 100644 --- a/tests/test_visualization/test_local_visualizer.py +++ b/tests/test_visualization/test_local_visualizer.py @@ -118,19 +118,14 @@ class TestSegLocalVisualizer(TestCase): [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 @@ class TestSegLocalVisualizer(TestCase): 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(