diff --git a/tensorflow/core/kernels/image/crop_and_resize_op.cc b/tensorflow/core/kernels/image/crop_and_resize_op.cc index 5c196df9cfe..4efc4ae8846 100644 --- a/tensorflow/core/kernels/image/crop_and_resize_op.cc +++ b/tensorflow/core/kernels/image/crop_and_resize_op.cc @@ -207,7 +207,7 @@ class CropAndResizeOp : public AsyncOpKernel { namespace functor { template struct CropAndResize { - bool operator()(const OpKernelContext* context, + bool operator()(OpKernelContext* context, typename TTypes::ConstTensor image, typename TTypes::ConstTensor boxes, typename TTypes::ConstTensor box_index, @@ -222,6 +222,17 @@ struct CropAndResize { const int crop_width = crops.dimension(2); const int depth = crops.dimension(3); + // Since `functor::CropAndResize` operates on float, we first validate + // that we don't overflow (since overflow causes undefined behavior which + // could result in segfault in this scenario). + const Eigen::Tensor only_finite_elements = + boxes.isfinite().all(); + if (!only_finite_elements()) { + context->SetStatus(errors::InvalidArgument( + "Boxes contains at least one element that is not finite")); + return false; + } + // Sharding across boxes. auto CropAndResizePerBox = [&](int64 start_box, int64 limit_box) { for (int b = start_box; b < limit_box; ++b) { diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py index a7c964d8a30..c77fe1a300c 100644 --- a/tensorflow/python/ops/image_ops_test.py +++ b/tensorflow/python/ops/image_ops_test.py @@ -5710,6 +5710,25 @@ class DecodeImageTest(test_util.TensorFlowTestCase, parameterized.TestCase): self.assertAllEqual(list(image2.shape), [12, 40, 20, 3]) self.assertAllEqual(image2, image3) + def testImageCropAndResize(self): + if test_util.is_gpu_available(): + op = image_ops_impl.crop_and_resize_v2( + image=array_ops.zeros((2, 1, 1, 1)), + boxes=[[1.0e+40, 0, 0, 0]], + box_indices=[1], + crop_size=[1, 1]) + self.evaluate(op) + else: + message = "Boxes contains at least one element that is not finite" + with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError), + message): + op = image_ops_impl.crop_and_resize_v2( + image=array_ops.zeros((2, 1, 1, 1)), + boxes=[[1.0e+40, 0, 0, 0]], + box_indices=[1], + crop_size=[1, 1]) + self.evaluate(op) + @parameterized.named_parameters( ("_jpeg", "JPEG", "jpeg_merge_test1.jpg"), ("_png", "PNG", "lena_rgba.png"),