tf-models-nightly 2.14.0.dev20231009__py2.py3-none-any.whl → 2.14.0.dev20231011__py2.py3-none-any.whl

official/vision/configs/semantic_segmentation.py

@@ -198,6 +198,7 @@ class SemanticSegmentationTask(cfg.TaskConfig):
   init_checkpoint_modules: Union[
       str, List[str]] = 'all'  # all, backbone, and/or decoder
   export_config: ExportConfig = dataclasses.field(default_factory=ExportConfig)
+  allow_image_summary: bool = True
 
 
 @exp_factory.register_config_factory('semantic_segmentation')

official/vision/tasks/semantic_segmentation.py

@@ -29,6 +29,7 @@ from official.vision.dataloaders import tfds_factory
 from official.vision.evaluation import segmentation_metrics
 from official.vision.losses import segmentation_losses
 from official.vision.modeling import factory
+from official.vision.utils.object_detection import visualization_utils
 
 
 @task_factory.register_task_cls(exp_cfg.SemanticSegmentationTask)
@@ -321,6 +322,14 @@ class SemanticSegmentationTask(base_task.Task):
     if metrics:
       self.process_metrics(metrics, labels, outputs)
 
+    if (
+        hasattr(self.task_config, 'allow_image_summary')
+        and self.task_config.allow_image_summary
+    ):
+      logs.update(
+          {'visualization': (tf.cast(features, dtype=tf.float32), outputs)}
+      )
+
     return logs
 
   def inference_step(self, inputs: tf.Tensor, model: tf.keras.Model):
@@ -330,17 +339,37 @@ class SemanticSegmentationTask(base_task.Task):
   def aggregate_logs(self, state=None, step_outputs=None):
     if state is None and self.iou_metric is not None:
       self.iou_metric.reset_states()
-      state = self.iou_metric
+
+    if 'visualization' in step_outputs:
+      # Update segmentation state for writing summary if there are artifacts for
+      # visualization.
+      if state is None:
+        state = {}
+      state.update(visualization_utils.update_segmentation_state(step_outputs))
+
+    if state is None:
+      # Create an arbitrary state to indicate it's not the first step in the
+      # following calls to this function.
+      state = True
+
     return state
 
   def reduce_aggregated_logs(self, aggregated_logs, global_step=None):
-    result = {}
+    logs = {}
     if self.iou_metric is not None:
       ious = self.iou_metric.result()
       # TODO(arashwan): support loading class name from a label map file.
       if self.task_config.evaluation.report_per_class_iou:
         for i, value in enumerate(ious.numpy()):
-          result.update({'iou/{}'.format(i): value})
+          logs.update({'iou/{}'.format(i): value})
       # Computes mean IoU
-      result.update({'mean_iou': tf.reduce_mean(ious)})
-    return result
+      logs.update({'mean_iou': tf.reduce_mean(ious)})
+
+    # Add visualization for summary.
+    if isinstance(aggregated_logs, dict) and 'image' in aggregated_logs:
+      validation_outputs = visualization_utils.visualize_segmentation_outputs(
+          logs=aggregated_logs, task_config=self.task_config
+      )
+      logs.update(validation_outputs)
+
+    return logs

official/vision/utils/object_detection/visualization_utils.py

@@ -894,3 +894,160 @@ def update_detection_state(step_outputs=None) -> Dict[str, Any]:
       state['detection_masks'] = tf.concat(detection_masks, axis=0)
 
   return state
+
+
+def update_segmentation_state(step_outputs=None) -> Dict[str, Any]:
+  """Updates segmentation state to optionally add input image and predictions."""
+  state = {}
+  if step_outputs:
+    state['image'] = tf.concat(step_outputs['visualization'][0], axis=0)
+    state['logits'] = tf.concat(
+        step_outputs['visualization'][1]['logits'], axis=0
+    )
+  return state
+
+
+def visualize_segmentation_outputs(
+    logs,
+    task_config,
+    original_image_spatial_shape=None,
+    true_image_shape=None,
+    image_mean: Optional[Union[float, List[float]]] = None,
+    image_std: Optional[Union[float, List[float]]] = None,
+    key: str = 'image/validation_outputs',
+) -> Dict[str, Any]:
+  """Visualizes the detection outputs.
+
+  It extracts images and predictions from logs and draws visualization on input
+  images. By default, it requires `detection_boxes`, `detection_classes` and
+  `detection_scores` in the prediction, and optionally accepts
+  `detection_keypoints` and `detection_masks`.
+
+  Args:
+    logs: A dictionaty of log that contains images and predictions.
+    task_config: A task config.
+    original_image_spatial_shape: A [N, 2] tensor containing the spatial size of
+      the original image.
+    true_image_shape: A [N, 3] tensor containing the spatial size of unpadded
+      original_image.
+    image_mean: An optional float or list of floats used as the mean pixel value
+      to normalize images.
+    image_std: An optional float or list of floats used as the std to normalize
+      images.
+    key: A string specifying the key of the returned dictionary.
+
+  Returns:
+    A dictionary of images with visualization drawn on it. Each key corresponds
+      to a 4D tensor with segments drawn on each image.
+  """
+  images = logs['image']
+  masks = np.argmax(logs['logits'], axis=-1)
+  num_classes = task_config.model.num_classes
+
+  def _denormalize_images(images: tf.Tensor) -> tf.Tensor:
+    if image_mean is None and image_std is None:
+      images *= tf.constant(
+          preprocess_ops.STDDEV_RGB, shape=[1, 1, 3], dtype=images.dtype
+      )
+      images += tf.constant(
+          preprocess_ops.MEAN_RGB, shape=[1, 1, 3], dtype=images.dtype
+      )
+    elif image_mean is not None and image_std is not None:
+      if isinstance(image_mean, float) and isinstance(image_std, float):
+        images = images * image_std + image_mean
+      elif isinstance(image_mean, list) and isinstance(image_std, list):
+        images *= tf.constant(image_std, shape=[1, 1, 3], dtype=images.dtype)
+        images += tf.constant(image_mean, shape=[1, 1, 3], dtype=images.dtype)
+      else:
+        raise ValueError(
+            '`image_mean` and `image_std` should be the same type.'
+        )
+    else:
+      raise ValueError(
+          'Both `image_mean` and `image_std` should be set or None at the same '
+          'time.'
+      )
+    return tf.cast(images, dtype=tf.uint8)
+
+  images = tf.nest.map_structure(
+      tf.identity,
+      tf.map_fn(
+          _denormalize_images,
+          elems=images,
+          fn_output_signature=tf.TensorSpec(
+              shape=images.shape.as_list()[1:], dtype=tf.uint8
+          ),
+          parallel_iterations=32,
+      ),
+  )
+
+  if images.shape[3] > 3:
+    images = images[:, :, :, 0:3]
+  elif images.shape[3] == 1:
+    images = tf.image.grayscale_to_rgb(images)
+  if true_image_shape is None:
+    true_shapes = tf.constant(-1, shape=[images.shape.as_list()[0], 3])
+  else:
+    true_shapes = true_image_shape
+  if original_image_spatial_shape is None:
+    original_shapes = tf.constant(-1, shape=[images.shape.as_list()[0], 2])
+  else:
+    original_shapes = original_image_spatial_shape
+
+  visualize_fn = functools.partial(_visualize_masks, num_classes=num_classes)
+  elems = [true_shapes, original_shapes, images, masks]
+
+  def draw_segments(image_and_segments):
+    """Draws boxes on image."""
+    true_shape = image_and_segments[0]
+    original_shape = image_and_segments[1]
+    if true_image_shape is not None:
+      image = shape_utils.pad_or_clip_nd(
+          image_and_segments[2], [true_shape[0], true_shape[1], 3]
+      )
+    if original_image_spatial_shape is not None:
+      image_and_segments[2] = _resize_original_image(image, original_shape)
+
+    image_with_boxes = tf.compat.v1.py_func(
+        visualize_fn, image_and_segments[2:], tf.uint8
+    )
+    return image_with_boxes
+
+  images_with_segments = tf.map_fn(
+      draw_segments, elems, dtype=tf.uint8, back_prop=False
+  )
+
+  outputs = {}
+  for i, image in enumerate(images_with_segments):
+    outputs[key + f'/{i}'] = image[None, ...]
+
+  return outputs
+
+
+def _visualize_masks(image, mask, num_classes, alpha=0.4):
+  """Visualizes semantic segmentation masks."""
+  solid_color = np.repeat(
+      np.expand_dims(np.zeros_like(mask), axis=2), 3, axis=2
+  )
+  for i in range(num_classes):
+    color = STANDARD_COLORS[i % len(STANDARD_COLORS)]
+    rgb = ImageColor.getrgb(color)
+    one_class_mask = np.where(mask == i, 1, 0)
+    solid_color = solid_color + np.expand_dims(
+        one_class_mask, axis=2
+    ) * np.reshape(list(rgb), [1, 1, 3])
+
+  pil_image = Image.fromarray(image)
+  pil_solid_color = (
+      Image.fromarray(np.uint8(solid_color))
+      .convert('RGBA')
+      .resize(pil_image.size)
+  )
+  pil_mask = (
+      Image.fromarray(np.uint8(255.0 * alpha * np.ones_like(mask)))
+      .convert('L')
+      .resize(pil_image.size)
+  )
+  pil_image = Image.composite(pil_solid_color, pil_image, pil_mask)
+  np.copyto(image, np.array(pil_image.convert('RGB')))
+  return image

{tf_models_nightly-2.14.0.dev20231009.dist-info → tf_models_nightly-2.14.0.dev20231011.dist-info}/METADATA

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

{tf_models_nightly-2.14.0.dev20231009.dist-info → tf_models_nightly-2.14.0.dev20231011.dist-info}/RECORD

@@ -874,7 +874,7 @@ official/vision/configs/maskrcnn.py,sha256=k5TT6AheXFJ1y_o-7Qxn0lsVyfkRngg0pERWc
 official/vision/configs/maskrcnn_test.py,sha256=t3Yx1GV4lCrQANwacC99NJf3Cy_guDIOii5qUvidG9M,1723
 official/vision/configs/retinanet.py,sha256=jYXROb2d3c1796V7AaAeQJQyfDLM2MQdD3I6NAODN9A,17719
 official/vision/configs/retinanet_test.py,sha256=yeNLguvsPCpxv1BQAJJCQrrDkveNA1gmZlDya9CIwck,1689
-official/vision/configs/semantic_segmentation.py,sha256=JkdKl_FKhSZq1wUcQUc_ersiOOgnPTytavanBX2ppis,30581
+official/vision/configs/semantic_segmentation.py,sha256=PJsmPt4bdORFFpVpswouOyK_yGiXbXLxRTxxklQBFzA,30616
 official/vision/configs/semantic_segmentation_test.py,sha256=EXXKMekGrik0uJ1O4zZ5AeN7BPDjm3a2uPDUTeDoL-A,1857
 official/vision/configs/video_classification.py,sha256=eLWT3ClzQAoJFxayR1AdGnFfiI0XYVOzrccYOGYx_FE,14513
 official/vision/configs/video_classification_test.py,sha256=tEprp-PLabvA0d7x52FtcCT6AmmRGPlP4dOEJ63fojU,1869
@@ -1059,7 +1059,7 @@ official/vision/tasks/__init__.py,sha256=pwY_FnD2dwCOuYUHfZoAvnONwDWalX5E1J27tSJ
 official/vision/tasks/image_classification.py,sha256=Q7TgcKEhHLP_64jlLbpu42Y8Hzdf3VhOm4Xb5dlEx1o,16689
 official/vision/tasks/maskrcnn.py,sha256=Gc5cu_W48bMz_5d8HL8UUR3uJTL0PMvKm3ZnyO2ALk4,25421
 official/vision/tasks/retinanet.py,sha256=UKiwdO7Id53nqpigW9Hmd-7xGtdEb0il_y5K3Zz3Fkw,18104
-official/vision/tasks/semantic_segmentation.py,sha256=rbUWzGJoI0ho_b0WYEGwpAidQ7vi-IDN7a1cN7V8eQ8,13246
+official/vision/tasks/semantic_segmentation.py,sha256=QA6qyxjdxptBRneGhECOtubvGo4ZyG17fFBV6Yz0hlY,14233
 official/vision/tasks/video_classification.py,sha256=E_AQwPvSzMx2MhnQctHKGKKshot-R9lwYLE6HFEux5M,14308
 official/vision/utils/__init__.py,sha256=1ToRMjre4mErL4Ek4_dMVxMjXNPossNXggV8fqbISao,609
 official/vision/utils/ops_test.py,sha256=g2a5Hij3N-nnA5vSBhootzeapYAMrZRvadXNUXtwnpc,4384
@@ -1078,7 +1078,7 @@ official/vision/utils/object_detection/preprocessor.py,sha256=Fklxck111TQEKsrWSM
 official/vision/utils/object_detection/region_similarity_calculator.py,sha256=yqQLLRT80IdAu3K_fILli_u1aL37lv0FpDtFcyRrPzs,4544
 official/vision/utils/object_detection/shape_utils.py,sha256=p3Q7e9gTTQNv1gnMAkuTkfXc6DYVB8mjk_Vjlq7bRlg,3608
 official/vision/utils/object_detection/target_assigner.py,sha256=fTkjedzhp_-RTUGR27tWbmnnLpW7F3lVCkZlr7Nv_9o,24198
-official/vision/utils/object_detection/visualization_utils.py,sha256=RzK96moeYTp90jn2uJ4l9NBI9-hSKuzM8vKhy4ZIl5Y,34757
+official/vision/utils/object_detection/visualization_utils.py,sha256=-g75oXcrcchAiJC4_63Z4CifJLN4Km6jUIyqIa4VY4Y,40256
 orbit/__init__.py,sha256=aQRo8zqIQ0Dw4JQReZeiB6MmuJLvvw4DbYHYti5AGys,1117
 orbit/controller.py,sha256=pQo060KNNWekzyyRuGBTNhVS0_9D4AJWylaW_Govhuc,25082
 orbit/controller_test.py,sha256=5b8JCx3s8mfzqYoOa1D3xbD-Qr5uIXcTXM3pLEGjRjo,31792
@@ -1112,9 +1112,9 @@ tensorflow_models/__init__.py,sha256=Ciz_YBke6teb6y42QyQTUBDdXJAiV7Qdu1zOoZvYiKw
 tensorflow_models/tensorflow_models_test.py,sha256=Kz2y4V-rtBhZFFfKD2soCq52hviSfJVV1L2ztqS-9oM,1385
 tensorflow_models/nlp/__init__.py,sha256=3dULDpUBpDi9vljpXadq6oJrWH4y6z42Bz2d3hopYZw,807
 tensorflow_models/vision/__init__.py,sha256=4y77XkHaH8qLls3-6ta4tMp3Xj8CLbB0ihH91HsQ9z4,833
-tf_models_nightly-2.14.0.dev20231009.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.14.0.dev20231009.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.14.0.dev20231009.dist-info/METADATA,sha256=VpWhHMvusYxGhF6UT1IFQsv5QhwBMma0iSVRvgRViig,1390
-tf_models_nightly-2.14.0.dev20231009.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.14.0.dev20231009.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.14.0.dev20231009.dist-info/RECORD,,
+tf_models_nightly-2.14.0.dev20231011.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.14.0.dev20231011.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.14.0.dev20231011.dist-info/METADATA,sha256=v1HZmRzGLdvqOpSYUi9Su-F1bbv4p1zqkuXEGmUXSCo,1390
+tf_models_nightly-2.14.0.dev20231011.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.14.0.dev20231011.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.14.0.dev20231011.dist-info/RECORD,,