tf-models-nightly 2.18.0.dev20240820__py2.py3-none-any.whl → 2.18.0.dev20240822__py2.py3-none-any.whl

official/projects/maskconver/tasks/multiscale_maskconver.py

@@ -0,0 +1,278 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Panoptic Multi-scale MaskConver task definition."""
+from typing import Any, Dict, List, Mapping, Optional, Tuple
+import tensorflow as tf, tf_keras
+
+from official.common import dataset_fn
+from official.core import task_factory
+from official.projects.maskconver.configs import multiscale_maskconver as exp_cfg
+from official.projects.maskconver.dataloaders import multiscale_maskconver_input
+from official.projects.maskconver.losses import maskconver_losses
+from official.projects.maskconver.modeling import factory
+from official.projects.maskconver.modeling.layers import copypaste
+from official.projects.maskconver.tasks import maskconver
+from official.projects.volumetric_models.losses import segmentation_losses as volumeteric_segmentation_losses
+from official.vision.dataloaders import input_reader_factory
+
+
+@task_factory.register_task_cls(exp_cfg.MultiScaleMaskConverTask)
+class PanopticMultiScaleMaskConverTask(maskconver.PanopticMaskRCNNTask):
+
+  """A single-replica view of training procedure.
+
+  Panoptic Mask R-CNN task provides artifacts for training/evalution procedures,
+  including loading/iterating over Datasets, initializing the model, calculating
+  the loss, post-processing, and customized metrics with reduction.
+  """
+
+  def build_model(self) -> tf_keras.Model:
+    """Build Panoptic Mask R-CNN model."""
+
+    tf_keras.utils.set_random_seed(0)
+    tf.config.experimental.enable_op_determinism()
+    input_specs = tf_keras.layers.InputSpec(
+        shape=[None] + self.task_config.model.input_size)
+
+    l2_weight_decay = self.task_config.losses.l2_weight_decay
+    # Divide weight decay by 2.0 to match the implementation of tf.nn.l2_loss.
+    # (https://www.tensorflow.org/api_docs/python/tf/keras/regularizers/l2)
+    # (https://www.tensorflow.org/api_docs/python/tf/nn/l2_loss)
+    l2_regularizer = (tf_keras.regularizers.l2(
+        l2_weight_decay / 2.0) if l2_weight_decay else None)
+
+    model = factory.build_multiscale_maskconver_model(
+        input_specs=input_specs,
+        model_config=self.task_config.model,
+        l2_regularizer=l2_regularizer)
+
+    # Get images and labels with batch size of 1.
+    images, labels = next(
+        iter(self.build_inputs(self.task_config.validation_data)))
+    images = tf.nest.map_structure(lambda x: x[0:1, ...], images)
+    labels = tf.nest.map_structure(lambda x: x[0:1, ...], labels)
+    _ = model(
+        images,
+        image_info=labels['image_info'],
+        training=False)
+    return model
+
+  def build_inputs(
+      self,
+      params: exp_cfg.DataConfig,
+      input_context: Optional[tf.distribute.InputContext] = None
+  ) -> tf.data.Dataset:
+    """Build input dataset."""
+    decoder_cfg = params.decoder.get()
+
+    if params.decoder.type == 'simple_decoder':
+      decoder = multiscale_maskconver_input.TfExampleDecoder(
+          regenerate_source_id=decoder_cfg.regenerate_source_id,
+          mask_binarize_threshold=decoder_cfg.mask_binarize_threshold,
+          include_panoptic_masks=decoder_cfg.include_panoptic_masks,
+          panoptic_category_mask_key=decoder_cfg.panoptic_category_mask_key,
+          panoptic_instance_mask_key=decoder_cfg.panoptic_instance_mask_key)
+    else:
+      raise ValueError('Unknown decoder type: {}!'.format(params.decoder.type))
+
+    if params.parser.copypaste:
+      sample_fn = copypaste.CopyPaste(
+          self.task_config.model.input_size[:2],
+          copypaste_frequency=params.parser.copypaste.copypaste_frequency,
+          copypaste_aug_scale_max=params.parser.copypaste.copypaste_aug_scale_max,
+          copypaste_aug_scale_min=params.parser.copypaste.copypaste_aug_scale_min,
+          aug_scale_min=params.parser.copypaste.aug_scale_min,
+          aug_scale_max=params.parser.copypaste.aug_scale_max,
+          random_flip=params.parser.aug_rand_hflip,
+          num_thing_classes=self.task_config.model.num_thing_classes)
+    else:
+      sample_fn = None
+
+    parser = multiscale_maskconver_input.Parser(
+        output_size=self.task_config.model.input_size[:2],
+        min_level=self.task_config.model.min_level,
+        max_level=self.task_config.model.max_level,
+        fpn_low_range=params.parser.fpn_low_range,
+        fpn_high_range=params.parser.fpn_high_range,
+        dtype=params.dtype,
+        aug_rand_hflip=params.parser.aug_rand_hflip,
+        aug_scale_min=params.parser.aug_scale_min,
+        aug_scale_max=params.parser.aug_scale_max,
+        max_num_instances=params.parser.max_num_instances,
+        segmentation_resize_eval_groundtruth=params.parser
+        .segmentation_resize_eval_groundtruth,
+        segmentation_groundtruth_padded_size=params.parser
+        .segmentation_groundtruth_padded_size,
+        segmentation_ignore_label=params.parser.segmentation_ignore_label,
+        panoptic_ignore_label=params.parser.panoptic_ignore_label,
+        num_panoptic_categories=self.task_config.model.num_classes,
+        num_thing_categories=self.task_config.model.num_thing_classes,
+        mask_target_level=params.parser.mask_target_level,
+        level=self.task_config.model.level,
+        gaussian_iou=params.parser.gaussaian_iou,
+        aug_type=params.parser.aug_type,)
+
+    reader = input_reader_factory.input_reader_generator(
+        params,
+        dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
+        sample_fn=sample_fn.copypaste_fn(
+            params.is_training) if sample_fn else None,
+        decoder_fn=decoder.decode,
+        parser_fn=parser.parse_fn(params.is_training))
+    dataset = reader.read(input_context=input_context)
+
+    return dataset
+
+  def build_losses(self,
+                   outputs: Mapping[str, Any],
+                   labels: Mapping[str, Any],
+                   iteration: Any,
+                   aux_losses: Optional[Any] = None,
+                   step=None) -> Dict[str, tf.Tensor]:
+    """Build Panoptic Mask R-CNN losses."""
+    # pylint: disable=line-too-long
+    loss_params = self._task_config.losses
+    center_loss_fn = maskconver_losses.PenaltyReducedLogisticFocalLoss(
+        alpha=loss_params.alpha, beta=loss_params.beta)
+
+    true_flattened_ct_heatmap = labels['panoptic_heatmaps']
+    true_flattened_ct_heatmap = tf.cast(true_flattened_ct_heatmap, tf.float32)
+
+    pred_flattened_ct_heatmap = outputs['class_heatmaps']
+    pred_flattened_ct_heatmap = tf.cast(pred_flattened_ct_heatmap, tf.float32)
+
+    center_loss = center_loss_fn(
+        target_tensor=true_flattened_ct_heatmap,
+        prediction_tensor=pred_flattened_ct_heatmap,
+        weights=1.0)
+
+    replica_context = tf.distribute.get_replica_context()
+    global_num_instances = replica_context.all_reduce(
+        tf.distribute.ReduceOp.SUM, labels['num_instances'])
+    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
+    num_instances = tf.cast(global_num_instances, tf.float32) / tf.cast(num_replicas, tf.float32) + 1.0
+
+    center_loss = tf.reduce_sum(center_loss) / num_instances
+
+    gt_masks = labels['panoptic_masks']
+    gt_mask_weights = labels['panoptic_mask_weights'][:, None, None, :] * tf.ones_like(gt_masks)
+    panoptic_padding_mask = labels['panoptic_padding_mask'][:, :, :, None] * tf.ones_like(gt_masks)
+
+    # gt_masks
+    _, h, w, q = gt_masks.get_shape().as_list()
+    predicted_masks = tf.cast(outputs['mask_proposal_logits'], tf.float32)
+    predicted_masks = tf.image.resize(
+        predicted_masks, tf.shape(gt_masks)[1:3], method='bilinear')
+
+    mask_loss_fn = tf_keras.losses.BinaryCrossentropy(
+        from_logits=True,
+        label_smoothing=0.0,
+        axis=-1,
+        reduction=tf_keras.losses.Reduction.NONE,
+        name='binary_crossentropy')
+
+    mask_weights = tf.cast(gt_masks >= 0, tf.float32) * gt_mask_weights  * (
+        1 - panoptic_padding_mask)  # b, h, w, # max inst
+    mask_loss = mask_loss_fn(
+        tf.expand_dims(gt_masks, -1),
+        tf.expand_dims(predicted_masks, -1),
+        sample_weight=tf.expand_dims(mask_weights, -1))
+
+    mask_loss = tf.reshape(mask_loss, [-1, h * w, q])
+    mask_loss = tf.reduce_sum(tf.reduce_mean(mask_loss, axis=1)) / num_instances
+
+    # Dice loss
+    masked_predictions = tf.sigmoid(predicted_masks) * tf.cast(
+        gt_mask_weights > 0, tf.float32) * (1 - panoptic_padding_mask)
+    masked_gt_masks = gt_masks * tf.cast(gt_mask_weights > 0, tf.float32) * (
+        1 - panoptic_padding_mask)
+
+    masked_predictions = tf.transpose(masked_predictions, [0, 3, 1, 2])
+    masked_predictions = tf.reshape(masked_predictions, [-1, h, w, 1])
+    masked_gt_masks = tf.transpose(masked_gt_masks, [0, 3, 1, 2])
+    masked_gt_masks = tf.reshape(masked_gt_masks, [-1, h, w, 1])
+
+    dice_loss_fn = volumeteric_segmentation_losses.SegmentationLossDiceScore(
+        metric_type='adaptive', axis=(2, 3))
+    dice_loss = dice_loss_fn(logits=masked_predictions, labels=masked_gt_masks)
+
+    total_loss = center_loss + loss_params.mask_weight * (mask_loss + dice_loss)
+    if aux_losses:
+      total_loss += tf.add_n(aux_losses)
+
+    total_loss = loss_params.loss_weight * total_loss
+
+    losses = {'total_loss': total_loss,
+              'mask_loss': mask_loss,
+              'center_loss': center_loss,
+              'dice_loss': dice_loss,}
+    return losses
+
+  def train_step(self,
+                 inputs: Tuple[Any, Any],
+                 model: tf_keras.Model,
+                 optimizer: tf_keras.optimizers.Optimizer,
+                 metrics: Optional[List[Any]] = None) -> Dict[str, Any]:
+    """Does forward and backward.
+
+    Args:
+      inputs: a dictionary of input tensors.
+      model: the model, forward pass definition.
+      optimizer: the optimizer for this training step.
+      metrics: a nested structure of metrics objects.
+
+    Returns:
+      A dictionary of logs.
+    """
+    images, labels = inputs
+    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
+
+    with tf.GradientTape() as tape:
+      outputs = model(
+          images,
+          box_indices=labels['panoptic_box_indices'],
+          classes=labels['panoptic_classes'],
+          training=True)
+      outputs = tf.nest.map_structure(
+          lambda x: tf.cast(x, tf.float32), outputs)
+
+      # Computes per-replica loss.
+      losses = self.build_losses(
+          outputs=outputs,
+          labels=labels,
+          aux_losses=model.losses,
+          iteration=optimizer.iterations,
+          step=optimizer.iterations)
+      scaled_loss = losses['total_loss'] / num_replicas
+
+      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
+      # scaled for numerical stability.
+      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
+        scaled_loss = optimizer.get_scaled_loss(scaled_loss)
+
+    tvars = model.trainable_variables
+    grads = tape.gradient(scaled_loss, tvars)
+    # Scales back gradient when LossScaleOptimizer is used.
+    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
+      grads = optimizer.get_unscaled_gradients(grads)
+    optimizer.apply_gradients(list(zip(grads, tvars)))
+
+    logs = {self.loss: losses['total_loss']}
+
+    if metrics:
+      for m in metrics:
+        m.update_state(losses[m.name])
+
+    return logs

official/projects/maskconver/train.py

@@ -0,0 +1,30 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Panoptic MaskRCNN trainer."""
+
+from absl import app
+
+from official.common import flags as tfm_flags
+from official.projects.maskconver.configs import maskconver as maskconver_cfg  # pylint: disable=unused-import
+from official.projects.maskconver.configs import multiscale_maskconver as multiscale_maskconver_cfg  # pylint: disable=unused-import
+from official.projects.maskconver.modeling import fpn  # pylint: disable=unused-import
+from official.projects.maskconver.tasks import maskconver as maskconver_task  # pylint: disable=unused-import
+from official.projects.maskconver.tasks import multiscale_maskconver as multiscale_maskconver_task  # pylint: disable=unused-import
+from official.vision import train
+
+
+if __name__ == '__main__':
+  tfm_flags.define_flags()
+  app.run(train.main)

{tf_models_nightly-2.18.0.dev20240820.dist-info → tf_models_nightly-2.18.0.dev20240822.dist-info}/METADATA

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

{tf_models_nightly-2.18.0.dev20240820.dist-info → tf_models_nightly-2.18.0.dev20240822.dist-info}/RECORD

@@ -163,7 +163,7 @@ official/legacy/transformer/optimizer.py,sha256=VbM_Dnblv4hmtaFugBh1cFK2_mnihxoW
 official/legacy/transformer/transformer.py,sha256=6ZCu61XxU1LRdMah1NKvEA-pF23gNr3n3d3rBybtv-Q,21761
 official/legacy/transformer/transformer_forward_test.py,sha256=pnvQ8g9k5FCQeKB523rtAxKfB4VbDtyeBrhy6xvfaeY,6070
 official/legacy/transformer/transformer_layers_test.py,sha256=0aCeywTKqSOTrZtkVBvy3FJhL4pwppeM4er1w01MAeI,3576
-official/legacy/transformer/transformer_main.py,sha256=fRFOvrbDjSPLmCFLhaPeD-xlDNm-8tqHiidV4Z8HokI,18193
+official/legacy/transformer/transformer_main.py,sha256=rn2E4MiwNL6uGwcl5Njnw-3h86HOq8sQfBp3gXITSRw,18110
 official/legacy/transformer/transformer_main_test.py,sha256=aZnbdbJCsDazU4tUhAs_xpuXZe0ZrZJgCxQy0ez92SY,6641
 official/legacy/transformer/transformer_test.py,sha256=2eD65briKm6pcEpveYcFuNh3wSWt3bcYETClbarewqc,3638
 official/legacy/transformer/translate.py,sha256=OO5B5qlSMu2hGMd2KYn1DyUaPDC4tuQm-2vA-xric8o,6961
@@ -531,6 +531,16 @@ official/projects/deepmac_maskrcnn/serving/detection_test.py,sha256=t5YdUlmPa3Ft
 official/projects/deepmac_maskrcnn/serving/export_saved_model.py,sha256=oXGf3Sd7umPWuBW3ZbOrTVIz8kSm23ki8N-OpHs__so,3991
 official/projects/deepmac_maskrcnn/tasks/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
 official/projects/deepmac_maskrcnn/tasks/deep_mask_head_rcnn.py,sha256=5GwDu4XVk1NRzGU9IA07Joqib7QLx6wZQ_h8fAfm-rY,9436
+official/projects/maskconver/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
+official/projects/maskconver/train.py,sha256=XDW0ArwYLcJpT137J3QgBEaoZHb-GE3rGbnKNsSVynI,1397
+official/projects/maskconver/configs/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
+official/projects/maskconver/configs/backbones.py,sha256=hbx64MddUeOQuviLZK4NsYnDMxaWTAJxLngUUUw8IIQ,1467
+official/projects/maskconver/configs/decoders.py,sha256=F1u9ef1YA0MphPYvzHWpIDTtY834E1zM-eLd3XJd-8c,1163
+official/projects/maskconver/configs/maskconver.py,sha256=mSjY3epff8zVG52YSNPcS1QMaKxgmQWkCfuV26SuD8Y,18925
+official/projects/maskconver/configs/multiscale_maskconver.py,sha256=yDZVLMSVNZyo23sTzfxzthDgHJh_VGGkYxyTEVEqjw8,8106
+official/projects/maskconver/tasks/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
+official/projects/maskconver/tasks/maskconver.py,sha256=AAoFojpD8r7MQdaDlYBjNs3-ryow1nvmVe1pUnWWwto,25256
+official/projects/maskconver/tasks/multiscale_maskconver.py,sha256=D_77WLT4zBBFdj-uo6yU3JfrxVvept2o87wg2sqbEz8,11730
 official/projects/maxvit/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
 official/projects/maxvit/registry_imports.py,sha256=anP9wyf9iahIxbm7wUXq32GXq5zUWQ1q6gSS7r1jY18,934
 official/projects/maxvit/train.py,sha256=03CS9XA8ag92YZpgAPED_DgPGMlnIwYNtPTAXFzuqdA,958
@@ -1212,9 +1222,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.18.0.dev20240820.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.18.0.dev20240820.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.18.0.dev20240820.dist-info/METADATA,sha256=4bsxNjZHn4lL4Bh4Fun2713UIwYgOXoccQHUAWlaydg,1432
-tf_models_nightly-2.18.0.dev20240820.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.18.0.dev20240820.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.18.0.dev20240820.dist-info/RECORD,,
+tf_models_nightly-2.18.0.dev20240822.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.18.0.dev20240822.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.18.0.dev20240822.dist-info/METADATA,sha256=IcMXUQ2bVpwSiwP4QbM0KZsqgPoLxn5WDJSbwMDa9oU,1432
+tf_models_nightly-2.18.0.dev20240822.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.18.0.dev20240822.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.18.0.dev20240822.dist-info/RECORD,,