Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix memory corruption issue with tf.raw_ops.DrawBoundingBoxesV2.
PiperOrigin-RevId: 372033910
Change-Id: I8a9f4efc1c8ddaacbc26ec1fbe4bfdd6791c226d
  • Loading branch information
Amit Patankar authored and tensorflower-gardener committed May 5, 2021
1 parent 31bd502 commit 79865b5
Showing 1 changed file with 6 additions and 0 deletions.
6 changes: 6 additions & 0 deletions tensorflow/core/kernels/image/draw_bounding_box_op.cc
Expand Up @@ -73,6 +73,12 @@ class DrawBoundingBoxesOp : public OpKernel {
errors::InvalidArgument("Channel depth should be either 1 (GRY), "
"3 (RGB), or 4 (RGBA)"));

OP_REQUIRES(
context, boxes.dim_size(2) == 4,
errors::InvalidArgument(
"The size of the third dimension of the box must be 4. Received: ",
boxes.dim_size(2)));

const int64 batch_size = images.dim_size(0);
const int64 height = images.dim_size(1);
const int64 width = images.dim_size(2);
Expand Down

0 comments on commit 79865b5

Please sign in to comment.