tf-models-nightly 2.17.0.dev20240530__py2.py3-none-any.whl → 2.17.0.dev20240601__py2.py3-none-any.whl

official/common/flags.py

@@ -112,3 +112,12 @@ def define_flags():
 
   flags.DEFINE_string(
       'tpu_platform', default=None, help='TPU platform type.')
+
+  flags.DEFINE_string(
+      'tfhub_handle',
+      None,
+      'TFHub handle for publishing the model to TFHub. The model '
+      'is published to TFHub if this flag is set. Please use a '
+      'TFHubPusher (go/tflex/standard_components/pusher) component if '
+      'running in TFleX.',
+  )

official/vision/serving/detection.py

@@ -41,27 +41,37 @@ class DetectionModule(export_base.ExportModule):
       return self._input_image_size
 
   def _build_model(self):
-
     nms_versions_supporting_dynamic_batch_size = {'batched', 'v2', 'v3'}
     nms_version = self.params.task.model.detection_generator.nms_version
-    if (self._batch_size is None and
-        nms_version not in nms_versions_supporting_dynamic_batch_size):
-      logging.info('nms_version is set to `batched` because `%s` '
-                   'does not support with dynamic batch size.', nms_version)
+    if (
+        self._batch_size is None
+        and nms_version not in nms_versions_supporting_dynamic_batch_size
+    ):
+      logging.info(
+          'nms_version is set to `batched` because `%s` '
+          'does not support with dynamic batch size.',
+          nms_version,
+      )
       self.params.task.model.detection_generator.nms_version = 'batched'
 
-    input_specs = tf_keras.layers.InputSpec(shape=[
-        self._batch_size, *self._padded_size, 3])
+    input_specs = tf_keras.layers.InputSpec(
+        shape=[self._batch_size, *self._padded_size, 3]
+    )
 
     if isinstance(self.params.task.model, configs.maskrcnn.MaskRCNN):
       model = factory.build_maskrcnn(
-          input_specs=input_specs, model_config=self.params.task.model)
+          input_specs=input_specs, model_config=self.params.task.model
+      )
     elif isinstance(self.params.task.model, configs.retinanet.RetinaNet):
       model = factory.build_retinanet(
-          input_specs=input_specs, model_config=self.params.task.model)
+          input_specs=input_specs, model_config=self.params.task.model
+      )
     else:
-      raise ValueError('Detection module not implemented for {} model.'.format(
-          type(self.params.task.model)))
+      raise ValueError(
+          'Detection module not implemented for {} model.'.format(
+              type(self.params.task.model)
+          )
+      )
 
     return model
 
@@ -73,7 +83,8 @@ class DetectionModule(export_base.ExportModule):
         max_level=model_params.max_level,
         num_scales=model_params.anchor.num_scales,
         aspect_ratios=model_params.anchor.aspect_ratios,
-        anchor_size=model_params.anchor.anchor_size)
+        anchor_size=model_params.anchor.anchor_size,
+    )
     return input_anchor(image_size=self._padded_size)
 
   def _build_inputs(self, image):
@@ -85,7 +96,8 @@ class DetectionModule(export_base.ExportModule):
 
     # Normalizes image with mean and std pixel values.
     image = preprocess_ops.normalize_image(
-        image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB)
+        image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB
+    )
 
     image, image_info = preprocess_ops.resize_and_crop_image(
         image,
@@ -131,20 +143,24 @@ class DetectionModule(export_base.ExportModule):
 
     Args:
       images: The images tensor.
+
     Returns:
       images: The images tensor cast to float.
       anchor_boxes: Dict mapping anchor levels to anchor boxes.
       image_info: Tensor containing the details of the image resizing.
-
     """
     model_params = self.params.task.model
     with tf.device('cpu:0'):
       # Tensor Specs for map_fn outputs (images, anchor_boxes, and image_info).
-      images_spec = tf.TensorSpec(shape=self._padded_size + [3],
-                                  dtype=tf.float32)
+      images_spec = tf.TensorSpec(
+          shape=self._padded_size + [3], dtype=tf.float32
+      )
 
-      num_anchors = model_params.anchor.num_scales * len(
-          model_params.anchor.aspect_ratios) * 4
+      num_anchors = (
+          model_params.anchor.num_scales
+          * len(model_params.anchor.aspect_ratios)
+          * 4
+      )
       anchor_shapes = []
       for level in range(model_params.min_level, model_params.max_level + 1):
         anchor_level_spec = tf.TensorSpec(
@@ -153,7 +169,8 @@ class DetectionModule(export_base.ExportModule):
                 math.ceil(self._padded_size[1] / 2**level),
                 num_anchors,
             ],
-            dtype=tf.float32)
+            dtype=tf.float32,
+        )
         anchor_shapes.append((str(level), anchor_level_spec))
 
       image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)
@@ -163,9 +180,14 @@ class DetectionModule(export_base.ExportModule):
           tf.map_fn(
               self._build_inputs,
               elems=images,
-              fn_output_signature=(images_spec, dict(anchor_shapes),
-                                   image_info_spec),
-              parallel_iterations=32))
+              fn_output_signature=(
+                  images_spec,
+                  dict(anchor_shapes),
+                  image_info_spec,
+              ),
+              parallel_iterations=32,
+          ),
+      )
 
       return images, anchor_boxes, image_info
 
@@ -174,6 +196,7 @@ class DetectionModule(export_base.ExportModule):
 
     Args:
       images: uint8 Tensor of shape [batch_size, None, None, 3]
+
     Returns:
       Tensor holding detection output logits.
     """
@@ -190,10 +213,15 @@ class DetectionModule(export_base.ExportModule):
         # [desired_height, desired_width], [y_scale, x_scale],
         # [y_offset, x_offset]]. When input_type is tflite, input image is
         # supposed to be preprocessed already.
-        image_info = tf.convert_to_tensor([[
-            self._input_image_size, self._input_image_size, [1.0, 1.0], [0, 0]
-        ]],
-                                          dtype=tf.float32)
+        image_info = tf.convert_to_tensor(
+            [[
+                self._input_image_size,
+                self._input_image_size,
+                [1.0, 1.0],
+                [0, 0],
+            ]],
+            dtype=tf.float32,
+        )
     input_image_shape = image_info[:, 1, :]
 
     # To overcome keras.Model extra limitation to save a model with layers that
@@ -226,20 +254,23 @@ class DetectionModule(export_base.ExportModule):
         # point outputs.
         if export_config.cast_num_detections_to_float:
           detections['num_detections'] = tf.cast(
-              detections['num_detections'], dtype=tf.float32)
+              detections['num_detections'], dtype=tf.float32
+          )
         if export_config.cast_detection_classes_to_float:
           detections['detection_classes'] = tf.cast(
-              detections['detection_classes'], dtype=tf.float32)
+              detections['detection_classes'], dtype=tf.float32
+          )
 
       final_outputs = {
           'detection_boxes': detections['detection_boxes'],
           'detection_scores': detections['detection_scores'],
           'detection_classes': detections['detection_classes'],
-          'num_detections': detections['num_detections']
+          'num_detections': detections['num_detections'],
       }
       if 'detection_outer_boxes' in detections:
-        final_outputs['detection_outer_boxes'] = (
-            detections['detection_outer_boxes'])
+        final_outputs['detection_outer_boxes'] = detections[
+            'detection_outer_boxes'
+        ]
     else:
       # For RetinaNet model, apply export_config.
       if isinstance(self.params.task.model, configs.retinanet.RetinaNet):
@@ -250,7 +281,7 @@ class DetectionModule(export_base.ExportModule):
           detections = self._normalize_coordinates(detections, keys, image_info)
       final_outputs = {
           'decoded_boxes': detections['decoded_boxes'],
-          'decoded_box_scores': detections['decoded_box_scores']
+          'decoded_box_scores': detections['decoded_box_scores'],
       }
 
     if 'detection_masks' in detections.keys():

official/vision/serving/export_saved_model_lib.py

@@ -15,7 +15,7 @@
 r"""Vision models export utility function for serving/inference."""
 
 import os
-from typing import Optional, List, Union, Text, Dict
+from typing import Dict, List, Optional, Union
 
 from absl import logging
 import tensorflow as tf, tf_keras
@@ -45,7 +45,7 @@ def export_inference_graph(
     log_model_flops_and_params: bool = False,
     checkpoint: Optional[tf.train.Checkpoint] = None,
     input_name: Optional[str] = None,
-    function_keys: Optional[Union[List[Text], Dict[Text, Text]]] = None,
+    function_keys: Optional[Union[List[str], Dict[str, str]]] = None,
     add_tpu_function_alias: Optional[bool] = False,
 ):
   """Exports inference graph for the model specified in the exp config.
@@ -83,57 +83,68 @@ def export_inference_graph(
 
   if export_checkpoint_subdir:
     output_checkpoint_directory = os.path.join(
-        export_dir, export_checkpoint_subdir)
+        export_dir, export_checkpoint_subdir
+    )
   else:
     output_checkpoint_directory = None
 
   if export_saved_model_subdir:
     output_saved_model_directory = os.path.join(
-        export_dir, export_saved_model_subdir)
+        export_dir, export_saved_model_subdir
+    )
   else:
     output_saved_model_directory = export_dir
 
   # TODO(arashwan): Offers a direct path to use ExportModule with Task objects.
   if not export_module:
-    if isinstance(params.task,
-                  configs.image_classification.ImageClassificationTask):
+    if isinstance(
+        params.task, configs.image_classification.ImageClassificationTask
+    ):
       export_module = image_classification.ClassificationModule(
           params=params,
           batch_size=batch_size,
           input_image_size=input_image_size,
           input_type=input_type,
           num_channels=num_channels,
-          input_name=input_name)
+          input_name=input_name,
+      )
     elif isinstance(params.task, configs.retinanet.RetinaNetTask) or isinstance(
-        params.task, configs.maskrcnn.MaskRCNNTask):
+        params.task, configs.maskrcnn.MaskRCNNTask
+    ):
       export_module = detection.DetectionModule(
           params=params,
           batch_size=batch_size,
           input_image_size=input_image_size,
           input_type=input_type,
           num_channels=num_channels,
-          input_name=input_name)
-    elif isinstance(params.task,
-                    configs.semantic_segmentation.SemanticSegmentationTask):
+          input_name=input_name,
+      )
+    elif isinstance(
+        params.task, configs.semantic_segmentation.SemanticSegmentationTask
+    ):
       export_module = semantic_segmentation.SegmentationModule(
           params=params,
           batch_size=batch_size,
           input_image_size=input_image_size,
           input_type=input_type,
           num_channels=num_channels,
-          input_name=input_name)
-    elif isinstance(params.task,
-                    configs.video_classification.VideoClassificationTask):
+          input_name=input_name,
+      )
+    elif isinstance(
+        params.task, configs.video_classification.VideoClassificationTask
+    ):
       export_module = video_classification.VideoClassificationModule(
           params=params,
           batch_size=batch_size,
           input_image_size=input_image_size,
           input_type=input_type,
           num_channels=num_channels,
-          input_name=input_name)
+          input_name=input_name,
+      )
     else:
-      raise ValueError('Export module not implemented for {} task.'.format(
-          type(params.task)))
+      raise ValueError(
+          'Export module not implemented for {} task.'.format(type(params.task))
+      )
 
   if add_tpu_function_alias:
     if input_type == 'image_tensor':
@@ -160,7 +171,8 @@ def export_inference_graph(
       checkpoint=checkpoint,
       checkpoint_path=checkpoint_path,
       timestamped=False,
-      save_options=save_options)
+      save_options=save_options,
+  )
 
   if output_checkpoint_directory:
     ckpt = tf.train.Checkpoint(model=export_module.model)
@@ -171,16 +183,16 @@ def export_inference_graph(
     inputs_kwargs = None
     if isinstance(
         params.task,
-        (configs.retinanet.RetinaNetTask, configs.maskrcnn.MaskRCNNTask)):
+        (configs.retinanet.RetinaNetTask, configs.maskrcnn.MaskRCNNTask),
+    ):
       # We need to create inputs_kwargs argument to specify the input shapes for
       # subclass model that overrides model.call to take multiple inputs,
       # e.g., RetinaNet model.
       inputs_kwargs = {
-          'images':
-              tf.TensorSpec([1] + input_image_size + [num_channels],
-                            tf.float32),
-          'image_shape':
-              tf.TensorSpec([1, 2], tf.float32)
+          'images': tf.TensorSpec(
+              [1] + input_image_size + [num_channels], tf.float32
+          ),
+          'image_shape': tf.TensorSpec([1, 2], tf.float32),
       }
       dummy_inputs = {
           k: tf.ones(v.shape.as_list(), tf.float32)
@@ -191,9 +203,14 @@ def export_inference_graph(
     else:
       logging.info(
           'Logging model flops and params not implemented for %s task.',
-          type(params.task))
+          type(params.task),
+      )
       return
-    train_utils.try_count_flops(export_module.model, inputs_kwargs,
-                                os.path.join(export_dir, 'model_flops.txt'))
-    train_utils.write_model_params(export_module.model,
-                                   os.path.join(export_dir, 'model_params.txt'))
+    train_utils.try_count_flops(
+        export_module.model,
+        inputs_kwargs,
+        os.path.join(export_dir, 'model_flops.txt'),
+    )
+    train_utils.write_model_params(
+        export_module.model, os.path.join(export_dir, 'model_params.txt')
+    )

{tf_models_nightly-2.17.0.dev20240530.dist-info → tf_models_nightly-2.17.0.dev20240601.dist-info}/METADATA

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

{tf_models_nightly-2.17.0.dev20240530.dist-info → tf_models_nightly-2.17.0.dev20240601.dist-info}/RECORD

@@ -3,7 +3,7 @@ official/common/__init__.py,sha256=Yqe5TjLrAR68lQ5G781dwlbbT6mELRByDipCtr4jQY4,6
 official/common/dataset_fn.py,sha256=MSRNhNGvMnZtjkuxFIZ90vGJgKlr2JeGaEYpz82Y3Ew,1858
 official/common/distribute_utils.py,sha256=FC4YgXohuUhp_-ZG5kNoerHdyOtGbRjCVdgKCKh5zyY,8572
 official/common/distribute_utils_test.py,sha256=ltR8wrIbfoX1miJElPQzZN2aioKX6Hx8-mfzoo-_WGc,4914
-official/common/flags.py,sha256=PT3F8-MMG0byEfu-PI5qqBUzl-Y4e2GAPDGLmYKTM-s,4286
+official/common/flags.py,sha256=weXgDeMK7tZ8eY8d5LgMLRQDJYuu2hrJF-XfpOqOP34,4579
 official/common/registry_imports.py,sha256=ez-j0sOvVyYBHHGM5hNjTF0jqHUbHI0zv-4SHe_AZ4s,843
 official/common/streamz_counters.py,sha256=OtdlnI5aosH_DWP8jx2MGEpjs5AhyMztr5E5qhrC4nE,1057
 official/core/__init__.py,sha256=84W8TGHlKgmzJjruHEJ0bGA4E4R_03tEeIkUcIGOP60,1265
@@ -1128,7 +1128,7 @@ 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=uZY4m3_KqkW5dLtFroN3TQuHPUQUXhP4iY9CfpW7UwY,10765
+official/vision/serving/detection.py,sha256=33aMP4oYpWAro0YyD_PGx3i0CmztQ84PmSQTa3HFynU,11022
 official/vision/serving/detection_test.py,sha256=s3ueCplPII8Am7tPSWcUv9KUcAqh1AWb6X91_B6qZIM,9165
 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
@@ -1136,7 +1136,7 @@ official/vision/serving/export_base_v2_test.py,sha256=khuY14W9Oi8LGlv_CvMwFiNnFb
 official/vision/serving/export_module_factory.py,sha256=TRnqqLkGPVhfDFzsUC6dQIfjGA3Ctv7vSLiHkyOc6fg,3549
 official/vision/serving/export_module_factory_test.py,sha256=kDTOmPhdvFXKC57UNo4RR5asLPP5D3hW83gULt4aWlU,4605
 official/vision/serving/export_saved_model.py,sha256=oeP99VFTOTUcxs2H8wFZtJy2sxiIBCn6vVcEEPjAIgQ,5758
-official/vision/serving/export_saved_model_lib.py,sha256=kuDsoWRvxI2sg3b0NgxJKNrMRN_PlpB4go7a6XXSYO8,8123
+official/vision/serving/export_saved_model_lib.py,sha256=KtaQc26BMHAyf6OllnyFFy-mmNcxjne8KVQqjSoMdH8,8119
 official/vision/serving/export_saved_model_lib_test.py,sha256=ggDN65ndXlab7cF0HefDCENDnAhJgs5C3ZCx0boZeeg,2409
 official/vision/serving/export_saved_model_lib_v2.py,sha256=8GQ0IrEB2r2OxsAW5tW3l076HvPSRQTwomSRhiwYddA,3710
 official/vision/serving/export_tfhub.py,sha256=2-LzrgbnxvdawL2v0s2RgClc6R6YidKK2vnbEaBwRSY,3500
@@ -1208,9 +1208,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=nc6A9K53OGqF25xN5St8EiWvdVbda
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
 tensorflow_models/uplift/__init__.py,sha256=mqfa55gweOdpKoaQyid4A_4u7xw__FcQeSIF0k_pYmI,999
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.17.0.dev20240530.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.17.0.dev20240530.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.17.0.dev20240530.dist-info/METADATA,sha256=_wjDgDF_kcyFxGJsy9ZsDCkVP-9a6ScSgVC109WtmVI,1432
-tf_models_nightly-2.17.0.dev20240530.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.17.0.dev20240530.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.17.0.dev20240530.dist-info/RECORD,,
+tf_models_nightly-2.17.0.dev20240601.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.17.0.dev20240601.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.17.0.dev20240601.dist-info/METADATA,sha256=zLRbvBHFR54SOb8VcP59SrXDleBc4QaNJ4rJP12KxzI,1432
+tf_models_nightly-2.17.0.dev20240601.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.17.0.dev20240601.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.17.0.dev20240601.dist-info/RECORD,,