# Path: meta_dataset/data/decoder.py
# def read_single_example(example_string):
# def read_example_and_parse_image(example_string):
#   def __init__(self,
#                image_size=None,
#                data_augmentation=None,
#                skip_tfexample_decoding=False):
#   def __call__(self, example_string):
#   def decode_with_label(self, example_string):
#   def __init__(self, feat_len):
#   def __call__(self, example_string):
#   def __init__(self):
#   def __call__(self, example_string):
# class ImageDecoder(object):
# class FeatureDecoder(object):
# class StringDecoder(object):
# 
# Path: meta_dataset/data/utils.py
# TRAIN_SUFFIX = 'train.tfrecords'
# TEST_SUFFIX = 'test.tfrecords'
# FILE_NAME_TEMPLATE = 'episode-{:04d}-{}'
# def get_label_counts(labels):
# def dump_as_tfrecord(path, images, labels):
# def get_file_path(folder, idx, split):
# def get_info_path(folder):


import functools
import json
import math
import os
import tensorflow.compat.v1 as tf
from meta_dataset.data import decoder
from meta_dataset.data import utils
from task_adaptation import data_loader
    "dmlab",
    "dsprites(predicted_attribute='label_orientation', num_classes=16)",
    "dsprites(predicted_attribute='label_x_position', num_classes=16)",
    "smallnorb(predicted_attribute='label_azimuth')",
    "smallnorb(predicted_attribute='label_elevation')",
    "kitti(task='closest_vehicle_distance')",
]
VTAB_DATASETS = VTAB_NATURAL + VTAB_SPECIALIZED + VTAB_STRUCTURED


def read_episode_as_dataset(episodes_dir,
                            episode_index,
                            split,
                            with_info=False):
  """This function reads a single episode from the directory.

  Args:
    episodes_dir: str, directory that has tf_record files.
    episode_index: int, of the episode to be loaded.
    split: str, `test` or `train`
    with_info: bool, if True reads the json file in the folder and returns
      number of

  Returns:
    decoded_dataset: tf.data.Dataset, with `image` and `label` fields.
    num_images_per_class_dict: tf.data.Dataset, if with_info=True.

  Raises:
    ValueError, when `split` is not one of {'train', 'test'}.
  """

# Path: meta_dataset/data/decoder.py
# def read_single_example(example_string):
# def read_example_and_parse_image(example_string):
#   def __init__(self,
#                image_size=None,
#                data_augmentation=None,
#                skip_tfexample_decoding=False):
#   def __call__(self, example_string):
#   def decode_with_label(self, example_string):
#   def __init__(self, feat_len):
#   def __call__(self, example_string):
#   def __init__(self):
#   def __call__(self, example_string):
# class ImageDecoder(object):
# class FeatureDecoder(object):
# class StringDecoder(object):
# 
# Path: meta_dataset/data/utils.py
# TRAIN_SUFFIX = 'train.tfrecords'
# TEST_SUFFIX = 'test.tfrecords'
# FILE_NAME_TEMPLATE = 'episode-{:04d}-{}'
# def get_label_counts(labels):
# def dump_as_tfrecord(path, images, labels):
# def get_file_path(folder, idx, split):
# def get_info_path(folder):


# Path: meta_dataset/data/read_episodes.py
import functools
import json
import math
import os
import tensorflow.compat.v1 as tf
from meta_dataset.data import decoder
from meta_dataset.data import utils
from task_adaptation import data_loader

    "dmlab",
    "dsprites(predicted_attribute='label_orientation', num_classes=16)",
    "dsprites(predicted_attribute='label_x_position', num_classes=16)",
    "smallnorb(predicted_attribute='label_azimuth')",
    "smallnorb(predicted_attribute='label_elevation')",
    "kitti(task='closest_vehicle_distance')",
]
VTAB_DATASETS = VTAB_NATURAL + VTAB_SPECIALIZED + VTAB_STRUCTURED


def read_episode_as_dataset(episodes_dir,
                            episode_index,
                            split,
                            with_info=False):
  """This function reads a single episode from the directory.

  Args:
    episodes_dir: str, directory that has tf_record files.
    episode_index: int, of the episode to be loaded.
    split: str, `test` or `train`
    with_info: bool, if True reads the json file in the folder and returns
      number of

  Returns:
    decoded_dataset: tf.data.Dataset, with `image` and `label` fields.
    num_images_per_class_dict: tf.data.Dataset, if with_info=True.

  Raises:
    ValueError, when `split` is not one of {'train', 'test'}.
  """

  episode_path = utils.get_file_path(episodes_dir, episode_index, split)