[3D Object Detection] Make iou_3d a function

keras_cv/layers/preprocessing3d/group_points_by_bounding_boxes.py

@@ -15,7 +15,6 @@
 import tensorflow as tf
 
 from keras_cv.layers.preprocessing3d import base_augmentation_layer_3d
-from keras_cv.ops import iou_3d
 from keras_cv.ops.point_cloud import is_within_box3d
 
 POINT_CLOUDS = base_augmentation_layer_3d.POINT_CLOUDS
@@ -76,7 +75,6 @@ def __init__(
         self._label_index = label_index
         self._min_points_per_bounding_boxes = min_points_per_bounding_boxes
         self._max_points_per_bounding_boxes = max_points_per_bounding_boxes
-        self._iou_3d = iou_3d.IoU3D
         self._auto_vectorize = False
 
     def augment_point_clouds_bounding_boxes(

keras_cv/ops/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from keras_cv.ops.iou_3d import IoU3D
+from keras_cv.ops.iou_3d import iou_3d
 from keras_cv.ops.point_cloud import _box_area
 from keras_cv.ops.point_cloud import _center_xyzWHD_to_corner_xyz
 from keras_cv.ops.point_cloud import _is_on_lefthand_side

keras_cv/ops/iou_3d.py

@@ -14,15 +14,16 @@
 # ============================================================================
 """IoU3D using a custom TF op."""
 
-from tensorflow.python.framework import load_library
-from tensorflow.python.platform import resource_loader
+from keras_cv.utils.resource_loader import LazySO
 
+keras_cv_custom_ops = LazySO("custom_ops/_keras_cv_custom_ops.so")
 
-class IoU3D:
+
+def iou_3d(y_true, y_pred):
     """Implements IoU computation for 3D upright rotated bounding boxes.
 
-    Note that this is implemented using a custom TensorFlow op. Initializing an
-    IoU3D object will attempt to load the binary for that op.
+    Note that this is implemented using a custom TensorFlow op. If you don't have
+    KerasCV installed with custom ops, calling this will fail.
 
     Boxes should have the format [center_x, center_y, center_z, dimension_x,
     dimension_y, dimension_z, heading (in radians)].
@@ -31,18 +32,8 @@ class IoU3D:
     ```python
     y_true = [[0, 0, 0, 2, 2, 2, 0], [1, 1, 1, 2, 2, 2, 3 * math.pi / 4]]
     y_pred = [[1, 1, 1, 2, 2, 2, math.pi / 4], [1, 1, 1, 2, 2, 2, 0]]
-    iou = IoU3D()
-    iou(y_true, y_pred)
+    iou_3d(y_true, y_pred)
     ```
     """
 
-    def __init__(self):
-        pairwise_iou_op = load_library.load_op_library(
-            resource_loader.get_path_to_datafile(
-                "../custom_ops/_keras_cv_custom_ops.so"
-            )
-        )
-        self.iou_3d = pairwise_iou_op.pairwise_iou3d
-
-    def __call__(self, y_true, y_pred):
-        return self.iou_3d(y_true, y_pred)
+    return keras_cv_custom_ops.ops.pairwise_iou3d(y_true, y_pred)

keras_cv/ops/iou_3d_test.py

@@ -19,7 +19,7 @@
 import pytest
 import tensorflow as tf
 
-from keras_cv.ops import IoU3D
+from keras_cv.ops import iou_3d
 
 
 class IoU3DTest(tf.test.TestCase):
@@ -45,8 +45,6 @@ def testOpCall(self):
         # the same offset of 135 degrees, which reduces to the square root of 0.5.
         expected_ious = [[1 / 15, 1 / 15], [1, 0.5**0.5]]
 
-        iou_3d = IoU3D()
-
         self.assertAllClose(iou_3d(box_preds, box_gt), expected_ious)
 
 

keras_cv/utils/resource_loader.py

@@ -0,0 +1,53 @@
+# Copyright 2022 The KerasCV Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities similar to tf.python.platform.resource_loader."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+import tensorflow as tf
+
+
+def get_project_root():
+    """Returns project root folder."""
+    return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+def get_path_to_datafile(path):
+    """Get the path to the specified file in the data dependencies.
+
+    The path is relative to keras_cv/
+
+    Args:
+      path: a string resource path relative to keras_cv/
+    Returns:
+      The path to the specified data file
+    """
+    root_dir = get_project_root()
+    return os.path.join(root_dir, path.replace("/", os.sep))
+
+
+class LazySO:
+    def __init__(self, relative_path):
+        self.relative_path = relative_path
+        self._ops = None
+
+    @property
+    def ops(self):
+        if self._ops is None:
+            self._ops = tf.load_op_library(get_path_to_datafile(self.relative_path))
+        return self._ops