No public description

PiperOrigin-RevId: 721434742

Author: tensorflower-gardener

official/projects/detr/tasks/detection.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 """DETR detection task definition."""
-
 from typing import Optional
 
 from absl import logging
@@ -48,25 +47,21 @@ class DetectionTask(base_task.Task):
   def build_model(self):
     """Build DETR model."""
 
-    input_specs = tf_keras.layers.InputSpec(
-        shape=[None] + self._task_config.model.input_size
-    )
+    input_specs = tf_keras.layers.InputSpec(shape=[None] +
+                                            self._task_config.model.input_size)
 
     backbone = backbones.factory.build_backbone(
         input_specs=input_specs,
         backbone_config=self._task_config.model.backbone,
-        norm_activation_config=self._task_config.model.norm_activation,
-    )
-
-    model = detr.DETR(
-        backbone,
-        self._task_config.model.backbone_endpoint_name,
-        self._task_config.model.num_queries,
-        self._task_config.model.hidden_size,
-        self._task_config.model.num_classes,
-        self._task_config.model.num_encoder_layers,
-        self._task_config.model.num_decoder_layers,
-    )
+        norm_activation_config=self._task_config.model.norm_activation)
+
+    model = detr.DETR(backbone,
+                      self._task_config.model.backbone_endpoint_name,
+                      self._task_config.model.num_queries,
+                      self._task_config.model.hidden_size,
+                      self._task_config.model.num_classes,
+                      self._task_config.model.num_encoder_layers,
+                      self._task_config.model.num_decoder_layers)
     return model
 
   def initialize(self, model: tf_keras.Model):
@@ -89,13 +84,12 @@ def initialize(self, model: tf_keras.Model):
       status = ckpt.restore(ckpt_dir_or_file)
       status.expect_partial().assert_existing_objects_matched()
 
-    logging.info(
-        'Finished loading pretrained checkpoint from %s', ckpt_dir_or_file
-    )
+    logging.info('Finished loading pretrained checkpoint from %s',
+                 ckpt_dir_or_file)
 
-  def build_inputs(
-      self, params, input_context: Optional[tf.distribute.InputContext] = None
-  ):
+  def build_inputs(self,
+                   params,
+                   input_context: Optional[tf.distribute.InputContext] = None):
     """Build input dataset."""
     if isinstance(params, coco.COCODataConfig):
       dataset = coco.COCODataLoader(params).load(input_context)
@@ -106,17 +100,14 @@ def build_inputs(
         decoder_cfg = params.decoder.get()
         if params.decoder.type == 'simple_decoder':
           decoder = tf_example_decoder.TfExampleDecoder(
-              regenerate_source_id=decoder_cfg.regenerate_source_id
-          )
+              regenerate_source_id=decoder_cfg.regenerate_source_id)
         elif params.decoder.type == 'label_map_decoder':
           decoder = tf_example_label_map_decoder.TfExampleDecoderLabelMap(
               label_map=decoder_cfg.label_map,
-              regenerate_source_id=decoder_cfg.regenerate_source_id,
-          )
+              regenerate_source_id=decoder_cfg.regenerate_source_id)
         else:
-          raise ValueError(
-              'Unknown decoder type: {}!'.format(params.decoder.type)
-          )
+          raise ValueError('Unknown decoder type: {}!'.format(
+              params.decoder.type))
 
       parser = detr_input.Parser(
           class_offset=self._task_config.losses.class_offset,
@@ -127,8 +118,7 @@ def build_inputs(
           params,
           dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
           decoder_fn=decoder.decode,
-          parser_fn=parser.parse_fn(params.is_training),
-      )
+          parser_fn=parser.parse_fn(params.is_training))
       dataset = reader.read(input_context=input_context)
 
     return dataset
@@ -187,8 +177,7 @@ def build_losses(self, outputs, labels, aux_losses=None):
     box_targets = labels['boxes']
 
     cost = self._compute_cost(
-        cls_outputs, box_outputs, cls_targets, box_targets
-    )
+        cls_outputs, box_outputs, cls_targets, box_targets)
 
     _, indices = matchers.hungarian_matching(cost)
     indices = tf.stop_gradient(indices)
@@ -199,53 +188,45 @@ def build_losses(self, outputs, labels, aux_losses=None):
 
     background = tf.equal(cls_targets, 0)
     num_boxes = tf.reduce_sum(
-        tf.cast(tf.logical_not(background), tf.float32), axis=-1
-    )
+        tf.cast(tf.logical_not(background), tf.float32), axis=-1)
 
     # Down-weight background to account for class imbalance.
     xentropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
-        labels=cls_targets, logits=cls_assigned
-    )
+        labels=cls_targets, logits=cls_assigned)
     cls_loss = self._task_config.losses.lambda_cls * tf.where(
-        background,
-        self._task_config.losses.background_cls_weight * xentropy,
-        xentropy,
-    )
+        background, self._task_config.losses.background_cls_weight * xentropy,
+        xentropy)
     cls_weights = tf.where(
         background,
         self._task_config.losses.background_cls_weight * tf.ones_like(cls_loss),
-        tf.ones_like(cls_loss),
-    )
+        tf.ones_like(cls_loss))
 
     # Box loss is only calculated on non-background class.
     l_1 = tf.reduce_sum(tf.abs(box_assigned - box_targets), axis=-1)
     box_loss = self._task_config.losses.lambda_box * tf.where(
-        background, tf.zeros_like(l_1), l_1
-    )
+        background, tf.zeros_like(l_1), l_1)
 
     # Giou loss is only calculated on non-background class.
-    giou = tf.linalg.diag_part(
-        1.0
-        - box_ops.bbox_generalized_overlap(
-            box_ops.cycxhw_to_yxyx(box_assigned),
-            box_ops.cycxhw_to_yxyx(box_targets),
-        )
-    )
+    giou = tf.linalg.diag_part(1.0 - box_ops.bbox_generalized_overlap(
+        box_ops.cycxhw_to_yxyx(box_assigned),
+        box_ops.cycxhw_to_yxyx(box_targets)
+        ))
     giou_loss = self._task_config.losses.lambda_giou * tf.where(
-        background, tf.zeros_like(giou), giou
-    )
+        background, tf.zeros_like(giou), giou)
 
     # Consider doing all reduce once in train_step to speed up.
     num_boxes_per_replica = tf.reduce_sum(num_boxes)
     cls_weights_per_replica = tf.reduce_sum(cls_weights)
     replica_context = tf.distribute.get_replica_context()
     num_boxes_sum, cls_weights_sum = replica_context.all_reduce(
         tf.distribute.ReduceOp.SUM,
-        [num_boxes_per_replica, cls_weights_per_replica],
-    )
-    cls_loss = tf.math.divide_no_nan(tf.reduce_sum(cls_loss), cls_weights_sum)
-    box_loss = tf.math.divide_no_nan(tf.reduce_sum(box_loss), num_boxes_sum)
-    giou_loss = tf.math.divide_no_nan(tf.reduce_sum(giou_loss), num_boxes_sum)
+        [num_boxes_per_replica, cls_weights_per_replica])
+    cls_loss = tf.math.divide_no_nan(
+        tf.reduce_sum(cls_loss), cls_weights_sum)
+    box_loss = tf.math.divide_no_nan(
+        tf.reduce_sum(box_loss), num_boxes_sum)
+    giou_loss = tf.math.divide_no_nan(
+        tf.reduce_sum(giou_loss), num_boxes_sum)
 
     aux_losses = tf.add_n(aux_losses) if aux_losses else 0.0
 
@@ -264,8 +245,7 @@ def build_metrics(self, training=True):
           annotation_file=self._task_config.annotation_file,
           include_mask=False,
           need_rescale_bboxes=True,
-          per_category_metrics=self._task_config.per_category_metrics,
-      )
+          per_category_metrics=self._task_config.per_category_metrics)
     return metrics
 
   def train_step(self, inputs, model, optimizer, metrics=None):
@@ -355,8 +335,7 @@ def validation_step(self, inputs, model, metrics=None):
 
     outputs = model(features, training=False)[-1]
     loss, cls_loss, box_loss, giou_loss = self.build_losses(
-        outputs=outputs, labels=labels, aux_losses=model.losses
-    )
+        outputs=outputs, labels=labels, aux_losses=model.losses)
 
     # Multiply for logging.
     # Since we expect the gradient replica sum to happen in the optimizer,
@@ -374,46 +353,35 @@ def validation_step(self, inputs, model, metrics=None):
     # This is for backward compatibility.
     if 'detection_boxes' not in outputs:
       detection_boxes = box_ops.cycxhw_to_yxyx(
-          outputs['box_outputs']
-      ) * tf.expand_dims(
-          tf.concat(
-              [
-                  labels['image_info'][:, 1:2, 0],
-                  labels['image_info'][:, 1:2, 1],
-                  labels['image_info'][:, 1:2, 0],
-                  labels['image_info'][:, 1:2, 1],
+          outputs['box_outputs']) * tf.expand_dims(
+              tf.concat([
+                  labels['image_info'][:, 1:2, 0], labels['image_info'][:, 1:2,
+                                                                        1],
+                  labels['image_info'][:, 1:2, 0], labels['image_info'][:, 1:2,
+                                                                        1]
               ],
-              axis=1,
-          ),
-          axis=1,
-      )
+                        axis=1),
+              axis=1)
     else:
       detection_boxes = outputs['detection_boxes']
 
-    if 'detection_scores' not in outputs:
-      detection_scores = tf.math.reduce_max(
-          tf.nn.softmax(outputs['cls_outputs'])[:, :, 1:], axis=-1
-      )
-    else:
-      detection_scores = outputs['detection_scores']
+    detection_scores = tf.math.reduce_max(
+        tf.nn.softmax(outputs['cls_outputs'])[:, :, 1:], axis=-1
+    ) if 'detection_scores' not in outputs else outputs['detection_scores']
 
     if 'detection_classes' not in outputs:
-      detection_classes = (
-          tf.math.argmax(outputs['cls_outputs'][:, :, 1:], axis=-1) + 1
-      )
+      detection_classes = tf.math.argmax(
+          outputs['cls_outputs'][:, :, 1:], axis=-1) + 1
     else:
       detection_classes = outputs['detection_classes']
 
     if 'num_detections' not in outputs:
       num_detections = tf.reduce_sum(
           tf.cast(
               tf.math.greater(
-                  tf.math.reduce_max(outputs['cls_outputs'], axis=-1), 0
-              ),
-              tf.int32,
-          ),
-          axis=-1,
-      )
+                  tf.math.reduce_max(outputs['cls_outputs'], axis=-1), 0),
+              tf.int32),
+          axis=-1)
     else:
       num_detections = outputs['num_detections']
 
@@ -423,21 +391,21 @@ def validation_step(self, inputs, model, metrics=None):
         'detection_classes': detection_classes,
         'num_detections': num_detections,
         'source_id': labels['id'],
-        'image_info': labels['image_info'],
+        'image_info': labels['image_info']
     }
 
     ground_truths = {
         'source_id': labels['id'],
         'height': labels['image_info'][:, 0:1, 0],
         'width': labels['image_info'][:, 0:1, 1],
         'num_detections': tf.reduce_sum(
-            tf.cast(tf.math.greater(labels['classes'], 0), tf.int32), axis=-1
-        ),
+            tf.cast(tf.math.greater(labels['classes'], 0), tf.int32), axis=-1),
         'boxes': labels['gt_boxes'],
         'classes': labels['classes'],
-        'is_crowds': labels['is_crowd'],
+        'is_crowds': labels['is_crowd']
     }
-    logs.update({'predictions': predictions, 'ground_truths': ground_truths})
+    logs.update({'predictions': predictions,
+                 'ground_truths': ground_truths})
 
     all_losses = {
         'cls_loss': cls_loss,
@@ -457,8 +425,8 @@ def aggregate_logs(self, state=None, step_outputs=None):
       state = self.coco_metric
 
     state.update_state(
-        step_outputs['ground_truths'], step_outputs['predictions']
-    )
+        step_outputs['ground_truths'],
+        step_outputs['predictions'])
     return state
 
   def reduce_aggregated_logs(self, aggregated_logs, global_step=None):