text: # Licensed to the Apache Software Foundation (ASF) under one
  # or more contributor license agreements.  See the NOTICE file
  # distributed with this work for additional information
  # regarding copyright ownership.  The ASF licenses this file
  # to you 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.
  
  import mxnet as mx
  import random
  from mxnet.io import DataBatch, DataIter
  import numpy as np
  
  def add_data_args(parser):
      data = parser.add_argument_group('Data', 'the input images')
      data.add_argument('--data-train', type=str, help='the training data')
      data.add_argument('--data-train-idx', type=str, default='', help='the index of training data')
      data.add_argument('--data-val', type=str, help='the validation data')
      data.add_argument('--data-val-idx', type=str, default='', help='the index of validation data')
      data.add_argument('--rgb-mean', type=str, default='123.68,116.779,103.939',
                        help='a tuple of size 3 for the mean rgb')
      data.add_argument('--pad-size', type=int, default=0,
                        help='padding the input image')
      data.add_argument('--image-shape', type=str,
                        help='the image shape feed into the network, e.g. (3,224,224)')
      data.add_argument('--num-classes', type=int, help='the number of classes')
      data.add_argument('--num-examples', type=int, help='the number of training examples')
      data.add_argument('--resize-train', type=int, help='resize size of train data')
      data.add_argument('--resize-val', type=int, help='rezize size of val data')
      data.add_argument('--data-nthreads', type=int, default=4,
                        help='number of threads for data decoding')
      data.add_argument('--benchmark', type=int, default=0,
                        help='if 1, then feed the network with synthetic data')
      return data
  
  def add_data_aug_args(parser):
      aug = parser.add_argument_group(
          'Image augmentations', 'implemented in src/io/image_aug_default.cc')
      aug.add_argument('--random-crop', type=int, default=1,
                       help='if or not randomly crop the image')
      aug.add_argument('--random-mirror', type=int, default=1,
                       help='if or not randomly flip horizontally')
      aug.add_argument('--max-random-h', type=int, default=0,
                       help='max change of hue, whose range is [0, 180]')
      aug.add_argument('--max-random-s', type=int, default=0,
                       help='max change of saturation, whose range is [0, 255]')
      aug.add_argument('--max-random-l', type=int, default=0,
                       help='max change of intensity, whose range is [0, 255]')
      aug.add_argument('--max-random-aspect-ratio', type=float, default=0,
                       help='max change of aspect ratio, whose range is [0, 1]')
      aug.add_argument('--max-random-rotate-angle', type=int, default=0,
                       help='max angle to rotate, whose range is [0, 360]')
      aug.add_argument('--max-random-shear-ratio', type=float, default=0,
                       help='max ratio to shear, whose range is [0, 1]')
      aug.add_argument('--max-random-scale', type=float, default=1,
                       help='max ratio to scale')
      aug.add_argument('--min-random-scale', type=float, default=1,
                       help='min ratio to scale, should >= img_size/input_shape. otherwise use --pad-size')
      return aug
  
  def set_data_aug_level(aug, level):
      if level >= 1:
          aug.set_defaults(random_crop=1, random_mirror=1)
      if level >= 2:
          aug.set_defaults(max_random_h=36, max_random_s=50, max_random_l=50)
      if level >= 3:
          aug.set_defaults(max_random_rotate_angle=10, max_random_shear_ratio=0.1, max_random_aspect_ratio=0.25)
  
  
  class SyntheticDataIter(DataIter):
      def __init__(self, num_classes, data_shape, max_iter, dtype):
          self.batch_size = data_shape[0]
          self.cur_iter = 0
          self.max_iter = max_iter
          self.dtype = dtype
          label = np.random.randint(0, num_classes, [self.batch_size,])
          data = np.random.uniform(-1, 1, data_shape)
          self.data = mx.nd.array(data, dtype=self.dtype, ctx=mx.Context('cpu_pinned', 0))
          self.label = mx.nd.array(label, dtype=self.dtype, ctx=mx.Context('cpu_pinned', 0))
      def __iter__(self):
          return self
      @property
      def provide_data(self):
          return [mx.io.DataDesc('data', self.data.shape, self.dtype)]
      @property
      def provide_label(self):
          return [mx.io.DataDesc('softmax_label', (self.batch_size,), self.dtype)]
      def next(self):
          self.cur_iter += 1
          if self.cur_iter <= self.max_iter:
              return DataBatch(data=(self.data,),
                               label=(self.label,),
                               pad=0,
                               index=None,
                               provide_data=self.provide_data,
                               provide_label=self.provide_label)
          else:
              raise StopIteration
      def __next__(self):
          return self.next()
      def reset(self):
          self.cur_iter = 0
  
  def get_rec_iter(args, kv=None):
      image_shape = tuple([int(l) for l in args.image_shape.split(',')])
      if 'benchmark' in args and args.benchmark:
          data_shape = (args.batch_size,) + image_shape
          train = SyntheticDataIter(args.num_classes, data_shape, 500, np.float32)
          return (train, None)
      if kv:
          (rank, nworker) = (kv.rank, kv.num_workers)
      else:
          (rank, nworker) = (0, 1)
      rgb_mean = [float(i) for i in args.rgb_mean.split(',')]
      train = mx.io.ImageRecordIter(
          path_imgrec         = args.data_train,
          path_imgidx         = args.data_train_idx,
          label_width         = 1,
          mean_r              = rgb_mean[0],
          mean_g              = rgb_mean[1],
          mean_b              = rgb_mean[2],
          data_name           = 'data',
          label_name          = 'softmax_label',
          data_shape          = image_shape,
          batch_size          = args.batch_size,
          rand_crop           = args.random_crop,
          max_random_scale    = args.max_random_scale,
          pad                 = args.pad_size,
          fill_value          = 127,
          resize              = args.resize_train,
          min_random_scale    = args.min_random_scale,
          max_aspect_ratio    = args.max_random_aspect_ratio,
          random_h            = args.max_random_h,
          random_s            = args.max_random_s,
          random_l            = args.max_random_l,
          max_rotate_angle    = args.max_random_rotate_angle,
          max_shear_ratio     = args.max_random_shear_ratio,
          rand_mirror         = args.random_mirror,
          preprocess_threads  = args.data_nthreads,
          shuffle             = True,
          num_parts           = nworker,
          part_index          = rank)
      if args.data_val is None:
          return (train, None)
      val = mx.io.ImageRecordIter(
          path_imgrec         = args.data_val,
          path_imgidx         = args.data_val_idx,
          label_width         = 1,
          mean_r              = rgb_mean[0],
          mean_g              = rgb_mean[1],
          mean_b              = rgb_mean[2],
          data_name           = 'data',
          label_name          = 'softmax_label',
          batch_size          = args.batch_size,
          data_shape          = image_shape,
          resize              = args.resize_val,
          preprocess_threads  = args.data_nthreads,
          rand_crop           = False,
          rand_mirror         = False,
          num_parts           = nworker,
          part_index          = rank)
      return (train, val)
  