Refactor visualisation util

PiperOrigin-RevId: 303010840

Author: Conchylicultor

tensorflow_datasets/core/visualization.py

@@ -0,0 +1,27 @@
+# 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.
+
+"""Visualizer utils."""
+
+from tensorflow_datasets.core.visualization.image_visualizer import ImageGridVisualizer
+from tensorflow_datasets.core.visualization.show_examples import show_examples
+from tensorflow_datasets.core.visualization.visualizer import Visualizer
+
+
+__all__ = [
+    "ImageGridVisualizer",
+    "show_examples",
+    "Visualizer",
+]

tensorflow_datasets/core/visualization/__init__.py

@@ -0,0 +1,150 @@
+# 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 absl import logging
+
+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: Figure to display.
+  """
+  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):
+    """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_info,
+      ds,
+      rows=3,
+      cols=3,
+      plot_scale=3.,
+      image_key=None,
+  ):
+    """Display the dataset.
+
+    Args:
+      ds_info: `tfds.core.DatasetInfo` object of the dataset to visualize.
+      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.
+      image_key: `string`, name of the feature that contains the image. If not
+         set, the system will try to auto-detect it.
+    """
+    # 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))
+
+    # Print the grid
+    fig = _make_grid(make_cell_fn, ds, rows, cols, plot_scale)

tensorflow_datasets/core/visualization/image_visualizer.py

@@ -0,0 +1,64 @@
+# 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
+"""Show example util.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow_datasets.core.visualization import image_visualizer
+
+_ALL_VISUALIZERS = [
+    image_visualizer.ImageGridVisualizer(),
+]
+
+
+def show_examples(ds_info, ds, **options_kwargs):
+  """Visualize images (and labels) from an image classification dataset.
+
+  This function is for interactive use (Colab, Jupyter). It displays and return
+  a plot of (rows*columns) images from a tf.data.Dataset.
+
+  Usage:
+  ```python
+  ds, ds_info = tfds.load('cifar10', split='train', with_info=True)
+  fig = tfds.show_examples(ds_info, ds)
+  ```
+
+  Args:
+    ds_info: The dataset info object to which extract the label and features
+      info. Available either through `tfds.load('mnist', with_info=True)` or
+      `tfds.builder('mnist').info`
+    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.
+    **options_kwargs: Additional display options, specific to the dataset type
+      to visualize. Are forwarded to `tfds.visualization.Visualizer.show`.
+      See the `tfds.visualization` for a list of available visualizers.
+
+  Returns:
+    fig: The `matplotlib.Figure` object
+  """
+  for visualizer in _ALL_VISUALIZERS:
+    if visualizer.match(ds_info):
+      visualizer.show(ds_info, ds, **options_kwargs)
+      break
+  else:
+    raise ValueError(
+        "Visualisation not supported for dataset `{}`".format(ds_info.name)
+    )