utils.py 文件源码

def load_kth_data(f_name, data_path, image_size, K, T): 
  flip = np.random.binomial(1,.5,1)[0]
  tokens = f_name.split()
  vid_path = data_path + tokens[0] + "_uncomp.avi"
  vid = imageio.get_reader(vid_path,"ffmpeg")
  low = int(tokens[1])
  high = np.min([int(tokens[2]),vid.get_length()])-K-T+1
  if low == high:
    stidx = 0 
  else:
    if low >= high: print(vid_path)
    stidx = np.random.randint(low=low, high=high)
  seq = np.zeros((image_size, image_size, K+T, 1), dtype="float32")
  for t in xrange(K+T):
    img = cv2.cvtColor(cv2.resize(vid.get_data(stidx+t),
                       (image_size,image_size)),
                       cv2.COLOR_RGB2GRAY)
    seq[:,:,t] = transform(img[:,:,None])

  if flip == 1:
    seq = seq[:,::-1]

  diff = np.zeros((image_size, image_size, K-1, 1), dtype="float32")
  for t in xrange(1,K):
    prev = inverse_transform(seq[:,:,t-1])
    next = inverse_transform(seq[:,:,t])
    diff[:,:,t-1] = next.astype("float32")-prev.astype("float32")

  return seq, diff