semantic_segmentation.py 3.01 KB
Newer Older
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
2
3
4
5
6
7
8
9
10
11
12
13
#
# 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.
Yeqing Li's avatar
Yeqing Li committed
14

Abdullah Rashwan's avatar
Abdullah Rashwan committed
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
"""Semantic segmentation input and model functions for serving/inference."""

import tensorflow as tf

from official.vision.beta.modeling import factory
from official.vision.beta.ops import preprocess_ops
from official.vision.beta.serving import export_base


MEAN_RGB = (0.485 * 255, 0.456 * 255, 0.406 * 255)
STDDEV_RGB = (0.229 * 255, 0.224 * 255, 0.225 * 255)


class SegmentationModule(export_base.ExportModule):
  """Segmentation Module."""

Hongkun Yu's avatar
Hongkun Yu committed
31
  def _build_model(self):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
32
33
34
    input_specs = tf.keras.layers.InputSpec(
        shape=[self._batch_size] + self._input_image_size + [3])

Hongkun Yu's avatar
Hongkun Yu committed
35
    return factory.build_segmentation_model(
Abdullah Rashwan's avatar
Abdullah Rashwan committed
36
        input_specs=input_specs,
Hongkun Yu's avatar
Hongkun Yu committed
37
        model_config=self.params.task.model,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
38
39
40
41
42
43
44
45
46
47
        l2_regularizer=None)

  def _build_inputs(self, image):
    """Builds classification model inputs for serving."""

    # Normalizes image with mean and std pixel values.
    image = preprocess_ops.normalize_image(image,
                                           offset=MEAN_RGB,
                                           scale=STDDEV_RGB)

Fan Yang's avatar
Fan Yang committed
48
    image, image_info = preprocess_ops.resize_and_crop_image(
Abdullah Rashwan's avatar
Abdullah Rashwan committed
49
50
51
52
53
        image,
        self._input_image_size,
        padded_size=self._input_image_size,
        aug_scale_min=1.0,
        aug_scale_max=1.0)
Fan Yang's avatar
Fan Yang committed
54
    return image, image_info
Abdullah Rashwan's avatar
Abdullah Rashwan committed
55

Hongkun Yu's avatar
Hongkun Yu committed
56
  def serve(self, images):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
57
58
59
60
61
62
63
    """Cast image to float and run inference.

    Args:
      images: uint8 Tensor of shape [batch_size, None, None, 3]
    Returns:
      Tensor holding classification output logits.
    """
Fan Yang's avatar
Fan Yang committed
64
65
    # Skip image preprocessing when input_type is tflite so it is compatible
    # with TFLite quantization.
Fan Yang's avatar
Fan Yang committed
66
    image_info = None
Fan Yang's avatar
Fan Yang committed
67
68
69
    if self._input_type != 'tflite':
      with tf.device('cpu:0'):
        images = tf.cast(images, dtype=tf.float32)
Fan Yang's avatar
Fan Yang committed
70
71
72
        images_spec = tf.TensorSpec(
            shape=self._input_image_size + [3], dtype=tf.float32)
        image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)
Fan Yang's avatar
Fan Yang committed
73

Fan Yang's avatar
Fan Yang committed
74
        images, image_info = tf.nest.map_structure(
Fan Yang's avatar
Fan Yang committed
75
76
77
78
            tf.identity,
            tf.map_fn(
                self._build_inputs,
                elems=images,
Fan Yang's avatar
Fan Yang committed
79
                fn_output_signature=(images_spec, image_info_spec),
Fan Yang's avatar
Fan Yang committed
80
                parallel_iterations=32))
Abdullah Rashwan's avatar
Abdullah Rashwan committed
81

Abdullah Rashwan's avatar
Abdullah Rashwan committed
82
83
84
    outputs = self.inference_step(images)
    outputs['logits'] = tf.image.resize(
        outputs['logits'], self._input_image_size, method='bilinear')
Abdullah Rashwan's avatar
Abdullah Rashwan committed
85

Fan Yang's avatar
Fan Yang committed
86
87
88
    if image_info is not None:
      outputs.update({'image_info': image_info})

Abdullah Rashwan's avatar
Abdullah Rashwan committed
89
    return outputs