text: import os
  import torch.utils.data as data
  import torchvision
  from torchvision import transforms
  import numpy as np
  
  from PIL import Image, ImageFile
  Image.MAX_IMAGE_PIXELS = None  # Disable DecompressionBombError
  ImageFile.LOAD_TRUNCATED_IMAGES = True  # Disable OSError: image file is truncated
  
  def train_transform():
      transform_list = [
          transforms.Resize(size=(512, 512)),
          transforms.RandomCrop(256),
          transforms.ToTensor()
      ]
      return transforms.Compose(transform_list)
  
  class FlatFolderDataset(data.Dataset):
      def __init__(self, root, transform):
          super(FlatFolderDataset, self).__init__()
          self.root = root
          self.paths = os.listdir(self.root)
          self.transform = transform
  
      def __getitem__(self, index):
          path = self.paths[index]
          img = Image.open(os.path.join(self.root, path)).convert('RGB')
          img = self.transform(img)
          return img
  
      def __len__(self):
          return len(self.paths)
  
      def name(self):
          return 'FlatFolderDataset'
  
  
          
  def test_transform(size, crop):
      transform_list = []
      if size != 0:
          transform_list.append(transforms.Resize(size))
      if crop:
          transform_list.append(transforms.CenterCrop(size))
      transform_list.append(transforms.ToTensor())
      transform = transforms.Compose(transform_list)
      return transform
  
  def InfiniteSampler(n):
      # i = 0
      i = n - 1
      order = np.random.permutation(n)
      while True:
          yield order[i]
          i += 1
          if i >= n:
              np.random.seed()
              order = np.random.permutation(n)
              i = 0
  
  
  class InfiniteSamplerWrapper(data.sampler.Sampler):
      def __init__(self, data_source):
          self.num_samples = len(data_source)
  
      def __iter__(self):
          return iter(InfiniteSampler(self.num_samples))
  
      def __len__(self):
          return 2 ** 31