Merged commit includes the following changes:

325897695  by Waymo Research:

    Internal change

--
322229379  by Waymo Research:

    Fix the initialization of the metrics variables tp support TPUEstimator.

--
320970608  by Waymo Research:

    Purely Google refactor

--
318516763  by Waymo Research:

    Add fields to the Open Dataset proto for third-party extensions.

--
317327330  by Waymo Research:

    Add BEV box utils

--
316303115  by Waymo Research:

    Fully qualify uses of namespaced int64 to guard against the situation where others have defined int64 to be a different primitive type. (On many platforms either long or long long are acceptable definitions.)

--
315382146  by Waymo Research:

    Minor comment update.

--

PiperOrigin-RevId: 325897695

waymo_open_dataset/dataset.proto

@@ -282,6 +282,11 @@ message Laser {
 }
 
 message Frame {
+  // The following field numbers are reserved for third-party extensions. Users
+  // may declare new fields in that range in their own .proto files without
+  // having to edit the original file.
+  extensions 1000 to max;
+
   // This context is the same for all frames belong to the same driving run
   // segment. Use context.name to identify frames belong to the same driving
   // segment. We do not store all frames from one driving segment in one proto

waymo_open_dataset/label.proto

@@ -81,6 +81,7 @@ message Label {
 
   // The total number of lidar points in this box.
   optional int32 num_lidar_points_in_box = 7;
+
 }
 
 // Non-self-intersecting 2d polygons. This polygon is not necessarily convex.

waymo_open_dataset/metrics/ops/detection_metrics_ops.cc

@@ -159,8 +159,8 @@ class DetectionMetricsOp final : public OpKernel {
     LOG(INFO) << "Parsing prediction "
               << input.prediction_bbox->shape().DebugString()
               << input.prediction_frame_id->shape();
-    absl::flat_hash_map<int64, std::vector<co::Object>> pds_map =
-        co::ParseObjectFromTensors(
+    absl::flat_hash_map<waymo::open_dataset::int64, std::vector<co::Object>>
+        pds_map = co::ParseObjectFromTensors(
             *input.prediction_bbox, *input.prediction_type,
             *input.prediction_frame_id, *input.prediction_score,
             *input.prediction_overlap_nlz, absl::nullopt, absl::nullopt,
@@ -169,8 +169,8 @@ class DetectionMetricsOp final : public OpKernel {
               << input.ground_truth_bbox->shape().DebugString()
               << input.ground_truth_frame_id->shape();
 
-    absl::flat_hash_map<int64, std::vector<co::Object>> gts_map =
-        co::ParseObjectFromTensors(
+    absl::flat_hash_map<waymo::open_dataset::int64, std::vector<co::Object>>
+        gts_map = co::ParseObjectFromTensors(
             *input.ground_truth_bbox, *input.ground_truth_type,
             *input.ground_truth_frame_id, absl::nullopt, absl::nullopt,
             *input.ground_truth_difficulty, absl::nullopt,

waymo_open_dataset/metrics/python/detection_metrics.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-""" tf.metrics implementation for detection metrics."""
+"""tf.metrics implementation for detection metrics."""
 
 import tensorflow as tf
 
@@ -36,11 +36,17 @@ def _update(name, update, init_shape, dtype):
   """
   with tf.compat.v1.variable_scope(
       'detection_metrics', reuse=tf.compat.v1.AUTO_REUSE):
+    initializer = lambda: tf.constant([], shape=init_shape, dtype=dtype)
     v = tf.compat.v1.get_local_variable(
         name,
         dtype=dtype,
+        collections=[
+            tf.compat.v1.GraphKeys.LOCAL_VARIABLES,
+            tf.compat.v1.GraphKeys.METRIC_VARIABLES
+        ],
         # init_shape is required to pass the shape inference check.
-        initializer=tf.constant([], shape=init_shape, dtype=dtype))
+        initializer=initializer,
+        validate_shape=False)
     shape = tf.concat([[-1], tf.shape(input=update)[1:]], axis=0)
     v_reshape = tf.reshape(v.value(), shape)
     v_assign = tf.compat.v1.assign(
@@ -89,10 +95,18 @@ def get_detection_metric_ops(
 
   Notation:
     * M: number of predicted boxes.
-    * D: number of box dimensions (4, 5 or 7).
+    * D: number of box dimensions. The number of box dimensions can be one of
+         the following:
+           4: Used for boxes with type TYPE_AA_2D (center_x, center_y, length,
+              width)
+           5: Used for boxes with type TYPE_2D (center_x, center_y, length,
+              width, heading).
+           7: Used for boxes with type TYPE_3D (center_x, center_y, center_z,
+              length, width, height, heading).
     * N: number of ground truth boxes.
 
   Args:
+    config: The metrics config defined in protos/metrics.proto.
     prediction_frame_id: [M] int64 tensor that identifies frame for each
       prediction.
     prediction_bbox: [M, D] tensor encoding the predicted bounding boxes.

waymo_open_dataset/utils/box_utils.py

@@ -23,7 +23,7 @@
 from waymo_open_dataset.utils import transform_utils
 
 __all__ = [
-    'is_within_box_3d', 'compute_num_points_in_box_3d',
+    'is_within_box_3d', 'compute_num_points_in_box_3d', 'is_within_box_2d',
     'get_upright_3d_box_corners', 'transform_point', 'transform_box'
 ]
 
@@ -70,6 +70,47 @@ def is_within_box_3d(point, box, name=None):
     return point_in_box
 
 
+def is_within_box_2d(point, box):
+  """Checks whether a point is in a BEV box given a set of points and boxes.
+
+  Args:
+    point: [N, 2] tensor. Inner dims are: [x, y].
+    box: [M, 5] tensor. Inner dims are: [center_x, center_y, length, width,
+      heading].
+
+  Returns:
+    point_in_box; [N, M] boolean tensor.
+  """
+
+  with tf.name_scope('IsWithinBox2D'):
+    center = box[:, 0:2]
+    dim = box[:, 2:4]
+    heading = box[:, 4]
+    # [M, 2, 2]
+    rotation = transform_utils.get_yaw_rotation_2d(heading)
+    # [M, 3, 3]
+    transform = transform_utils.get_transform(rotation, center)
+    # [M, 3, 3]
+    transform = tf.linalg.inv(transform)
+    # [M, 2, 2]
+    rotation = transform[:, 0:2, 0:2]
+    # [M, 2]
+    translation = transform[:, 0:2, 2]
+
+    # [N, M, 2]
+    point_in_box_frame = tf.einsum('nj,mij->nmi', point, rotation) + translation
+    # [N, M, 2]
+    point_in_box = tf.logical_and(point_in_box_frame <= dim * 0.5,
+                                  point_in_box_frame >= -dim * 0.5)
+    # [N, M]
+    point_in_box = tf.cast(
+        tf.reduce_prod(
+            input_tensor=tf.cast(point_in_box, dtype=tf.int32), axis=-1),
+        dtype=tf.bool)
+
+    return point_in_box
+
+
 def compute_num_points_in_box_3d(point, box, name=None):
   """Computes the number of points in each box given a set of points and boxes.
 

waymo_open_dataset/utils/transform_utils.py

@@ -20,7 +20,10 @@
 
 import tensorflow as tf
 
-__all__ = ['get_yaw_rotation', 'get_rotation_matrix', 'get_transform']
+__all__ = [
+    'get_yaw_rotation', 'get_yaw_rotation_2d', 'get_rotation_matrix',
+    'get_transform'
+]
 
 
 def get_yaw_rotation(yaw, name=None):
@@ -49,6 +52,28 @@ def get_yaw_rotation(yaw, name=None):
                     axis=-2)
 
 
+def get_yaw_rotation_2d(yaw):
+  """Gets a rotation matrix given yaw only for 2d.
+
+  Args:
+    yaw: x-rotation in radians. This tensor can be any shape except an empty
+      one.
+
+  Returns:
+    A rotation tensor with the same data type of the input. Its shape is
+      [input_shape, 2, 2].
+  """
+  with tf.name_scope('GetYawRotation2D'):
+    cos_yaw = tf.cos(yaw)
+    sin_yaw = tf.sin(yaw)
+
+    return tf.stack([
+        tf.stack([cos_yaw, -1.0 * sin_yaw], axis=-1),
+        tf.stack([sin_yaw, cos_yaw], axis=-1),
+    ],
+                    axis=-2)
+
+
 def get_rotation_matrix(roll, pitch, yaw, name=None):
   """Gets a rotation matrix given roll, pitch, yaw.
 
@@ -101,27 +126,27 @@ def get_rotation_matrix(roll, pitch, yaw, name=None):
     return tf.matmul(r_yaw, tf.matmul(r_pitch, r_roll))
 
 
-def get_transform(rotation, translation, name=None):
-  """Combines 3x3 rotation and 3x1 translation to 4x4 transform.
+def get_transform(rotation, translation):
+  """Combines NxN rotation and Nx1 translation to (N+1)x(N+1) transform.
 
   Args:
-    rotation: [..., 3, 3] rotation tensor.
-    translation: [..., 3] translation tensor. This must have the same type as
+    rotation: [..., N, N] rotation tensor.
+    translation: [..., N] translation tensor. This must have the same type as
       rotation.
-    name: the op name.
 
   Returns:
-    transform: [..., 4, 4] transform tensor. This has the same type as rotation.
+    transform: [..., (N+1), (N+1)] transform tensor. This has the same type as
+      rotation.
   """
-  with tf.compat.v1.name_scope(name, 'GetTransform', [rotation, translation]):
-    # [..., 3, 1]
-    translation_3_1 = translation[..., tf.newaxis]
-    # [..., 3, 4]
-    transform = tf.concat([rotation, translation_3_1], axis=-1)
-    # [..., 3]
+  with tf.name_scope('GetTransform'):
+    # [..., N, 1]
+    translation_n_1 = translation[..., tf.newaxis]
+    # [..., N, N+1]
+    transform = tf.concat([rotation, translation_n_1], axis=-1)
+    # [..., N]
     last_row = tf.zeros_like(translation)
-    # [..., 4]
+    # [..., N+1]
     last_row = tf.concat([last_row, tf.ones_like(last_row[..., 0:1])], axis=-1)
-    # [..., 4, 4]
+    # [..., N+1, N+1]
     transform = tf.concat([transform, last_row[..., tf.newaxis, :]], axis=-2)
     return transform