code: def labels_to_features_wrapper(data_fn):
      """Add wrapper to data_fn that add labels to features."""
  
      def wrapper(params, *args, **kwargs):
          """Wrapper for input_fn that adds contents of labels to features.labels."""
          # Workaround for Estimator API that forces 'labels' to be None when in
          # predict mode.
          # https://github.com/tensorflow/tensorflow/issues/17824
          # See also train_util.create_estimator
          assert not args
          dataset = data_fn(params=params, **kwargs)
          features_with_labels_type = collections.namedtuple(
              type(dataset.output_shapes[0]).__name__ + "WithLabels",
              dataset.output_shapes[0]._fields + ("labels",),
          )
  
          def add_labels_to_features(features, labels):
              features_dict = features._asdict()
              features_dict.update(labels=labels)
              return features_with_labels_type(**features_dict), labels
  
          return dataset.map(add_labels_to_features)
  
      return wrapper
  
file_path: magenta/magenta/models/onsets_frames_transcription/infer_util.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from magenta.models.onsets_frames_transcription import data
  from note_seq import sequences_lib
  import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: labels_to_features_wrapper(data_fn) and returns: wrapper, dataset, features_with_labels_type,  labels