content: #!/usr/bin/env python
  
  """
  A browser for the NORB and small NORB datasets. Navigate the images by
  choosing the values for the label vector. Note that for the 'big' NORB
  dataset, you can only set the first 5 label dimensions. You can then cycle
  through the 3-12 images that fit those labels.
  """
  
  import sys
  import os
  import argparse
  import numpy
  import warnings
  
  try:
      import matplotlib
      from matplotlib import pyplot
  except ImportError as import_error:
      warnings.warn("Can't use this script without matplotlib.")
      matplotlib = None
      pyplot = None
  
  from pylearn2.datasets.new_norb import NORB
  from pylearn2.utils import safe_zip, serial
  
  
  def _parse_args():
      parser = argparse.ArgumentParser(
          description="Browser for NORB dataset.")
  
      parser.add_argument('--which_norb',
                          type=str,
                          required=False,
                          choices=('big', 'small'),
                          help="'Selects the (big) NORB, or the Small NORB.")
  
      parser.add_argument('--which_set',
                          type=str,
                          required=False,
                          choices=('train', 'test', 'both'),
                          help="'train', or 'test'")
  
      parser.add_argument('--pkl',
                          type=str,
                          required=False,
                          help=".pkl file of NORB dataset")
  
      parser.add_argument('--stereo_viewer',
                          action='store_true',
                          help="Swaps left and right stereo images, so you "
                          "can see them in 3D by crossing your eyes.")
  
      parser.add_argument('--no_norm',
                          action='store_true',
                          help="Don't normalize pixel values")
  
      result = parser.parse_args()
  
      if (result.pkl is not None) == (result.which_norb is not None or
                                      result.which_set is not None):
          print("Must supply either --pkl, or both --which_norb and "
                "--which_set.")
          sys.exit(1)
  
      if (result.which_norb is None) != (result.which_set is None):
          print("When not supplying --pkl, you must supply both "
                "--which_norb and --which_set.")
          sys.exit(1)
  
      if result.pkl is not None:
          if not result.pkl.endswith('.pkl'):
              print("--pkl must be a filename that ends in .pkl")
              sys.exit(1)
  
          if not os.path.isfile(result.pkl):
              print("couldn't find --pkl file '%s'" % result.pkl)
              sys.exit(1)
  
      return result
  
  
  def _make_grid_to_short_label(dataset):
      """
      Returns an array x such that x[a][b] gives label index a's b'th unique
      value. In other words, it maps label grid indices a, b to the
      corresponding label value.
      """
      unique_values = [sorted(list(frozenset(column)))
                       for column
                       in dataset.y[:, :5].transpose()]
  
      # If dataset contains blank images, removes the '-1' labels
      # corresponding to blank images, since they aren't contained in the
      # label grid.
      category_index = dataset.label_name_to_index['category']
      unique_categories = unique_values[category_index]
      category_to_name = dataset.label_to_value_funcs[category_index]
      if any(category_to_name(category) == 'blank'
             for category in unique_categories):
          for d in range(1, len(unique_values)):
              assert unique_values[d][0] == -1, ("unique_values: %s" %
                                                 str(unique_values))
              unique_values[d] = unique_values[d][1:]
  
      return unique_values
  
  
  def _get_blank_label(dataset):
      """
      Returns the label vector associated with blank images.
  
      If dataset is a Small NORB (i.e. it has no blank images), this returns
      None.
      """
  
      category_index = dataset.label_name_to_index['category']
      category_to_name = dataset.label_to_value_funcs[category_index]
      blank_label = 5
  
      try:
          blank_name = category_to_name(blank_label)
      except ValueError:
          # Returns None if there is no 'blank' category (e.g. if we're using
          # the small NORB dataset.
          return None
  
      assert blank_name == 'blank'
  
      blank_rowmask = dataset.y[:, category_index] == blank_label
      blank_labels = dataset.y[blank_rowmask, :]
  
      if not blank_rowmask.any():
          return None
  
      if not numpy.all(blank_labels[0, :] == blank_labels[1:, :]):
          raise ValueError("Expected all labels of category 'blank' to have "
                           "the same value, but they differed.")
  
      return blank_labels[0, :].copy()
  
  
  def _make_label_to_row_indices(labels):
      """
      Returns a map from short labels (the first 5 elements of the label
      vector) to the list of row indices of rows in the dense design matrix
      with that label.
  
      For Small NORB, all unique short labels have exactly one row index.
  
      For big NORB, a short label can have 0-N row indices.
      """
      result = {}
  
      for row_index, label in enumerate(labels):
  
          short_label = tuple(label[:5])
          if result.get(short_label, None) is None:
              result[short_label] = []
  
          result[short_label].append(row_index)
  
      return result
  
  
  def main():
      """Top-level function."""
  
      args = _parse_args()
  
      if args.pkl is not None:
          dataset = serial.load(args.pkl)
      else:
          dataset = NORB(args.which_norb, args.which_set)
  
      # Indexes into the first 5 labels, which live on a 5-D grid.
      grid_indices = [0, ] * 5
  
      grid_to_short_label = _make_grid_to_short_label(dataset)
  
      # Maps 5-D label vector to a list of row indices for dataset.X, dataset.y
      # that have those labels.
      label_to_row_indices = _make_label_to_row_indices(dataset.y)
  
      # Indexes into the row index lists returned by label_to_row_indices.
      object_image_index = [0, ]
      blank_image_index = [0, ]
      blank_label = _get_blank_label(dataset)
  
      # Index into grid_indices currently being edited
      grid_dimension = [0, ]
  
      dataset_is_stereo = 's' in dataset.view_converter.axes
      figure, all_axes = pyplot.subplots(1,
                                         3 if dataset_is_stereo else 2,
                                         squeeze=True,
                                         figsize=(10, 3.5))
  
      set_name = (os.path.split(args.pkl)[1] if args.which_set is None
                  else "%sing set" % args.which_set)
      figure.canvas.set_window_title("NORB dataset (%s)" % set_name)
  
      label_text = figure.suptitle('Up/down arrows choose label, '
                                   'left/right arrows change it',
                                   x=0.1,
                                   horizontalalignment="left")
  
      # Hides axes' tick marks
      for axes in all_axes:
          axes.get_xaxis().set_visible(False)
          axes.get_yaxis().set_visible(False)
  
      text_axes, image_axes = (all_axes[0], all_axes[1:])
      image_captions = (('left', 'right') if dataset_is_stereo
                        else ('mono image', ))
  
      if args.stereo_viewer:
          image_captions = tuple(reversed(image_captions))
  
      for image_ax, caption in safe_zip(image_axes, image_captions):
          image_ax.set_title(caption)
  
      text_axes.set_frame_on(False)  # Hides background of text_axes
  
      def is_blank(grid_indices):
          assert len(grid_indices) == 5
          assert all(x >= 0 for x in grid_indices)
  
          ci = dataset.label_name_to_index['category']  # category index
          category = grid_to_short_label[ci][grid_indices[ci]]
          category_name = dataset.label_to_value_funcs[ci](category)
          return category_name == 'blank'
  
      def get_short_label(grid_indices):
          """
          Returns the first 5 elements of the label vector pointed to by
          grid_indices. We use the first 5, since they're the labels used by
          both the 'big' and Small NORB datasets.
          """
  
          # Need to special-case the 'blank' category, since it lies outside of
          # the grid.
          if is_blank(grid_indices):   # won't happen with SmallNORB
              return tuple(blank_label[:5])
          else:
              return tuple(grid_to_short_label[i][g]
                           for i, g in enumerate(grid_indices))
  
      def get_row_indices(grid_indices):
          short_label = get_short_label(grid_indices)
          return label_to_row_indices.get(short_label, None)
  
      axes_to_pixels = {}
  
      def redraw(redraw_text, redraw_images):
          row_indices = get_row_indices(grid_indices)
  
          if row_indices is None:
              row_index = None
              image_index = 0
              num_images = 0
          else:
              image_index = (blank_image_index
                             if is_blank(grid_indices)
                             else object_image_index)[0]
              row_index = row_indices[image_index]
              num_images = len(row_indices)
  
          def draw_text():
              if row_indices is None:
                  padding_length = dataset.y.shape[1] - len(grid_indices)
                  current_label = (tuple(get_short_label(grid_indices)) +
                                   (0, ) * padding_length)
              else:
                  current_label = dataset.y[row_index, :]
  
              label_names = dataset.label_index_to_name
  
              label_values = [label_to_value(label) for label_to_value, label
                              in safe_zip(dataset.label_to_value_funcs,
                                          current_label)]
  
              lines = ['%s: %s' % (t, v)
                       for t, v
                       in safe_zip(label_names, label_values)]
  
              if dataset.y.shape[1] > 5:
                  # Inserts image number & blank line between editable and
                  # fixed labels.
                  lines = (lines[:5] +
                           ['No such image' if num_images == 0
                            else 'image: %d of %d' % (image_index + 1,
                                                      num_images),
                            '\n'] +
                           lines[5:])
  
              # prepends the current index's line with an arrow.
              lines[grid_dimension[0]] = '==> ' + lines[grid_dimension[0]]
  
              text_axes.clear()
  
              # "transAxes": 0, 0 = bottom-left, 1, 1 at upper-right.
              text_axes.text(0, 0.5,  # coords
                             '\n'.join(lines),
                             verticalalignment='center',
                             transform=text_axes.transAxes)
  
          def draw_images():
              if row_indices is None:
                  for axis in image_axes:
                      axis.clear()
              else:
                  data_row = dataset.X[row_index:row_index + 1, :]
  
                  axes_names = dataset.view_converter.axes
                  assert len(axes_names) in (4, 5)
                  assert axes_names[0] == 'b'
                  assert axes_names[-3] == 0
                  assert axes_names[-2] == 1
                  assert axes_names[-1] == 'c'
  
                  def draw_image(image, axes):
                      assert len(image.shape) == 2
                      norm = matplotlib.colors.NoNorm() if args.no_norm else None
                      axes_to_pixels[axes] = image
                      axes.imshow(image, norm=norm, cmap='gray')
  
                  if 's' in axes_names:
                      image_pair = \
                          dataset.get_topological_view(mat=data_row,
                                                       single_tensor=True)
                      # Shaves off the singleton dimensions
                      # (batch # and channel #), leaving just 's', 0, and 1.
                      image_pair = tuple(image_pair[0, :, :, :, 0])
  
                      if args.stereo_viewer:
                          image_pair = tuple(reversed(image_pair))
  
                      for axis, image in safe_zip(image_axes, image_pair):
                          draw_image(image, axis)
                  else:
                      image = dataset.get_topological_view(mat=data_row)
                      image = image[0, :, :, 0]
                      draw_image(image, image_axes[0])
  
          if redraw_text:
              draw_text()
  
          if redraw_images:
              draw_images()
  
          figure.canvas.draw()
  
      default_status_text = ("mouseover image%s for pixel values" %
                             ("" if len(image_axes) == 1 else "s"))
      status_text = figure.text(0.5, 0.1, default_status_text)
  
      def on_mouse_motion(event):
          original_text = status_text.get_text()
  
          if event.inaxes not in image_axes:
              status_text.set_text(default_status_text)
          else:
              pixels = axes_to_pixels[event.inaxes]
              row = int(event.ydata + .5)
              col = int(event.xdata + .5)
              status_text.set_text("Pixel value: %g" % pixels[row, col])
  
          if status_text.get_text != original_text:
              figure.canvas.draw()
  
      def on_key_press(event):
  
          def add_mod(arg, step, size):
              return (arg + size + step) % size
  
          def incr_index_type(step):
              num_dimensions = len(grid_indices)
              if dataset.y.shape[1] > 5:
                  # If dataset is big NORB, add one for the image index
                  num_dimensions += 1
  
              grid_dimension[0] = add_mod(grid_dimension[0],
                                          step,
                                          num_dimensions)
  
          def incr_index(step):
              assert step in (0, -1, 1), ("Step was %d" % step)
  
              image_index = (blank_image_index
                             if is_blank(grid_indices)
                             else object_image_index)
  
              if grid_dimension[0] == 5:  # i.e. the image index
                  row_indices = get_row_indices(grid_indices)
                  if row_indices is None:
                      image_index[0] = 0
                  else:
                      # increment the image index
                      image_index[0] = add_mod(image_index[0],
                                               step,
                                               len(row_indices))
              else:
                  # increment one of the grid indices
                  gd = grid_dimension[0]
                  grid_indices[gd] = add_mod(grid_indices[gd],
                                             step,
                                             len(grid_to_short_label[gd]))
  
                  row_indices = get_row_indices(grid_indices)
                  if row_indices is None:
                      image_index[0] = 0
                  else:
                      # some grid indices have 2 images instead of 3.
                      image_index[0] = min(image_index[0], len(row_indices))
  
          # Disables left/right key if we're currently showing a blank,
          # and the current index type is neither 'category' (0) nor
          # 'image number' (5)
          disable_left_right = (is_blank(grid_indices) and
                                not (grid_dimension[0] in (0, 5)))
  
          if event.key == 'up':
              incr_index_type(-1)
              redraw(True, False)
          elif event.key == 'down':
              incr_index_type(1)
              redraw(True, False)
          elif event.key == 'q':
              sys.exit(0)
          elif not disable_left_right:
              if event.key == 'left':
                  incr_index(-1)
                  redraw(True, True)
              elif event.key == 'right':
                  incr_index(1)
                  redraw(True, True)
  
      figure.canvas.mpl_connect('key_press_event', on_key_press)
      figure.canvas.mpl_connect('motion_notify_event', on_mouse_motion)
      redraw(True, True)
  
      pyplot.show()
  
  
  if __name__ == '__main__':
      main()
  