code: # This code is modified from https://github.com/facebookresearch/low-shot-shrink-hallucinate
  
  import torch
  from PIL import Image
  import numpy as np
  import torchvision.transforms as transforms
  import additional_transforms as add_transforms
  from abc import abstractmethod
  from torchvision.datasets import CIFAR100, CIFAR10
  
  identity = lambda x:x
  class SimpleDataset:
      def __init__(self, mode, dataset, transform, target_transform=identity):
          self.transform = transform
          self.dataset = dataset
          self.target_transform = target_transform
  
          self.meta = {}
  
          self.meta['image_names'] = []
          self.meta['image_labels'] = []
          if self.dataset == "CIFAR100":
  
              d = CIFAR100("./", train=True, download=True)
              for i, (data, label) in enumerate(d):
                  if mode == "base":
                      if label % 3 == 0:
                          self.meta['image_names'].append(data)
                          self.meta['image_labels'].append(label)
                  elif mode == "val":
                      if label % 3 == 1:
                          self.meta['image_names'].append(data)
                          self.meta['image_labels'].append(label)           
                  else:
                      if label % 3 == 2:
                          self.meta['image_names'].append(data)
                          self.meta['image_labels'].append(label)  
  
          elif self.dataset == "CIFAR10":
              d = CIFAR10("./", train=True, download=True)
              for i, (data, label) in enumerate(d):
                  if mode == "novel":
                      self.meta['image_names'].append(data)
                      self.meta['image_labels'].append(label)  
  
      def __getitem__(self, i):
  
          img = self.transform(self.meta['image_names'][i])
          target = self.target_transform(self.meta['image_labels'][i])
  
          return img, target
  
      def __len__(self):
          return len(self.meta['image_names'])
  
  
  class SetDataset:
      def __init__(self, mode, dataset, batch_size, transform):
  
          self.sub_meta = {}
          self.cl_list = range(100)
          self.dataset = dataset
  
          if mode == "base":
              type_ = 0
          elif mode == "val":
              type_ = 1
          else:
              type_ = 2
  
          for cl in self.cl_list:
              if cl % 3 == type_:
                  self.sub_meta[cl] = []
  
          if self.dataset == "CIFAR100":
              d = CIFAR100("./", train=True, download=True)
          elif self.dataset == "CIFAR10":
              d = CIFAR10("./", train=True, download=True)
  
  
          for i, (data, label) in enumerate(d):
              if label % 3 == type_:
                  self.sub_meta[label].append(data)
      
          self.sub_dataloader = [] 
          sub_data_loader_params = dict(batch_size = batch_size,
                                    shuffle = True,
                                    num_workers = 0, #use main thread only or may receive multiple batches
                                    pin_memory = False)        
          for cl in self.cl_list:
              if cl % 3 == type_:
                  sub_dataset = SubDataset(self.sub_meta[cl], cl, transform = transform )
                  self.sub_dataloader.append( torch.utils.data.DataLoader(sub_dataset, **sub_data_loader_params) )
  
      def __getitem__(self,i):
          return next(iter(self.sub_dataloader[i]))
  
      def __len__(self):
          return len(self.sub_dataloader)
  
  class SubDataset:
      def __init__(self, sub_meta, cl, transform=transforms.ToTensor(), target_transform=identity):
          self.sub_meta = sub_meta
          self.cl = cl 
          self.transform = transform
          self.target_transform = target_transform
  
      def __getitem__(self,i):
  
          img = self.transform(self.sub_meta[i])
          target = self.target_transform(self.cl)
          return img, target
  
      def __len__(self):
          return len(self.sub_meta)
  
  class EpisodicBatchSampler(object):
      def __init__(self, n_classes, n_way, n_episodes):
          self.n_classes = n_classes
          self.n_way = n_way
          self.n_episodes = n_episodes
  
      def __len__(self):
          return self.n_episodes
  
      def __iter__(self):
          for i in range(self.n_episodes):
              yield torch.randperm(self.n_classes)[:self.n_way]
  
  class TransformLoader:
      def __init__(self, image_size, 
                   normalize_param    = dict(mean= [0.485, 0.456, 0.406] , std=[0.229, 0.224, 0.225]),
                   jitter_param       = dict(Brightness=0.4, Contrast=0.4, Color=0.4)):
          self.image_size = image_size
          self.normalize_param = normalize_param
          self.jitter_param = jitter_param
      
      def parse_transform(self, transform_type):
          if transform_type=='ImageJitter':
              method = add_transforms.ImageJitter( self.jitter_param )
              return method
          method = getattr(transforms, transform_type)
          if transform_type=='RandomSizedCrop':
              return method(self.image_size) 
          elif transform_type=='CenterCrop':
              return method(self.image_size) 
          elif transform_type=='Scale':
              return method([int(self.image_size*1.15), int(self.image_size*1.15)])
          elif transform_type=='Normalize':
              return method(**self.normalize_param )
          else:
              return method()
  
      def get_composed_transform(self, aug = False):
          if aug:
              transform_list = ['RandomSizedCrop', 'ImageJitter', 'RandomHorizontalFlip', 'ToTensor', 'Normalize']
          else:
              transform_list = ['Scale','CenterCrop', 'ToTensor', 'Normalize']
  
          transform_funcs = [ self.parse_transform(x) for x in transform_list]
          transform = transforms.Compose(transform_funcs)
          return transform
  
  class DataManager(object):
      @abstractmethod
      def get_data_loader(self, data_file, aug):
          pass 
  
  class SimpleDataManager(DataManager):
      def __init__(self, dataset, image_size, batch_size):        
          super(SimpleDataManager, self).__init__()
          self.batch_size = batch_size
          self.trans_loader = TransformLoader(image_size)
          self.dataset = dataset
  
      def get_data_loader(self, mode, aug): #parameters that would change on train/val set
          transform = self.trans_loader.get_composed_transform(aug)
          dataset = SimpleDataset(mode, self.dataset, transform)
  
          data_loader_params = dict(batch_size = self.batch_size, shuffle = True, num_workers = 12, pin_memory = True)       
          data_loader = torch.utils.data.DataLoader(dataset, **data_loader_params)
  
          return data_loader
  
  class SetDataManager(DataManager):
      def __init__(self, mode, dataset, image_size, n_way=5, n_support=5, n_query=16, n_eposide = 100):        
          super(SetDataManager, self).__init__()
          self.image_size = image_size
          self.n_way = n_way
          self.batch_size = n_support + n_query
          self.n_eposide = n_eposide
          self.mode = mode
          self.dataset = dataset
  
          self.trans_loader = TransformLoader(image_size)
  
      def get_data_loader(self, aug): #parameters that would change on train/val set
          transform = self.trans_loader.get_composed_transform(aug)
          dataset = SetDataset(self.mode, self.dataset, self.batch_size, transform)
          sampler = EpisodicBatchSampler(len(dataset), self.n_way, self.n_eposide )  
          data_loader_params = dict(batch_sampler = sampler,  num_workers = 12, pin_memory = True)       
          data_loader = torch.utils.data.DataLoader(dataset, **data_loader_params)
          return data_loader
  
  if __name__ == '__main__':
      pass
apis: torch.utils.data.DataLoader torch.randperm