From b84bff476514c1a2ee80d9f1bc31a9cb5dcc2ee5 Mon Sep 17 00:00:00 2001 From: Yong Tang Date: Mon, 14 May 2018 17:57:23 -0700 Subject: Improve shape function of `tf.image.draw_bounding_boxes` (#19237) * Improve shape function of `tf.image.draw_bounding_boxes` The `tf.image.draw_bounding_boxes` requires `boxes` to be 3-D shape though there was no check on shape function. This fix improves the shape function by restricting the boxes to 3-D. Signed-off-by: Yong Tang * Add additional check to make sure boxes shape ends with 4 ([batch, num_bounding_boxes, 4]) Signed-off-by: Yong Tang * Address review feedback with addtional shape checks. Signed-off-by: Yong Tang * Add unit tests Signed-off-by: Yong Tang --- tensorflow/core/ops/image_ops_test.cc | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) (limited to 'tensorflow/core/ops/image_ops_test.cc') diff --git a/tensorflow/core/ops/image_ops_test.cc b/tensorflow/core/ops/image_ops_test.cc index 5f0b391b0d..517af26b44 100644 --- a/tensorflow/core/ops/image_ops_test.cc +++ b/tensorflow/core/ops/image_ops_test.cc @@ -312,4 +312,23 @@ TEST(ImageOpsTest, QuantizedResizeBilinear_ShapeFn) { INFER_OK(op, "[1,?,3,?];[2];[];[]", "[d0_0,20,30,d0_3];[];[]"); } +TEST(ImageOpsTest, DrawBoundingBoxes_ShapeFn) { + ShapeInferenceTestOp op("DrawBoundingBoxes"); + op.input_tensors.resize(2); + + // Check images. + INFER_ERROR("must be rank 4", op, "[1,?,3];?"); + INFER_ERROR("should be either 1 (GRY), 3 (RGB), or 4 (RGBA)", + op, "[1,?,?,5];?"); + + // Check boxes. + INFER_ERROR("must be rank 3", op, "[1,?,?,4];[1,4]"); + INFER_ERROR("Dimension must be 4", op, "[1,?,?,4];[1,2,2]"); + + // OK shapes. + INFER_OK(op, "[4,?,?,4];?", "in0"); + INFER_OK(op, "[?,?,?,?];[?,?,?]", "in0"); + INFER_OK(op, "[4,?,?,4];[?,?,?]", "in0"); + INFER_OK(op, "[4,?,?,4];[?,?,4]", "in0"); +} } // end namespace tensorflow -- cgit v1.2.3