Internal change

PiperOrigin-RevId: 489040948

official/projects/detr/tasks/detection.py

@@ -338,29 +338,50 @@ class DetectionTask(base_task.Task):
     # Evaluator class handles loss metric for you.
     logs = {self.loss: loss}
 
+    # 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]
+              ],
+                        axis=1),
+              axis=1)
+    else:
+      detection_boxes = outputs['detection_boxes']
+
+    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
+    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)
+    else:
+      num_detections = outputs['num_detections']
+
     predictions = {
-        'detection_boxes':
-            outputs['detection_boxes'] * 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),
-        'detection_scores':
-            outputs['detection_scores'],
-        'detection_classes':
-            outputs['detection_classes'],
-        # Fix this. It's not being used at the moment.
-        'num_detections':
-            outputs['num_detections'],
-        'source_id':
-            labels['id'],
-        'image_info':
-            labels['image_info']
+        'detection_boxes': detection_boxes,
+        'detection_scores': detection_scores,
+        'detection_classes': detection_classes,
+        'num_detections': num_detections,
+        'source_id': labels['id'],
+        'image_info': labels['image_info']
     }
+
     ground_truths = {
         'source_id': labels['id'],
         'height': labels['image_info'][:, 0:1, 0],