# 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/read_episodes.py
# VTAB_NATURAL = [
#     "caltech101",
#     "cifar(num_classes=100)",
#     "dtd",
#     "oxford_flowers102",
#     "oxford_iiit_pet",
#     "sun397",
#     "svhn"
# ]
# VTAB_SPECIALIZED = [
#     "diabetic_retinopathy(config='btgraham-300')",
#     "eurosat",
#     "resisc45",
#     "patch_camelyon",
# ]
# VTAB_STRUCTURED = [
#     "clevr(task='closest_object_distance')",
#     "clevr(task='count_all')",
#     "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):
# def read_episodes_from_records(episodes_dir,
#                                train_suffix=utils.TRAIN_SUFFIX,
#                                test_suffix=utils.TEST_SUFFIX):
#   def get_total_img_count(file_name):
#   def _load_and_batch(file_path, n_images):
# def read_episodes_from_records_multiple_sources(episodes_dir_list,
#                                                 train_suffix=utils.TRAIN_SUFFIX,
#                                                 test_suffix=utils.TEST_SUFFIX):
# def read_vtab_as_episode(vtab_key,
#                          image_size=224,
#                          query_size_limit=500,
#                          data_dir=None):
#   def cast_label_fn(ex):
# 
# 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 collections
import json
import os
import random
import absl.testing.parameterized as parameterized
import tensorflow.compat.v1 as tf
from meta_dataset.data import decoder
from meta_dataset.data import read_episodes
from meta_dataset.data import utils
from task_adaptation import data_loader
        self.assertAllEqual(label1, label2)
      test_loader = read_episodes.read_episode_as_dataset(record_dir, i, 'test')
      for (example, (img2, label2)) in zip(test_loader, test_ds):
        img1, label1 = example['image'], example['label']
        self.assertAllEqual(img1, img2)
        self.assertAllEqual(label1, label2)

  def test_single_dataset_with_info(self):
    record_dir, episodes = self.dump_episodes(2)
    for i, (train_ds, test_ds) in enumerate(episodes):
      _, num_imgs_per_class = read_episodes.read_episode_as_dataset(
          record_dir, i, 'train', with_info=True)
      all_train_labels = [label2.numpy() for _, label2 in train_ds]
      expected_counts_train = dict(collections.Counter(all_train_labels))
      self.assertDictEqual(expected_counts_train, num_imgs_per_class)
      _, num_imgs_per_class = read_episodes.read_episode_as_dataset(
          record_dir, i, 'test', with_info=True)
      all_test_labels = [label2.numpy() for _, label2 in test_ds]
      expected_counts_test = dict(collections.Counter(all_test_labels))
      self.assertDictEqual(expected_counts_test, num_imgs_per_class)

  def test_read_as_meta_dataset(self):
    record_dir, episodes = self.dump_episodes(2)
    episode_ds, n_episodes = read_episodes.read_episodes_from_records(
        record_dir)
    total_episodes = 0
    for episode, (train_ds, test_ds) in zip(episode_ds, episodes):
      total_episodes += 1
      train_episode, test_episode = episode
      for (example_string, (img2, label2)) in zip(train_episode, train_ds):

# 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/read_episodes.py
# VTAB_NATURAL = [
#     "caltech101",
#     "cifar(num_classes=100)",
#     "dtd",
#     "oxford_flowers102",
#     "oxford_iiit_pet",
#     "sun397",
#     "svhn"
# ]
# VTAB_SPECIALIZED = [
#     "diabetic_retinopathy(config='btgraham-300')",
#     "eurosat",
#     "resisc45",
#     "patch_camelyon",
# ]
# VTAB_STRUCTURED = [
#     "clevr(task='closest_object_distance')",
#     "clevr(task='count_all')",
#     "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):
# def read_episodes_from_records(episodes_dir,
#                                train_suffix=utils.TRAIN_SUFFIX,
#                                test_suffix=utils.TEST_SUFFIX):
#   def get_total_img_count(file_name):
#   def _load_and_batch(file_path, n_images):
# def read_episodes_from_records_multiple_sources(episodes_dir_list,
#                                                 train_suffix=utils.TRAIN_SUFFIX,
#                                                 test_suffix=utils.TEST_SUFFIX):
# def read_vtab_as_episode(vtab_key,
#                          image_size=224,
#                          query_size_limit=500,
#                          data_dir=None):
#   def cast_label_fn(ex):
# 
# 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/dump_and_read_episodes_test.py
import collections
import json
import os
import random
import absl.testing.parameterized as parameterized
import tensorflow.compat.v1 as tf
from meta_dataset.data import decoder
from meta_dataset.data import read_episodes
from meta_dataset.data import utils
from task_adaptation import data_loader

        self.assertAllEqual(label1, label2)
      test_loader = read_episodes.read_episode_as_dataset(record_dir, i, 'test')
      for (example, (img2, label2)) in zip(test_loader, test_ds):
        img1, label1 = example['image'], example['label']
        self.assertAllEqual(img1, img2)
        self.assertAllEqual(label1, label2)

  def test_single_dataset_with_info(self):
    record_dir, episodes = self.dump_episodes(2)
    for i, (train_ds, test_ds) in enumerate(episodes):
      _, num_imgs_per_class = read_episodes.read_episode_as_dataset(
          record_dir, i, 'train', with_info=True)
      all_train_labels = [label2.numpy() for _, label2 in train_ds]
      expected_counts_train = dict(collections.Counter(all_train_labels))
      self.assertDictEqual(expected_counts_train, num_imgs_per_class)
      _, num_imgs_per_class = read_episodes.read_episode_as_dataset(
          record_dir, i, 'test', with_info=True)
      all_test_labels = [label2.numpy() for _, label2 in test_ds]
      expected_counts_test = dict(collections.Counter(all_test_labels))
      self.assertDictEqual(expected_counts_test, num_imgs_per_class)

  def test_read_as_meta_dataset(self):
    record_dir, episodes = self.dump_episodes(2)
    episode_ds, n_episodes = read_episodes.read_episodes_from_records(
        record_dir)
    total_episodes = 0
    for episode, (train_ds, test_ds) in zip(episode_ds, episodes):
      total_episodes += 1
      train_episode, test_episode = episode
      for (example_string, (img2, label2)) in zip(train_episode, train_ds):

        example = decoder.read_example_and_parse_image(example_string)