text: # Code for "TDN: Temporal Difference Networks for Efficient Action Recognition"
  # arXiv: 2012.10071
  # Limin Wang, Zhan Tong, Bin Ji, Gangshan Wu
  # tongzhan@smail.nju.edu.cn
  
  import os
  
  
  def return_ucf101(modality):
      filename_categories = 101
      if modality == 'RGB':
          filename_imglist_train = "ucfTrainTestlist/ucf101_rgb_train_split_3.txt"
          filename_imglist_val = "ucfTrainTestlist/ucf101_rgb_val_split_3.txt"
          prefix = 'image_{:05d}.jpg'
      elif modality == 'Flow':
          filename_imglist_train = 'file_list/ucf101_flow_train_split_3.txt'
          filename_imglist_val = 'file_list/ucf101_flow_val_split_3.txt'
          prefix = 'flow_{}_{:05d}.jpg'
      else:
          raise NotImplementedError('no such modality:' + modality)
      return filename_categories, filename_imglist_train, filename_imglist_val, prefix
  
  
  def return_hmdb51(modality):
      filename_categories = 51
      if modality == 'RGB':
          filename_imglist_train = "splits/hmdb51_rgb_train_split_3.txt"
          filename_imglist_val = "splits/hmdb51_rgb_val_split_3.txt"
          prefix = '{:05d}.jpg'
      elif modality == 'Flow':
          filename_imglist_train = 'HMDB51/splits/hmdb51_flow_train_split_1.txt'
          filename_imglist_val = 'HMDB51/splits/hmdb51_flow_val_split_1.txt'
          prefix = 'flow_{}_{:05d}.jpg'
      else:
          raise NotImplementedError('no such modality:' + modality)
      return filename_categories, filename_imglist_train, filename_imglist_val, prefix
  
  
  def return_something(modality):
      filename_categories = 174
      if modality == 'RGB' or modality == 'RGBDiff':
          filename_imglist_train = "/your_path_to/something-something-v1-train_list.txt"
          filename_imglist_val = "/your_path_to/something-something-v1-val_list.txt"
          prefix = '{:05d}.jpg'
      elif modality == 'Flow':
          filename_imglist_train = '/your_path_to/something/v1/train_videofolder_flow.txt'
          filename_imglist_val = '/your_path_to/something/v1/val_videofolder_flow.txt'
          prefix = '{:06d}-{}_{:05d}.jpg'
      else:
          print('no such modality:' + modality)
          raise NotImplementedError
      return filename_categories, filename_imglist_train, filename_imglist_val, prefix
  
  
  def return_somethingv2(modality):
      filename_categories = 174
      if modality == 'RGB':
          filename_imglist_train = "/your_path_to/something-something-v2-train_list.txt"
          filename_imglist_val = "/your_path_to/something-something-v2-val_list.txt"
          prefix = '{:06d}.jpg'
      elif modality == 'Flow':
          filename_imglist_train = 'something/v2/train_videofolder_flow.txt'
          filename_imglist_val = 'something/v2/val_videofolder_flow.txt'
          prefix = '{:06d}.jpg'
      else:
          raise NotImplementedError('no such modality:' + modality)
      return filename_categories, filename_imglist_train, filename_imglist_val, prefix
  
  
  def return_jester(modality):
      filename_categories = 'jester/category.txt'
      if modality == 'RGB':
          prefix = '{:05d}.jpg'
          filename_imglist_train = 'jester/train_videofolder.txt'
          filename_imglist_val = 'jester/val_videofolder.txt'
      else:
          raise NotImplementedError('no such modality:' + modality)
      return filename_categories, filename_imglist_train, filename_imglist_val, prefix
  
  
  def return_kinetics(modality):
      filename_categories = 400
      if modality == 'RGB':
          filename_imglist_train = "/your_path_to/list_kinetics-400/train.csv"
          filename_imglist_val = "/your_path_to/list_kinetics-400/val.csv"
          prefix = 'image_{:05d}.jpg'
      else:
          raise NotImplementedError('no such modality:' + modality)
      return filename_categories, filename_imglist_train, filename_imglist_val, prefix
  
  
  def return_dataset(root_path, dataset, modality):
      dict_single = {'jester': return_jester, 'something': return_something, 'somethingv2': return_somethingv2,
                     'ucf101': return_ucf101, 'hmdb51': return_hmdb51,
                     'kinetics': return_kinetics}
      if dataset in dict_single:
          file_categories, file_imglist_train, file_imglist_val, prefix = dict_single[dataset](modality)
      else:
          raise ValueError('Unknown dataset ' + dataset)
  
      file_imglist_train = os.path.join(root_path, file_imglist_train)
      file_imglist_val = os.path.join(root_path, file_imglist_val)
      if isinstance(file_categories, str):
          # file_categories = os.path.join(ROOT_DATASET, file_categories)
          with open(file_categories) as f:
              lines = f.readlines()
          categories = [item.rstrip() for item in lines]
      else:
          categories = [None] * file_categories
      n_class = len(categories)
      print('{}: {} classes'.format(dataset, n_class))
      return n_class, file_imglist_train, file_imglist_val, prefix
  