def visualize_image_retrieval(self,
                                  image: np.ndarray,
                                  data_sample: DataSample,
                                  prototype_dataset: BaseDataset,
                                  topk: int = 1,
                                  draw_score: bool = True,
                                  resize: Optional[int] = None,
                                  text_cfg: dict = dict(),
                                  show: bool = False,
                                  wait_time: float = 0,
                                  out_file: Optional[str] = None,
                                  name: Optional[str] = '',
                                  step: int = 0) -> None:
        """Visualize image retrieval result.

        This method will draw the input image and the images retrieved from the
        prototype dataset.

        Args:
            image (np.ndarray): The image to draw. The format should be RGB.
            data_sample (:obj:`DataSample`): The annotation of the image.
            prototype_dataset (:obj:`BaseDataset`): The prototype dataset.
                It should have `get_data_info` method and return a dict
                includes `img_path`.
            draw_score (bool): Whether to draw the match scores of the
                retrieved images. Defaults to True.
            resize (int, optional): Resize the long edge of the image to the
                specified length before visualization. Defaults to None.
            text_cfg (dict): Extra text setting, which accepts arguments of
                :func:`plt.text`. Defaults to an empty dict.
            show (bool): Whether to display the drawn image in a window, please
                confirm your are able to access the graphical interface.
                Defaults to False.
            wait_time (float): The display time (s). Defaults to 0, which means
                "forever".
            out_file (str, optional): Extra path to save the visualization
                result. If specified, the visualizer will only save the result
                image to the out_file and ignore its storage backends.
                Defaults to None.
            name (str): The image identifier. It's useful when using the
                storage backends of the visualizer to save or display the
                image. Defaults to an empty string.
            step (int): The global step value. It's useful to record a
                series of visualization results for the same image with the
                storage backends. Defaults to 0.

        Returns:
            np.ndarray: The visualization image.
        """
        text_cfg = {**self.DEFAULT_TEXT_CFG, **text_cfg}
        if resize is not None:
            image = mmcv.imrescale(image, (resize, resize))

        match_scores, indices = torch.topk(data_sample.pred_score, k=topk)

        figure = create_figure(margin=True)
        gs = figure.add_gridspec(2, topk)
        query_plot = figure.add_subplot(gs[0, :])
        query_plot.axis(False)
        query_plot.imshow(image)

        for k, (score, sample_idx) in enumerate(zip(match_scores, indices)):
            sample = prototype_dataset.get_data_info(sample_idx.item())
            value_image = mmcv.imread(sample['img_path'])[..., ::-1]
            value_plot = figure.add_subplot(gs[1, k])
            value_plot.axis(False)
            value_plot.imshow(value_image)
            if draw_score:
                value_plot.text(
                    5,
                    5,
                    f'{score:.2f}',
                    **text_cfg,
                )
        drawn_img = img_from_canvas(figure.canvas)
        self.set_image(drawn_img)

        if show:
            self.show(drawn_img, win_name=name, wait_time=wait_time)

        if out_file is not None:
            # save the image to the target file instead of vis_backends
            mmcv.imwrite(drawn_img[..., ::-1], out_file)
        else:
            self.add_image(name, drawn_img, step=step)

        return drawn_img
Visualize image retrieval result.

        This method will draw the input image and the images retrieved from the
        prototype dataset.

        Args:
            image (np.ndarray): The image to draw. The format should be RGB.
            data_sample (:obj:`DataSample`): The annotation of the image.
            prototype_dataset (:obj:`BaseDataset`): The prototype dataset.
                It should have `get_data_info` method and return a dict
                includes `img_path`.
            draw_score (bool): Whether to draw the match scores of the
                retrieved images. Defaults to True.
            resize (int, optional): Resize the long edge of the image to the
                specified length before visualization. Defaults to None.
            text_cfg (dict): Extra text setting, which accepts arguments of
                :func:`plt.text`. Defaults to an empty dict.
            show (bool): Whether to display the drawn image in a window, please
                confirm your are able to access the graphical interface.
                Defaults to False.
            wait_time (float): The display time (s). Defaults to 0, which means
                "forever".
            out_file (str, optional): Extra path to save the visualization
                result. If specified, the visualizer will only save the result
                image to the out_file and ignore its storage backends.
                Defaults to None.
            name (str): The image identifier. It's useful when using the
                storage backends of the visualizer to save or display the
                image. Defaults to an empty string.
            step (int): The global step value. It's useful to record a
                series of visualization results for the same image with the
                storage backends. Defaults to 0.

        Returns:
            np.ndarray: The visualization image.
        
mmpretrain/visualization/visualizer.py