text: # Lint as: python2, python3
  # Copyright 2018 The TensorFlow Authors All Rights Reserved.
  #
  # 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.
  # ==============================================================================
  
  """Converts ADE20K data to TFRecord file format with Example protos."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  import os
  import random
  import sys
  import build_data
  from six.moves import range
  import tensorflow as tf
  
  FLAGS = tf.app.flags.FLAGS
  
  tf.app.flags.DEFINE_string(
      'train_image_folder',
      './ADE20K/ADEChallengeData2016/images/training',
      'Folder containing trainng images')
  tf.app.flags.DEFINE_string(
      'train_image_label_folder',
      './ADE20K/ADEChallengeData2016/annotations/training',
      'Folder containing annotations for trainng images')
  
  tf.app.flags.DEFINE_string(
      'val_image_folder',
      './ADE20K/ADEChallengeData2016/images/validation',
      'Folder containing validation images')
  
  tf.app.flags.DEFINE_string(
      'val_image_label_folder',
      './ADE20K/ADEChallengeData2016/annotations/validation',
      'Folder containing annotations for validation')
  
  tf.app.flags.DEFINE_string(
      'output_dir', './ADE20K/tfrecord',
      'Path to save converted tfrecord of Tensorflow example')
  
  _NUM_SHARDS = 4
  
  
  def _convert_dataset(dataset_split, dataset_dir, dataset_label_dir):
    """Converts the ADE20k dataset into into tfrecord format.
  
    Args:
      dataset_split: Dataset split (e.g., train, val).
      dataset_dir: Dir in which the dataset locates.
      dataset_label_dir: Dir in which the annotations locates.
  
    Raises:
      RuntimeError: If loaded image and label have different shape.
    """
  
    img_names = tf.gfile.Glob(os.path.join(dataset_dir, '*.jpg'))
    random.shuffle(img_names)
    seg_names = []
    for f in img_names:
      # get the filename without the extension
      basename = os.path.basename(f).split('.')[0]
      # cover its corresponding *_seg.png
      seg = os.path.join(dataset_label_dir, basename+'.png')
      seg_names.append(seg)
  
    num_images = len(img_names)
    num_per_shard = int(math.ceil(num_images / _NUM_SHARDS))
  
    image_reader = build_data.ImageReader('jpeg', channels=3)
    label_reader = build_data.ImageReader('png', channels=1)
  
    for shard_id in range(_NUM_SHARDS):
      output_filename = os.path.join(
          FLAGS.output_dir,
          '%s-%05d-of-%05d.tfrecord' % (dataset_split, shard_id, _NUM_SHARDS))
      with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
        start_idx = shard_id * num_per_shard
        end_idx = min((shard_id + 1) * num_per_shard, num_images)
        for i in range(start_idx, end_idx):
          sys.stdout.write('\r>> Converting image %d/%d shard %d' % (
              i + 1, num_images, shard_id))
          sys.stdout.flush()
          # Read the image.
          image_filename = img_names[i]
          image_data = tf.gfile.FastGFile(image_filename, 'rb').read()
          height, width = image_reader.read_image_dims(image_data)
          # Read the semantic segmentation annotation.
          seg_filename = seg_names[i]
          seg_data = tf.gfile.FastGFile(seg_filename, 'rb').read()
          seg_height, seg_width = label_reader.read_image_dims(seg_data)
          if height != seg_height or width != seg_width:
            raise RuntimeError('Shape mismatched between image and label.')
          # Convert to tf example.
          example = build_data.image_seg_to_tfexample(
              image_data, img_names[i], height, width, seg_data)
          tfrecord_writer.write(example.SerializeToString())
      sys.stdout.write('\n')
      sys.stdout.flush()
  
  
  def main(unused_argv):
    tf.gfile.MakeDirs(FLAGS.output_dir)
    _convert_dataset(
        'train', FLAGS.train_image_folder, FLAGS.train_image_label_folder)
    #_convert_dataset('val', FLAGS.val_image_folder, FLAGS.val_image_label_folder)
  
  
  if __name__ == '__main__':
    tf.app.run()
  