data_path = "gs://sidewalks-tfx-lowres/sidewalks-tfrecords"
local_data_path = "data"
model_file = "modules/model.py"
model_fn = "modules.model.run_fn"
preprocessing_file = "modules/preprocessing.py"
preprocessing_fn = "modules.preprocessing.preprocessing_fn"import tfx
tfx.__version__from tfx import v1 as tfx
from tfx.orchestration.experimental.interactive.interactive_context import InteractiveContext
from tfx.components import ImportExampleGen
from tfx.components import SchemaGen
from tfx.components import StatisticsGen
from tfx.components import Trainer
from tfx.components import Transform
from tfx.components import Evaluator
from tfx.proto import example_gen_pb2

from tfx.types import Channel
from tfx.types.standard_artifacts import Model
from tfx.types.standard_artifacts import ModelBlessing
from tfx.dsl.components.common import resolver
from tfx.dsl.experimental.latest_blessed_model_resolver import LatestBlessedModelResolver

import tensorflow_model_analysis as tfmacontext = InteractiveContext()input_config = example_gen_pb2.Input(
    splits=[
        example_gen_pb2.Input.Split(name="train", pattern="train-*.tfrec"),
        example_gen_pb2.Input.Split(name="eval", pattern="val-*.tfrec"),
    ]
)
example_gen = ImportExampleGen(
    input_base=local_data_path, 
    input_config=input_config
)context.run(example_gen)statistics_gen = StatisticsGen(examples=example_gen.outputs["examples"])context.run(statistics_gen)context.show(statistics_gen.outputs['statistics'])schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'])context.run(schema_gen)%%writefile {preprocessing_file}

import tensorflow as tf
from tensorflow.keras.applications import mobilenet_v2

_INPUT_IMG_SIZE = 128

_IMAGE_KEY = "image"
_IMAGE_SHAPE_KEY = "image_shape"
_LABEL_KEY = "label"
_LABEL_SHAPE_KEY = "label_shape"


def _transformed_name(key: str) -> str:
    return key + "_xf"


# output should have the same keys as inputs
def preprocess(inputs):
    image_shape = inputs[_IMAGE_SHAPE_KEY]
    label_shape = inputs[_LABEL_SHAPE_KEY]
    
    images = tf.reshape(inputs[_IMAGE_KEY], [image_shape[0], image_shape[1], 3])
    labels = tf.reshape(inputs[_LABEL_KEY], [label_shape[0], label_shape[1], 1])

    return {
        _IMAGE_KEY: images,
        _IMAGE_SHAPE_KEY: inputs[_IMAGE_SHAPE_KEY],
        _LABEL_KEY: labels,
        _LABEL_SHAPE_KEY: inputs[_LABEL_SHAPE_KEY],
    }


def preprocessing_fn(inputs):
    """tf.transform's callback function for preprocessing inputs.
    Args:
      inputs: map from feature keys to raw not-yet-transformed features.
    Returns:
      Map from string feature key to transformed feature operations.
    """
    # print(inputs)
    outputs = {}

    features = tf.map_fn(preprocess, inputs)

    features[_IMAGE_KEY] = tf.image.resize(features[_IMAGE_KEY], [_INPUT_IMG_SIZE, _INPUT_IMG_SIZE])
    features[_LABEL_KEY] = tf.image.resize(features[_LABEL_KEY], [_INPUT_IMG_SIZE, _INPUT_IMG_SIZE])

    image_features = mobilenet_v2.preprocess_input(features[_IMAGE_KEY])

    outputs[_transformed_name(_IMAGE_KEY)] = image_features
    outputs[_transformed_name(_LABEL_KEY)] = features[_LABEL_KEY]

    return outputstransform = Transform(
    examples=example_gen.outputs["examples"],
    schema=schema_gen.outputs["schema"],
    preprocessing_fn=preprocessing_fn,
)context.run(transform)%%writefile {model_file}

from typing import List, Dict, Tuple

import absl
import tensorflow as tf
import tensorflow_transform as tft
from tensorflow.keras.optimizers import Adam
from tfx_bsl.tfxio import dataset_options
from tfx.components.trainer.fn_args_utils import FnArgs
from tfx.components.trainer.fn_args_utils import DataAccessor

_CONCRETE_INPUT = "pixel_values"
_RAW_IMG_SIZE = 256
_INPUT_IMG_SIZE = 128
_TRAIN_LENGTH = 800
_EVAL_LENGTH = 200
_TRAIN_BATCH_SIZE = 64
_EVAL_BATCH_SIZE = 64
_EPOCHS = 1
_LR = 0.00006

_IMAGE_KEY = "image"
_LABEL_KEY = "label"

def INFO(text: str):
    absl.logging.info(text)

def _transformed_name(key: str) -> str:
    return key + "_xf"

"""
    _serving_preprocess, _serving_preprocess_fn, and 
    _model_exporter functions are defined to provide pre-
    processing capabilities when the model is served.
"""


def _serving_preprocess(string_input):
    decoded_input = tf.io.decode_base64(string_input)
    decoded = tf.io.decode_jpeg(decoded_input, channels=3)
    decoded = decoded / 255
    resized = tf.image.resize(decoded, size=(_INPUT_IMG_SIZE, _INPUT_IMG_SIZE))
    return resized


@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])
def _serving_preprocess_fn(string_input):
    decoded_images = tf.map_fn(
        _serving_preprocess, string_input, dtype=tf.float32, back_prop=False
    )
    return {_CONCRETE_INPUT: decoded_images}


def _model_exporter(model: tf.keras.Model):
    m_call = tf.function(model.call).get_concrete_function(
        tf.TensorSpec(shape=[None, _INPUT_IMG_SIZE, _INPUT_IMG_SIZE, 3], dtype=tf.float32, name=_CONCRETE_INPUT)
    )

    @tf.function(input_signature=[tf.TensorSpec([None], tf.string)])
    def serving_fn(string_input):
        images = _serving_preprocess_fn(string_input)
        logits = m_call(**images)
        seg_mask = tf.math.argmax(logits, -1)
        return {"seg_mask": seg_mask}

    return serving_fn

def _get_transform_features_signature(model, tf_transform_output):
    # the layer is added as an attribute to the model in order to make sure that
    # the model assets are handled correctly when exporting.    
    model.tft_layer = tf_transform_output.transform_features_layer()

    @tf.function(
        input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string, name="examples")]
    )
    def serve_tf_examples_fn(serialized_tf_examples):
        """Returns the output to be used in the serving signature."""
        feature_spec = tf_transform_output.raw_feature_spec()
        parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)

        transformed_features = model.tft_layer(parsed_features)

        return transformed_features
        
    return serve_tf_examples_fn

def _get_tf_examples_serving_signature(model, tf_transform_output):
    """
    Returns a serving signature that accepts `tensorflow.Example`.

    This signature will be used for evaluation or bulk inference.
    """
    
    model.tft_layer = tf_transform_output.transform_features_layer()

    @tf.function(
        input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string, name="examples")]
    )
    def serve_tf_examples_fn(
        serialized_tf_example: tf.Tensor,
    ) -> Dict[str, tf.Tensor]:
        """Returns the output to be used in the serving signature."""
        # Load the schema of raw examples.
        raw_feature_spec = tf_transform_output.raw_feature_spec()
        # Remove label feature since these will not be present at serving time.
        # Parse the examples using schema into raw features
        raw_features = tf.io.parse_example(serialized_tf_example, raw_feature_spec)

        # Preprocess the raw features
        transformed_features = model.tft_layer(raw_features)

        # Run preprocessed inputs through the model to get the prediction
        outputs = model(transformed_features)

        return {
            _transformed_name(_LABEL_KEY): outputs
        }

    return serve_tf_examples_fn


"""
    _input_fn reads TFRecord files passed from the upstream 
    TFX component, Transform. Assume the dataset is already 
    transformed appropriately. 
"""

def _input_fn(
    file_pattern: List[str],
    data_accessor: DataAccessor,
    tf_transform_output: tft.TFTransformOutput,
    is_train: bool = False,
    batch_size: int = 200,
) -> tf.data.Dataset:
    dataset = data_accessor.tf_dataset_factory(
        file_pattern,
        dataset_options.TensorFlowDatasetOptions(
            batch_size=batch_size, label_key=_transformed_name(_LABEL_KEY)
        ),
        tf_transform_output.transformed_metadata.schema,
    )

    return dataset


def _build_model(num_labels) -> tf.keras.Model:
    base_model = tf.keras.applications.MobileNetV2(
        input_shape=[128, 128, 3], include_top=False
    )

    # Use the activations of these layers
    layer_names = [
        "block_1_expand_relu",  # 64x64
        "block_3_expand_relu",  # 32x32
        "block_6_expand_relu",  # 16x16
        "block_13_expand_relu",  # 8x8
        "block_16_project",  # 4x4
    ]
    base_model_outputs = [base_model.get_layer(name).output for name in layer_names]

    # Create the feature extraction model
    down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs)
    down_stack.trainable = False

    up_stack = [
        upsample(512, 3),  # 4x4 -> 8x8
        upsample(256, 3),  # 8x8 -> 16x16
        upsample(128, 3),  # 16x16 -> 32x32
        upsample(64, 3),  # 32x32 -> 64x64
    ]

    inputs = tf.keras.layers.Input(
        shape=[128, 128, 3], name=_transformed_name(_IMAGE_KEY)
    )

    # Downsampling through the model
    skips = down_stack(inputs)
    x = skips[-1]
    skips = reversed(skips[:-1])

    # Upsampling and establishing the skip connections
    for up, skip in zip(up_stack, skips):
        x = up(x)
        concat = tf.keras.layers.Concatenate()
        x = concat([x, skip])

    # This is the last layer of the model
    last = tf.keras.layers.Conv2DTranspose(
        filters=num_labels, kernel_size=3, strides=2, padding="same", name=_transformed_name(_LABEL_KEY)
    )  # 64x64 -> 128x128

    x = last(x)

    model = tf.keras.Model(inputs=inputs, outputs=x)
    model.compile(
        optimizer=Adam(learning_rate=_LR),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=["sparse_categorical_accuracy"],
    )
    return model


"""
    InstanceNormalization class and upsample function are
    borrowed from pix2pix in [TensorFlow Example repository](
    https://github.com/tensorflow/examples/tree/master/tensorflow_examples/models/pix2pix)
"""


class InstanceNormalization(tf.keras.layers.Layer):
    """Instance Normalization Layer (https://arxiv.org/abs/1607.08022)."""

    def __init__(self, epsilon=1e-5):
        super(InstanceNormalization, self).__init__()
        self.epsilon = epsilon

    def build(self, input_shape):
        self.scale = self.add_weight(
            name="scale",
            shape=input_shape[-1:],
            initializer=tf.random_normal_initializer(1.0, 0.02),
            trainable=True,
        )

        self.offset = self.add_weight(
            name="offset", shape=input_shape[-1:], initializer="zeros", trainable=True
        )

    def call(self, x):
        mean, variance = tf.nn.moments(x, axes=[1, 2], keepdims=True)
        inv = tf.math.rsqrt(variance + self.epsilon)
        normalized = (x - mean) * inv
        return self.scale * normalized + self.offset


def upsample(filters, size, norm_type="batchnorm", apply_dropout=False):
    """Upsamples an input.
    Conv2DTranspose => Batchnorm => Dropout => Relu
    Args:
      filters: number of filters
      size: filter size
      norm_type: Normalization type; either 'batchnorm' or 'instancenorm'.
      apply_dropout: If True, adds the dropout layer
    Returns:
      Upsample Sequential Model
    """

    initializer = tf.random_normal_initializer(0.0, 0.02)

    result = tf.keras.Sequential()
    result.add(
        tf.keras.layers.Conv2DTranspose(
            filters,
            size,
            strides=2,
            padding="same",
            kernel_initializer=initializer,
            use_bias=False,
        )
    )

    if norm_type.lower() == "batchnorm":
        result.add(tf.keras.layers.BatchNormalization())
    elif norm_type.lower() == "instancenorm":
        result.add(InstanceNormalization())

    if apply_dropout:
        result.add(tf.keras.layers.Dropout(0.5))

    result.add(tf.keras.layers.ReLU())

    return result


def run_fn(fn_args: FnArgs):
    tf_transform_output = tft.TFTransformOutput(fn_args.transform_output)

    train_dataset = _input_fn(
        fn_args.train_files,
        fn_args.data_accessor,
        tf_transform_output,
        is_train=True,
        batch_size=_TRAIN_BATCH_SIZE,
    )

    eval_dataset = _input_fn(
        fn_args.eval_files,
        fn_args.data_accessor,
        tf_transform_output,        
        is_train=False,
        batch_size=_EVAL_BATCH_SIZE,
    )

    num_labels = 35
    model = _build_model(num_labels)

    model.fit(
        train_dataset,
        steps_per_epoch=_TRAIN_LENGTH // _TRAIN_BATCH_SIZE,
        validation_data=eval_dataset,
        validation_steps=_EVAL_LENGTH // _TRAIN_BATCH_SIZE,
        epochs=_EPOCHS,
    )

    model.save(
        fn_args.serving_model_dir, save_format="tf", signatures={
            "serving_default": _model_exporter(model),
            "transform_features": _get_transform_features_signature(
                model, tf_transform_output
            ),
            "from_examples": _get_tf_examples_serving_signature(
                model, tf_transform_output
            ),            
        }
    )
trainer = Trainer(
    run_fn=model_fn,
    transformed_examples=transform.outputs["transformed_examples"],
    transform_graph=transform.outputs["transform_graph"],
    schema=schema_gen.outputs["schema"],
)context.run(trainer)