# 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.
r"""Beam pipeline to create TFRecord files from JPEG files stored on GCS.

These are the TFRecord format expected by  the resnet and amoebanet models.
Example usage:
python -m jpeg_to_tf_record.py \
       --train_csv gs://cloud-ml-data/img/flower_photos/train_set.csv \
       --validation_csv gs://cloud-ml-data/img/flower_photos/eval_set.csv \
       --labels_file /tmp/labels.txt \
       --project_id $PROJECT \
       --output_dir gs://${BUCKET}/tpu/imgclass/data

The format of the CSV files is:
    URL-of-image,label
And the format of the labels_file is simply a list of strings one-per-line.
"""

from __future__ import print_function

import argparse
import datetime
import os
import shutil
import subprocess
import sys
import apache_beam as beam
import tensorflow.compat.v1 as tf


def _int64_feature(value):
  """Wrapper for inserting int64 features into Example proto."""
  if not isinstance(value, list):
    value = [value]
  return tf.train.Feature(int64_list=tf.train.Int64List(value=value))


def _bytes_feature(value):
  """Wrapper for inserting bytes features into Example proto."""
  return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def _convert_to_example(filename, image_buffer, label_int, label_str, height,
                        width):
  """Build an Example proto for an example.

  Args:
    filename: string, path to an image file, e.g., '/path/to/example.JPG'
    image_buffer: string, JPEG encoding of RGB image
    label_int: integer, identifier for ground truth (0-based)
    label_str: string, identifier for ground truth, e.g., 'daisy'
    height: integer, image height in pixels
    width: integer, image width in pixels
  Returns:
    Example proto
  """
  colorspace = 'RGB'
  channels = 3
  image_format = 'JPEG'

  example = tf.train.Example(
      features=tf.train.Features(
          feature={
              'image/height': _int64_feature(height),
              'image/width': _int64_feature(width),
              'image/colorspace': _bytes_feature(colorspace),
              'image/channels': _int64_feature(channels),
              'image/class/label': _int64_feature(label_int +
                                                  1),  # model expects 1-based
              'image/class/synset': _bytes_feature(label_str),
              'image/format': _bytes_feature(image_format),
              'image/filename': _bytes_feature(os.path.basename(filename)),
              'image/encoded': _bytes_feature(image_buffer)
          }))
  return example


class ImageCoder(object):
  """Helper class that provides TensorFlow image coding utilities."""

  def __init__(self):
    # Create a single Session to run all image coding calls.
    self._sess = tf.Session()

    # Initializes function that decodes RGB JPEG data.
    self._decode_jpeg_data = tf.placeholder(dtype=tf.string)
    self._decode_jpeg = tf.image.decode_jpeg(self._decode_jpeg_data, channels=3)

  def decode_jpeg(self, image_data):
    image = self._sess.run(
        self._decode_jpeg, feed_dict={self._decode_jpeg_data: image_data})
    assert len(image.shape) == 3
    assert image.shape[2] == 3
    return image

  def __del__(self):
    self._sess.close()


def _get_image_data(filename, coder):
  """Process a single image file.

  Args:
    filename: string, path to an image file e.g., '/path/to/example.JPG'.
    coder: instance of ImageCoder to provide TensorFlow image coding utils.
  Returns:
    image_buffer: string, JPEG encoding of RGB image.
    height: integer, image height in pixels.
    width: integer, image width in pixels.
  """
  # Read the image file.
  with tf.gfile.FastGFile(filename, 'r') as ifp:
    image_data = ifp.read()

  # Decode the RGB JPEG.
  image = coder.decode_jpeg(image_data)

  # Check that image converted to RGB
  assert len(image.shape) == 3
  height = image.shape[0]
  width = image.shape[1]
  assert image.shape[2] == 3

  return image_data, height, width


def convert_to_example(csvline, categories):
  """Parse a line of CSV file and convert to TF Record.

  Args:
    csvline: line from input CSV file
    categories: list of labels
  Yields:
    serialized TF example if the label is in categories
  """
  filename, label = csvline.encode('ascii', 'ignore').split(',')
  if label in categories:
    # ignore labels not in categories list
    coder = ImageCoder()
    image_buffer, height, width = _get_image_data(filename, coder)
    del coder
    example = _convert_to_example(filename, image_buffer,
                                  categories.index(label), label, height, width)
    yield example.SerializeToString()


if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--train_csv',
      # pylint: disable=line-too-long
      help=
      'Path to input.  Each line of input has two fields  image-file-name and label separated by a comma',
      required=True)
  parser.add_argument(
      '--validation_csv',
      # pylint: disable=line-too-long
      help=
      'Path to input.  Each line of input has two fields  image-file-name and label separated by a comma',
      required=True)
  parser.add_argument(
      '--labels_file',
      help='Path to file containing list of labels, one per line',
      required=True)
  parser.add_argument(
      '--project_id',
      help='ID (not name) of your project. Ignored by DirectRunner',
      required=True)
  parser.add_argument(
      '--runner',
      help='If omitted, uses DataFlowRunner if output_dir starts with gs://',
      default=None)
  parser.add_argument(
      '--output_dir', help='Top-level directory for TF Records', required=True)

  args = parser.parse_args()
  arguments = args.__dict__

  JOBNAME = (
      'preprocess-images-' + datetime.datetime.now().strftime('%y%m%d-%H%M%S'))

  PROJECT = arguments['project_id']
  OUTPUT_DIR = arguments['output_dir']

  # set RUNNER using command-line arg or based on output_dir path
  on_cloud = OUTPUT_DIR.startswith('gs://')
  if arguments['runner']:
    RUNNER = arguments['runner']
  else:
    RUNNER = 'DataflowRunner' if on_cloud else 'DirectRunner'

  # clean-up output directory since Beam will name files 0000-of-0004 etc.
  # and this could cause confusion if earlier run has 0000-of-0005, for eg
  if on_cloud:
    try:
      subprocess.check_call('gsutil -m rm -r {}'.format(OUTPUT_DIR).split())
    except subprocess.CalledProcessError:
      pass
  else:
    shutil.rmtree(OUTPUT_DIR, ignore_errors=True)
    os.makedirs(OUTPUT_DIR)

  # read list of labels
  with tf.gfile.FastGFile(arguments['labels_file'], 'r') as f:
    LABELS = [line.rstrip() for line in f]
  print('Read in {} labels, from {} to {}'.format(
      len(LABELS), LABELS[0], LABELS[-1]))
  if len(LABELS) < 2:
    print('Require at least two labels')
    sys.exit(-1)

  # set up Beam pipeline to convert images to TF Records
  options = {
      'staging_location': os.path.join(OUTPUT_DIR, 'tmp', 'staging'),
      'temp_location': os.path.join(OUTPUT_DIR, 'tmp'),
      'job_name': JOBNAME,
      'project': PROJECT,
      'teardown_policy': 'TEARDOWN_ALWAYS',
      'save_main_session': True
  }
  opts = beam.pipeline.PipelineOptions(flags=[], **options)

  with beam.Pipeline(RUNNER, options=opts) as p:
    # BEAM tasks
    for step in ['train', 'validation']:
      _ = (
          p
          | '{}_read_csv'.format(step) >> beam.io.ReadFromText(
              arguments['{}_csv'.format(step)])
          | '{}_convert'.format(step) >>
          beam.FlatMap(lambda line: convert_to_example(line, LABELS))
          | '{}_write_tfr'.format(step) >> beam.io.tfrecordio.WriteToTFRecord(
              os.path.join(OUTPUT_DIR, step)))
