06_preprocessing/jpeg_to_tfrecord_tft.py

#!/usr/bin/env python3

# Copyright 2020 Google Inc. 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"""

Apache Beam pipeline to create TFRecord files from JPEG files stored on GCS.
This pipeline will split the data 80:10:10,
convert the images to lie in [-1, 1] range and resize them.

Modify the constants and TF Record format as needed.

Example usage:
python3 -m jpeg_to_tfrecord_tft \
       --all_data gs://cloud-ml-data/img/flower_photos/all_data.csv \
       --labels_file gs://cloud-ml-data/img/flower_photos/dict.txt \
       --project_id $PROJECT \
       --output_dir gs://${BUCKET}/data/flower_tfrecords \
       --resize 448,448

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.
"""

import argparse
import datetime
import os
import shutil
import subprocess
import sys
import tempfile
import apache_beam as beam
import tensorflow as tf
import numpy as np
import tensorflow_transform as tft
import tensorflow_transform.beam as tft_beam
from tfx_bsl.public import tfxio

IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS = 448, 448, 3
LABELS = []

def _string_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value.encode('utf-8')]))

def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))

def _float_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))

def decode_image(img_bytes):
    IMG_CHANNELS = 3
    return tf.image.decode_jpeg(img_bytes, channels=IMG_CHANNELS)

def assign_record_to_split(rec):
    rnd = np.random.rand()
    if rnd < 0.8:
        return ('train', rec)
    if rnd < 0.9:
        return ('valid', rec)
    return ('test', rec)

def yield_records_for_split(x, desired_split):
    split, rec = x
    # print(split, desired_split, split == desired_split)
    if split == desired_split:
        yield rec

def write_records(OUTPUT_DIR, splits, split):
    # same 80:10:10 split
    # The flowers dataset takes about 1GB, so 20 files means 50MB each
    nshards = 16 if (split == 'train') else 2
    _ = (splits
         | 'only_{}'.format(split) >> beam.FlatMap(
             lambda x: yield_records_for_split(x, split))
         | 'write_{}'.format(split) >> beam.io.tfrecordio.WriteToTFRecord(
             os.path.join(OUTPUT_DIR, split),
             file_name_suffix='.gz', num_shards=nshards)
        )

def decode_image(img_bytes):
    img = tf.image.decode_jpeg(img_bytes, channels=IMG_CHANNELS)
    return img

def tft_preprocess(img_record): 
    # tft_preprocess gets a batch, but decode_jpeg can only read individual files
    img = tf.map_fn(decode_image, img_record['img_bytes'],
                    fn_output_signature=tf.float32)
    img = tf.image.convert_image_dtype(img, tf.float32) # [0,1]
    img = tf.image.resize_with_pad(img, IMG_HEIGHT, IMG_WIDTH)
    return {
        'image': img,
        'label': img_record['label'],
        'label_int': img_record['label_int']
    }

def create_input_record(filename, label):
    label_list = label.to_pylist()
    filename_list = filename.to_pylist()
    assert len(filename_list) == 1 and len(filename_list[0]) == 1
    assert len(label_list) == 1 and len(label_list[0]) == 1
    contents = tf.io.read_file(filename_list[0][0]).numpy()
    return {
        'img_bytes': contents,
        'label': label_list[0][0],
        'label_int': LABELS.index(label_list[0][0].decode())
    }

def run_main(arguments):
    global IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS, LABELS
    
    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']
        on_cloud = (RUNNER == 'DataflowRunner')
    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)
   
    # tf.config.run_functions_eagerly(not on_cloud)

    # read list of labels
    with tf.io.gfile.GFile(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)

    # resize the input images
    ht, wd = arguments['resize'].split(',')
    IMG_HEIGHT = int(ht)
    IMG_WIDTH = int(wd)
    print("Will resize input images to {}x{}".format(IMG_HEIGHT, IMG_WIDTH))
        
    # make it repeatable
    np.random.seed(10)

    # 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,
        'max_num_workers': 20, # autoscale up to 20
        'region': arguments['region'],
        'teardown_policy': 'TEARDOWN_ALWAYS',
        'save_main_session': True,
        'requirements_file': 'requirements.txt'
    }
    opts = beam.pipeline.PipelineOptions(flags=[], **options)

    RAW_DATA_SCHEMA = tft.tf_metadata.dataset_schema.schema_utils.schema_from_feature_spec({
            'filename': tf.io.FixedLenFeature([], tf.string),
            'label': tf.io.FixedLenFeature([], tf.string),
        })
    IMG_BYTES_METADATA = tft.tf_metadata.dataset_metadata.DatasetMetadata(
        tft.tf_metadata.dataset_schema.schema_utils.schema_from_feature_spec({
            'img_bytes': tf.io.FixedLenFeature([], tf.string),
            'label': tf.io.FixedLenFeature([], tf.string),
            'label_int': tf.io.FixedLenFeature([], tf.int64)
        })
    )
    csv_tfxio = tfxio.CsvTFXIO(file_pattern=arguments['all_data'],
                               column_names=['filename', 'label'],
                               schema=RAW_DATA_SCHEMA,
                               telemetry_descriptors=['standalone_tft'])
    with beam.Pipeline(RUNNER, options=opts) as p:
        with tft_beam.Context(temp_dir=os.path.join(OUTPUT_DIR, 'tmp', 'beam_context')):
            img_records = (p
                      | 'read_csv' >> csv_tfxio.BeamSource(batch_size=1)
                      | 'img_record' >> beam.Map(
                          lambda x: create_input_record(x[0], x[1])))

            # tf.transform preprocessing
            # note that our preprocessing is simply to resize the images
            # so there is no need to be careful to run analysis only on training data

            # Ideally, we could have done csv_tfxio.TensorAdapterConfig()
            # but here, we are processing bytes, not the filenames we read from CSV
            raw_dataset = (img_records, IMG_BYTES_METADATA)

            transformed_dataset, transform_fn = (
                raw_dataset | 'tft_img' >> tft_beam.AnalyzeAndTransformDataset(tft_preprocess)
            )
            transformed_data, transformed_metadata = transformed_dataset
            transformed_data_coder = tft.coders.ExampleProtoCoder(transformed_metadata.schema)

            # write the cropped images
            splits = (transformed_data
                      | 'create_tfr' >> beam.Map(transformed_data_coder.encode)
                      | 'assign_ds' >> beam.Map(assign_record_to_split)
                      )

            for split in ['train', 'valid', 'test']:
                write_records(OUTPUT_DIR, splits, split)

            # make sure to write out a SavedModel with the tf transforms that were carried out
            _ = (
                transform_fn | 'write_tft' >> tft_beam.WriteTransformFn(
                    os.path.join(OUTPUT_DIR, 'tft'))
            )

            if on_cloud:
                print("Submitting {} job: {}".format(RUNNER, JOBNAME))
                print("Monitor at https://console.cloud.google.com/dataflow/jobs")
            else:
                print("Running on DirectRunner. Please hold on ...")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--all_data',
        # 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(
        '--region',
        help='Cloud Region to run in. Ignored for DirectRunner',
        default='us-central1')
    parser.add_argument(
        '--resize',
        help='Specify the img_height,img_width that you want images resized.',
        default='448,448')
    parser.add_argument(
        '--output_dir', help='Top-level directory for TF Records', required=True)

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