text: import torch
  import torch.utils.data as data
  
  import os, math, random
  from os.path import *
  import numpy as np
  from tqdm import tqdm
  
  from glob import glob
  import utils.frame_utils as frame_utils
  from  utils.tools import find_le
  
  #from scipy.misc import imread, imresize
  import cv2
  
  def has_file_allowed_extension(filename, extensions):
      """Checks if a file is an allowed extension.
  
      Args:
          filename (string): path to a file
          extensions (iterable of strings): extensions to consider (lowercase)
  
      Returns:
          bool: True if the filename ends with one of given extensions
      """
      filename_lower = filename.lower()
      return any(filename_lower.endswith(ext) for ext in extensions)
  
  
  def is_image_file(filename):
      """Checks if a file is an allowed image extension.
  
      Args:
          filename (string): path to a file
  
      Returns:
          bool: True if the filename ends with a known image extension
      """
      return has_file_allowed_extension(filename, IMG_EXTENSIONS)
  
  def list_directories(root, full_path = False):
      if full_path:
          return [join(root, d) for d in os.listdir(root) if exists(join(root, d))]
      else:
          return [d for d in os.listdir(root) if exists(join(root, d))]
  
  def list_images(root, full_path = True):
      if full_path:
          return [join(root, i) for i in os.listdir(root) if exists(join(root, i))]
      else:
          return [i for i in os.listdir(root) if is_image_file(join(root, i))]
  
  IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']
  VIDEO_EXTENSION = ['.mp4', '.mov']
  
  class StaticRandomCrop(object):
      def __init__(self, image_size, crop_size):
          self.th, self.tw = crop_size
          h, w = image_size
          self.h1 = random.randint(0, h - self.th)
          self.w1 = random.randint(0, w - self.tw)
  
      def __call__(self, img):
          return img[self.h1:(self.h1+self.th), self.w1:(self.w1+self.tw),:]
  
  class StaticCenterCrop(object):
      def __init__(self, image_size, crop_size):
          self.th, self.tw = crop_size
          self.h, self.w = image_size
      def __call__(self, img):
          return img[(self.h-self.th)//2:(self.h+self.th)//2, (self.w-self.tw)//2:(self.w+self.tw)//2,:]
  
  class TVQA(data.Dataset):
      def __init__(self, args, is_cropped = False, root = '', dstype = '', replicates = 1):
          self.args = args
          self.is_cropped = is_cropped
          self.crop_size = args.crop_size
          self.render_size = args.inference_size
          self.replicates = replicates
  
          self.loc = join('./preprocessed/lists', 'tvqa')
          if not isdir(self.loc):
              os.makedirs(self.loc)
  
          self.root = root
  
          self.make_dataset()
          self._integrate_lengths()
  
      def make_dataset(self):
          video_file = join(self.loc, 'video_list.npy')
          video_len_file = join(self.loc, 'video_lengths.npy')
          video_image_file = join(self.loc, 'image_list.npy')
  
          if exists(video_file):
              self.video_list = np.load(video_file, allow_pickle = True)
              self.video_lengths = np.load(video_len_file, allow_pickle = True)
              self.image_list = np.load(video_image_file, allow_pickle = True)
              return
  
          self.video_list = []
          self.video_lengths = []
          self.image_list = []
          tvshows = list_directories(self.root, full_path = True)
          for show in tqdm(tvshows):
              clips = list_directories(show, full_path = True)
              for clip in clips:
                  self.video_list.append(clip)
                  image_list = list_images(join(self.root, show, clip), full_path = True)
                  image_list = sorted(image_list)
                  self.image_list.append(image_list)
                  self.video_lengths.append(len(image_list))
  
          np.save(video_file, self.video_list)
          np.save(video_len_file, self.video_lengths)
          np.save(video_image_file, self.image_list)
  
      def _integrate_lengths(self):
          # Save video_lengths as an integral array for faster lookups
          total = 0
          for idx in range(len(self.video_lengths)):
              total += self.video_lengths[idx] - 1
              self.video_lengths[idx] = total
  
      def __len__(self):
          return self.video_lengths[-1]
  
      def _find_index(self, index):
          if index >= self.video_lengths[0]:
              video_idx = find_le(self.video_lengths, index)
              image_idx = index - self.video_lengths[video_idx]
              video_idx += 1
          else:
              video_idx = 0
              image_idx = index
  
          return video_idx, image_idx
  
      def __getitem__(self, index):
          video_idx, image_idx = self._find_index(index)
          video_path = self.video_list[video_idx]
  
          img1 = cv2.imread(self.image_list[video_idx][image_idx])
          img2 = cv2.imread(self.image_list[video_idx][image_idx + 1])
  
          if img1 is None:
              np.zeros((240, 360, 3))
  
          if img2 is None:
              img2 = np.zeros_like(img1)
  
          if (img1.shape[0] < self.render_size[0]):
              scale = self.render_size[0] / img1.shape[0]
              img1 = cv2.resize(img1, None, fx = scale, fy = scale)
              img2 = cv2.resize(img2, None, fx = scale, fy = scale)
  
          if (img1.shape[1] < self.render_size[1]):
              scale = self.render_size[1] / img1.shape[1]
              img1 = cv2.resize(img1, None, fx = scale, fy = scale)
              img2 = cv2.resize(img2, None, fx = scale, fy = scale)
  
          self.frame_size = img1.shape
          image_size = img1.shape[:2]
  
          if (self.render_size[0] < 0) or (self.render_size[1] < 0) \
             or (self.render_size[0] % 64) or (self.render_size[1] % 64):
              self.render_size[0] = ( (self.frame_size[0]) // 64) * 64
              self.render_size[1] = ( (self.frame_size[1]) // 64) * 64
  
          self.args.inference_size = self.render_size
  
          images = [img1, img2]
          flow = np.zeros((img1.shape[0], img1.shape[1], 2))
  
          if self.is_cropped:
              cropper = StaticRandomCrom(image_size, self.crop_size)
          else:
              cropper = StaticCenterCrop(image_size, self.render_size)
  
          images = list(map(cropper, images))
          flow = cropper(flow)
  
          images = np.array(images).transpose(3, 0, 1, 2)
          flow = flow.transpose(2, 0, 1)
  
          images = torch.from_numpy(images.astype(np.float32))
          flow = torch.from_numpy(flow.astype(np.float32))
  
          pre_path, post_path = video_path.split('frames_hq/')
          post_path = os.path.splitext(post_path)[0]
          out_video_path = join(pre_path, 'flow', post_path)
          if not os.path.exists(out_video_path):
              os.makedirs(out_video_path)
  
          flo_filename = join(out_video_path, "{:05d}.flo".format(image_idx))
          flo_rgb_filename = join(out_video_path, "{:05d}.jpg".format(image_idx))
          return [images], [flow], [flo_filename], [flo_rgb_filename]
  
  
  class Kinetics(data.Dataset):
      def __init__(self, args, is_cropped = False, root = '', dstype = '', replicates = 1):
          self.args = args
          self.is_cropped = is_cropped
          self.crop_size = args.crop_size
          self.render_size = args.inference_size
          self.replicates = replicates
  
          self.loc = join('./preprocessed/lists', 'kinetics-600')
          if not isdir(self.loc):
              os.makedirs(self.loc)
  
          self.root = root
  
          self.videos, self.video_lengths = self._parallelly_make_dataset()
          self._integrate_lengths()
  
      def _parallelly_make_dataset(self):
          import multiprocessing
          from multiprocessing import Process
          from multiprocessing import JoinableQueue as Queue
  
          name_file = '{}/video_list.npy'.format(self.loc)
          len_file = '{}/video_lengths.npy'.format(self.loc)
  
          if isfile(name_file):
              video_list = np.load(name_file)
              video_lengths = np.load(len_file)
              return video_list, video_lengths
  
          q = Queue()
          qvideo_list = Queue()
  
          fnames_list = []
          for root, _, fnames in tqdm(os.walk(self.root)):
              for fname in sorted(fnames):
                  fnames_list.append(os.path.join(root, fname))
  
          def parallel_worker(fnames_chunk):
              item = q.get()
              for fname in tqdm(fnames_chunk):
                  if has_file_allowed_extension(fname, VIDEO_EXTENSION):
                      video_path = fname
                      vc = cv2.VideoCapture(video_path)
                      length = int(vc.get(cv2.CAP_PROP_FRAME_COUNT))
                      if length > 0 and vc.isOpened():
                          qvideo_list.put((video_path, length))
                          qvideo_list.task_done()
                      vc.release()
              q.task_done()
  
          processes = 32
          n = len(fnames_list)
          chunk = int(n / processes)
          if chunk == 0:
              chunk = 1
          fnames_chunks = [fnames_list[i*chunk:(i+1)*chunk] \
                          for i in range((n + chunk - 1) // chunk)]
          for i in range(processes):
              q.put(i)
              multiprocessing.Process(target = parallel_worker,
                                      args = (fnames_chunks[i],)).start()
  
          q.join()
          qvideo_list.join()
  
          video_list = []
          video_lengths = []
  
          while qvideo_list.qsize() != 0:
              video, length = qvideo_list.get()
              video_list.append(video)
              video_lengths.append(length)
  
          np.save(name_file, video_list)
          np.save(len_file, video_lengths)
  
          return video_list, video_lengths
  
      def _integrate_lengths(self):
          # Save video_lengths as an integral array for faster lookups
          total = 0
          for idx in range(len(self.video_lengths)):
              total += self.video_lengths[idx]
              self.video_lengths[idx] = total
  
      def __len__(self):
          return self.video_lengths[-1]
  
      def _find_index(self, index):
          if index >= self.video_lengths[0]:
              video_idx = find_le(self.video_lengths, index)
              image_idx = index - self.video_lengths[video_idx]
              video_idx += 1
          else:
              video_idx = 0
              image_idx = index
  
          return video_idx, image_idx
  
      def __getitem__(self, index):
          video_idx, image_idx = self._find_index(index)
          video_path = self.videos[video_idx]
          vc = cv2.VideoCapture(self.videos[video_idx])
          length = int(vc.get(cv2.CAP_PROP_FRAME_COUNT))
  
          start_frame = 0
          return_val = True
          while start_frame < image_idx and return_val:
              return_val, frame = vc.read()
              start_frame += 1
  
          count = 0
          return_val = True
          ret, img1 = vc.read()
          ret, img2 = vc.read()
  
          if img1 is None:
              np.zeros((240, 360, 3))
  
          if img2 is None:
              img2 = np.zeros_like(img1)
  
          if (img1.shape[0] < self.render_size[0]):
              scale = self.render_size[0] / img1.shape[0]
              img1 = cv2.resize(img1, None, fx = scale, fy = scale)
              img2 = cv2.resize(img2, None, fx = scale, fy = scale)
  
          if (img1.shape[1] < self.render_size[1]):
              scale = self.render_size[1] / img1.shape[1]
              img1 = cv2.resize(img1, None, fx = scale, fy = scale)
              img2 = cv2.resize(img2, None, fx = scale, fy = scale)
  
          self.frame_size = img1.shape
          image_size = img1.shape[:2]
  
          if (self.render_size[0] < 0) or (self.render_size[1] < 0) \
             or (self.render_size[0] % 64) or (self.render_size[1] % 64):
              self.render_size[0] = ( (self.frame_size[0]) // 64) * 64
              self.render_size[1] = ( (self.frame_size[1]) // 64) * 64
  
          self.args.inference_size = self.render_size
  
          images = [img1, img2]
          flow = np.zeros((img1.shape[0], img1.shape[1], 2))
  
          if self.is_cropped:
              cropper = StaticRandomCrom(image_size, self.crop_size)
          else:
              cropper = StaticCenterCrop(image_size, self.render_size)
  
          images = list(map(cropper, images))
          flow = cropper(flow)
  
          images = np.array(images).transpose(3, 0, 1, 2)
          flow = flow.transpose(2, 0, 1)
  
          images = torch.from_numpy(images.astype(np.float32))
          flow = torch.from_numpy(flow.astype(np.float32))
  
          pre_path, post_path = video_path.split('train/')
          post_path = os.path.splitext(post_path)[0]
          out_video_path = join(pre_path, 'flow', post_path)
          if not os.path.exists(out_video_path):
              os.makedirs(out_video_path)
  
          flo_filename = join(out_video_path, "{:06d}.flo".format(image_idx))
          flo_rgb_filename = join(out_video_path, "{:06d}.png".format(image_idx))
          return [images], [flow], [flo_filename], [flo_rgb_filename]
  
  class MpiSintel(data.Dataset):
      def __init__(self, args, is_cropped = False, root = '', dstype = 'clean', replicates = 1):
          self.args = args
          self.is_cropped = is_cropped
          self.crop_size = args.crop_size
          self.render_size = args.inference_size
          self.replicates = replicates
  
          flow_root = join(root, 'flow')
          image_root = join(root, dstype)
  
          file_list = sorted(glob(join(flow_root, '*/*.flo')))
  
          self.flow_list = []
          self.image_list = []
  
          for file in file_list:
              if 'test' in file:
                  # print file
                  continue
  
              fbase = file[len(flow_root)+1:]
              fprefix = fbase[:-8]
              fnum = int(fbase[-8:-4])
  
              img1 = join(image_root, fprefix + "%04d"%(fnum+0) + '.png')
              img2 = join(image_root, fprefix + "%04d"%(fnum+1) + '.png')
  
              if not isfile(img1) or not isfile(img2) or not isfile(file):
                  continue
  
              self.image_list += [[img1, img2]]
              self.flow_list += [file]
  
          self.size = len(self.image_list)
  
          self.frame_size = frame_utils.read_gen(self.image_list[0][0]).shape
  
          if (self.render_size[0] < 0) or (self.render_size[1] < 0) or (self.frame_size[0]%64) or (self.frame_size[1]%64):
              self.render_size[0] = ( (self.frame_size[0])//64 ) * 64
              self.render_size[1] = ( (self.frame_size[1])//64 ) * 64
  
          args.inference_size = self.render_size
  
          assert (len(self.image_list) == len(self.flow_list))
  
      def __getitem__(self, index):
  
          index = index % self.size
  
          img1 = frame_utils.read_gen(self.image_list[index][0])
          img2 = frame_utils.read_gen(self.image_list[index][1])
  
          flow = frame_utils.read_gen(self.flow_list[index])
  
          images = [img1, img2]
          image_size = img1.shape[:2]
  
          if self.is_cropped:
              cropper = StaticRandomCrop(image_size, self.crop_size)
          else:
              cropper = StaticCenterCrop(image_size, self.render_size)
          images = list(map(cropper, images))
          flow = cropper(flow)
  
          images = np.array(images).transpose(3,0,1,2)
          flow = flow.transpose(2,0,1)
  
          images = torch.from_numpy(images.astype(np.float32))
          flow = torch.from_numpy(flow.astype(np.float32))
  
          return [images], [flow]
  
      def __len__(self):
          return self.size * self.replicates
  
  class MpiSintelClean(MpiSintel):
      def __init__(self, args, is_cropped = False, root = '', replicates = 1):
          super(MpiSintelClean, self).__init__(args, is_cropped = is_cropped, root = root, dstype = 'clean', replicates = replicates)
  
  class MpiSintelFinal(MpiSintel):
      def __init__(self, args, is_cropped = False, root = '', replicates = 1):
          super(MpiSintelFinal, self).__init__(args, is_cropped = is_cropped, root = root, dstype = 'final', replicates = replicates)
  
  class FlyingChairs(data.Dataset):
    def __init__(self, args, is_cropped, root = '/path/to/FlyingChairs_release/data', replicates = 1):
      self.args = args
      self.is_cropped = is_cropped
      self.crop_size = args.crop_size
      self.render_size = args.inference_size
      self.replicates = replicates
  
      images = sorted( glob( join(root, '*.ppm') ) )
  
      self.flow_list = sorted( glob( join(root, '*.flo') ) )
  
      assert (len(images)//2 == len(self.flow_list))
  
      self.image_list = []
      for i in range(len(self.flow_list)):
          im1 = images[2*i]
          im2 = images[2*i + 1]
          self.image_list += [ [ im1, im2 ] ]
  
      assert len(self.image_list) == len(self.flow_list)
  
      self.size = len(self.image_list)
  
      self.frame_size = frame_utils.read_gen(self.image_list[0][0]).shape
  
      if (self.render_size[0] < 0) or (self.render_size[1] < 0) or (self.frame_size[0]%64) or (self.frame_size[1]%64):
          self.render_size[0] = ( (self.frame_size[0])//64 ) * 64
          self.render_size[1] = ( (self.frame_size[1])//64 ) * 64
  
      args.inference_size = self.render_size
  
    def __getitem__(self, index):
      index = index % self.size
  
      img1 = frame_utils.read_gen(self.image_list[index][0])
      img2 = frame_utils.read_gen(self.image_list[index][1])
  
      flow = frame_utils.read_gen(self.flow_list[index])
  
      images = [img1, img2]
      image_size = img1.shape[:2]
      if self.is_cropped:
          cropper = StaticRandomCrop(image_size, self.crop_size)
      else:
          cropper = StaticCenterCrop(image_size, self.render_size)
      images = list(map(cropper, images))
      flow = cropper(flow)
  
  
      images = np.array(images).transpose(3,0,1,2)
      flow = flow.transpose(2,0,1)
  
      images = torch.from_numpy(images.astype(np.float32))
      flow = torch.from_numpy(flow.astype(np.float32))
  
      return [images], [flow]
  
    def __len__(self):
      return self.size * self.replicates
  
  class FlyingThings(data.Dataset):
    def __init__(self, args, is_cropped, root = '/path/to/flyingthings3d', dstype = 'frames_cleanpass', replicates = 1):
      self.args = args
      self.is_cropped = is_cropped
      self.crop_size = args.crop_size
      self.render_size = args.inference_size
      self.replicates = replicates
  
      image_dirs = sorted(glob(join(root, dstype, 'TRAIN/*/*')))
      image_dirs = sorted([join(f, 'left') for f in image_dirs] + [join(f, 'right') for f in image_dirs])
  
      flow_dirs = sorted(glob(join(root, 'optical_flow_flo_format/TRAIN/*/*')))
      flow_dirs = sorted([join(f, 'into_future/left') for f in flow_dirs] + [join(f, 'into_future/right') for f in flow_dirs])
  
      assert (len(image_dirs) == len(flow_dirs))
  
      self.image_list = []
      self.flow_list = []
  
      for idir, fdir in zip(image_dirs, flow_dirs):
          images = sorted( glob(join(idir, '*.png')) )
          flows = sorted( glob(join(fdir, '*.flo')) )
          for i in range(len(flows)):
              self.image_list += [ [ images[i], images[i+1] ] ]
              self.flow_list += [flows[i]]
  
      assert len(self.image_list) == len(self.flow_list)
  
      self.size = len(self.image_list)
  
      self.frame_size = frame_utils.read_gen(self.image_list[0][0]).shape
  
      if (self.render_size[0] < 0) or (self.render_size[1] < 0) or (self.frame_size[0]%64) or (self.frame_size[1]%64):
          self.render_size[0] = ( (self.frame_size[0])//64 ) * 64
          self.render_size[1] = ( (self.frame_size[1])//64 ) * 64
  
      args.inference_size = self.render_size
  
    def __getitem__(self, index):
      index = index % self.size
  
      img1 = frame_utils.read_gen(self.image_list[index][0])
      img2 = frame_utils.read_gen(self.image_list[index][1])
  
      flow = frame_utils.read_gen(self.flow_list[index])
  
      images = [img1, img2]
      image_size = img1.shape[:2]
      if self.is_cropped:
          cropper = StaticRandomCrop(image_size, self.crop_size)
      else:
          cropper = StaticCenterCrop(image_size, self.render_size)
      images = list(map(cropper, images))
      flow = cropper(flow)
  
  
      images = np.array(images).transpose(3,0,1,2)
      flow = flow.transpose(2,0,1)
  
      images = torch.from_numpy(images.astype(np.float32))
      flow = torch.from_numpy(flow.astype(np.float32))
  
      return [images], [flow]
  
    def __len__(self):
      return self.size * self.replicates
  
  class FlyingThingsClean(FlyingThings):
      def __init__(self, args, is_cropped = False, root = '', replicates = 1):
          super(FlyingThingsClean, self).__init__(args, is_cropped = is_cropped, root = root, dstype = 'frames_cleanpass', replicates = replicates)
  
  class FlyingThingsFinal(FlyingThings):
      def __init__(self, args, is_cropped = False, root = '', replicates = 1):
          super(FlyingThingsFinal, self).__init__(args, is_cropped = is_cropped, root = root, dstype = 'frames_finalpass', replicates = replicates)
  
  class ChairsSDHom(data.Dataset):
    def __init__(self, args, is_cropped, root = '/path/to/chairssdhom/data', dstype = 'train', replicates = 1):
      self.args = args
      self.is_cropped = is_cropped
      self.crop_size = args.crop_size
      self.render_size = args.inference_size
      self.replicates = replicates
  
      image1 = sorted( glob( join(root, dstype, 't0/*.png') ) )
      image2 = sorted( glob( join(root, dstype, 't1/*.png') ) )
      self.flow_list = sorted( glob( join(root, dstype, 'flow/*.flo') ) )
  
      assert (len(image1) == len(self.flow_list))
  
      self.image_list = []
      for i in range(len(self.flow_list)):
          im1 = image1[i]
          im2 = image2[i]
          self.image_list += [ [ im1, im2 ] ]
  
      assert len(self.image_list) == len(self.flow_list)
  
      self.size = len(self.image_list)
  
      self.frame_size = frame_utils.read_gen(self.image_list[0][0]).shape
  
      if (self.render_size[0] < 0) or (self.render_size[1] < 0) or (self.frame_size[0]%64) or (self.frame_size[1]%64):
          self.render_size[0] = ( (self.frame_size[0])//64 ) * 64
          self.render_size[1] = ( (self.frame_size[1])//64 ) * 64
  
      args.inference_size = self.render_size
  
    def __getitem__(self, index):
      index = index % self.size
  
      img1 = frame_utils.read_gen(self.image_list[index][0])
      img2 = frame_utils.read_gen(self.image_list[index][1])
  
      flow = frame_utils.read_gen(self.flow_list[index])
      flow = flow[::-1,:,:]
  
      images = [img1, img2]
      image_size = img1.shape[:2]
      if self.is_cropped:
          cropper = StaticRandomCrop(image_size, self.crop_size)
      else:
          cropper = StaticCenterCrop(image_size, self.render_size)
      images = list(map(cropper, images))
      flow = cropper(flow)
  
  
      images = np.array(images).transpose(3,0,1,2)
      flow = flow.transpose(2,0,1)
  
      images = torch.from_numpy(images.astype(np.float32))
      flow = torch.from_numpy(flow.astype(np.float32))
  
      return [images], [flow]
  
    def __len__(self):
      return self.size * self.replicates
  
  class ChairsSDHomTrain(ChairsSDHom):
      def __init__(self, args, is_cropped = False, root = '', replicates = 1):
          super(ChairsSDHomTrain, self).__init__(args, is_cropped = is_cropped, root = root, dstype = 'train', replicates = replicates)
  
  class ChairsSDHomTest(ChairsSDHom):
      def __init__(self, args, is_cropped = False, root = '', replicates = 1):
          super(ChairsSDHomTest, self).__init__(args, is_cropped = is_cropped, root = root, dstype = 'test', replicates = replicates)
  
  class ImagesFromFolder(data.Dataset):
    def __init__(self, args, is_cropped, root = '/path/to/frames/only/folder', iext = 'png', replicates = 1):
      self.args = args
      self.is_cropped = is_cropped
      self.crop_size = args.crop_size
      self.render_size = args.inference_size
      self.replicates = replicates
  
      images = sorted( glob( join(root, '*.' + iext) ) )
      self.image_list = []
      for i in range(len(images)-1):
          im1 = images[i]
          im2 = images[i+1]
          self.image_list += [ [ im1, im2 ] ]
  
      self.size = len(self.image_list)
  
      self.frame_size = frame_utils.read_gen(self.image_list[0][0]).shape
  
      if (self.render_size[0] < 0) or (self.render_size[1] < 0) or (self.frame_size[0]%64) or (self.frame_size[1]%64):
          self.render_size[0] = ( (self.frame_size[0])//64 ) * 64
          self.render_size[1] = ( (self.frame_size[1])//64 ) * 64
  
      args.inference_size = self.render_size
  
    def __getitem__(self, index):
      index = index % self.size
  
      img1 = frame_utils.read_gen(self.image_list[index][0])
      img2 = frame_utils.read_gen(self.image_list[index][1])
  
      images = [img1, img2]
      image_size = img1.shape[:2]
      if self.is_cropped:
          cropper = StaticRandomCrop(image_size, self.crop_size)
      else:
          cropper = StaticCenterCrop(image_size, self.render_size)
      images = list(map(cropper, images))
  
      images = np.array(images).transpose(3,0,1,2)
      images = torch.from_numpy(images.astype(np.float32))
  
      return [images], [torch.zeros(images.size()[0:1] + (2,) + images.size()[-2:])]
  
    def __len__(self):
      return self.size * self.replicates
  
  '''
  import argparse
  import sys, os
  import importlib
  from scipy.misc import imsave
  import numpy as np
  
  import datasets
  reload(datasets)
  
  parser = argparse.ArgumentParser()
  args = parser.parse_args()
  args.inference_size = [1080, 1920]
  args.crop_size = [384, 512]
  args.effective_batch_size = 1
  
  index = 500
  v_dataset = datasets.MpiSintelClean(args, True, root='../MPI-Sintel/flow/training')
  a, b = v_dataset[index]
  im1 = a[0].numpy()[:,0,:,:].transpose(1,2,0)
  im2 = a[0].numpy()[:,1,:,:].transpose(1,2,0)
  imsave('./img1.png', im1)
  imsave('./img2.png', im2)
  flow_utils.writeFlow('./flow.flo', b[0].numpy().transpose(1,2,0))
  
  '''
  