No public description

PiperOrigin-RevId: 671550909

Author: tensorflower-gardener

official/vision/serving/semantic_segmentation.py

@@ -26,7 +26,8 @@ class SegmentationModule(export_base.ExportModule):
 
   def _build_model(self):
     input_specs = tf_keras.layers.InputSpec(
-        shape=[self._batch_size] + self._input_image_size + [3])
+        shape=[self._batch_size] + self._input_image_size + [self._num_channels]
+    )
 
     return factory.build_segmentation_model(
         input_specs=input_specs,
@@ -72,7 +73,9 @@ def serve(self, images):
     if self._input_type != 'tflite':
       with tf.device('cpu:0'):
         images_spec = tf.TensorSpec(
-            shape=self._input_image_size + [3], dtype=tf.float32)
+            shape=self._input_image_size + [self._num_channels],
+            dtype=tf.float32,
+        )
         image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)
 
         images, image_info = tf.nest.map_structure(

official/vision/serving/semantic_segmentation_test.py

@@ -49,20 +49,22 @@ def _export_from_module(self, module, input_type, save_directory):
         {input_type: 'serving_default'})
     tf.saved_model.save(module, save_directory, signatures=signatures)
 
-  def _get_dummy_input(self, input_type, input_image_size):
+  def _get_dummy_input(self, input_type, input_image_size, num_channels):
     """Get dummy input for the given input type."""
 
     height = input_image_size[0]
     width = input_image_size[1]
     if input_type == 'image_tensor':
-      return tf.zeros((1, height, width, 3), dtype=np.uint8)
+      return tf.zeros((1, height, width, num_channels), dtype=np.uint8)
     elif input_type == 'image_bytes':
-      image = Image.fromarray(np.zeros((height, width, 3), dtype=np.uint8))
+      image = Image.fromarray(
+          np.zeros((height, width, num_channels), dtype=np.uint8)
+      )
       byte_io = io.BytesIO()
       image.save(byte_io, 'PNG')
       return [byte_io.getvalue()]
     elif input_type == 'tf_example':
-      image_tensor = tf.zeros((height, width, 3), dtype=tf.uint8)
+      image_tensor = tf.zeros((height, width, num_channels), dtype=tf.uint8)
       encoded_jpeg = tf.image.encode_jpeg(tf.constant(image_tensor)).numpy()
       example = tf.train.Example(
           features=tf.train.Features(
@@ -73,7 +75,7 @@ def _get_dummy_input(self, input_type, input_image_size):
               })).SerializeToString()
       return [example]
     elif input_type == 'tflite':
-      return tf.zeros((1, height, width, 3), dtype=np.float32)
+      return tf.zeros((1, height, width, num_channels), dtype=np.float32)
 
   @parameterized.parameters(
       ('image_tensor', False, [112, 112], False),
@@ -105,7 +107,7 @@ def test_export(self, input_type, rescale_output, input_image_size,
     imported = tf.saved_model.load(tmp_dir)
     segmentation_fn = imported.signatures['serving_default']
 
-    images = self._get_dummy_input(input_type, input_image_size)
+    images = self._get_dummy_input(input_type, input_image_size, num_channels=3)
     if input_type != 'tflite':
       processed_images, _ = tf.nest.map_structure(
           tf.stop_gradient,
@@ -128,6 +130,68 @@ def test_export(self, input_type, rescale_output, input_image_size,
     out = segmentation_fn(tf.constant(images))
     self.assertAllClose(out['logits'].numpy(), expected_output.numpy())
 
+  @parameterized.parameters(
+      ('image_tensor',),
+      ('tflite',),
+  )
+  def test_export_with_extra_input_channels(self, input_type):
+    tmp_dir = self.get_temp_dir()
+    num_channels = 6
+    params = exp_factory.get_exp_config('mnv2_deeplabv3_pascal')
+    params.task.init_checkpoint = None
+    params.task.model.input_size = [112, 112, num_channels]
+    params.task.export_config.rescale_output = False
+    params.task.train_data.preserve_aspect_ratio = False
+    params.task.train_data.image_feature.mean = [0.5] * num_channels
+    params.task.train_data.image_feature.stddev = [0.5] * num_channels
+    params.task.train_data.image_feature.num_channels = num_channels
+    module = semantic_segmentation.SegmentationModule(
+        params,
+        batch_size=1,
+        input_image_size=[112, 112],
+        input_type=input_type,
+        num_channels=num_channels,
+    )
+
+    self._export_from_module(module, input_type, tmp_dir)
+
+    self.assertTrue(os.path.exists(os.path.join(tmp_dir, 'saved_model.pb')))
+    self.assertTrue(
+        os.path.exists(os.path.join(tmp_dir, 'variables', 'variables.index'))
+    )
+    self.assertTrue(
+        os.path.exists(
+            os.path.join(tmp_dir, 'variables', 'variables.data-00000-of-00001')
+        )
+    )
+
+    imported = tf.saved_model.load(tmp_dir)
+    segmentation_fn = imported.signatures['serving_default']
+
+    images = self._get_dummy_input(input_type, [112, 112], num_channels)
+
+    if input_type != 'tflite':
+      processed_images, _ = tf.nest.map_structure(
+          tf.stop_gradient,
+          tf.map_fn(
+              module._build_inputs,
+              elems=tf.zeros((1, 112, 112, num_channels), dtype=tf.uint8),
+              fn_output_signature=(
+                  tf.TensorSpec(
+                      shape=[112, 112, num_channels], dtype=tf.float32
+                  ),
+                  tf.TensorSpec(shape=[4, 2], dtype=tf.float32),
+              ),
+          ),
+      )
+    else:
+      processed_images = images
+
+    logits = module.model(processed_images, training=False)['logits']
+    expected_output = tf.image.resize(logits, [112, 112], method='bilinear')
+    out = segmentation_fn(tf.constant(images))
+    self.assertAllClose(out['logits'].numpy(), expected_output.numpy())
+
   def test_export_invalid_batch_size(self):
     batch_size = 3
     tmp_dir = self.get_temp_dir()