code: import os
  import glob
  import random
  
  import numpy as np
  import pandas as pd
  from imageio import mimread
  from skimage.color import gray2rgb
  from skimage import io, img_as_float32
  from sklearn.model_selection import train_test_split
  
  from torch.utils.data import Dataset
  from data.augmentation import AllAugmentationTransform
  
  
  def read_video(name, frame_shape):
      """
      Read video which can be:
        - an image of concatenated frames
        - '.mp4' and'.gif'
        - folder with videos
      """
  
      if os.path.isdir(name):
          frames = sorted(os.listdir(name))
          num_frames = len(frames)
          video_array = np.array(
              [img_as_float32(io.imread(os.path.join(name, frames[idx]))) for idx in range(num_frames)])
      elif name.lower().endswith('.png') or name.lower().endswith('.jpg'):
          image = io.imread(name)
  
          if len(image.shape) == 2 or image.shape[2] == 1:
              image = gray2rgb(image)
  
          if image.shape[2] == 4:
              image = image[..., :3]
  
          image = img_as_float32(image)
  
          video_array = np.moveaxis(image, 1, 0)
  
          video_array = video_array.reshape((-1,) + frame_shape)
          video_array = np.moveaxis(video_array, 1, 2)
      elif name.lower().endswith('.gif') or name.lower().endswith('.mp4') or name.lower().endswith('.mov'):
          video = np.array(mimread(name))
          if len(video.shape) == 3:
              video = np.array([gray2rgb(frame) for frame in video])
          if video.shape[-1] == 4:
              video = video[..., :3]
          video_array = img_as_float32(video)
      else:
          raise Exception("Unknown file extensions  %s" % name)
  
      return video_array
  
  
  class FramesDataset(Dataset):
      """
      Dataset of videos, each video can be represented as:
        - an image of concatenated frames
        - '.mp4' or '.gif'
        - folder with all frames
      """
  
      def __init__(self, root_dir, frame_shape=(256, 256, 3), id_sampling=False, is_train=True,
                   random_seed=0, pairs_list=None, augmentation_params=None):
          self.root_dir = root_dir
          self.videos = os.listdir(root_dir)
          self.frame_shape = tuple(frame_shape)
          self.pairs_list = pairs_list
          self.id_sampling = id_sampling
          if os.path.exists(os.path.join(root_dir, 'train')):
              assert os.path.exists(os.path.join(root_dir, 'test'))
              print("Use predefined train-test split.")
              if id_sampling:
                  train_videos = {os.path.basename(video).split('#')[0] for video in
                                  os.listdir(os.path.join(root_dir, 'train'))}
                  train_videos = list(train_videos)
              else:
                  train_videos = os.listdir(os.path.join(root_dir, 'train'))
              test_videos = os.listdir(os.path.join(root_dir, 'test'))
              self.root_dir = os.path.join(self.root_dir, 'train' if is_train else 'test')
          else:
              print("Use random train-test split.")
              train_videos, test_videos = train_test_split(self.videos, random_state=random_seed, test_size=0.2)
  
          if is_train:
              self.videos = train_videos
          else:
              self.videos = test_videos
  
          self.is_train = is_train
  
          if self.is_train:
              self.transform = AllAugmentationTransform(**augmentation_params)
          else:
              self.transform = None
  
      def __len__(self):
          return len(self.videos)
  
      def __getitem__(self, idx):
          if self.is_train and self.id_sampling:
              name = self.videos[idx]
              path = np.random.choice(glob.glob(os.path.join(self.root_dir, name + '*.mp4')))
          else:
              name = self.videos[idx]
              path = os.path.join(self.root_dir, name)
  
          video_name = os.path.basename(path)
  
          if self.is_train and os.path.isdir(path):
              frames = os.listdir(path)
              num_frames = len(frames)
              frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2))
              video_array = [img_as_float32(io.imread(os.path.join(path, frames[idx]))) for idx in frame_idx]
          else:
              video_array = read_video(path, frame_shape=self.frame_shape)
              num_frames = len(video_array)
              frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2)) if self.is_train else range(
                  num_frames)
              video_array = list(video_array[frame_idx])
  
          if self.transform is not None:
              video_array = self.transform(video_array)
  
          out = dict()
          if self.is_train:
              source = np.array(video_array[0], dtype='float32')
              driving = np.array(video_array[1], dtype='float32')
  
              out['source'] = source.transpose((2, 0, 1))
              out['driving'] = driving.transpose((2, 0, 1))
          else:
              video = np.array(video_array, dtype='float32')
              out['video'] = video.transpose((3, 0, 1, 2))
  
          out['name'] = video_name
  
          return out
  
  
  class PairedDataset(Dataset):
      """Dataset of pairs for animation."""
  
      def __init__(self, initial_dataset, number_of_pairs, seed=0):
          self.initial_dataset = initial_dataset
          pairs_list = self.initial_dataset.pairs_list
  
          np.random.seed(seed)
  
          if pairs_list is None:
              max_idx = min(number_of_pairs, len(initial_dataset))
              nx, ny = max_idx, max_idx
              xy = np.mgrid[:nx, :ny].reshape(2, -1).T
              number_of_pairs = min(xy.shape[0], number_of_pairs)
              self.pairs = xy.take(np.random.choice(xy.shape[0], number_of_pairs, replace=False), axis=0)
          else:
              videos = self.initial_dataset.videos
              name_to_index = {name: index for index, name in enumerate(videos)}
              pairs = pd.read_csv(pairs_list)
              pairs = pairs[np.logical_and(pairs['source'].isin(videos), pairs['driving'].isin(videos))]
  
              number_of_pairs = min(pairs.shape[0], number_of_pairs)
              self.pairs = []
              self.start_frames = []
              for ind in range(number_of_pairs):
                  self.pairs.append(
                      (name_to_index[pairs['driving'].iloc[ind]], name_to_index[pairs['source'].iloc[ind]]))
  
      def __len__(self):
          return len(self.pairs)
  
      def __getitem__(self, idx):
          pair = self.pairs[idx]
          first = self.initial_dataset[pair[0]]
          second = self.initial_dataset[pair[1]]
          first = {'driving_' + key: value for key, value in first.items()}
          second = {'source_' + key: value for key, value in second.items()}
  
          return {**first, **second}
  
apis: pandas.read_csv numpy.random.seed numpy.moveaxis numpy.random.choice numpy.array sklearn.model_selection.train_test_split