box_ops update

official/vision/beta/projects/yolo/ops/box_ops.py

@@ -38,32 +38,7 @@ def yxyx_to_xcycwh(box: tf.Tensor):
   return box
 
 
-@tf.custom_gradient
-def _xcycwh_to_yxyx(box: tf.Tensor, scale):
-  """Private function to allow custom gradients with defaults."""
-  with tf.name_scope('xcycwh_to_yxyx'):
-    xy, wh = tf.split(box, 2, axis=-1)
-    xy_min = xy - wh / 2
-    xy_max = xy + wh / 2
-    x_min, y_min = tf.split(xy_min, 2, axis=-1)
-    x_max, y_max = tf.split(xy_max, 2, axis=-1)
-    box = tf.concat([y_min, x_min, y_max, x_max], axis=-1)
-
-    def delta(dbox):
-      # y_min = top, x_min = left, y_max = bottom, x_max = right
-      dt, dl, db, dr = tf.split(dbox, 4, axis=-1)
-      dx = dl + dr
-      dy = dt + db
-      dw = (dr - dl) / scale
-      dh = (db - dt) / scale
-
-      dbox = tf.concat([dx, dy, dw, dh], axis=-1)
-      return dbox, 0.0
-
-  return box, delta
-
-
-def xcycwh_to_yxyx(box: tf.Tensor, darknet=False):
+def xcycwh_to_yxyx(box: tf.Tensor):
   """Converts boxes from x_center, y_center, width, height to yxyx format.
 
   Args:
@@ -74,11 +49,13 @@ def xcycwh_to_yxyx(box: tf.Tensor, darknet=False):
   Returns:
     box: a `Tensor` whose shape is the same as `box` in new format.
   """
-  if darknet:
-    scale = 1.0
-  else:
-    scale = 2.0
-  box = _xcycwh_to_yxyx(box, scale)
+  with tf.name_scope('xcycwh_to_yxyx'):
+    xy, wh = tf.split(box, 2, axis=-1)
+    xy_min = xy - wh / 2
+    xy_max = xy + wh / 2
+    x_min, y_min = tf.split(xy_min, 2, axis=-1)
+    x_max, y_max = tf.split(xy_max, 2, axis=-1)
+    box = tf.concat([y_min, x_min, y_max, x_max], axis=-1)
   return box
 
 
@@ -98,8 +75,9 @@ def intersect_and_union(box1, box2, yxyx=False):
     intersection: a `Tensor` who represents the intersection.
     union: a `Tensor` who represents the union.
   """
-
   if not yxyx:
+    box1_area = tf.reduce_prod(tf.split(box1, 2, axis=-1)[-1], axis=-1)
+    box2_area = tf.reduce_prod(tf.split(box2, 2, axis=-1)[-1], axis=-1)
     box1 = xcycwh_to_yxyx(box1)
     box2 = xcycwh_to_yxyx(box2)
 
@@ -110,13 +88,14 @@ def intersect_and_union(box1, box2, yxyx=False):
   intersect_wh = tf.math.maximum(intersect_maxes - intersect_mins, 0.0)
   intersection = tf.reduce_prod(intersect_wh, axis=-1)
 
+  if yxyx:
     box1_area = tf.reduce_prod(b1ma - b1mi, axis=-1)
     box2_area = tf.reduce_prod(b2ma - b2mi, axis=-1)
   union = box1_area + box2_area - intersection
   return intersection, union
 
 
-def smallest_encompassing_box(box1, box2, yxyx=False):
+def smallest_encompassing_box(box1, box2, yxyx=False, clip=False):
   """Calculates the smallest box that encompasses box1 and box2.
 
   Args:
@@ -141,15 +120,15 @@ def smallest_encompassing_box(box1, box2, yxyx=False):
 
   bcmi = tf.math.minimum(b1mi, b2mi)
   bcma = tf.math.maximum(b1ma, b2ma)
-
-  bca = tf.reduce_prod(bcma - bcmi, keepdims=True, axis=-1)
   box_c = tf.concat([bcmi, bcma], axis=-1)
 
   if not yxyx:
     box_c = yxyx_to_xcycwh(box_c)
 
-  box_c = tf.where(bca == 0.0, tf.zeros_like(box_c), box_c)
-  return box_c
+  if clip:
+    bca = tf.reduce_prod(bcma - bcmi, keepdims=True, axis=-1)
+    box_c = tf.where(bca <= 0.0, tf.zeros_like(box_c), box_c)
+  return bcmi, bcma, box_c
 
 
 def compute_iou(box1, box2, yxyx=False):
@@ -166,15 +145,13 @@ def compute_iou(box1, box2, yxyx=False):
   Returns:
     iou: a `Tensor` who represents the intersection over union.
   """
-  # get box corners
   with tf.name_scope('iou'):
     intersection, union = intersect_and_union(box1, box2, yxyx=yxyx)
     iou = math_ops.divide_no_nan(intersection, union)
-    iou = math_ops.rm_nan_inf(iou, val=0.0)
   return iou
 
 
-def compute_giou(box1, box2, yxyx=False, darknet=False):
+def compute_giou(box1, box2, yxyx=False):
   """Calculates the General intersection over union between box1 and box2.
 
   Args:
@@ -191,27 +168,21 @@ def compute_giou(box1, box2, yxyx=False, darknet=False):
     giou: a `Tensor` who represents the General intersection over union.
   """
   with tf.name_scope('giou'):
-    # get IOU
     if not yxyx:
-      box1 = xcycwh_to_yxyx(box1, darknet=darknet)
-      box2 = xcycwh_to_yxyx(box2, darknet=darknet)
-      yxyx = True
+      yxyx1 = xcycwh_to_yxyx(box1)
+      yxyx2 = xcycwh_to_yxyx(box2)
+    else:
+      yxyx1, yxyx2 = box1, box2
 
-    intersection, union = intersect_and_union(box1, box2, yxyx=yxyx)
+    cmi, cma, _ = smallest_encompassing_box(yxyx1, yxyx2, yxyx=True)
+    intersection, union = intersect_and_union(yxyx1, yxyx2, yxyx=True)
     iou = math_ops.divide_no_nan(intersection, union)
-    iou = math_ops.rm_nan_inf(iou, val=0.0)
 
-    # find the smallest box to encompase both box1 and box2
-    boxc = smallest_encompassing_box(box1, box2, yxyx=yxyx)
-    if yxyx:
-      boxc = yxyx_to_xcycwh(boxc)
-    _, cwch = tf.split(boxc, 2, axis=-1)
-    c = tf.math.reduce_prod(cwch, axis=-1)
+    bcwh = cma - cmi
+    c = tf.math.reduce_prod(bcwh, axis=-1)
 
-    # compute giou
     regularization = math_ops.divide_no_nan((c - union), c)
     giou = iou - regularization
-    giou = tf.clip_by_value(giou, clip_value_min=-1.0, clip_value_max=1.0)
   return iou, giou
 
 
@@ -236,30 +207,27 @@ def compute_diou(box1, box2, beta=1.0, yxyx=False, darknet=False):
   with tf.name_scope('diou'):
     # compute center distance
     if not yxyx:
-      box1 = xcycwh_to_yxyx(box1, darknet=darknet)
-      box2 = xcycwh_to_yxyx(box2, darknet=darknet)
-      yxyx = True
-
-    intersection, union = intersect_and_union(box1, box2, yxyx=yxyx)
-    boxc = smallest_encompassing_box(box1, box2, yxyx=yxyx)
+      xycc1, xycc2 = box1, box2
+      yxyx1 = xcycwh_to_yxyx(box1)
+      yxyx2 = xcycwh_to_yxyx(box2)
+    else:
+      yxyx1, yxyx2 = box1, box2
+      xycc1 = yxyx_to_xcycwh(box1)
+      xycc2 = yxyx_to_xcycwh(box2)
 
+    cmi, cma, _ = smallest_encompassing_box(yxyx1, yxyx2, yxyx=True)
+    intersection, union = intersect_and_union(yxyx1, yxyx2, yxyx=True)
     iou = math_ops.divide_no_nan(intersection, union)
-    iou = math_ops.rm_nan_inf(iou, val=0.0)
-    if yxyx:
-      boxc = yxyx_to_xcycwh(boxc)
-      box1 = yxyx_to_xcycwh(box1)
-      box2 = yxyx_to_xcycwh(box2)
 
-    b1xy, _ = tf.split(box1, 2, axis=-1)
-    b2xy, _ = tf.split(box2, 2, axis=-1)
-    _, bcwh = tf.split(boxc, 2, axis=-1)
+    b1xy, _ = tf.split(xycc1, 2, axis=-1)
+    b2xy, _ = tf.split(xycc2, 2, axis=-1)
+    bcwh = cma - cmi
 
     center_dist = tf.reduce_sum((b1xy - b2xy)**2, axis=-1)
     c_diag = tf.reduce_sum(bcwh**2, axis=-1)
 
     regularization = math_ops.divide_no_nan(center_dist, c_diag)
     diou = iou - regularization**beta
-    diou = tf.clip_by_value(diou, clip_value_min=-1.0, clip_value_max=1.0)
   return iou, diou
 
 
@@ -280,33 +248,48 @@ def compute_ciou(box1, box2, yxyx=False, darknet=False):
     ciou: a `Tensor` who represents the complete intersection over union.
   """
   with tf.name_scope('ciou'):
-    # compute DIOU and IOU
-    iou, diou = compute_diou(box1, box2, yxyx=yxyx, darknet=darknet)
+    if not yxyx:
+      xycc1, xycc2 = box1, box2
+      yxyx1 = xcycwh_to_yxyx(box1)
+      yxyx2 = xcycwh_to_yxyx(box2)
+    else:
+      yxyx1, yxyx2 = box1, box2
+      xycc1 = yxyx_to_xcycwh(box1)
+      xycc2 = yxyx_to_xcycwh(box2)
 
-    if yxyx:
-      box1 = yxyx_to_xcycwh(box1)
-      box2 = yxyx_to_xcycwh(box2)
-
-    _, _, b1w, b1h = tf.split(box1, 4, axis=-1)
-    _, _, b2w, b2h = tf.split(box1, 4, axis=-1)
-
-    # computer aspect ratio consistency
-    terma = tf.cast(math_ops.divide_no_nan(b1w, b1h), tf.float32)
-    termb = tf.cast(math_ops.divide_no_nan(b2w, b2h), tf.float32)
-    arcterm = tf.square(tf.math.atan(terma) - tf.math.atan(termb))
-    v = tf.squeeze(4 * arcterm / (math.pi**2), axis=-1)
-    v = tf.cast(v, b1w.dtype)
-
-    a = tf.stop_gradient(math_ops.divide_no_nan(v, ((1 - iou) + v)))
-    ciou = diou - (v * a)
-    ciou = tf.clip_by_value(ciou, clip_value_min=-1.0, clip_value_max=1.0)
+    # Build the smallest encomapssing box.
+    cmi, cma, _ = smallest_encompassing_box(yxyx1, yxyx2, yxyx=True)
+    intersection, union = intersect_and_union(yxyx1, yxyx2, yxyx=True)
+    iou = math_ops.divide_no_nan(intersection, union)
+
+    b1xy, b1w, b1h = tf.split(xycc1, [2, 1, 1], axis=-1)
+    b2xy, b2w, b2h = tf.split(xycc2, [2, 1, 1], axis=-1)
+    bchw = cma - cmi
+
+    # Center regularization
+    center_dist = tf.reduce_sum((b1xy - b2xy)**2, axis=-1)
+    c_diag = tf.reduce_sum(bchw**2, axis=-1)
+    regularization = math_ops.divide_no_nan(center_dist, c_diag)
+
+    # Computer aspect ratio consistency
+    terma = math_ops.divide_no_nan(b1w, b1h)  # gt
+    termb = math_ops.divide_no_nan(b2w, b2h)  # pred
+    arcterm = tf.squeeze(
+        tf.math.pow(tf.math.atan(termb) - tf.math.atan(terma), 2), axis=-1)
+    v = (4 / math.pi**2) * arcterm
+
+    # Compute the aspect ratio weight, should be treated as a constant
+    a = tf.stop_gradient(math_ops.divide_no_nan(v, 1 - iou + v))
+
+    if darknet:
+      grad_scale = tf.stop_gradient(tf.square(b2w) + tf.square(b2h))
+      v *= tf.squeeze(grad_scale, axis=-1)
+
+    ciou = iou - regularization - (v * a)
   return iou, ciou
 
 
-def aggregated_comparitive_iou(boxes1,
-                               boxes2=None,
-                               iou_type=0,
-                               beta=0.6):
+def aggregated_comparitive_iou(boxes1, boxes2=None, iou_type=0, beta=0.6):
   """Calculates the IOU between two set of boxes.
 
   Similar to bbox_overlap but far more versitile.
@@ -333,11 +316,11 @@ def aggregated_comparitive_iou(boxes1,
   else:
     boxes2 = tf.transpose(boxes1, perm=(0, 2, 1, 3))
 
-  if iou_type == 0:  # diou
+  if iou_type == 0 or iou_type == "diou":  #diou
     _, iou = compute_diou(boxes1, boxes2, beta=beta, yxyx=True)
-  elif iou_type == 1:  # giou
+  elif iou_type == 1 or iou_type == "giou":  #giou
     _, iou = compute_giou(boxes1, boxes2, yxyx=True)
-  elif iou_type == 2:  # ciou
+  elif iou_type == 2 or iou_type == "ciou":  #ciou
     _, iou = compute_ciou(boxes1, boxes2, yxyx=True)
   else:
     iou = compute_iou(boxes1, boxes2, yxyx=True)