code: def parse_function(image_size, raw_image_key_name):
      """Generate parse function for parsing the TFRecord training dataset.
  
      Read the image example and resize it to desired size.
  
      Args:
        image_size: int, target size to resize the image to
        raw_image_key_name: str, name of the JPEG image in each TFRecord entry
  
      Returns:
        A map function to use with tf.data.Dataset.map() .
      """
  
      def func(example_proto):
          """A generator to be used as representative_dataset for TFLiteConverter."""
          image_raw = tf.io.parse_single_example(
              example_proto,
              features={
                  raw_image_key_name: tf.FixedLenFeature([], tf.string)
              },
          )
          image = tf.image.decode_jpeg(image_raw[raw_image_key_name])
          image = tf.expand_dims(image, axis=0)
          image = tf.image.resize_bilinear(
              image, (image_size, image_size)
          )
          image = tf.squeeze(image, axis=0)
          image = image / 255.0
          return image
  
      return func
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import tempfile
  from magenta.models.arbitrary_image_stylization import arbitrary_image_stylization_build_mobilenet_model
  import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 5 for calling a global function with example usage: parse_function(image_size, raw_image_key_name) and returns: func, image