code: def get_augment(
      train: core.DataSet,
      extra: Optional[List[Type[AugmentPool]]] = None,
  ):
      pool, a = FLAGS.augment.split("(")
      a = a[:-1].split(",")
      a, kwargs = a[:3], {
          k: v for k, v in (x.split("=") for x in a[3:])
      }
      kwargs["tfops"] = tuple(
          get_tf_augment(ai, size=train.image_shape[0]) for ai in a
      )
      for x in list(globals().values()) + (extra or []):
          if inspect.isclass(x) and issubclass(x, AugmentPool):
              if x.NAME == pool:
                  pool = x
      return functools.partial(pool, **kwargs)
  
file_path: objax/examples/fixmatch/libml/augment/augment.py
imports: import functools
  import inspect
  import multiprocessing
  from typing import Tuple, Callable, Union, List, Type, Optional
  import numpy
  import tensorflow
  from absl import flags
  import objax
  from examples.fixmatch.libml.augment import ctaugment
  from examples.fixmatch.libml.augment.core import get_tf_augment
  from examples.fixmatch.libml.data import core
label_desc: Use image node_id 1 for calling a global function with example usage: get_augment(train, extra) and returns: functools