karolmajek/maxdeeplab

1# coding=utf-82# Copyright 2021 The Deeplab2 Authors.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16r"""Script to export deeplab model to saved model."""17 18import functools19from typing import Any, MutableMapping, Sequence, Text20 21from absl import app22from absl import flags23import tensorflow as tf24 25from google.protobuf import text_format26from deeplab2 import config_pb227from deeplab2.data import dataset28from deeplab2.data.preprocessing import input_preprocessing29from deeplab2.model import utils30from deeplab2.trainer import train_lib31 32 33_FLAGS_EXPERIMENT_OPTION_PATH = flags.DEFINE_string(34    'experiment_option_path',35    default='',36    help='Path to the experiment option text proto.')37 38_FLAGS_CKPT_PATH = flags.DEFINE_string(39    'checkpoint_path',40    default='',41    help='Path to the saved checkpoint.')42 43_FLAGS_OUTPUT_PATH = flags.DEFINE_string(44    'output_path',45    default='',46    help='Output directory path for the exported saved model.')47 48_FLAGS_MERGE_WITH_TF_OP = flags.DEFINE_boolean(49    'merge_with_tf_op',50    default=False,51    help='Whether to use customized TF op for merge semantic and instance '52    'predictions. Set it to True to reproduce the numbers as reported in '53    'paper, but the saved model would require specifically compiled TensorFlow '54    'to run.')55 56 57class DeepLabModule(tf.Module):58  """Class that runs DeepLab inference end-to-end."""59 60  def __init__(self, config: config_pb2.ExperimentOptions, ckpt_path: Text,61               use_tf_op: bool = False):62    super().__init__(name='DeepLabModule')63 64    dataset_options = config.eval_dataset_options65    dataset_name = dataset_options.dataset66    crop_height, crop_width = dataset_options.crop_size67 68    config.evaluator_options.merge_semantic_and_instance_with_tf_op = use_tf_op69    # Disable drop path and recompute grad as they are only used in training.70    config.model_options.backbone.drop_path_keep_prob = 1.071 72    deeplab_model = train_lib.create_deeplab_model(73        config,74        dataset.MAP_NAME_TO_DATASET_INFO[dataset_name])75    self._is_motion_deeplab = (76        config.model_options.WhichOneof('meta_architecture') ==77        'motion_deeplab')78 79    # For now we only support batch size of 1 for saved model.80    input_shape = train_lib.build_deeplab_model(81        deeplab_model, (crop_height, crop_width), batch_size=1)82    self._input_depth = input_shape[-1]83 84    checkpoint = tf.train.Checkpoint(**deeplab_model.checkpoint_items)85    # Not all saved variables (e.g. variables from optimizer) will be restored.86    # `expect_partial()` to suppress the warning.87    checkpoint.restore(ckpt_path).expect_partial()88    self._model = deeplab_model89 90    self._preprocess_fn = functools.partial(91        input_preprocessing.preprocess_image_and_label,92        label=None,93        crop_height=crop_height,94        crop_width=crop_width,95        prev_label=None,96        min_resize_value=dataset_options.min_resize_value,97        max_resize_value=dataset_options.max_resize_value,98        resize_factor=dataset_options.resize_factor,99        is_training=False)100 101  def get_input_spec(self):102    """Returns TensorSpec of input tensor needed for inference."""103    # We expect a single 3D, uint8 tensor with shape [height, width, channels].104    return tf.TensorSpec(shape=[None, None, self._input_depth], dtype=tf.uint8)105 106  @tf.function107  def __call__(self, input_tensor: tf.Tensor) -> MutableMapping[Text, Any]:108    """Performs a forward pass.109 110    Args:111      input_tensor: An uint8 input tensor of type tf.Tensor with shape [height,112        width, channels].113 114    Returns:115      A dictionary containing the results of the specified DeepLab architecture.116      The results are bilinearly upsampled to input size before returning.117    """118    input_size = [tf.shape(input_tensor)[0], tf.shape(input_tensor)[1]]119 120    if self._is_motion_deeplab:121      # For motion deeplab, split the input tensor to current and previous122      # frame before preprocessing, and re-assemble them.123      image, prev_image = tf.split(input_tensor, 2, axis=2)124      (resized_image, processed_image, _, processed_prev_image,125       _) = self._preprocess_fn(image=image, prev_image=prev_image)126      processed_image = tf.concat(127          [processed_image, processed_prev_image], axis=2)128    else:129      (resized_image, processed_image, _, _, _) = self._preprocess_fn(130          image=input_tensor)131 132    resized_size = tf.shape(resized_image)[0:2]133    # Making input tensor to 4D to fit model input requirements.134    outputs = self._model(tf.expand_dims(processed_image, 0), training=False)135    # We only undo-preprocess for those defined in tuples in model/utils.py.136    return utils.undo_preprocessing(outputs, resized_size,137                                    input_size)138 139 140def main(argv: Sequence[str]) -> None:141  if len(argv) > 1:142    raise app.UsageError('Too many command-line arguments.')143 144  config = config_pb2.ExperimentOptions()145  with tf.io.gfile.GFile(_FLAGS_EXPERIMENT_OPTION_PATH.value, 'r') as f:146    text_format.Parse(f.read(), config)147 148  module = DeepLabModule(149      config, _FLAGS_CKPT_PATH.value, _FLAGS_MERGE_WITH_TF_OP.value)150 151  signatures = module.__call__.get_concrete_function(module.get_input_spec())152  tf.saved_model.save(153      module, _FLAGS_OUTPUT_PATH.value, signatures=signatures)154 155 156if __name__ == '__main__':157  app.run(main)158