def get_i3d_features(data_shape: Tuple[int]):
  """Gets the i3d features for videos.

  Args:
    data_shape: data shape.
  Returns:
    A callable function to extract features.
  """
  batch_size = 256
  with tf.Graph().as_default():
    input_videos = tf.placeholder(tf.float32,
                                  shape=(batch_size,) + data_shape[1:])
    features = create_i3d_embedding(
        normalize_videos(input_videos, (224, 224)), batch_size)
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())
    sess.run(tf.tables_initializer())
    def feature_extractor(data, sess, input_videos, features):
      output_features = []
      num_batches = int(np.ceil(data.shape[0] / batch_size))
      for i in tqdm.tqdm(range(num_batches),
                         desc='extract features', unit='batch'):
        batch = data[i * batch_size:(i + 1) * batch_size]
        pad_length = batch_size - batch.shape[0]
        if pad_length > 0:
          batch = np.pad(
              batch, ((0, pad_length), (0, 0), (0, 0), (0, 0), (0, 0)),
              mode='constant',
              constant_values=0)
        output = sess.run(features, {input_videos: batch})
        output_features.append(output[:batch_size - pad_length])
      output_features = np.concatenate(output_features, axis=0)
      return output_features
    return functools.partial(
        feature_extractor,
        sess=sess,
        input_videos=input_videos,
        features=features)
Gets the i3d features for videos.

  Args:
    data_shape: data shape.
  Returns:
    A callable function to extract features.
  
videogvt/tools/compute_fvd.py