tf-models-nightly 2.16.0.dev20240216__py2.py3-none-any.whl → 2.16.0.dev20240218__py2.py3-none-any.whl

official/vision/serving/detection.py

@@ -78,6 +78,11 @@ class DetectionModule(export_base.ExportModule):
 
   def _build_inputs(self, image):
     """Builds detection model inputs for serving."""
+
+    if isinstance(image, tf.RaggedTensor):
+      image = image.to_tensor()
+    image = tf.cast(image, dtype=tf.float32)
+
     # Normalizes image with mean and std pixel values.
     image = preprocess_ops.normalize_image(
         image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB)
@@ -134,8 +139,6 @@ class DetectionModule(export_base.ExportModule):
     """
     model_params = self.params.task.model
     with tf.device('cpu:0'):
-      images = tf.cast(images, dtype=tf.float32)
-
       # Tensor Specs for map_fn outputs (images, anchor_boxes, and image_info).
       images_spec = tf.TensorSpec(shape=self._padded_size + [3],
                                   dtype=tf.float32)

official/vision/serving/export_base.py

@@ -18,6 +18,7 @@ import abc
 from typing import Dict, List, Mapping, Optional, Text
 
 import tensorflow as tf, tf_keras
+
 from official.core import config_definitions as cfg
 from official.core import export_base
 
@@ -25,15 +26,17 @@ from official.core import export_base
 class ExportModule(export_base.ExportModule, metaclass=abc.ABCMeta):
   """Base Export Module."""
 
-  def __init__(self,
-               params: cfg.ExperimentConfig,
-               *,
-               batch_size: int,
-               input_image_size: List[int],
-               input_type: str = 'image_tensor',
-               num_channels: int = 3,
-               model: Optional[tf_keras.Model] = None,
-               input_name: Optional[str] = None):
+  def __init__(
+      self,
+      params: cfg.ExperimentConfig,
+      *,
+      batch_size: int,
+      input_image_size: List[int],
+      input_type: str = 'image_tensor',
+      num_channels: int = 3,
+      model: Optional[tf_keras.Model] = None,
+      input_name: Optional[str] = None
+  ):
     """Initializes a module for export.
 
     Args:
@@ -72,17 +75,20 @@ class ExportModule(export_base.ExportModule, metaclass=abc.ABCMeta):
     if len(self._input_image_size) == 2:
       # Decode an image if 2D input is expected.
       image_tensor = tf.image.decode_image(
-          encoded_image_bytes, channels=self._num_channels)
+          encoded_image_bytes, channels=self._num_channels
+      )
       image_tensor.set_shape((None, None, self._num_channels))
     else:
       # Convert raw bytes into a tensor and reshape it, if not 2D input.
       image_tensor = tf.io.decode_raw(encoded_image_bytes, out_type=tf.uint8)
-      image_tensor = tf.reshape(image_tensor,
-                                self._input_image_size + [self._num_channels])
+      image_tensor = tf.reshape(
+          image_tensor, self._input_image_size + [self._num_channels]
+      )
     return image_tensor
 
   def _decode_tf_example(
-      self, tf_example_string_tensor: tf.train.Example) -> tf.Tensor:
+      self, tf_example_string_tensor: tf.train.Example
+  ) -> tf.Tensor:
     """Decodes a TF Example to an image tensor.
 
     Args:
@@ -94,7 +100,8 @@ class ExportModule(export_base.ExportModule, metaclass=abc.ABCMeta):
     """
     keys_to_features = {'image/encoded': tf.io.FixedLenFeature((), tf.string)}
     parsed_tensors = tf.io.parse_single_example(
-        serialized=tf_example_string_tensor, features=keys_to_features)
+        serialized=tf_example_string_tensor, features=keys_to_features
+    )
     image_tensor = self._decode_image(parsed_tensors['image/encoded'])
     image_tensor.set_shape(
         [None] * len(self._input_image_size) + [self._num_channels]
@@ -107,7 +114,8 @@ class ExportModule(export_base.ExportModule, metaclass=abc.ABCMeta):
 
   @tf.function
   def inference_from_image_tensors(
-      self, inputs: tf.Tensor) -> Mapping[str, tf.Tensor]:
+      self, inputs: tf.Tensor
+  ) -> Mapping[str, tf.Tensor]:
     return self.serve(inputs)
 
   @tf.function
@@ -117,39 +125,57 @@ class ExportModule(export_base.ExportModule, metaclass=abc.ABCMeta):
   @tf.function
   def inference_from_image_bytes(self, inputs: tf.Tensor):
     with tf.device('cpu:0'):
-      images = tf.nest.map_structure(
-          tf.identity,
-          tf.map_fn(
-              self._decode_image,
-              elems=inputs,
-              fn_output_signature=tf.TensorSpec(
-                  shape=[None] * len(self._input_image_size) +
-                  [self._num_channels],
-                  dtype=tf.uint8),
-              parallel_iterations=32))
-      images = tf.stack(images)
+      if len(inputs) == 1:
+        images = tf.nest.map_structure(
+            tf.identity,
+            tf.map_fn(
+                self._decode_image,
+                elems=inputs,
+                fn_output_signature=tf.TensorSpec(
+                    shape=[None] * len(self._input_image_size)
+                    + [self._num_channels],
+                    dtype=tf.uint8,
+                ),
+                parallel_iterations=32,
+            ),
+        )
+        images = tf.stack(images)
+      else:
+        images = []
+        # Need to use for loop instead of enumerate as enumerate is not well
+        # supported in the exported savedmodel.
+        for i in range(len(inputs)):  # pylint: disable=range-len-with-index-access
+          images.append(tf.cast(self._decode_image(inputs[i]), tf.uint8))
+        images = tf.ragged.stack(images)
     return self.serve(images)
 
   @tf.function
-  def inference_from_tf_example(self,
-                                inputs: tf.Tensor) -> Mapping[str, tf.Tensor]:
+  def inference_from_tf_example(
+      self, inputs: tf.Tensor
+  ) -> Mapping[str, tf.Tensor]:
     with tf.device('cpu:0'):
-      images = tf.nest.map_structure(
-          tf.identity,
-          tf.map_fn(
-              self._decode_tf_example,
-              elems=inputs,
-              # Height/width of the shape of input images is unspecified (None)
-              # at the time of decoding the example, but the shape will
-              # be adjusted to conform to the input layer of the model,
-              # by _run_inference_on_image_tensors() below.
-              fn_output_signature=tf.TensorSpec(
-                  shape=[None] * len(self._input_image_size) +
-                  [self._num_channels],
-                  dtype=tf.uint8),
-              dtype=tf.uint8,
-              parallel_iterations=32))
-      images = tf.stack(images)
+      if len(inputs) == 1:
+        images = tf.nest.map_structure(
+            tf.identity,
+            tf.map_fn(
+                self._decode_tf_example,
+                elems=inputs,
+                fn_output_signature=tf.TensorSpec(
+                    shape=[None] * len(self._input_image_size)
+                    + [self._num_channels],
+                    dtype=tf.uint8,
+                ),
+                parallel_iterations=32,
+            ),
+        )
+        images = tf.stack(images)
+      else:
+        images = []
+        # Need to use for loop instead of enumerate as enumerate is not well
+        # supported in the exported savedmodel.
+        for i in range(len(inputs)):  # pylint: disable=range-len-with-index-access
+          images.append(tf.cast(self._decode_tf_example(inputs[i]), tf.uint8))
+        images = tf.ragged.stack(images)
     return self.serve(images)
 
   def get_inference_signatures(self, function_keys: Dict[Text, Text]):
@@ -167,33 +193,46 @@ class ExportModule(export_base.ExportModule, metaclass=abc.ABCMeta):
     for key, def_name in function_keys.items():
       if key == 'image_tensor':
         input_signature = tf.TensorSpec(
-            shape=[self._batch_size] + [None] * len(self._input_image_size) +
-            [self._num_channels],
+            shape=[self._batch_size]
+            + [None] * len(self._input_image_size)
+            + [self._num_channels],
             dtype=tf.uint8,
-            name=self._input_name)
-        signatures[
-            def_name] = self.inference_from_image_tensors.get_concrete_function(
-                input_signature)
+            name=self._input_name,
+        )
+        signatures[def_name] = (
+            self.inference_from_image_tensors.get_concrete_function(
+                input_signature
+            )
+        )
       elif key == 'image_bytes':
         input_signature = tf.TensorSpec(
-            shape=[self._batch_size], dtype=tf.string, name=self._input_name)
-        signatures[
-            def_name] = self.inference_from_image_bytes.get_concrete_function(
-                input_signature)
+            shape=[self._batch_size], dtype=tf.string, name=self._input_name
+        )
+        signatures[def_name] = (
+            self.inference_from_image_bytes.get_concrete_function(
+                input_signature
+            )
+        )
       elif key == 'serve_examples' or key == 'tf_example':
         input_signature = tf.TensorSpec(
-            shape=[self._batch_size], dtype=tf.string, name=self._input_name)
-        signatures[
-            def_name] = self.inference_from_tf_example.get_concrete_function(
-                input_signature)
+            shape=[self._batch_size], dtype=tf.string, name=self._input_name
+        )
+        signatures[def_name] = (
+            self.inference_from_tf_example.get_concrete_function(
+                input_signature
+            )
+        )
       elif key == 'tflite':
         input_signature = tf.TensorSpec(
-            shape=[self._batch_size] + self._input_image_size +
-            [self._num_channels],
+            shape=[self._batch_size]
+            + self._input_image_size
+            + [self._num_channels],
             dtype=tf.float32,
-            name=self._input_name)
+            name=self._input_name,
+        )
         signatures[def_name] = self.inference_for_tflite.get_concrete_function(
-            input_signature)
+            input_signature
+        )
       else:
         raise ValueError('Unrecognized `input_type`')
     return signatures

official/vision/serving/image_classification.py

@@ -36,6 +36,9 @@ class ClassificationModule(export_base.ExportModule):
   def _build_inputs(self, image):
     """Builds classification model inputs for serving."""
     # Center crops and resizes image.
+    if isinstance(image, tf.RaggedTensor):
+      image = image.to_tensor()
+    image = tf.cast(image, dtype=tf.float32)
     if self.params.task.train_data.aug_crop:
       image = preprocess_ops.center_crop_image(image)
 
@@ -62,8 +65,6 @@ class ClassificationModule(export_base.ExportModule):
     # with TFLite quantization.
     if self._input_type != 'tflite':
       with tf.device('cpu:0'):
-        images = tf.cast(images, dtype=tf.float32)
-
         images = tf.nest.map_structure(
             tf.identity,
             tf.map_fn(

official/vision/serving/semantic_segmentation.py

@@ -35,6 +35,9 @@ class SegmentationModule(export_base.ExportModule):
 
   def _build_inputs(self, image):
     """Builds classification model inputs for serving."""
+    if isinstance(image, tf.RaggedTensor):
+      image = image.to_tensor()
+    image = tf.cast(image, dtype=tf.float32)
 
     # Normalizes image with mean and std pixel values.
     image_feature = self.params.task.train_data.image_feature
@@ -68,7 +71,6 @@ class SegmentationModule(export_base.ExportModule):
     image_info = None
     if self._input_type != 'tflite':
       with tf.device('cpu:0'):
-        images = tf.cast(images, dtype=tf.float32)
         images_spec = tf.TensorSpec(
             shape=self._input_image_size + [3], dtype=tf.float32)
         image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)

{tf_models_nightly-2.16.0.dev20240216.dist-info → tf_models_nightly-2.16.0.dev20240218.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.16.0.dev20240216
+Version: 2.16.0.dev20240218
 Summary: TensorFlow Official Models
 Home-page: https://github.com/tensorflow/models
 Author: Google Inc.

{tf_models_nightly-2.16.0.dev20240216.dist-info → tf_models_nightly-2.16.0.dev20240218.dist-info}/RECORD

@@ -1079,9 +1079,9 @@ official/vision/ops/spatial_transform_ops.py,sha256=PVEJGAn0ygtsrid84vD5GgV0jsjy
 official/vision/ops/target_gather.py,sha256=Ir3X76yXYEVFSYX5h-yfS8SMkY37GYuypBP2C8ykggo,3965
 official/vision/ops/target_gather_test.py,sha256=yiTZg7w4HYh19pw9yIDppvenscp8MSBGlhDko180l08,2561
 official/vision/serving/__init__.py,sha256=YlRnCfAvBWlB7gn-Dz32dxVIDBCzxH8PhdM14kHxPgs,702
-official/vision/serving/detection.py,sha256=5N5bvkWGoBKjP7K2gYxzY3wnL8E5XXEIxKsZXxebq2Y,10693
+official/vision/serving/detection.py,sha256=uZY4m3_KqkW5dLtFroN3TQuHPUQUXhP4iY9CfpW7UwY,10765
 official/vision/serving/detection_test.py,sha256=s3ueCplPII8Am7tPSWcUv9KUcAqh1AWb6X91_B6qZIM,9165
-official/vision/serving/export_base.py,sha256=NnLU1MaLNZYCaZFhd_KqFoFScxfGET4q6rj5QS7bXRY,7541
+official/vision/serving/export_base.py,sha256=6HS9Tj1iHhYvkmg7qqOk_l0oGD3uot9hyI4EFJDklbo,8190
 official/vision/serving/export_base_v2.py,sha256=GYIqt-xaOv4UztAKjx-acD-9i2pjftMw46DWRMy7Bsk,2741
 official/vision/serving/export_base_v2_test.py,sha256=khuY14W9Oi8LGlv_CvMwFiNnFbgpPVzvvD_hugf1_lk,2880
 official/vision/serving/export_module_factory.py,sha256=TRnqqLkGPVhfDFzsUC6dQIfjGA3Ctv7vSLiHkyOc6fg,3549
@@ -1095,9 +1095,9 @@ official/vision/serving/export_tfhub_lib.py,sha256=BZC4XSbly0DLDSlWnphjWaD0Q6SOt
 official/vision/serving/export_tflite.py,sha256=OhIVHrcDvB83p4WroheWcTmeYArMV9TJcgUVJWytWIc,5115
 official/vision/serving/export_tflite_lib.py,sha256=2AWkyEsLvMBE19m2WAa0IpyFsHV8sIR4Gvcv5ZQWbmg,7272
 official/vision/serving/export_utils.py,sha256=8mJb1MF_6kk3lbqZOZq2Lwu4A2L1KWxjnWnV_ZpYlVI,4881
-official/vision/serving/image_classification.py,sha256=a27sm77w-tqFsfm0zF0cA8RaXhR1HwBrVBbaWRJLK4I,2868
+official/vision/serving/image_classification.py,sha256=ikvRyxMwGnFj7ptk6TT9H2YEJmxoZUcAHAihwghXgUk,2936
 official/vision/serving/image_classification_test.py,sha256=y1Mv8UDxdCfs2rRIUBgnLFE9f-4V6uOn592ueNMOc4s,4622
-official/vision/serving/semantic_segmentation.py,sha256=qTXfGJgz6dkygYq9B1YF2HOdzoBc7VjeY0DsCgYDrIU,4091
+official/vision/serving/semantic_segmentation.py,sha256=1RmEOpYz-HjB-VKa3k9p3gZj3h_3ob5d-3RO9_cXfH0,4160
 official/vision/serving/semantic_segmentation_test.py,sha256=QVx12ciMA6T20d3wuKu8pe6ekEyPj8aNbKbRKxOgRik,5638
 official/vision/serving/video_classification.py,sha256=s37SdFoASmX0b3MubTpMdebPOvegx0Nj7yFogu5rYXE,6884
 official/vision/serving/video_classification_test.py,sha256=vx-o4y_mkgLnCrLfw-uWm6S5nrZKIcDcv_1wMEicXrE,4232
@@ -1158,9 +1158,9 @@ tensorflow_models/__init__.py,sha256=etxw45SHxuwFCRX5qGxGMP83II0JfJulzNl5GSNJvhw
 tensorflow_models/tensorflow_models_test.py,sha256=AxUYUdiQn416UR7jg0h6rmv688esvlKDfpyDCIQkF18,1395
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.16.0.dev20240216.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.16.0.dev20240216.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.16.0.dev20240216.dist-info/METADATA,sha256=NUFGQjonQ-9RXFYzDiy6pGG1BjAK00v3l2cA-nD9E84,1432
-tf_models_nightly-2.16.0.dev20240216.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.16.0.dev20240216.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.16.0.dev20240216.dist-info/RECORD,,
+tf_models_nightly-2.16.0.dev20240218.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.16.0.dev20240218.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.16.0.dev20240218.dist-info/METADATA,sha256=B9tEBf5OIvvaK4sZY0DGKxIh_U47mLKz1EplgQqNOv0,1432
+tf_models_nightly-2.16.0.dev20240218.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.16.0.dev20240218.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.16.0.dev20240218.dist-info/RECORD,,