# 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. # ============================================================================== """Wrapper for providing semantic segmentation video data.""" import tensorflow as tf from feelvos import input_preprocess from feelvos import model from feelvos.utils import mask_damaging from feelvos.utils import train_utils slim = tf.contrib.slim dataset_data_provider = slim.dataset_data_provider MIN_LABEL_COUNT = 10 def decode_image_sequence(tensor, image_format='jpeg', shape=None, channels=3, raw_dtype=tf.uint8): """Decodes a sequence of images. Args: tensor: the tensor of strings to decode, shape: [num_images] image_format: a string (possibly tensor) with the format of the image. Options include 'jpeg', 'png', and 'raw'. shape: a list or tensor of the decoded image shape for a single image. channels: if 'shape' is None, the third dimension of the image is set to this value. raw_dtype: if the image is encoded as raw bytes, this is the method of decoding the bytes into values. Returns: The decoded images with shape [time, height, width, channels]. """ handler = slim.tfexample_decoder.Image( shape=shape, channels=channels, dtype=raw_dtype, repeated=True) return handler.tensors_to_item({'image/encoded': tensor, 'image/format': image_format}) def _get_data(data_provider, dataset_split, video_frames_are_decoded): """Gets data from data provider. Args: data_provider: An object of slim.data_provider. dataset_split: Dataset split. video_frames_are_decoded: Boolean, whether the video frames are already decoded Returns: image: Image Tensor. label: Label Tensor storing segmentation annotations. object_label: An integer refers to object_label according to labelmap. If the example has more than one object_label, take the first one. image_name: Image name. height: Image height. width: Image width. video_id: String tensor representing the name of the video. Raises: ValueError: Failed to find label. """ if video_frames_are_decoded: image, = data_provider.get(['image']) else: image, = data_provider.get(['image/encoded']) # Some datasets do not contain image_name. if 'image_name' in data_provider.list_items(): image_name, = data_provider.get(['image_name']) else: image_name = tf.constant('') height, width = data_provider.get(['height', 'width']) label = None if dataset_split != 'test': if video_frames_are_decoded: if 'labels_class' not in data_provider.list_items(): raise ValueError('Failed to find labels.') label, = data_provider.get(['labels_class']) else: key = 'segmentation/object/encoded' if key not in data_provider.list_items(): raise ValueError('Failed to find labels.') label, = data_provider.get([key]) object_label = None video_id, = data_provider.get(['video_id']) return image, label, object_label, image_name, height, width, video_id def _has_foreground_and_background_in_first_frame(label, subsampling_factor): """Checks if the labels have foreground and background in the first frame. Args: label: Label tensor of shape [num_frames, height, width, 1]. subsampling_factor: Integer, the subsampling factor. Returns: Boolean, whether the labels have foreground and background in the first frame. """ h, w = train_utils.resolve_shape(label)[1:3] label_downscaled = tf.squeeze( tf.image.resize_nearest_neighbor(label[0, tf.newaxis], [h // subsampling_factor, w // subsampling_factor], align_corners=True), axis=0) is_bg = tf.equal(label_downscaled, 0) is_fg = tf.logical_not(is_bg) # Just using reduce_any was not robust enough, so lets make sure the count # is above MIN_LABEL_COUNT. fg_count = tf.reduce_sum(tf.cast(is_fg, tf.int32)) bg_count = tf.reduce_sum(tf.cast(is_bg, tf.int32)) has_bg = tf.greater_equal(fg_count, MIN_LABEL_COUNT) has_fg = tf.greater_equal(bg_count, MIN_LABEL_COUNT) return tf.logical_and(has_bg, has_fg) def _has_foreground_and_background_in_first_frame_2(label, decoder_output_stride): """Checks if the labels have foreground and background in the first frame. Second attempt, this time we use the actual output dimension for resizing. Args: label: Label tensor of shape [num_frames, height, width, 1]. decoder_output_stride: Integer, the stride of the decoder output. Returns: Boolean, whether the labels have foreground and background in the first frame. """ h, w = train_utils.resolve_shape(label)[1:3] h_sub = model.scale_dimension(h, 1.0 / decoder_output_stride) w_sub = model.scale_dimension(w, 1.0 / decoder_output_stride) label_downscaled = tf.squeeze( tf.image.resize_nearest_neighbor(label[0, tf.newaxis], [h_sub, w_sub], align_corners=True), axis=0) is_bg = tf.equal(label_downscaled, 0) is_fg = tf.logical_not(is_bg) # Just using reduce_any was not robust enough, so lets make sure the count # is above MIN_LABEL_COUNT. fg_count = tf.reduce_sum(tf.cast(is_fg, tf.int32)) bg_count = tf.reduce_sum(tf.cast(is_bg, tf.int32)) has_bg = tf.greater_equal(fg_count, MIN_LABEL_COUNT) has_fg = tf.greater_equal(bg_count, MIN_LABEL_COUNT) return tf.logical_and(has_bg, has_fg) def _has_enough_pixels_of_each_object_in_first_frame( label, decoder_output_stride): """Checks if for each object (incl. background) enough pixels are visible. During test time, we will usually not see a reference frame in which only very few pixels of one object are visible. These cases can be problematic during training, especially if more than the 1-nearest neighbor is used. That's why this function can be used to detect and filter these cases. Args: label: Label tensor of shape [num_frames, height, width, 1]. decoder_output_stride: Integer, the stride of the decoder output. Returns: Boolean, whether the labels have enough pixels of each object in the first frame. """ h, w = train_utils.resolve_shape(label)[1:3] h_sub = model.scale_dimension(h, 1.0 / decoder_output_stride) w_sub = model.scale_dimension(w, 1.0 / decoder_output_stride) label_downscaled = tf.squeeze( tf.image.resize_nearest_neighbor(label[0, tf.newaxis], [h_sub, w_sub], align_corners=True), axis=0) _, _, counts = tf.unique_with_counts( tf.reshape(label_downscaled, [-1])) has_enough_pixels_per_object = tf.reduce_all( tf.greater_equal(counts, MIN_LABEL_COUNT)) return has_enough_pixels_per_object def get(dataset, num_frames_per_video, crop_size, batch_size, min_resize_value=None, max_resize_value=None, resize_factor=None, min_scale_factor=1., max_scale_factor=1., scale_factor_step_size=0, preprocess_image_and_label=True, num_readers=1, num_threads=1, dataset_split=None, is_training=True, model_variant=None, batch_capacity_factor=32, video_frames_are_decoded=False, decoder_output_stride=None, first_frame_finetuning=False, sample_only_first_frame_for_finetuning=False, sample_adjacent_and_consistent_query_frames=False, remap_labels_to_reference_frame=True, generate_prev_frame_mask_by_mask_damaging=False, three_frame_dataset=False, add_prev_frame_label=True): """Gets the dataset split for semantic segmentation. This functions gets the dataset split for semantic segmentation. In particular, it is a wrapper of (1) dataset_data_provider which returns the raw dataset split, (2) input_preprcess which preprocess the raw data, and (3) the Tensorflow operation of batching the preprocessed data. Then, the output could be directly used by training, evaluation or visualization. Args: dataset: An instance of slim Dataset. num_frames_per_video: The number of frames used per video crop_size: Image crop size [height, width]. batch_size: Batch size. min_resize_value: Desired size of the smaller image side. max_resize_value: Maximum allowed size of the larger image side. resize_factor: Resized dimensions are multiple of factor plus one. min_scale_factor: Minimum scale factor value. max_scale_factor: Maximum scale factor value. scale_factor_step_size: The step size from min scale factor to max scale factor. The input is randomly scaled based on the value of (min_scale_factor, max_scale_factor, scale_factor_step_size). preprocess_image_and_label: Boolean variable specifies if preprocessing of image and label will be performed or not. num_readers: Number of readers for data provider. num_threads: Number of threads for batching data. dataset_split: Dataset split. is_training: Is training or not. model_variant: Model variant (string) for choosing how to mean-subtract the images. See feature_extractor.network_map for supported model variants. batch_capacity_factor: Batch capacity factor affecting the training queue batch capacity. video_frames_are_decoded: Boolean, whether the video frames are already decoded decoder_output_stride: Integer, the stride of the decoder output. first_frame_finetuning: Boolean, whether to only sample the first frame for fine-tuning. sample_only_first_frame_for_finetuning: Boolean, whether to only sample the first frame during fine-tuning. This should be False when using lucid or wonderland data, but true when fine-tuning on the first frame only. Only has an effect if first_frame_finetuning is True. sample_adjacent_and_consistent_query_frames: Boolean, if true, the query frames (all but the first frame which is the reference frame) will be sampled such that they are adjacent video frames and have the same crop coordinates and flip augmentation. remap_labels_to_reference_frame: Boolean, whether to remap the labels of the query frames to match the labels of the (downscaled) reference frame. If a query frame contains a label which is not present in the reference, it will be mapped to background. generate_prev_frame_mask_by_mask_damaging: Boolean, whether to generate the masks used as guidance from the previous frame by damaging the ground truth mask. three_frame_dataset: Boolean, whether the dataset has exactly three frames per video of which the first is to be used as reference and the two others are consecutive frames to be used as query frames. add_prev_frame_label: Boolean, whether to sample one more frame before the first query frame to obtain a previous frame label. Only has an effect, if sample_adjacent_and_consistent_query_frames is True and generate_prev_frame_mask_by_mask_damaging is False. Returns: A dictionary of batched Tensors for semantic segmentation. Raises: ValueError: dataset_split is None, or Failed to find labels. """ if dataset_split is None: raise ValueError('Unknown dataset split.') if model_variant is None: tf.logging.warning('Please specify a model_variant. See ' 'feature_extractor.network_map for supported model ' 'variants.') data_provider = dataset_data_provider.DatasetDataProvider( dataset, num_readers=num_readers, num_epochs=None if is_training else 1, shuffle=is_training) image, label, object_label, image_name, height, width, video_id = _get_data( data_provider, dataset_split, video_frames_are_decoded) sampling_is_valid = tf.constant(True) if num_frames_per_video is not None: total_num_frames = tf.shape(image)[0] if first_frame_finetuning or three_frame_dataset: if sample_only_first_frame_for_finetuning: assert not sample_adjacent_and_consistent_query_frames, ( 'this option does not make sense for sampling only first frame.') # Sample the first frame num_frames_per_video times. sel_indices = tf.tile(tf.constant(0, dtype=tf.int32)[tf.newaxis], multiples=[num_frames_per_video]) else: if sample_adjacent_and_consistent_query_frames: if add_prev_frame_label: num_frames_per_video += 1 # Since this is first frame fine-tuning, we'll for now assume that # each sequence has exactly 3 images: the ref frame and 2 adjacent # query frames. assert num_frames_per_video == 3 with tf.control_dependencies([tf.assert_equal(total_num_frames, 3)]): sel_indices = tf.constant([1, 2], dtype=tf.int32) else: # Sample num_frames_per_video - 1 query frames which are not the # first frame. sel_indices = tf.random_shuffle( tf.range(1, total_num_frames))[:(num_frames_per_video - 1)] # Concat first frame as reference frame to the front. sel_indices = tf.concat([tf.constant(0, dtype=tf.int32)[tf.newaxis], sel_indices], axis=0) else: if sample_adjacent_and_consistent_query_frames: if add_prev_frame_label: # Sample one more frame which we can use to provide initial softmax # feedback. num_frames_per_video += 1 ref_idx = tf.random_shuffle(tf.range(total_num_frames))[0] sampling_is_valid = tf.greater_equal(total_num_frames, num_frames_per_video) def sample_query_start_idx(): return tf.random_shuffle( tf.range(total_num_frames - num_frames_per_video + 1))[0] query_start_idx = tf.cond(sampling_is_valid, sample_query_start_idx, lambda: tf.constant(0, dtype=tf.int32)) def sample_sel_indices(): return tf.concat( [ref_idx[tf.newaxis], tf.range( query_start_idx, query_start_idx + (num_frames_per_video - 1))], axis=0) sel_indices = tf.cond( sampling_is_valid, sample_sel_indices, lambda: tf.zeros((num_frames_per_video,), dtype=tf.int32)) else: # Randomly sample some frames from the video. sel_indices = tf.random_shuffle( tf.range(total_num_frames))[:num_frames_per_video] image = tf.gather(image, sel_indices, axis=0) if not video_frames_are_decoded: image = decode_image_sequence(image) if label is not None: if num_frames_per_video is not None: label = tf.gather(label, sel_indices, axis=0) if not video_frames_are_decoded: label = decode_image_sequence(label, image_format='png', channels=1) # Sometimes, label is saved as [num_frames_per_video, height, width] or # [num_frames_per_video, height, width, 1]. We change it to be # [num_frames_per_video, height, width, 1]. if label.shape.ndims == 3: label = tf.expand_dims(label, 3) elif label.shape.ndims == 4 and label.shape.dims[3] == 1: pass else: raise ValueError('Input label shape must be ' '[num_frames_per_video, height, width],' ' or [num_frames, height, width, 1]. ' 'Got {}'.format(label.shape.ndims)) label.set_shape([None, None, None, 1]) # Add size of first dimension since tf can't figure it out automatically. image.set_shape((num_frames_per_video, None, None, None)) if label is not None: label.set_shape((num_frames_per_video, None, None, None)) preceding_frame_label = None if preprocess_image_and_label: if num_frames_per_video is None: raise ValueError('num_frame_per_video must be specified for preproc.') original_images = [] images = [] labels = [] if sample_adjacent_and_consistent_query_frames: num_frames_individual_preproc = 1 else: num_frames_individual_preproc = num_frames_per_video for frame_idx in range(num_frames_individual_preproc): original_image_t, image_t, label_t = ( input_preprocess.preprocess_image_and_label( image[frame_idx], label[frame_idx], crop_height=crop_size[0] if crop_size is not None else None, crop_width=crop_size[1] if crop_size is not None else None, min_resize_value=min_resize_value, max_resize_value=max_resize_value, resize_factor=resize_factor, min_scale_factor=min_scale_factor, max_scale_factor=max_scale_factor, scale_factor_step_size=scale_factor_step_size, ignore_label=dataset.ignore_label, is_training=is_training, model_variant=model_variant)) original_images.append(original_image_t) images.append(image_t) labels.append(label_t) if sample_adjacent_and_consistent_query_frames: imgs_for_preproc = [image[frame_idx] for frame_idx in range(1, num_frames_per_video)] labels_for_preproc = [label[frame_idx] for frame_idx in range(1, num_frames_per_video)] original_image_rest, image_rest, label_rest = ( input_preprocess.preprocess_images_and_labels_consistently( imgs_for_preproc, labels_for_preproc, crop_height=crop_size[0] if crop_size is not None else None, crop_width=crop_size[1] if crop_size is not None else None, min_resize_value=min_resize_value, max_resize_value=max_resize_value, resize_factor=resize_factor, min_scale_factor=min_scale_factor, max_scale_factor=max_scale_factor, scale_factor_step_size=scale_factor_step_size, ignore_label=dataset.ignore_label, is_training=is_training, model_variant=model_variant)) original_images.extend(original_image_rest) images.extend(image_rest) labels.extend(label_rest) assert len(original_images) == num_frames_per_video assert len(images) == num_frames_per_video assert len(labels) == num_frames_per_video if remap_labels_to_reference_frame: # Remap labels to indices into the labels of the (downscaled) reference # frame, or 0, i.e. background, for labels which are not present # in the reference. reference_labels = labels[0][tf.newaxis] h, w = train_utils.resolve_shape(reference_labels)[1:3] embedding_height = model.scale_dimension( h, 1.0 / decoder_output_stride) embedding_width = model.scale_dimension( w, 1.0 / decoder_output_stride) reference_labels_embedding_size = tf.squeeze( tf.image.resize_nearest_neighbor( reference_labels, tf.stack([embedding_height, embedding_width]), align_corners=True), axis=0) # Get sorted unique labels in the reference frame. labels_in_ref_frame, _ = tf.unique( tf.reshape(reference_labels_embedding_size, [-1])) labels_in_ref_frame = tf.contrib.framework.sort(labels_in_ref_frame) for idx in range(1, len(labels)): ref_label_mask = tf.equal( labels[idx], labels_in_ref_frame[tf.newaxis, tf.newaxis, :]) remapped = tf.argmax(tf.cast(ref_label_mask, tf.uint8), axis=-1, output_type=tf.int32) # Set to 0 if label is not present is_in_ref = tf.reduce_any(ref_label_mask, axis=-1) remapped *= tf.cast(is_in_ref, tf.int32) labels[idx] = remapped[..., tf.newaxis] if sample_adjacent_and_consistent_query_frames: if first_frame_finetuning and generate_prev_frame_mask_by_mask_damaging: preceding_frame_label = mask_damaging.damage_masks(labels[1]) elif add_prev_frame_label: # Discard the image of the additional frame and take the label as # initialization for softmax feedback. original_images = [original_images[0]] + original_images[2:] preceding_frame_label = labels[1] images = [images[0]] + images[2:] labels = [labels[0]] + labels[2:] num_frames_per_video -= 1 original_image = tf.stack(original_images, axis=0) image = tf.stack(images, axis=0) label = tf.stack(labels, axis=0) else: if label is not None: # Need to set label shape due to batching. label.set_shape([num_frames_per_video, None if crop_size is None else crop_size[0], None if crop_size is None else crop_size[1], 1]) original_image = tf.to_float(tf.zeros_like(label)) if crop_size is None: height = tf.shape(image)[1] width = tf.shape(image)[2] else: height = crop_size[0] width = crop_size[1] sample = {'image': image, 'image_name': image_name, 'height': height, 'width': width, 'video_id': video_id} if label is not None: sample['label'] = label if object_label is not None: sample['object_label'] = object_label if preceding_frame_label is not None: sample['preceding_frame_label'] = preceding_frame_label if not is_training: # Original image is only used during visualization. sample['original_image'] = original_image if is_training: if first_frame_finetuning: keep_input = tf.constant(True) else: keep_input = tf.logical_and(sampling_is_valid, tf.logical_and( _has_enough_pixels_of_each_object_in_first_frame( label, decoder_output_stride), _has_foreground_and_background_in_first_frame_2( label, decoder_output_stride))) batched = tf.train.maybe_batch(sample, keep_input=keep_input, batch_size=batch_size, num_threads=num_threads, capacity=batch_capacity_factor * batch_size, dynamic_pad=True) else: batched = tf.train.batch(sample, batch_size=batch_size, num_threads=num_threads, capacity=batch_capacity_factor * batch_size, dynamic_pad=True) # Flatten from [batch, num_frames_per_video, ...] to # batch * num_frames_per_video, ...]. cropped_height = train_utils.resolve_shape(batched['image'])[2] cropped_width = train_utils.resolve_shape(batched['image'])[3] if num_frames_per_video is None: first_dim = -1 else: first_dim = batch_size * num_frames_per_video batched['image'] = tf.reshape(batched['image'], [first_dim, cropped_height, cropped_width, 3]) if label is not None: batched['label'] = tf.reshape(batched['label'], [first_dim, cropped_height, cropped_width, 1]) return batched