text: # coding=utf-8
  # Copyright 2020 The TensorFlow Datasets Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  # Lint as: python3
  """Image visualizer."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from typing import Optional
  
  from absl import logging
  import tensorflow.compat.v2 as tf
  
  from tensorflow_datasets.core import dataset_info
  from tensorflow_datasets.core import dataset_utils
  from tensorflow_datasets.core import features as features_lib
  from tensorflow_datasets.core import lazy_imports_lib
  from tensorflow_datasets.core.visualization import visualizer
  
  
  def _make_grid(plot_single_ex_fn, ds, rows, cols, plot_scale):
    """Plot each individual example in a grid.
  
    Args:
      plot_single_ex_fn: Function with fill a single cell of the grid, with
        signature `fn(ax: matplotlib.axes.Axes, ex: Nested[np.array]) -> None`
      ds: `tf.data.Dataset`. The tf.data.Dataset object to visualize. Examples
        should not be batched. Examples will be consumed in order until
        (rows * cols) are read or the dataset is consumed.
      rows: `int`, number of rows of the display grid.
      cols: `int`, number of columns of the display grid.
      plot_scale: `float`, controls the plot size of the images. Keep this
        value around 3 to get a good plot. High and low values may cause
        the labels to get overlapped.
  
    Returns:
      fig: The `matplotlib.Figure` object.
    """
    plt = lazy_imports_lib.lazy_imports.matplotlib.pyplot
  
    num_examples = rows * cols
    examples = list(dataset_utils.as_numpy(ds.take(num_examples)))
  
    fig = plt.figure(figsize=(plot_scale * cols, plot_scale * rows))
    fig.subplots_adjust(hspace=1 / plot_scale, wspace=1 / plot_scale)
  
    for i, ex in enumerate(examples):
      ax = fig.add_subplot(rows, cols, i+1)
      plot_single_ex_fn(ax, ex)
  
    plt.show()
    return fig
  
  
  def _add_image(ax, image):
    """Add the image to the given `matplotlib.axes.Axes`."""
    plt = lazy_imports_lib.lazy_imports.matplotlib.pyplot
  
    if len(image.shape) != 3:
      raise ValueError(
          'Image dimension should be 3. tfds.show_examples does not support '
          'batched examples or video.')
    _, _, c = image.shape
    if c == 1:
      image = image.reshape(image.shape[:2])
    ax.imshow(image, cmap='gray')
    ax.grid(False)
    plt.xticks([], [])
    plt.yticks([], [])
  
  
  class ImageGridVisualizer(visualizer.Visualizer):
    """Visualizer for supervised image datasets."""
  
    def match(self, ds_info: dataset_info.DatasetInfo) -> bool:
      """See base class."""
      # Supervised required a single image key
      image_keys = visualizer.extract_keys(ds_info.features, features_lib.Image)
      return len(image_keys) >= 1
  
    def show(
        self,
        ds: tf.data.Dataset,
        ds_info: dataset_info.DatasetInfo,
        rows: int = 3,
        cols: int = 3,
        plot_scale: float = 3.,
        image_key: Optional[str] = None,
    ):
      """Display the dataset.
  
      Args:
        ds: `tf.data.Dataset`. The tf.data.Dataset object to visualize. Examples
          should not be batched. Examples will be consumed in order until
          (rows * cols) are read or the dataset is consumed.
        ds_info: `tfds.core.DatasetInfo` object of the dataset to visualize.
        rows: `int`, number of rows of the display grid.
        cols: `int`, number of columns of the display grid.
        plot_scale: `float`, controls the plot size of the images. Keep this
          value around 3 to get a good plot. High and low values may cause
          the labels to get overlapped.
        image_key: `string`, name of the feature that contains the image. If not
           set, the system will try to auto-detect it.
  
      Returns:
        fig: The pyplot figure.
      """
      # Extract the image key
      if not image_key:
        image_keys = visualizer.extract_keys(ds_info.features, features_lib.Image)
        if len(image_keys) > 1:
          raise ValueError(
              'Multiple image features detected in the dataset. '
              'Use `image_key` argument to override. Images detected: {}'.format(
                  image_keys))
        image_key = image_keys[0]
  
      # Optionally extract the label key
      label_keys = visualizer.extract_keys(
          ds_info.features, features_lib.ClassLabel)
      label_key = label_keys[0] if len(label_keys) == 1 else None
      if not label_key:
        logging.info('Was not able to auto-infer label.')
  
      # Single image display
      def make_cell_fn(ax, ex):
        plt = lazy_imports_lib.lazy_imports.matplotlib.pyplot
  
        if not isinstance(ex, dict):
          raise ValueError(
              '{} requires examples as `dict`, with the same '
              'structure as `ds_info.features`. It is currently not compatible '
              'with `as_supervised=True`. Received: {}'.format(
                  type(self).__name__, type(ex)))
  
        _add_image(ax, ex[image_key])
        if label_key:
          label = ex[label_key]
          label_str = ds_info.features[label_key].int2str(label)
          plt.xlabel('{} ({})'.format(label_str, label))
  
      # Returns the grid.
      fig = _make_grid(make_cell_fn, ds, rows, cols, plot_scale)
      return fig
  