tf-models-nightly 2.16.0.dev20240223__py2.py3-none-any.whl → 2.16.0.dev20240225__py2.py3-none-any.whl

official/vision/ops/augment.py

@@ -665,6 +665,7 @@ def _fill_rectangle_video(image,
   image_time = tf.shape(image)[0]
   image_height = tf.shape(image)[1]
   image_width = tf.shape(image)[2]
+  image_channels = tf.shape(image)[3]
 
   lower_pad = tf.maximum(0, center_height - half_height)
   upper_pad = tf.maximum(0, image_height - center_height - half_height)
@@ -681,7 +682,7 @@ def _fill_rectangle_video(image,
       padding_dims,
       constant_values=1)
   mask = tf.expand_dims(mask, -1)
-  mask = tf.tile(mask, [1, 1, 1, 3])
+  mask = tf.tile(mask, [1, 1, 1, image_channels])
 
   if replace is None:
     fill = tf.random.normal(tf.shape(image), dtype=image.dtype)

official/vision/ops/augment_test.py

@@ -479,12 +479,14 @@ class MixupAndCutmixTest(tf.test.TestCase, parameterized.TestCase):
     self.assertAllGreaterEqual(aug_labels, label_smoothing / num_classes -
                                1e4)  # With tolerance
 
-  def test_mixup_changes_video(self):
+  @parameterized.product(num_channels=[3, 4])
+  def test_mixup_changes_video(self, num_channels: int):
     batch_size = 12
     num_classes = 1000
     label_smoothing = 0.1
 
-    images = tf.random.normal((batch_size, 8, 224, 224, 3), dtype=tf.float32)
+    images = tf.random.normal(
+        (batch_size, 8, 224, 224, num_channels), dtype=tf.float32)
     labels = tf.range(batch_size)
     augmenter = augment.MixupAndCutmix(
         mixup_alpha=1., cutmix_alpha=0., num_classes=num_classes)
@@ -500,12 +502,14 @@ class MixupAndCutmixTest(tf.test.TestCase, parameterized.TestCase):
                                1e4)  # With tolerance
     self.assertFalse(tf.math.reduce_all(images == aug_images))
 
-  def test_cutmix_changes_video(self):
+  @parameterized.product(num_channels=[3, 4])
+  def test_cutmix_changes_video(self, num_channels: int):
     batch_size = 12
     num_classes = 1000
     label_smoothing = 0.1
 
-    images = tf.random.normal((batch_size, 8, 224, 224, 3), dtype=tf.float32)
+    images = tf.random.normal(
+        (batch_size, 8, 224, 224, num_channels), dtype=tf.float32)
     labels = tf.range(batch_size)
     augmenter = augment.MixupAndCutmix(
         mixup_alpha=0., cutmix_alpha=1., num_classes=num_classes)

official/vision/serving/export_base.py

@@ -125,28 +125,20 @@ class ExportModule(export_base.ExportModule, metaclass=abc.ABCMeta):
   @tf.function
   def inference_from_image_bytes(self, inputs: tf.Tensor):
     with tf.device('cpu:0'):
-      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)
+      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)
     return self.serve(images)
 
   @tf.function
@@ -154,28 +146,20 @@ class ExportModule(export_base.ExportModule, metaclass=abc.ABCMeta):
       self, inputs: tf.Tensor
   ) -> Mapping[str, tf.Tensor]:
     with tf.device('cpu:0'):
-      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)
+      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)
     return self.serve(images)
 
   def get_inference_signatures(self, function_keys: Dict[Text, Text]):

official/vision/serving/image_classification.py

@@ -53,6 +53,41 @@ class ClassificationModule(export_base.ExportModule):
         image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB)
     return image
 
+  def _decode_image(self, encoded_image_bytes: str) -> tf.Tensor:
+    """Decodes an image bytes to an image tensor.
+
+    Use `tf.image.decode_image` to decode an image if input is expected to be 2D
+    image; otherwise use `tf.io.decode_raw` to convert the raw bytes to tensor
+    and reshape it to desire shape.
+
+    Args:
+      encoded_image_bytes: An encoded image string to be decoded.
+
+    Returns:
+      A decoded image tensor.
+    """
+    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
+      )
+      image_tensor.set_shape((None, None, self._num_channels))
+      # Resize image to input_size to support varible image resolutions in a
+      # batch for tf_example and image_bytes input type.
+      image_tensor = tf.image.resize(
+          tf.cast(image_tensor, tf.float32),
+          self._input_image_size,
+          method=tf.image.ResizeMethod.BILINEAR,
+      )
+      image_tensor = tf.cast(image_tensor, tf.uint8)
+    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]
+      )
+    return image_tensor
+
   def serve(self, images):
     """Cast image to float and run inference.
 

official/vision/serving/image_classification_test.py

@@ -29,46 +29,50 @@ from official.vision.serving import image_classification
 
 class ImageClassificationExportTest(tf.test.TestCase, parameterized.TestCase):
 
-  def _get_classification_module(self, input_type):
+  def _get_classification_module(self, input_type, batch_size=1):
     params = exp_factory.get_exp_config('resnet_imagenet')
     params.task.model.backbone.resnet.model_id = 18
     classification_module = image_classification.ClassificationModule(
         params,
-        batch_size=1,
+        batch_size=batch_size,
         input_image_size=[224, 224],
-        input_type=input_type)
+        input_type=input_type,
+    )
     return classification_module
 
   def _export_from_module(self, module, input_type, save_directory):
     signatures = module.get_inference_signatures(
-        {input_type: 'serving_default'})
-    tf.saved_model.save(module,
-                        save_directory,
-                        signatures=signatures)
+        {input_type: 'serving_default'}
+    )
+    tf.saved_model.save(module, save_directory, signatures=signatures)
 
-  def _get_dummy_input(self, input_type):
+  def _get_dummy_input(self, input_type, image_size=224):
     """Get dummy input for the given input type."""
 
     if input_type == 'image_tensor':
-      return tf.zeros((1, 224, 224, 3), dtype=np.uint8)
+      return tf.zeros((1, image_size, image_size, 3), dtype=np.uint8)
     elif input_type == 'image_bytes':
-      image = Image.fromarray(np.zeros((224, 224, 3), dtype=np.uint8))
+      image = Image.fromarray(
+          np.zeros((image_size, image_size, 3), 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((224, 224, 3), dtype=tf.uint8)
+      image_tensor = tf.zeros((image_size, image_size, 3), dtype=tf.uint8)
       encoded_jpeg = tf.image.encode_jpeg(tf.constant(image_tensor)).numpy()
       example = tf.train.Example(
           features=tf.train.Features(
               feature={
-                  'image/encoded':
-                      tf.train.Feature(
-                          bytes_list=tf.train.BytesList(value=[encoded_jpeg])),
-              })).SerializeToString()
+                  'image/encoded': tf.train.Feature(
+                      bytes_list=tf.train.BytesList(value=[encoded_jpeg])
+                  ),
+              }
+          )
+      ).SerializeToString()
       return [example]
     elif input_type == 'tflite':
-      return tf.zeros((1, 224, 224, 3), dtype=np.float32)
+      return tf.zeros((1, image_size, image_size, 3), dtype=np.float32)
 
   @parameterized.parameters(
       {'input_type': 'image_tensor'},
@@ -86,10 +90,14 @@ class ImageClassificationExportTest(tf.test.TestCase, parameterized.TestCase):
     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')))
+    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)
     classification_fn = imported.signatures['serving_default']
@@ -101,6 +109,58 @@ class ImageClassificationExportTest(tf.test.TestCase, parameterized.TestCase):
           tf.map_fn(
               module._build_inputs,
               elems=tf.zeros((1, 224, 224, 3), dtype=tf.uint8),
+              fn_output_signature=tf.TensorSpec(
+                  shape=[224, 224, 3], dtype=tf.float32
+              ),
+          ),
+      )
+    else:
+      processed_images = images
+    expected_logits = module.model(processed_images, training=False)
+    expected_prob = tf.nn.softmax(expected_logits)
+    out = classification_fn(tf.constant(images))
+
+    # The imported model should contain any trackable attrs that the original
+    # model had.
+    self.assertTrue(hasattr(imported.model, 'test_trackable'))
+    self.assertAllClose(out['logits'].numpy(), expected_logits.numpy())
+    self.assertAllClose(out['probs'].numpy(), expected_prob.numpy())
+
+  @parameterized.parameters(
+      {'input_type': 'image_bytes'},
+      {'input_type': 'tf_example'},
+  )
+  def test_multi_size_images_inference(self, input_type='image_tensor'):
+    tmp_dir = self.get_temp_dir()
+    module = self._get_classification_module(input_type, batch_size=None)
+    # Test that the model restores any attrs that are trackable objects
+    # (eg: tables, resource variables, keras models/layers, tf.hub modules).
+    module.model.test_trackable = tf_keras.layers.InputLayer(input_shape=(4,))
+
+    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)
+    classification_fn = imported.signatures['serving_default']
+
+    images = self._get_dummy_input(input_type, image_size=224)
+    images += self._get_dummy_input(input_type, image_size=196)
+
+    if input_type != 'tflite':
+      processed_images = tf.nest.map_structure(
+          tf.stop_gradient,
+          tf.map_fn(
+              module._build_inputs,
+              elems=tf.zeros((2, 224, 224, 3), dtype=tf.uint8),
               fn_output_signature=tf.TensorSpec(
                   shape=[224, 224, 3], dtype=tf.float32)))
     else:

{tf_models_nightly-2.16.0.dev20240223.dist-info → tf_models_nightly-2.16.0.dev20240225.dist-info}/METADATA

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

{tf_models_nightly-2.16.0.dev20240223.dist-info → tf_models_nightly-2.16.0.dev20240225.dist-info}/RECORD

@@ -1060,8 +1060,8 @@ official/vision/ops/anchor.py,sha256=ruGKslsQid6A0yfiJuG6sf_JEbJuHHpWL77D9s3SNx4
 official/vision/ops/anchor_generator.py,sha256=293RH42qBk8zMtfiJovtBJ6Hn9yi4Ln-Az-sGXNAOR0,7274
 official/vision/ops/anchor_generator_test.py,sha256=grVbHuxlu1W7xbjMErV3q0ARxCesHN6q-7LvLcCi-_4,5296
 official/vision/ops/anchor_test.py,sha256=OQPiWHNUkQi3KaMWzOYd9HuqJa9vQ6Za4_6tf79X0qY,7633
-official/vision/ops/augment.py,sha256=cL_g4gHRabtZrEXyDUyjBD-1d0JBib2Wilf2AOTlub0,108239
-official/vision/ops/augment_test.py,sha256=R4A5UZwVc_LCTgLCoTRf25BS8v21kATjUMAFrrc-g28,25711
+official/vision/ops/augment.py,sha256=20K-OWV_GLbpzv8sY5rBe6x3H5MIRbTlDWOPTfNlGeY,108290
+official/vision/ops/augment_test.py,sha256=9QISeqJ79JgkI2-Hn4pMxPRRin3mbOd7mSJ4v7S8W6E,25881
 official/vision/ops/box_matcher.py,sha256=AvZd7CUUZnT4FwETLyVz3Uxb3gO-o94OwlZrvm7CtX0,9067
 official/vision/ops/box_matcher_test.py,sha256=YdxaTITaUyKBA27Pss5MZKF2ibBfSu879222c7lKngU,2438
 official/vision/ops/box_ops.py,sha256=FRKLYzVYd3zFqEYX6Ql4Pr39csvR2Vc-OQv-WxhegFM,34618
@@ -1081,7 +1081,7 @@ official/vision/ops/target_gather_test.py,sha256=yiTZg7w4HYh19pw9yIDppvenscp8MSB
 official/vision/serving/__init__.py,sha256=YlRnCfAvBWlB7gn-Dz32dxVIDBCzxH8PhdM14kHxPgs,702
 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=6HS9Tj1iHhYvkmg7qqOk_l0oGD3uot9hyI4EFJDklbo,8190
+official/vision/serving/export_base.py,sha256=ah8Cuh_cxpptCpRAjptvA00a-yWgd-KokLk3UBLdVt4,7363
 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,8 +1095,8 @@ 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=ikvRyxMwGnFj7ptk6TT9H2YEJmxoZUcAHAihwghXgUk,2936
-official/vision/serving/image_classification_test.py,sha256=y1Mv8UDxdCfs2rRIUBgnLFE9f-4V6uOn592ueNMOc4s,4622
+official/vision/serving/image_classification.py,sha256=Pl2TG7Kjq3-DXh-MywjXQ1MvUmQQLWIer6NgA66iAMM,4300
+official/vision/serving/image_classification_test.py,sha256=qZmuiQewptSQdY2iQEkx8rHjULybgADuXsQ84SjLCok,6759
 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
@@ -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.dev20240223.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.16.0.dev20240223.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.16.0.dev20240223.dist-info/METADATA,sha256=jx-KXEQueFDytdSV3RXmLPWtny7on21QimoauNUnf3s,1432
-tf_models_nightly-2.16.0.dev20240223.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.16.0.dev20240223.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.16.0.dev20240223.dist-info/RECORD,,
+tf_models_nightly-2.16.0.dev20240225.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.16.0.dev20240225.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.16.0.dev20240225.dist-info/METADATA,sha256=G4ySKJnIDr-Te1AAJO0ABym3d7WfOagwR4n0y-GMAJI,1432
+tf_models_nightly-2.16.0.dev20240225.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.16.0.dev20240225.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.16.0.dev20240225.dist-info/RECORD,,