models/research/feelvos/datasets/video_dataset.py

# Copyright 2018 The TensorFlow Authors All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Provides data from video object segmentation datasets.

This file provides both images and annotations (instance segmentations) for
TensorFlow. Currently, we support the following datasets:

1. DAVIS 2017 (https://davischallenge.org/davis2017/code.html).

2. DAVIS 2016 (https://davischallenge.org/davis2016/code.html).

3. YouTube-VOS (https://youtube-vos.org/dataset/download).
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import os.path
import tensorflow as tf
from feelvos.datasets import tfsequence_example_decoder

slim = tf.contrib.slim
dataset = slim.dataset
tfexample_decoder = slim.tfexample_decoder


_ITEMS_TO_DESCRIPTIONS = {
    'image': 'A color image of varying height and width.',
    'labels_class': ('A semantic segmentation label whose size matches image.'
                     'Its values range from 0 (background) to num_classes.'),
}

# Named tuple to describe the dataset properties.
DatasetDescriptor = collections.namedtuple(
    'DatasetDescriptor',
    ['splits_to_sizes',   # Splits of the dataset into training, val, and test.
     'num_classes',   # Number of semantic classes.
     'ignore_label',  # Ignore label value.
    ]
)

_DAVIS_2016_INFORMATION = DatasetDescriptor(
    splits_to_sizes={'train': [30, 1830],
                     'val': [20, 1376]},
    num_classes=2,
    ignore_label=255,
)

_DAVIS_2017_INFORMATION = DatasetDescriptor(
    splits_to_sizes={'train': [60, 4219],
                     'val': [30, 2023],
                     'test-dev': [30, 2037]},
    num_classes=None,  # Number of instances per videos differ.
    ignore_label=255,
)

_YOUTUBE_VOS_2018_INFORMATION = DatasetDescriptor(
    # Leave these sizes as None to allow for different splits into
    # training and validation sets.
    splits_to_sizes={'train': [None, None],
                     'val': [None, None]},
    num_classes=None,  # Number of instances per video differs.
    ignore_label=255,
)

_DATASETS_INFORMATION = {
    'davis_2016': _DAVIS_2016_INFORMATION,
    'davis_2017': _DAVIS_2017_INFORMATION,
    'youtube_vos_2018': _YOUTUBE_VOS_2018_INFORMATION,
}

# Default file pattern of SSTable. Note we include '-' to avoid the confusion
# between `train-` and `trainval-` sets.
_FILE_PATTERN = '%s-*'


def get_dataset(dataset_name,
                split_name,
                dataset_dir,
                file_pattern=None,
                data_type='tf_sequence_example',
                decode_video_frames=False):
  """Gets an instance of slim Dataset.

  Args:
    dataset_name: String, dataset name.
    split_name: String, the train/val Split name.
    dataset_dir: String, the directory of the dataset sources.
    file_pattern: String, file pattern of SSTable.
    data_type: String, data type. Currently supports 'tf_example' and
      'annotated_image'.
    decode_video_frames: Boolean, decode the images or not. Not decoding it here
        is useful if we subsample later

  Returns:
    An instance of slim Dataset.

  Raises:
    ValueError: If the dataset_name or split_name is not recognized, or if
      the dataset_type is not supported.
  """
  if dataset_name not in _DATASETS_INFORMATION:
    raise ValueError('The specified dataset is not supported yet.')

  splits_to_sizes = _DATASETS_INFORMATION[dataset_name].splits_to_sizes

  if split_name not in splits_to_sizes:
    raise ValueError('data split name %s not recognized' % split_name)

  # Prepare the variables for different datasets.
  num_classes = _DATASETS_INFORMATION[dataset_name].num_classes
  ignore_label = _DATASETS_INFORMATION[dataset_name].ignore_label

  if file_pattern is None:
    file_pattern = _FILE_PATTERN
  file_pattern = os.path.join(dataset_dir, file_pattern % split_name)
  if data_type == 'tf_sequence_example':
    keys_to_context_features = {
        'image/format': tf.FixedLenFeature((), tf.string, default_value='jpeg'),
        'image/height': tf.FixedLenFeature((), tf.int64, default_value=0),
        'image/width': tf.FixedLenFeature((), tf.int64, default_value=0),
        'segmentation/object/format': tf.FixedLenFeature(
            (), tf.string, default_value='png'),
        'video_id': tf.FixedLenFeature((), tf.string, default_value='unknown')
    }
    label_name = 'class' if dataset_name == 'davis_2016' else 'object'
    keys_to_sequence_features = {
        'image/encoded': tf.FixedLenSequenceFeature((), dtype=tf.string),
        'segmentation/{}/encoded'.format(label_name):
            tf.FixedLenSequenceFeature((), tf.string),
        'segmentation/{}/encoded'.format(label_name):
            tf.FixedLenSequenceFeature((), tf.string),
    }
    items_to_handlers = {
        'height': tfexample_decoder.Tensor('image/height'),
        'width': tfexample_decoder.Tensor('image/width'),
        'video_id': tfexample_decoder.Tensor('video_id')
    }
    if decode_video_frames:
      decode_image_handler = tfexample_decoder.Image(
          image_key='image/encoded',
          format_key='image/format',
          channels=3,
          repeated=True)
      items_to_handlers['image'] = decode_image_handler
      decode_label_handler = tfexample_decoder.Image(
          image_key='segmentation/{}/encoded'.format(label_name),
          format_key='segmentation/{}/format'.format(label_name),
          channels=1,
          repeated=True)
      items_to_handlers['labels_class'] = decode_label_handler
    else:
      items_to_handlers['image/encoded'] = tfexample_decoder.Tensor(
          'image/encoded')
      items_to_handlers[
          'segmentation/object/encoded'] = tfexample_decoder.Tensor(
              'segmentation/{}/encoded'.format(label_name))
    decoder = tfsequence_example_decoder.TFSequenceExampleDecoder(
        keys_to_context_features, keys_to_sequence_features, items_to_handlers)
  else:
    raise ValueError('Unknown data type.')

  size = splits_to_sizes[split_name]
  if isinstance(size, collections.Sequence):
    num_videos = size[0]
    num_samples = size[1]
  else:
    num_videos = 0
    num_samples = size

  return dataset.Dataset(
      data_sources=file_pattern,
      reader=tf.TFRecordReader,
      decoder=decoder,
      num_samples=num_samples,
      num_videos=num_videos,
      items_to_descriptions=_ITEMS_TO_DESCRIPTIONS,
      ignore_label=ignore_label,
      num_classes=num_classes,
      name=dataset_name,
      multi_label=True)