|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
r"""Beam pipeline that generates Vimeo-90K (train or test) triplet TFRecords. |
|
|
|
Vimeo-90K dataset is built upon 5,846 videos downloaded from vimeo.com. The list |
|
of the original video links are available here: |
|
https://github.com/anchen1011/toflow/blob/master/data/original_vimeo_links.txt. |
|
Each video is further cropped into a fixed spatial size of (448 x 256) to create |
|
89,000 video clips. |
|
|
|
The Vimeo-90K dataset is designed for four video processing tasks. This script |
|
creates the TFRecords of frame triplets for frame interpolation task. |
|
|
|
Temporal frame interpolation triplet dataset: |
|
- 73,171 triplets of size (448x256) extracted from 15K subsets of Vimeo-90K. |
|
- The triplets are pre-split into (train,test) = (51313,3782) |
|
- Download links: |
|
Test-set: http://data.csail.mit.edu/tofu/testset/vimeo_interp_test.zip |
|
Train+test-set: http://data.csail.mit.edu/tofu/dataset/vimeo_triplet.zip |
|
|
|
For more information, see the arXiv paper, project page or the GitHub link. |
|
@article{xue17toflow, |
|
author = {Xue, Tianfan and |
|
Chen, Baian and |
|
Wu, Jiajun and |
|
Wei, Donglai and |
|
Freeman, William T}, |
|
title = {Video Enhancement with Task-Oriented Flow}, |
|
journal = {arXiv}, |
|
year = {2017} |
|
} |
|
Project: http://toflow.csail.mit.edu/ |
|
GitHub: https://github.com/anchen1011/toflow |
|
|
|
Inputs to the script are (1) the directory to the downloaded and unzipped folder |
|
(2) the filepath of the text-file that lists the subfolders of the triplets. |
|
|
|
Output TFRecord is a tf.train.Example proto of each image triplet. |
|
The feature_map takes the form: |
|
feature_map { |
|
'frame_0/encoded': |
|
tf.io.FixedLenFeature((), tf.string, default_value=''), |
|
'frame_0/format': |
|
tf.io.FixedLenFeature((), tf.string, default_value='jpg'), |
|
'frame_0/height': |
|
tf.io.FixedLenFeature((), tf.int64, default_value=0), |
|
'frame_0/width': |
|
tf.io.FixedLenFeature((), tf.int64, default_value=0), |
|
'frame_1/encoded': |
|
tf.io.FixedLenFeature((), tf.string, default_value=''), |
|
'frame_1/format': |
|
tf.io.FixedLenFeature((), tf.string, default_value='jpg'), |
|
'frame_1/height': |
|
tf.io.FixedLenFeature((), tf.int64, default_value=0), |
|
'frame_1/width': |
|
tf.io.FixedLenFeature((), tf.int64, default_value=0), |
|
'frame_2/encoded': |
|
tf.io.FixedLenFeature((), tf.string, default_value=''), |
|
'frame_2/format': |
|
tf.io.FixedLenFeature((), tf.string, default_value='jpg'), |
|
'frame_2/height': |
|
tf.io.FixedLenFeature((), tf.int64, default_value=0), |
|
'frame_2/width': |
|
tf.io.FixedLenFeature((), tf.int64, default_value=0) |
|
'path': |
|
tf.io.FixedLenFeature((), tf.string, default_value='') |
|
} |
|
|
|
Usage example: |
|
python3 -m frame_interpolation.datasets.create_vimeo90K_tfrecord \ |
|
--input_dir=<root folder of vimeo90K dataset> \ |
|
--input_triplet_list_filepath=<filepath of tri_{test|train}list.txt> \ |
|
--output_tfrecord_filepath=<output tfrecord filepath> |
|
""" |
|
import os |
|
|
|
from . import util |
|
from absl import app |
|
from absl import flags |
|
from absl import logging |
|
import apache_beam as beam |
|
import numpy as np |
|
import tensorflow as tf |
|
|
|
|
|
_INPUT_DIR = flags.DEFINE_string( |
|
'input_dir', |
|
default='/path/to/raw_vimeo_interp/sequences', |
|
help='Path to the root directory of the vimeo frame interpolation dataset. ' |
|
'We expect the data to have been downloaded and unzipped.\n' |
|
'Folder structures:\n' |
|
'| raw_vimeo_dataset/\n' |
|
'| sequences/\n' |
|
'| | 00001\n' |
|
'| | | 0389/\n' |
|
'| | | | im1.png\n' |
|
'| | | | im2.png\n' |
|
'| | | | im3.png\n' |
|
'| | | ...\n' |
|
'| | 00002/\n' |
|
'| | ...\n' |
|
'| readme.txt\n' |
|
'| tri_trainlist.txt\n' |
|
'| tri_testlist.txt \n') |
|
|
|
_INTPUT_TRIPLET_LIST_FILEPATH = flags.DEFINE_string( |
|
'input_triplet_list_filepath', |
|
default='/path/to/raw_vimeo_dataset/tri_{test|train}list.txt', |
|
help='Text file containing a list of sub-directories of input triplets.') |
|
|
|
_OUTPUT_TFRECORD_FILEPATH = flags.DEFINE_string( |
|
'output_tfrecord_filepath', |
|
default=None, |
|
help='Filepath to the output TFRecord file.') |
|
|
|
_NUM_SHARDS = flags.DEFINE_integer('num_shards', |
|
default=200, |
|
help='Number of shards used for the output.') |
|
|
|
|
|
_INTERPOLATOR_IMAGES_MAP = { |
|
'frame_0': 'im1.png', |
|
'frame_1': 'im2.png', |
|
'frame_2': 'im3.png', |
|
} |
|
|
|
|
|
def main(unused_argv): |
|
"""Creates and runs a Beam pipeline to write frame triplets as a TFRecord.""" |
|
with tf.io.gfile.GFile(_INTPUT_TRIPLET_LIST_FILEPATH.value, 'r') as fid: |
|
triplets_list = np.loadtxt(fid, dtype=str) |
|
|
|
triplet_dicts = [] |
|
for triplet in triplets_list: |
|
triplet_dict = { |
|
image_key: os.path.join(_INPUT_DIR.value, triplet, image_basename) |
|
for image_key, image_basename in _INTERPOLATOR_IMAGES_MAP.items() |
|
} |
|
triplet_dicts.append(triplet_dict) |
|
p = beam.Pipeline('DirectRunner') |
|
(p | 'ReadInputTripletDicts' >> beam.Create(triplet_dicts) |
|
| 'GenerateSingleExample' >> beam.ParDo( |
|
util.ExampleGenerator(_INTERPOLATOR_IMAGES_MAP)) |
|
| 'WriteToTFRecord' >> beam.io.tfrecordio.WriteToTFRecord( |
|
file_path_prefix=_OUTPUT_TFRECORD_FILEPATH.value, |
|
num_shards=_NUM_SHARDS.value, |
|
coder=beam.coders.BytesCoder())) |
|
result = p.run() |
|
result.wait_until_finish() |
|
|
|
logging.info('Succeeded in creating the output TFRecord file: \'%s@%s\'.', |
|
_OUTPUT_TFRECORD_FILEPATH.value, str(_NUM_SHARDS.value)) |
|
|
|
if __name__ == '__main__': |
|
app.run(main) |
|
|