Rewrite the GIOU loss op in matrix form to avoid the map_fn. map_fn could be...

Rewrite the GIOU loss op in matrix form to avoid the map_fn. map_fn could be prohibitively costly when dealing with large number of bounding boxes.

PiperOrigin-RevId: 366854232

research/object_detection/utils/ops.py

@@ -1149,12 +1149,8 @@ def giou(boxes1, boxes2):
     a tensor of shape [num_boxes] containing GIoUs
 
   """
-  def _two_boxes_giou(boxes):
-    """Compute giou between two boxes."""
-    boxes1, boxes2 = boxes
-
-    pred_ymin, pred_xmin, pred_ymax, pred_xmax = tf.unstack(boxes1)
-    gt_ymin, gt_xmin, gt_ymax, gt_xmax = tf.unstack(boxes2)
+  pred_ymin, pred_xmin, pred_ymax, pred_xmax = tf.unstack(boxes1, axis=1)
+  gt_ymin, gt_xmin, gt_ymax, gt_xmax = tf.unstack(boxes2, axis=1)
 
   gt_area = (gt_ymax - gt_ymin) * (gt_xmax - gt_xmin)
   pred_area = (pred_ymax - pred_ymin) * (pred_xmax - pred_xmin)
@@ -1173,15 +1169,12 @@ def giou(boxes1, boxes2):
   hull_area = (y2_c - y1_c) * (x2_c - x1_c)
 
   union_area = gt_area + pred_area - intersection_area
-    iou = tf.where(
-        tf.equal(union_area, 0.0), 0.0, intersection_area / union_area)
+  iou = tf.where(tf.equal(union_area, 0.0),
+                 tf.zeros_like(union_area), intersection_area / union_area)
   giou_ = iou - tf.where(hull_area > 0.0,
                          (hull_area - union_area) / hull_area, iou)
-
   return giou_
 
-  return shape_utils.static_or_dynamic_map_fn(_two_boxes_giou, [boxes1, boxes2])
-
 
 def center_to_corner_coordinate(input_tensor):
   """Converts input boxes from center to corner representation."""