title: TensorFlow - Read video frames from TFRecords file
body: <p><strong>TLDR;</strong> <em>my question is on how to load compressed video frames from TFRecords.</em></p>
  
  <p>I am setting up a data pipeline for training deep learning models on a large video dataset (<a href="https://deepmind.com/research/open-source/open-source-datasets/kinetics/" rel="noreferrer">Kinetics</a>). For this I am using TensorFlow, more specifically the <code>tf.data.Dataset</code> and <code>TFRecordDataset</code> structures. As the dataset contains ~300k videos of 10 seconds, there is a large amount of data to deal with. During training, I want to randomly sample 64 consecutive frames from a video, therefore fast random sampling is important. For achieving this there are a number of data loading scenarios possible during training:</p>
  
  <ol>
  <li><strong>Sample from Video.</strong> Load the videos using <code>ffmpeg</code> or <code>OpenCV</code> and sample frames. Not ideal as seeking in videos is tricky, and decoding video streams is much slower than decoding JPG.</li>
  <li><strong>JPG Images.</strong> Preprocess the dataset by extracting all video frames as JPG. This generates a huge amount of files, which is probably not going to be fast due to random access.</li>
  <li><strong>Data Containers.</strong> Preprocess the dataset to <code>TFRecords</code> or <code>HDF5</code> files. Requires more work getting the pipeline ready, but most likely to be the fastest of those options.</li>
  </ol>
  
  <p>I have decided to go for option (3) and use <code>TFRecord</code> files to store a preprocessed version of the dataset. However, this is also not as straightforward as it seems, for example:</p>
  
  <ol>
  <li><strong>Compression.</strong> Storing the video frames as uncompressed byte data in TFRecords will require a huge amount of disk space. Therefore, I extract all the video frames, apply JPG compression and store the compressed bytes as TFRecords. </li>
  <li><strong>Video Data.</strong> We are dealing with video, so each example in the TFRecords file will be quite large and contains several video frames (typically 250-300 for 10 seconds of video, depending on the frame rate). </li>
  </ol>
  
  <p>I have wrote the following code to preprocess the video dataset and write the video frames as TFRecord files (each of ~5GB in size):</p>
  
  <pre><code>def _int64_feature(value):
      """Wrapper for inserting int64 features into Example proto."""
      if not isinstance(value, list):
          value = [value]
      return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
  
  def _bytes_feature(value):
      """Wrapper for inserting bytes features into Example proto."""
      return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
  
  
  with tf.python_io.TFRecordWriter(output_file) as writer:
  
    # Read and resize all video frames, np.uint8 of size [N,H,W,3]
    frames = ... 
  
    features = {}
    features['num_frames']  = _int64_feature(frames.shape[0])
    features['height']      = _int64_feature(frames.shape[1])
    features['width']       = _int64_feature(frames.shape[2])
    features['channels']    = _int64_feature(frames.shape[3])
    features['class_label'] = _int64_feature(example['class_id'])
    features['class_text']  = _bytes_feature(tf.compat.as_bytes(example['class_label']))
    features['filename']    = _bytes_feature(tf.compat.as_bytes(example['video_id']))
  
    # Compress the frames using JPG and store in as bytes in:
    # 'frames/000001', 'frames/000002', ...
    for i in range(len(frames)):
        ret, buffer = cv2.imencode(".jpg", frames[i])
        features["frames/{:04d}".format(i)] = _bytes_feature(tf.compat.as_bytes(buffer.tobytes()))
  
    tfrecord_example = tf.train.Example(features=tf.train.Features(feature=features))
    writer.write(tfrecord_example.SerializeToString())
  </code></pre>
  
  <p>This works fine; the dataset is nicely written as TFRecord files with the frames as compressed JPG bytes. My question regards, how to read the TFRecord files during training, randomly sample 64 frames from a video and decode the JPG images. </p>
  
  <p>According to <a href="https://www.tensorflow.org/programmers_guide/datasets" rel="noreferrer">TensorFlow's documentation</a> on <code>tf.Data</code> we need to do something like: </p>
  
  <pre><code>filenames = tf.placeholder(tf.string, shape=[None])
  dataset = tf.data.TFRecordDataset(filenames)
  dataset = dataset.map(...)  # Parse the record into tensors.
  dataset = dataset.repeat()  # Repeat the input indefinitely.
  dataset = dataset.batch(32)
  iterator = dataset.make_initializable_iterator()
  training_filenames = ["/var/data/file1.tfrecord", "/var/data/file2.tfrecord"]
  sess.run(iterator.initializer, feed_dict={filenames: training_filenames})
  </code></pre>
  
  <p>There are many example on how to do this with images, and that is quite straightforward. However, for video and random sampling of frames I am stuck. The <code>tf.train.Features</code> object stores the frames as <code>frame/00001</code>, <code>frame/000002</code> etc. My first question is how to randomly sample a set of consecutive frames from this inside the <code>dataset.map()</code> function? Considerations are that each frame has a variable number of bytes due to JPG compression and need to be decoded using <code>tf.image.decode_jpeg</code>.</p>
  
  <p>Any help how to best setup reading video sampels from TFRecord files would be appreciated! </p>