Spaces:
Runtime error
Runtime error
| from typing import Optional, Sequence | |
| import mmcv | |
| import numpy as np | |
| from mmengine.dist import master_only | |
| from mmseg.structures import SegDataSample | |
| from mmseg.visualization import SegLocalVisualizer | |
| from opencd.registry import VISUALIZERS | |
| class CDLocalVisualizer(SegLocalVisualizer): | |
| """Change Detection Local Visualizer. """ | |
| def add_datasample( | |
| self, | |
| name: str, | |
| image: np.ndarray, | |
| image_from_to: Sequence[np.array], | |
| 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, | |
| with_labels: Optional[bool] = False) -> None: | |
| """Draw datasample and save to all backends. | |
| - If GT and prediction are plotted at the same time, they are | |
| displayed in a stitched image where the left image is the | |
| 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. | |
| image (np.ndarray): The image to draw. | |
| image_from_to (Sequence[np.array]): The image pairs to draw. | |
| gt_sample (:obj:`SegDataSample`, optional): GT SegDataSample. | |
| Defaults to None. | |
| pred_sample (:obj:`SegDataSample`, optional): Prediction | |
| SegDataSample. Defaults to None. | |
| draw_gt (bool): Whether to draw GT SegDataSample. Default to True. | |
| draw_pred (bool): Whether to draw Prediction SegDataSample. | |
| 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. | |
| with_labels(bool, optional): Add semantic labels in visualization | |
| result, Defaults to True. | |
| """ | |
| exist_img_from_to = True if len(image_from_to) > 0 else False | |
| if exist_img_from_to: | |
| assert len(image_from_to) == 2, '`image_from_to` contains `from` ' \ | |
| 'and `to` images' | |
| classes = self.dataset_meta.get('classes', None) | |
| palette = self.dataset_meta.get('palette', None) | |
| semantic_classes = self.dataset_meta.get('semantic_classes', None) | |
| semantic_palette = self.dataset_meta.get('semantic_palette', None) | |
| gt_img_data = None | |
| gt_img_data_from = None | |
| gt_img_data_to = None | |
| pred_img_data = None | |
| pred_img_data_from = None | |
| pred_img_data_to = None | |
| drawn_img_from = None | |
| drawn_img_to = None | |
| if draw_gt and data_sample is not None and 'gt_sem_seg' in data_sample: | |
| gt_img_data = image | |
| assert classes is not None, 'class information is ' \ | |
| 'not provided when ' \ | |
| 'visualizing change ' \ | |
| 'deteaction results.' | |
| gt_img_data = self._draw_sem_seg(gt_img_data, data_sample.gt_sem_seg, | |
| classes, palette, with_labels) | |
| if draw_gt and data_sample is not None and 'gt_sem_seg_from' in data_sample \ | |
| and 'gt_sem_seg_to' in data_sample: | |
| if exist_img_from_to: | |
| gt_img_data_from = image_from_to[0] | |
| gt_img_data_to = image_from_to[1] | |
| else: | |
| gt_img_data_from = np.zeros_like(image) | |
| gt_img_data_to = np.zeros_like(image) | |
| assert semantic_classes is not None, 'class information is ' \ | |
| 'not provided when ' \ | |
| 'visualizing change ' \ | |
| 'deteaction results.' | |
| gt_img_data_from = self._draw_sem_seg(gt_img_data_from, | |
| data_sample.gt_sem_seg_from, semantic_classes, | |
| semantic_palette, with_labels) | |
| gt_img_data_to = self._draw_sem_seg(gt_img_data_to, | |
| data_sample.gt_sem_seg_to, semantic_classes, | |
| semantic_palette, with_labels) | |
| if (draw_pred and data_sample is not None | |
| and 'pred_sem_seg' in data_sample): | |
| pred_img_data = image | |
| assert classes is not None, 'class information is ' \ | |
| 'not provided when ' \ | |
| 'visualizing semantic ' \ | |
| 'segmentation results.' | |
| pred_img_data = self._draw_sem_seg(pred_img_data, | |
| data_sample.pred_sem_seg, | |
| classes, palette, | |
| with_labels) | |
| if (draw_pred and data_sample is not None and 'pred_sem_seg_from' in data_sample \ | |
| and 'pred_sem_seg_to' in data_sample): | |
| if exist_img_from_to: | |
| pred_img_data_from = image_from_to[0] | |
| pred_img_data_to = image_from_to[1] | |
| else: | |
| pred_img_data_from = np.zeros_like(image) | |
| pred_img_data_to = np.zeros_like(image) | |
| assert semantic_classes is not None, 'class information is ' \ | |
| 'not provided when ' \ | |
| 'visualizing change ' \ | |
| 'deteaction results.' | |
| pred_img_data_from = self._draw_sem_seg(pred_img_data_from, | |
| data_sample.pred_sem_seg_from, semantic_classes, | |
| semantic_palette, with_labels) | |
| pred_img_data_to = self._draw_sem_seg(pred_img_data_to, | |
| data_sample.pred_sem_seg_to, semantic_classes, | |
| semantic_palette, with_labels) | |
| if gt_img_data is not None and pred_img_data is not None: | |
| drawn_img = np.concatenate((gt_img_data, pred_img_data), axis=1) | |
| elif gt_img_data is not None: | |
| drawn_img = gt_img_data | |
| else: | |
| drawn_img = pred_img_data | |
| if gt_img_data_from is not None and pred_img_data_from is not None: | |
| drawn_img_from = np.concatenate((gt_img_data_from, pred_img_data_from), axis=1) | |
| elif gt_img_data_from is not None: | |
| drawn_img_from = gt_img_data_from | |
| else: | |
| drawn_img_from = pred_img_data_from | |
| if gt_img_data_to is not None and pred_img_data_to is not None: | |
| drawn_img_to = np.concatenate((gt_img_data_to, pred_img_data_to), axis=1) | |
| elif gt_img_data_to is not None: | |
| drawn_img_to = gt_img_data_to | |
| else: | |
| drawn_img_to = pred_img_data_to | |
| if show: | |
| if drawn_img_from is not None and drawn_img_to is not None: | |
| drawn_img_cat = np.concatenate((drawn_img, drawn_img_from, drawn_img_to), axis=0) | |
| self.show(drawn_img_cat, win_name=name, wait_time=wait_time) | |
| else: | |
| self.show(drawn_img, win_name=name, wait_time=wait_time) | |
| if out_file is not None: | |
| if drawn_img_from is not None and drawn_img_to is not None: | |
| drawn_img_cat = np.concatenate((drawn_img, drawn_img_from, drawn_img_to), axis=0) | |
| mmcv.imwrite(mmcv.bgr2rgb(drawn_img_cat), out_file) | |
| else: | |
| mmcv.imwrite(mmcv.bgr2rgb(drawn_img), out_file) | |
| else: | |
| self.add_image(name, drawn_img, drawn_img_from, drawn_img_to, step) | |
| def add_image(self, name: str, | |
| image: np.ndarray, | |
| image_from: np.ndarray = None, | |
| image_to: np.ndarray = None, | |
| step: int = 0) -> None: | |
| """Record the image. | |
| Args: | |
| name (str): The image identifier. | |
| image (np.ndarray, optional): The image to be saved. The format | |
| should be RGB. Defaults to None. | |
| step (int): Global step value to record. Defaults to 0. | |
| """ | |
| for vis_backend in self._vis_backends.values(): | |
| vis_backend.add_image(name, image, image_from, image_to, step) # type: ignore | |
| def set_image(self, image: np.ndarray) -> None: | |
| """Set the image to draw. | |
| Args: | |
| image (np.ndarray): The image to draw. | |
| """ | |
| assert image is not None | |
| image = image.astype('uint8') | |
| self._image = image | |
| self.width, self.height = image.shape[1], image.shape[0] | |
| # print(image.shape) | |
| self._default_font_size = max( | |
| np.sqrt(self.height * self.width) // 90, 10) | |
| self.fig_save.set_size_inches( # type: ignore | |
| self.width / self.dpi, self.height / self.dpi) | |
| # self.canvas = mpl.backends.backend_cairo.FigureCanvasCairo(fig) | |
| self.ax_save.cla() | |
| self.ax_save.axis(False) | |
| self.ax_save.imshow( | |
| image, | |
| extent=(0, self.width, self.height, 0), | |
| interpolation='none') | |