# Copyright 2018 Google LLC
#
# 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
#
#     https://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.

"""Load MNIST dataset, and generate its n-digit version."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

import numpy as np
from PIL import Image


def main():
  output_dir = 'data'
  num_digits = 2
  domain_gap = 'number'
  mnist_split = 'train'
  num_visualize = 10

  # Construct dataset directory.
  dataset_dir = os.path.join(output_dir,
                             'dataset_mnist_%d_%s' % (num_digits,
                                                      domain_gap))

  path = os.path.join(dataset_dir, '%s.npz' % mnist_split)
  with open(path, 'r') as f:
    data = np.load(path)
    labels = data['labels']
    images = data['images']

  visualize_dir = os.path.join(dataset_dir, 'visualization')
  if not os.path.exists(visualize_dir):
    os.makedirs(visualize_dir)

  visualize_indices = np.random.choice(range(len(labels)),
                                       num_visualize, replace=False)
  
  for i in visualize_indices:
    im = Image.fromarray(images[i])
    im.save(os.path.join(visualize_dir,
                         'sample_image_%d_label_%d.jpg' % (i, labels[i])))

if __name__ == '__main__':
  main()


@@ -812,59 +812,349 @@
 :%0A  
-output_dir = 'data'%0A  num_digits = 2%0A  domain_gap =
+parser = argparse.ArgumentParser(%0A    description='Visualize created n-digit MNIST dataset.')%0A%0A  parser.add_argument('--num_digits', default=1, type=int,%0A                      help='Number of concatenated digits per data point.')%0A%0A  parser.add_argument('--domain_gap', default='instance', type=str,%0A                      choices=%5B'instance',
  'nu
@@ -1162,53 +1162,595 @@
 ber'
+%5D,
 %0A  
-mnist_split = 'train'%0A  num_visualize = 10
+                    help='How to split training and test sets'%0A                           'of the n-digit mnist.')%0A%0A  parser.add_argument('--num_visualize', default=10, type=int,%0A                      help='How many samples to visualize.')%0A%0A  parser.add_argument('--output_dir', default='data', type=str,%0A                      help='Directory to write the dataset.')%0A%0A  parser.add_argument('--mnist_split', default='data', type=str,%0A                      help='Which MNIST split was used to generate data?'%0A                           '(train or test)')%0A%0A  args = parser.parse_args()
 %0A%0A  
@@ -1809,16 +1809,21 @@
 th.join(
+args.
 output_d
@@ -1880,16 +1880,21 @@
 _%25s' %25 (
+args.
 num_digi
@@ -1951,16 +1951,21 @@
         
+args.
 domain_g
@@ -1984,32 +1984,37 @@
  = os.path.join(
+args.
 dataset_dir, '%25s
@@ -2021,16 +2021,21 @@
 .npz' %25 
+args.
 mnist_sp
@@ -2181,16 +2181,21 @@
 th.join(
+args.
 dataset_
@@ -2386,16 +2386,21 @@
         
+args.
 num_visu
@@ -2515,16 +2515,21 @@
 th.join(
+args.
 visualiz

Update example_visualization.py