data_path = "gs://beans-lowres/tfrecords"
local_data_path = "data"

model_file = "modules/model.py"
model_fn = "modules.model.run_fn"

proprocessing_file = "modules/preprocessing.py"
preprocessing_fn = "modules.preprocessing.preprocessing_fn"

schema_file = "schema.pbtxt"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 StatisticsGen
from tfx.components import SchemaGen
from tfx.components import ExampleValidator
from tfx.components import Transform
from tfx.components import Trainer
from tfx.components import Pusher
from tfx.proto import example_gen_pb2context = 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'])%%writefile {schema_file}

feature {
  name: "image"
  type: FLOAT

  presence {
    min_fraction: 1.0
  }

  float_domain {
    min: 0
    max: 255
  }

  shape {
    dim { size: 256 }
    dim { size: 256 }
    dim { size: 3 }
  }
}

feature {
  name: "image_shape"
  type: INT

  presence {
    min_fraction: 1.0
  }

  shape {
    dim { size: 3 }
  }
}

feature {
  name: "label"
  type: INT

  presence {
    min_fraction: 1.0
  }

  int_domain {
    min: 0
    max: 2
  }

  shape {
    dim { size: 1 }
  }  
}schema_gen = tfx.components.ImportSchemaGen(
    schema_file=schema_file)

context.run(schema_gen)example_validator = ExampleValidator(
  statistics=statistics_gen.outputs['statistics'],
  schema=schema_gen.outputs['schema']   
)

context.run(example_validator)context.show(example_validator.outputs['anomalies'])%%writefile {proprocessing_file}

import tensorflow as tf

IMAGE_KEY = "image"
LABEL_KEY = "label"

MODEL_INPUT_IMAGE_KEY = "pixel_values"
MODEL_INPUT_LABEL_KEY = "labels"

INPUT_IMG_SIZE = 224

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 = {}

    inputs[IMAGE_KEY] = tf.image.resize(
        inputs[IMAGE_KEY], [INPUT_IMG_SIZE, INPUT_IMG_SIZE]
    )

    inputs[IMAGE_KEY] = inputs[IMAGE_KEY] / 255.0
    inputs[IMAGE_KEY] = tf.transpose(inputs[IMAGE_KEY], [0, 3, 1, 2])

    outputs[MODEL_INPUT_IMAGE_KEY] = inputs[IMAGE_KEY]
    outputs[MODEL_INPUT_LABEL_KEY] = inputs[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 transformers import ViTFeatureExtractor, TFViTForImageClassification
from tfx.components.trainer.fn_args_utils import FnArgs

from tfx_bsl.tfxio import dataset_options
from tfx.components.trainer.fn_args_utils import DataAccessor

feature_extractor = ViTFeatureExtractor()

_TRAIN_LENGTH = 128
_EVAL_LENGTH = 128
_TRAIN_BATCH_SIZE = 8
_EVAL_BATCH_SIZE = 8

_EPOCHS = 1

_LABELS = ['angular_leaf_spot', 'bean_rust', 'healthy']

_CONCRETE_INPUT = "pixel_values"

_MODEL_INPUT_LABEL_KEY = "labels"

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

def _normalize_img(
    img, mean=feature_extractor.image_mean, std=feature_extractor.image_std
):
    img = img / 255
    mean = tf.constant(mean)
    std = tf.constant(std)
    return (img - mean) / std


def _preprocess_serving(string_input):
    decoded_input = tf.io.decode_base64(string_input)
    decoded = tf.io.decode_jpeg(decoded_input, channels=3)
    resized = tf.image.resize(decoded, size=(224, 224))
    normalized = _normalize_img(resized)
    normalized = tf.transpose(
        normalized, (2, 0, 1)
    )  # Since HF models are channel-first.
    return normalized


@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])
def _preprocess_fn(string_input):
    decoded_images = tf.map_fn(
        _preprocess_serving, 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, 3, 224, 224], dtype=tf.float32, name=_CONCRETE_INPUT
        )
    )

    @tf.function(input_signature=[tf.TensorSpec([None], tf.string)])
    def serving_fn(string_input):
        labels = tf.constant(list(model.config.id2label.values()), dtype=tf.string)
        images = _preprocess_fn(string_input)

        predictions = m_call(**images)
        indices = tf.argmax(predictions.logits, axis=1)
        pred_source = tf.gather(params=labels, indices=indices)
        probs = tf.nn.softmax(predictions.logits, axis=1)
        pred_confidence = tf.reduce_max(probs, axis=1)
        return {"label": pred_source, "confidence": pred_confidence}

    return serving_fn


def _input_fn(
    file_pattern: List[str],
    data_accessor: DataAccessor,
    tf_transform_output: tft.TFTransformOutput,
    is_train: bool = False,
    batch_size: int = 32,
) -> tf.data.Dataset:
    INFO(f"Reading data from: {file_pattern}")

    dataset = data_accessor.tf_dataset_factory(
        file_pattern,
        dataset_options.TensorFlowDatasetOptions(
            batch_size=batch_size, label_key=_MODEL_INPUT_LABEL_KEY
        ),
        tf_transform_output.transformed_metadata.schema,
    )

    return dataset

def _build_model():
  id2label={str(i): c for i, c in enumerate(_LABELS)}
  label2id={c: str(i) for i, c in enumerate(_LABELS)}

  model = TFViTForImageClassification.from_pretrained(
    "google/vit-base-patch16-224-in21k",
    num_labels=len(_LABELS),
    label2id=label2id,
    id2label=id2label,
  )

  model.layers[0].trainable=False

  loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
  model.compile(optimizer='adam', loss=loss, metrics=["accuracy"])

  return model  

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,
    )

    model = _build_model()

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

    model.save(
        fn_args.serving_model_dir, 
        save_format="tf", 
        signatures=_model_exporter(model)
    )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)