diff --git a/tensorflow/core/kernels/image/crop_and_resize_op.cc b/tensorflow/core/kernels/image/crop_and_resize_op.cc index 41e30e543c1085..571e54a5a15700 100644 --- a/tensorflow/core/kernels/image/crop_and_resize_op.cc +++ b/tensorflow/core/kernels/image/crop_and_resize_op.cc @@ -170,14 +170,15 @@ class CropAndResizeOp : public AsyncOpKernel { context, crop_height > 0 && crop_width > 0, errors::InvalidArgument("crop dimensions must be positive"), done); + TensorShape shape; + OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(num_boxes), done); + OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(crop_height), done); + OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(crop_width), done); + OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(depth), done); // Allocate output tensor. Tensor* output = nullptr; - OP_REQUIRES_OK_ASYNC( - context, - context->allocate_output( - 0, TensorShape({num_boxes, crop_height, crop_width, depth}), - &output), - done); + OP_REQUIRES_OK_ASYNC(context, context->allocate_output(0, shape, &output), + done); auto compute_callback = [this, context, output]() { const Tensor& image = context->input(0); @@ -417,14 +418,15 @@ class CropAndResizeGradImageOp : public AsyncOpKernel { done); } + TensorShape shape; + OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(batch_size), done); + OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(image_height), done); + OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(image_width), done); + OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(depth), done); // Allocate output tensor. Tensor* output = nullptr; - OP_REQUIRES_OK_ASYNC( - context, - context->allocate_output( - 0, TensorShape({batch_size, image_height, image_width, depth}), - &output), - done); + OP_REQUIRES_OK_ASYNC(context, context->allocate_output(0, shape, &output), + done); auto compute_callback = [this, context, output]() { const Tensor& grads = context->input(0); diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py index bfe4671328f460..86ac731a3080fe 100644 --- a/tensorflow/python/ops/image_ops_test.py +++ b/tensorflow/python/ops/image_ops_test.py @@ -6051,6 +6051,16 @@ def testImageCropAndResize(self): crop_size=[1, 1]) self.evaluate(op) + def testImageCropAndResizeWithInvalidInput(self): + with self.session(): + with self.assertRaises((errors.InternalError, ValueError)): + op = image_ops_impl.crop_and_resize_v2( + image=np.ones((1, 1, 1, 1)), + boxes=np.ones((11, 4)), + box_indices=np.ones((11)), + crop_size=[2065374891, 1145309325]) + self.evaluate(op) + @parameterized.named_parameters( ("_jpeg", "JPEG", "jpeg_merge_test1.jpg"), ("_png", "PNG", "lena_rgba.png"),