Repository: keras-team/keras-cv

Issue:
RandomZoom raises error when input is RaggedTensor
Standalone script to reproduce
```python
import tensorflow as tf

from keras_cv.layers.preprocessing import RandomZoom

if __name__ == "__main__":
    layer = RandomZoom(height_factor=0.5, width_factor=0.5)
    inputs = tf.ragged.stack(
        [
            tf.ones((5, 5, 3)),
            tf.ones((8, 8, 3)),
        ]
    )
    outputs = layer(inputs)  # error occurs!
    assert isinstance(outputs, tf.RaggedTensor)
    print(type(outputs))

```

Error message
```bash
tensorflow.python.framework.errors_impl.InvalidArgumentError: Exception encountered when calling layer 'random_zoom' (type RandomZoom).

in user code:

    File ".../keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py", line 253, in _unwrap_ragged_image_call  *
        images = self.augment_ragged_image(
    File ".../keras_cv/layers/preprocessing/random_zoom.py", line 172, in augment_ragged_image  *
        return self.augment_images(
    File ".../keras_cv/layers/preprocessing/random_zoom.py", line 184, in augment_images  *
        zooms = tf.cast(

    InvalidArgumentError: {{function_node __wrapped__ConcatV2_N_2_device_/job:localhost/replica:0/task:0/device:GPU:0}} ConcatOp : Expected concatenating dimensions in the range [-1, 1), but got 1 [Op:ConcatV2] name: concat
```

`RandomZoom` is not included in `keras_cv/layers/preprocessing/ragged_image_test.py`.
I can fix it once you approve.

Discussion:
LukeWood: /gcbrun

Patch:
diff --git a/keras_cv/custom_ops/box_util.cc b/keras_cv/custom_ops/box_util.cc
index 22457c0..a0d5a63 100644
--- a/keras_cv/custom_ops/box_util.cc
+++ b/keras_cv/custom_ops/box_util.cc
@@ -269,7 +269,8 @@ double RotatedBox2D::IoU(const RotatedBox2D& other) const {
   return intersection_area / union_area;
 }
 
-bool RotatedBox2D::left_hand_side(const Vertex& point, const Vertex& v1, const Vertex& v2) const {
+bool RotatedBox2D::left_hand_side(const Vertex& point, const Vertex& v1,
+                                  const Vertex& v2) const {
   double d1 = (point.y - v1.y) * (v2.x - v1.x);
   double d2 = (point.x - v1.x) * (v2.y - v1.y);
   return d1 >= d2;
@@ -330,7 +331,8 @@ std::vector<Vertex> ParseVerticesFromTensor(const Tensor& points_tensor) {
   return points3d;
 }
 
-std::vector<int> GetMinXIndexFromBoxes(std::vector<Upright3DBox>& boxes, std::vector<double>& points) {
+std::vector<int> GetMinXIndexFromBoxes(std::vector<Upright3DBox>& boxes,
+                                       std::vector<double>& points) {
   std::vector<int> res;
   res.reserve(boxes.size());
   auto p_begin = points.begin();
@@ -346,7 +348,8 @@ std::vector<int> GetMinXIndexFromBoxes(std::vector<Upright3DBox>& boxes, std::ve
   return res;
 }
 
-std::vector<int> GetMaxXIndexFromBoxes(std::vector<Upright3DBox>& boxes, std::vector<double>& points) {
+std::vector<int> GetMaxXIndexFromBoxes(std::vector<Upright3DBox>& boxes,
+                                       std::vector<double>& points) {
   std::vector<int> res;
   res.reserve(boxes.size());
   auto p_begin = points.begin();
@@ -362,7 +365,8 @@ std::vector<int> GetMaxXIndexFromBoxes(std::vector<Upright3DBox>& boxes, std::ve
   return res;
 }
 
-std::vector<int> GetMinYIndexFromBoxes(std::vector<Upright3DBox>& boxes, std::vector<double>& points) {
+std::vector<int> GetMinYIndexFromBoxes(std::vector<Upright3DBox>& boxes,
+                                       std::vector<double>& points) {
   std::vector<int> res;
   res.reserve(boxes.size());
   auto p_begin = points.begin();
@@ -378,7 +382,8 @@ std::vector<int> GetMinYIndexFromBoxes(std::vector<Upright3DBox>& boxes, std::ve
   return res;
 }
 
-std::vector<int> GetMaxYIndexFromBoxes(std::vector<Upright3DBox>& boxes, std::vector<double>& points) {
+std::vector<int> GetMaxYIndexFromBoxes(std::vector<Upright3DBox>& boxes,
+                                       std::vector<double>& points) {
   std::vector<int> res;
   res.reserve(boxes.size());
   auto p_begin = points.begin();
diff --git a/keras_cv/custom_ops/kernels/withinbox_op.cc b/keras_cv/custom_ops/kernels/withinbox_op.cc
index 2d74400..14d57a8 100644
--- a/keras_cv/custom_ops/kernels/withinbox_op.cc
+++ b/keras_cv/custom_ops/kernels/withinbox_op.cc
@@ -38,7 +38,8 @@ class WithinBoxOp : public OpKernel {
     const int num_boxes = boxes.dim_size(0);
     Tensor* box_indices = nullptr;
     OP_REQUIRES_OK(
-        ctx, ctx->allocate_output("box_indices", TensorShape({num_points}), &box_indices));
+        ctx, ctx->allocate_output("box_indices", TensorShape({num_points}),
+                                  &box_indices));
     auto boxes_indices_t = box_indices->flat<int>();
     for (auto i = 0; i < num_points; ++i) boxes_indices_t(i) = -1;
 
@@ -52,13 +53,13 @@ class WithinBoxOp : public OpKernel {
 
     // sort, return sorted value and indices
     std::sort(p_indices_x.begin(), p_indices_x.end(),
-      [&points_vec](const int& a, const int& b) -> bool {
-        return points_vec[a].x < points_vec[b].x;
-      });
+              [&points_vec](const int& a, const int& b) -> bool {
+                return points_vec[a].x < points_vec[b].x;
+              });
     std::sort(p_indices_y.begin(), p_indices_y.end(),
-      [&points_vec](const int& a, const int& b) -> bool {
-        return points_vec[a].y < points_vec[b].y;
-      });
+              [&points_vec](const int& a, const int& b) -> bool {
+                return points_vec[a].y < points_vec[b].y;
+              });
     std::vector<double> sorted_points_x;
     sorted_points_x.reserve(num_points);
     std::vector<double> sorted_points_y;
@@ -68,12 +69,16 @@ class WithinBoxOp : public OpKernel {
       sorted_points_y.emplace_back(points_vec[p_indices_y[i]].y);
     }
 
-    // for each box, find all point indices whose x values are within box boundaries
-    // when the box is rotated, the box boundary is the minimum and maximum x for all vertices
-    std::vector<int> points_x_min = box::GetMinXIndexFromBoxes(boxes_vec, sorted_points_x);
-    std::vector<int> points_x_max = box::GetMaxXIndexFromBoxes(boxes_vec, sorted_points_x);
+    // for each box, find all point indices whose x values are within box
+    // boundaries when the box is rotated, the box boundary is the minimum and
+    // maximum x for all vertices
+    std::vector<int> points_x_min =
+        box::GetMinXIndexFromBoxes(boxes_vec, sorted_points_x);
+    std::vector<int> points_x_max =
+        box::GetMaxXIndexFromBoxes(boxes_vec, sorted_points_x);
     std::vector<std::unordered_set<int>> points_x_indices(num_boxes);
-    auto set_fn_x = [&points_x_min, &points_x_max, &p_indices_x, &points_x_indices](int64_t begin, int64_t end) {
+    auto set_fn_x = [&points_x_min, &points_x_max, &p_indices_x,
+                     &points_x_indices](int64_t begin, int64_t end) {
       for (int64_t idx = begin; idx < end; ++idx) {
         std::unordered_set<int> p_set;
         int p_start = points_x_min[idx];
@@ -88,12 +93,16 @@ class WithinBoxOp : public OpKernel {
     const Eigen::TensorOpCost cost(num_points, num_boxes, 3);
     device.parallelFor(num_boxes, cost, set_fn_x);
 
-    // for each box, find all point indices whose y values are within box boundaries
-    // when the box is rotated, the box boundary is the minimum and maximum x for all vertices
-    std::vector<int> points_y_min = box::GetMinYIndexFromBoxes(boxes_vec, sorted_points_y);
-    std::vector<int> points_y_max = box::GetMaxYIndexFromBoxes(boxes_vec, sorted_points_y);
+    // for each box, find all point indices whose y values are within box
+    // boundaries when the box is rotated, the box boundary is the minimum and
+    // maximum x for all vertices
+    std::vector<int> points_y_min =
+        box::GetMinYIndexFromBoxes(boxes_vec, sorted_points_y);
+    std::vector<int> points_y_max =
+        box::GetMaxYIndexFromBoxes(boxes_vec, sorted_points_y);
     std::vector<std::unordered_set<int>> points_y_indices(num_boxes);
-    auto set_fn_y = [&points_y_min, &points_y_max, &p_indices_y, &points_y_indices](int64_t begin, int64_t end) {
+    auto set_fn_y = [&points_y_min, &points_y_max, &p_indices_y,
+                     &points_y_indices](int64_t begin, int64_t end) {
       for (int64_t idx = begin; idx < end; ++idx) {
         std::unordered_set<int> p_set;
         int p_start = points_y_min[idx];
@@ -108,7 +117,9 @@ class WithinBoxOp : public OpKernel {
 
     // for the intersection of x indices set and y indices set, check if
     // those points are within the box
-    auto within_fn = [&points_x_indices, &points_y_indices, &boxes_vec, &points_vec, &boxes_indices_t](int64_t begin, int64_t end) {
+    auto within_fn = [&points_x_indices, &points_y_indices, &boxes_vec,
+                      &points_vec,
+                      &boxes_indices_t](int64_t begin, int64_t end) {
       for (int64_t idx = begin; idx < end; ++idx) {
         std::unordered_set<int>& set_a = points_x_indices[idx];
         std::unordered_set<int>& set_b = points_y_indices[idx];
@@ -128,12 +139,10 @@ class WithinBoxOp : public OpKernel {
       }
     };
     device.parallelFor(num_boxes, cost, within_fn);
-
   }
 };
 
-REGISTER_KERNEL_BUILDER(Name("KcvWithinBox").Device(DEVICE_CPU),
-                        WithinBoxOp);
+REGISTER_KERNEL_BUILDER(Name("KcvWithinBox").Device(DEVICE_CPU), WithinBoxOp);
 
 }  // namespace kerascv
 }  // namespace tensorflow
diff --git a/keras_cv/custom_ops/ops/withinbox_op.cc b/keras_cv/custom_ops/ops/withinbox_op.cc
index cdc206d..b949f62 100644
--- a/keras_cv/custom_ops/ops/withinbox_op.cc
+++ b/keras_cv/custom_ops/ops/withinbox_op.cc
@@ -23,7 +23,6 @@ REGISTER_OP("KcvWithinBox")
     .Input("boxes: float")
     .Output("box_indices: int32")
     .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
-      c->set_output(
-          0, c->MakeShape({c->Dim(c->input(0), 0)}));
+      c->set_output(0, c->MakeShape({c->Dim(c->input(0), 0)}));
       return tensorflow::Status();
     });
diff --git a/keras_cv/layers/preprocessing/random_zoom.py b/keras_cv/layers/preprocessing/random_zoom.py
index d10e502..a347a5e 100644
--- a/keras_cv/layers/preprocessing/random_zoom.py
+++ b/keras_cv/layers/preprocessing/random_zoom.py
@@ -156,9 +156,17 @@ class RandomZoom(VectorizedBaseImageAugmentationLayer):
         return {"height_zooms": height_zooms, "width_zooms": width_zooms}
 
     def augment_ragged_image(self, image, transformation, **kwargs):
-        return self.augment_images(
+        image = tf.expand_dims(image, axis=0)
+        width_zooms = transformation["width_zooms"]
+        height_zooms = transformation["height_zooms"]
+        transformation = {
+            "height_zooms": tf.expand_dims(height_zooms, axis=0),
+            "width_zooms": tf.expand_dims(width_zooms, axis=0),
+        }
+        image = self.augment_images(
             images=image, transformations=transformation, **kwargs
         )
+        return tf.squeeze(image, axis=0)
 
     def augment_images(self, images, transformations, **kwargs):
         images = preprocessing_utils.ensure_tensor(images, self.compute_dtype)


Test patch:
diff --git a/keras_cv/layers/preprocessing/ragged_image_test.py b/keras_cv/layers/preprocessing/ragged_image_test.py
index 1e0987f..e8713dc 100644
--- a/keras_cv/layers/preprocessing/ragged_image_test.py
+++ b/keras_cv/layers/preprocessing/ragged_image_test.py
@@ -78,6 +78,11 @@ CONSISTENT_OUTPUT_TEST_CONFIGURATIONS = [
         layers.RandomTranslation,
         {"height_factor": 0.5, "width_factor": 0.5},
     ),
+    (
+        "RandomZoom",
+        layers.RandomZoom,
+        {"height_factor": 0.2, "width_factor": 0.5},
+    ),
     ("Solarization", layers.Solarization, {"value_range": (0, 255)}),
 ]