git.sesse.net Git - voxel-flow/blob - dataset.py

   1 """Implements a dataset class for handling image data"""
   2 from utils.image_utils import imread, imsave
   3
   4 DATA_PATH_BASE = '/home/VoxelFlow/dataset/ucf101_triplets/'
   5
   6 class Dataset(object):
   7   def __init__(self, data_list_file=None, process_func=None):
   8     """
   9       Args:
  10     """
  11     self.data_list_file = data_list_file
  12     if process_func:
  13       self.process_func = process_func
  14     else:
  15       self.process_func = self.process_func
  16
  17   def read_data_list_file(self):
  18     """Reads the data list_file into python list
  19     """
  20     f = open(self.data_list_file)
  21     data_list =  [DATA_PATH_BASE+line.rstrip() for line in f]
  22     self.data_list = data_list
  23     return data_list
  24
  25   def process_func(self, example_line):
  26     """Process the single example line and return data
  27       Default behavior, assumes each line is the path to a single image.
  28       This is used to train a VAE.
  29     """
  30     return imread(example_line)