code: def _get_infer_batches(tf_dataset, model, batch_size):
      """Provides batches of pileup images from this dataset.
  
      This instantiates an iterator on the dataset, and returns the
      image, variant, alt_allele_indices, features in batches. It calls
      model.preprocess_images on the images (but note that we will be moving
      that step into model_fn for the Estimator api).
  
      Args:
        tf_dataset: DeepVariantInput.
        model: DeepVariantModel.
        batch_size: int.  The batch size.
  
      Returns:
        (image, variant, alt_allele_indices)
  
      Raises:
        ValueError: if the dataset has the wrong mode.
      """
      if tf_dataset.mode != tf_estimator.ModeKeys.PREDICT:
          raise ValueError(
              "tf_dataset.mode is {} but must be PREDICT.".format(
                  tf_dataset.mode
              )
          )
  
      params = dict(batch_size=batch_size)
      features = tf.compat.v1.data.make_one_shot_iterator(
          tf_dataset(params)
      ).get_next()
  
      images = features["image"]
      if tf_dataset.tensor_shape:
          # tensor_shape will be None if the input was an empty file.
          images = model.preprocess_images(images)
      variant = features["variant"]
      alt_allele_indices = features["alt_allele_indices"]
  
      return images, variant, alt_allele_indices
  
imports: import collections
  import errno
  import sys
  from unittest import mock
  from absl import flags
  from absl import logging
  from absl.testing import absltest
  from absl.testing import flagsaver
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tensorflow import estimator
  from deepvariant import call_variants_slim
  from deepvariant import dv_utils
  from deepvariant import modeling
  from deepvariant import testdata
  from deepvariant.protos import deepvariant_pb2
  from third_party.nucleus.io import tfrecord
  from third_party.nucleus.testing import test_utils
  from third_party.nucleus.util import variant_utils
label_desc: Use image node_id 2 for calling a global function with example usage: _get_infer_batches(tf_dataset, model, batch_size) and returns: images,  variant,  alt_allele_indices