code: def calculate_style_bottleneck(
      session,
      predict_saved_model,
      style_dataset,
      batch_size=10,
      min_sample_count=30,
  ):
      """Calculate style bottleneck using style predict SavedModel.
  
      Args:
        session: tf.Session, the session that contains subgraph to load the training
          dataset.
        predict_saved_model: str, path to the style predict SavedModel.
        style_dataset: tf.data.Dataset that contains training style images.
        batch_size: int, Batch size used when calculating style bottlenecks.
        min_sample_count: int, the minimum number of sample needed.
  
      Returns:
        [[[int]]] Value of style bottleneck for the style images being processed.
      """
  
      styles_bottleneck = None
  
      tf.logging.info("Started calculating style bottlenecks.")
  
      # Load and run the style predict SavedModel
      loaded = tf.saved_model.loader.load(
          session, ["serve"], predict_saved_model
      )
      bottleneck_tensor_name = (
          loaded.signature_def["serving_default"]
          .outputs["style_bottleneck"]
          .name
      )
  
      # Load the style dataset
      dataset = style_dataset.batch(batch_size)
      iterator = tf.data.make_initializable_iterator(dataset)
      session.run(iterator.initializer)
      next_element = iterator.get_next()
  
      for _ in range(np.math.ceil(min_sample_count / batch_size)):
          style_image = session.run(next_element)
  
          style_bottleneck_batch = session.run(
              bottleneck_tensor_name,
              feed_dict={"style_image:0": style_image},
          )
          if styles_bottleneck is None:
              styles_bottleneck = style_bottleneck_batch
          else:
              styles_bottleneck = np.append(
                  styles_bottleneck, style_bottleneck_batch, axis=0
              )
  
      tf.logging.info("Finished calculating style bottlenecks.")
  
      return styles_bottleneck
  
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 9 for calling a global function with example usage: calculate_style_bottleneck(session, predict_saved_model, style_dataset, batch_size, min_sample_count) and returns: styles_bottleneck