[MaskRCNN] accelerate numpy iou computation

examples/FasterRCNN/data.py

@@ -16,13 +16,32 @@ from tensorpack.utils import logger
 
 from coco import COCODetection
 from utils.generate_anchors import generate_anchors
-from utils.np_box_ops import iou as np_iou
 from utils.np_box_ops import area as np_area
 from common import (
     DataFromListOfDict, CustomResize, filter_boxes_inside_shape,
     box_to_point8, point8_to_box, segmentation_to_mask)
 from config import config as cfg
 
+try:
+    import pycocotools.mask as cocomask
+
+    # Much faster than utils/np_box_ops
+    def np_iou(A, B):
+        def to_xywh(box):
+            box = box.copy()
+            box[:, 2] -= box[:, 0]
+            box[:, 3] -= box[:, 1]
+            return box
+
+        ret = cocomask.iou(
+            to_xywh(A), to_xywh(B),
+            np.zeros((len(B),), dtype=np.bool))
+        # can accelerate even more, if using float32
+        return ret.astype('float32')
+
+except ImportError:
+    from utils.np_box_ops import iou as np_iou
+
 
 class MalformedData(BaseException):
     pass

examples/FasterRCNN/utils/np_box_ops.py

@@ -50,12 +50,12 @@ def intersection(boxes1, boxes2):
   all_pairs_min_ymax = np.minimum(y_max1, np.transpose(y_max2))
   all_pairs_max_ymin = np.maximum(y_min1, np.transpose(y_min2))
   intersect_heights = np.maximum(
-      np.zeros(all_pairs_max_ymin.shape),
+      np.zeros(all_pairs_max_ymin.shape, dtype='f4'),
       all_pairs_min_ymax - all_pairs_max_ymin)
   all_pairs_min_xmax = np.minimum(x_max1, np.transpose(x_max2))
   all_pairs_max_xmin = np.maximum(x_min1, np.transpose(x_min2))
   intersect_widths = np.maximum(
-      np.zeros(all_pairs_max_xmin.shape),
+      np.zeros(all_pairs_max_xmin.shape, dtype='f4'),
       all_pairs_min_xmax - all_pairs_max_xmin)
   return intersect_heights * intersect_widths
 
@@ -93,5 +93,5 @@ def ioa(boxes1, boxes2):
     a numpy array with shape [N, M] representing pairwise ioa scores.
   """
   intersect = intersection(boxes1, boxes2)
-  areas = np.expand_dims(area(boxes2), axis=0)
-  return intersect / areas
+  inv_areas = np.expand_dims(1.0 / area(boxes2), axis=0)
+  return intersect * inv_areas