tf-models-nightly 2.17.0.dev20240528__py2.py3-none-any.whl → 2.20.0.dev20251205__py2.py3-none-any.whl

official/vision/serving/detection.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -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,
@@ -124,6 +136,14 @@ class DetectionModule(export_base.ExportModule):
 
     return detections_dict
 
+  def _flatten_output(self, feature_map, feature_size=4):
+    flatten_outputs = []
+    for level_output in feature_map.values():
+      flatten_outputs.append(
+          tf.reshape(level_output, (self._batch_size, -1, feature_size))
+      )
+    return tf.concat(flatten_outputs, axis=1)
+
   def preprocess(
       self, images: tf.Tensor
   ) -> Tuple[tf.Tensor, Mapping[str, tf.Tensor], tf.Tensor]:
@@ -131,20 +151,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 +177,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 +188,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 +204,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 +221,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 +262,35 @@ 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'
+        ]
+    elif (
+        isinstance(self.params.task.model, configs.retinanet.RetinaNet)
+        and not self.params.task.model.detection_generator.decode_boxes
+    ):
+      final_outputs = {
+          'raw_boxes': self._flatten_output(detections['box_outputs'], 4),
+          'raw_scores': tf.sigmoid(
+              self._flatten_output(
+                  detections['cls_outputs'], self.params.task.model.num_classes
+              )
+          ),
+      }
     else:
       # For RetinaNet model, apply export_config.
       if isinstance(self.params.task.model, configs.retinanet.RetinaNet):
@@ -250,7 +301,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/detection_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -38,6 +38,7 @@ class DetectionExportTest(tf.test.TestCase, parameterized.TestCase):
       normalized_coordinates=False,
       nms_version='batched',
       output_intermediate_features=False,
+      decode_boxes=True,
   ):
     params = exp_factory.get_exp_config(experiment_name)
     params.task.model.outer_boxes_scale = outer_boxes_scale
@@ -48,6 +49,8 @@ class DetectionExportTest(tf.test.TestCase, parameterized.TestCase):
     params.task.model.detection_generator.nms_version = nms_version
     if output_intermediate_features:
       params.task.export_config.output_intermediate_features = True
+    if not decode_boxes:
+      params.task.model.detection_generator.decode_boxes = False
     detection_module = detection.DetectionModule(
         params,
         batch_size=1,
@@ -232,6 +235,41 @@ class DetectionExportTest(tf.test.TestCase, parameterized.TestCase):
         max_values.numpy(), tf.ones_like(max_values).numpy()
     )
 
+  @parameterized.parameters(
+      'retinanet_mobile_coco',
+      'retinanet_spinenet_coco',
+  )
+  def test_export_without_decoding_boxes(
+      self,
+      experiment_name,
+  ):
+    input_type = 'tflite'
+    tmp_dir = self.get_temp_dir()
+    module = self._get_detection_module(
+        experiment_name,
+        input_type=input_type,
+        apply_nms=False,
+        decode_boxes=False,
+    )
+
+    self._export_from_module(module, input_type, tmp_dir)
+
+    imported = tf.saved_model.load(tmp_dir)
+    detection_fn = imported.signatures['serving_default']
+
+    images = self._get_dummy_input(
+        input_type, batch_size=1, image_size=(640, 640)
+    )
+    outputs = detection_fn(tf.constant(images))
+
+    self.assertContainsSubset(
+        {
+            'raw_boxes',
+            'raw_scores',
+        },
+        outputs.keys(),
+    )
+
 
 if __name__ == '__main__':
   tf.test.main()

official/vision/serving/export_base.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/export_base_v2.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/export_base_v2_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/export_module_factory.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/export_module_factory_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/export_saved_model.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/export_saved_model_lib.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -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')
+    )

official/vision/serving/export_saved_model_lib_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/export_saved_model_lib_v2.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/export_tfhub.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,7 +14,6 @@
 
 """A script to export the image classification as a TF-Hub SavedModel."""
 
-# Import libraries
 from absl import app
 from absl import flags
 

official/vision/serving/export_tfhub_lib.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,8 +15,6 @@
 """A script to export a TF-Hub SavedModel."""
 from typing import List, Optional
 
-# Import libraries
-
 import tensorflow as tf, tf_keras
 
 from official.core import config_definitions as cfg

official/vision/serving/export_tflite.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/export_tflite_lib.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/export_utils.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/image_classification.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/image_classification_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/semantic_segmentation.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -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 @@ class SegmentationModule(export_base.ExportModule):
     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(