experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #24282 from yongtang:15692-draw_bounding_boxes-colors

Browse Source 
PiperOrigin-RevId: 242188454
This commit is contained in:
TensorFlower Gardener 2019-04-05 14:14:31 -07:00
commit e3eec7a4da
11 changed files with 210 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
tensorflow/core/api_def/base_api/api_def_DrawBoundingBoxesV2.pbtxt Normal file
View File

@ -0,0 +1,43 @@
op {
graph_op_name: "DrawBoundingBoxesV2"
in_arg {
name: "images"
description: <<END
4-D with shape `[batch, height, width, depth]`. A batch of images.
END
}
in_arg {
name: "boxes"
description: <<END
3-D with shape `[batch, num_bounding_boxes, 4]` containing bounding
boxes.
END
}
in_arg {
name: "colors"
description: <<END
2-D. A list of RGBA colors to cycle through for the boxes.
END
}
out_arg {
name: "output"
description: <<END
4-D with the same shape as `images`. The batch of input images with
bounding boxes drawn on the images.
END
}
summary: "Draw bounding boxes on a batch of images."
description: <<END
Outputs a copy of `images` but draws on top of the pixels zero or more bounding
boxes specified by the locations in `boxes`. The coordinates of the each
bounding box in `boxes` are encoded as `[y_min, x_min, y_max, x_max]`. The
bounding box coordinates are floats in `[0.0, 1.0]` relative to the width and
height of the underlying image.
For example, if an image is 100 x 200 pixels (height x width) and the bounding
box is `[0.1, 0.2, 0.5, 0.9]`, the upper-left and bottom-right coordinates of
the bounding box will be `(40, 10)` to `(100, 50)` (in (x,y) coordinates).
Parts of the bounding box may fall outside the image.
END
}

4
tensorflow/core/api_def/python_api/api_def_DrawBoundingBoxes.pbtxt
View File

@ -1,6 +1,4 @@
op {
graph_op_name: "DrawBoundingBoxes"
endpoint {
name: "image.draw_bounding_boxes"
}
visibility: HIDDEN
}

4
tensorflow/core/api_def/python_api/api_def_DrawBoundingBoxesV2.pbtxt Normal file
View File

@ -0,0 +1,4 @@
op {
graph_op_name: "DrawBoundingBoxesV2"
visibility: HIDDEN
}

74
tensorflow/core/kernels/draw_bounding_box_op.cc
View File

@ -25,6 +25,30 @@ limitations under the License.
namespace tensorflow {
namespace {
std::vector<std::vector<float>> DefaultColorTable(int depth) {
std::vector<std::vector<float>> color_table;
color_table.emplace_back(std::vector<float>({1, 1, 0, 1})); // 0: yellow
color_table.emplace_back(std::vector<float>({0, 0, 1, 1})); // 1: blue
color_table.emplace_back(std::vector<float>({1, 0, 0, 1})); // 2: red
color_table.emplace_back(std::vector<float>({0, 1, 0, 1})); // 3: lime
color_table.emplace_back(std::vector<float>({0.5, 0, 0.5, 1})); // 4: purple
color_table.emplace_back(std::vector<float>({0.5, 0.5, 0, 1})); // 5: olive
color_table.emplace_back(std::vector<float>({0.5, 0, 0, 1})); // 6: maroon
color_table.emplace_back(std::vector<float>({0, 0, 0.5, 1})); // 7: navy blue
color_table.emplace_back(std::vector<float>({0, 1, 1, 1})); // 8: aqua
color_table.emplace_back(std::vector<float>({1, 0, 1, 1})); // 9: fuchsia
if (depth == 1) {
for (int64 i = 0; i < color_table.size(); i++) {
color_table[i][0] = 1;
}
}
return color_table;
}
} // namespace
template <class T>
class DrawBoundingBoxesOp : public OpKernel {
public:
@ -52,30 +76,31 @@ class DrawBoundingBoxesOp : public OpKernel {
const int64 batch_size = images.dim_size(0);
const int64 height = images.dim_size(1);
const int64 width = images.dim_size(2);
const int64 color_table_length = 10;
std::vector<std::vector<float>> color_table;
if (context->num_inputs() == 3) {
const Tensor& colors_tensor = context->input(2);
OP_REQUIRES(context, colors_tensor.shape().dims() == 2,
errors::InvalidArgument("colors must be a 2-D matrix",
colors_tensor.shape().DebugString()));
OP_REQUIRES(context, colors_tensor.shape().dim_size(1) >= depth,
errors::InvalidArgument("colors must have equal or more ",
"channels than the image provided: ",
colors_tensor.shape().DebugString()));
if (colors_tensor.NumElements() != 0) {
color_table.clear();
// 0: yellow
// 1: blue
// 2: red
// 3: lime
// 4: purple
// 5: olive
// 6: maroon
// 7: navy blue
// 8: aqua
// 9: fuchsia
float color_table[color_table_length][4] = {
{1, 1, 0, 1}, {0, 0, 1, 1}, {1, 0, 0, 1}, {0, 1, 0, 1},
{0.5, 0, 0.5, 1}, {0.5, 0.5, 0, 1}, {0.5, 0, 0, 1}, {0, 0, 0.5, 1},
{0, 1, 1, 1}, {1, 0, 1, 1},
};
// Reset first color channel to 1 if image is GRY.
// For GRY images, this means all bounding boxes will be white.
if (depth == 1) {
for (int64 i = 0; i < color_table_length; i++) {
color_table[i][0] = 1;
auto colors = colors_tensor.matrix<float>();
for (int64 i = 0; i < colors.dimension(0); i++) {
std::vector<float> color_value(4);
for (int64 j = 0; j < 4; j++) {
color_value[j] = colors(i, j);
}
color_table.emplace_back(color_value);
}
}
}
if (color_table.empty()) {
color_table = DefaultColorTable(depth);
}
Tensor* output;
OP_REQUIRES_OK(
@ -90,7 +115,7 @@ class DrawBoundingBoxesOp : public OpKernel {
const int64 num_boxes = boxes.dim_size(1);
const auto tboxes = boxes.tensor<T, 3>();
for (int64 bb = 0; bb < num_boxes; ++bb) {
int64 color_index = bb % color_table_length;
int64 color_index = bb % color_table.size();
const int64 min_box_row =
static_cast<float>(tboxes(b, bb, 0)) * (height - 1);
const int64 min_box_row_clamp = std::max<int64>(min_box_row, int64{0});
@ -179,6 +204,9 @@ class DrawBoundingBoxesOp : public OpKernel {
#define REGISTER_CPU_KERNEL(T) \
REGISTER_KERNEL_BUILDER( \
Name("DrawBoundingBoxes").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
DrawBoundingBoxesOp<T>); \
REGISTER_KERNEL_BUILDER( \
Name("DrawBoundingBoxesV2").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
DrawBoundingBoxesOp<T>);
TF_CALL_half(REGISTER_CPU_KERNEL);
TF_CALL_float(REGISTER_CPU_KERNEL);

11
tensorflow/core/ops/image_ops.cc
View File

@ -601,6 +601,17 @@ REGISTER_OP("DrawBoundingBoxes")
return shape_inference::UnchangedShape(c);
});
// --------------------------------------------------------------------------
REGISTER_OP("DrawBoundingBoxesV2")
.Input("images: T")
.Input("boxes: float")
.Input("colors: float")
.Output("output: T")
.Attr("T: {float, half} = DT_FLOAT")
.SetShapeFn([](InferenceContext* c) {
return shape_inference::UnchangedShapeWithRankAtLeast(c, 3);
});
// --------------------------------------------------------------------------
REGISTER_OP("SampleDistortedBoundingBox")
.Input("image_size: T")

25
tensorflow/python/kernel_tests/draw_bounding_box_op_test.py
View File

@ -54,15 +54,18 @@ class DrawBoundingBoxOpTest(test.TestCase):
image[height - 1, 0:width, 0:depth] = color
return image
def _testDrawBoundingBoxColorCycling(self, img):
def _testDrawBoundingBoxColorCycling(self, img, colors=None):
"""Tests if cycling works appropriately.
Args:
img: 3-D numpy image on which to draw.
"""
color_table = colors
if colors is None:
# THIS TABLE MUST MATCH draw_bounding_box_op.cc
color_table = np.asarray([[1, 1, 0, 1], [0, 0, 1, 1], [1, 0, 0, 1],
[0, 1, 0, 1], [0.5, 0, 0.5, 1], [0.5, 0.5, 0, 1],
[0, 1, 0, 1], [0.5, 0, 0.5,
1], [0.5, 0.5, 0, 1],
[0.5, 0, 0, 1], [0, 0, 0.5, 1], [0, 1, 1, 1],
[1, 0, 1, 1]])
assert len(img.shape) == 3
@ -85,9 +88,9 @@ class DrawBoundingBoxOpTest(test.TestCase):
image = ops.convert_to_tensor(image)
image = image_ops_impl.convert_image_dtype(image, dtypes.float32)
image = array_ops.expand_dims(image, 0)
image = image_ops.draw_bounding_boxes(image, bboxes)
image = image_ops.draw_bounding_boxes(image, bboxes, colors=colors)
with self.cached_session(use_gpu=False) as sess:
op_drawn_image = np.squeeze(self.evaluate(image), 0)
op_drawn_image = np.squeeze(sess.run(image), 0)
self.assertAllEqual(test_drawn_image, op_drawn_image)
def testDrawBoundingBoxRGBColorCycling(self):
@ -105,6 +108,20 @@ class DrawBoundingBoxOpTest(test.TestCase):
image = np.zeros([4, 4, 1], "float32")
self._testDrawBoundingBoxColorCycling(image)
def testDrawBoundingBoxRGBColorCyclingWithColors(self):
"""Test if RGB color cycling works correctly with provided colors."""
image = np.zeros([10, 10, 3], "float32")
colors = np.asarray([[1, 1, 0, 1], [0, 0, 1, 1], [0.5, 0, 0.5, 1],
[0.5, 0.5, 0, 1], [0, 1, 1, 1], [1, 0, 1, 1]])
self._testDrawBoundingBoxColorCycling(image, colors=colors)
def testDrawBoundingBoxRGBAColorCyclingWithColors(self):
"""Test if RGBA color cycling works correctly with provided colors."""
image = np.zeros([10, 10, 4], "float32")
colors = np.asarray([[0.5, 0, 0.5, 1], [0.5, 0.5, 0, 1], [0.5, 0, 0, 1],
[0, 0, 0.5, 1]])
self._testDrawBoundingBoxColorCycling(image, colors=colors)
if __name__ == "__main__":
test.main()

63
tensorflow/python/ops/image_ops_impl.py
View File

@ -3355,7 +3355,6 @@ def crop_and_resize_v1( # pylint: disable=missing-docstring
crop_and_resize_v1.__doc__ = gen_image_ops.crop_and_resize.__doc__
@tf_export(v1=['image.extract_glimpse'])
def extract_glimpse(
input, # pylint: disable=redefined-builtin
@ -3555,3 +3554,65 @@ def combined_non_max_suppression(boxes,
return gen_image_ops.combined_non_max_suppression(
boxes, scores, max_output_size_per_class, max_total_size, iou_threshold,
score_threshold, pad_per_class)
@tf_export('image.draw_bounding_boxes', v1=[])
def draw_bounding_boxes_v2(images, boxes, colors, name=None):
"""Draw bounding boxes on a batch of images.
Outputs a copy of `images` but draws on top of the pixels zero or more
bounding boxes specified by the locations in `boxes`. The coordinates of the
each bounding box in `boxes` are encoded as `[y_min, x_min, y_max, x_max]`.
The bounding box coordinates are floats in `[0.0, 1.0]` relative to the width
and height of the underlying image.
For example, if an image is 100 x 200 pixels (height x width) and the bounding
box is `[0.1, 0.2, 0.5, 0.9]`, the upper-left and bottom-right coordinates of
the bounding box will be `(40, 10)` to `(180, 50)` (in (x,y) coordinates).
Parts of the bounding box may fall outside the image.
Args:
images: A `Tensor`. Must be one of the following types: `float32`, `half`.
4-D with shape `[batch, height, width, depth]`. A batch of images.
boxes: A `Tensor` of type `float32`. 3-D with shape `[batch,
num_bounding_boxes, 4]` containing bounding boxes.
colors: A `Tensor` of type `float32`. 2-D. A list of RGBA colors to cycle
through for the boxes.
name: A name for the operation (optional).
Returns:
A `Tensor`. Has the same type as `images`.
"""
if colors is None and not compat.forward_compatible(2019, 5, 1):
return gen_image_ops.draw_bounding_boxes(images, boxes, name)
return gen_image_ops.draw_bounding_boxes_v2(images, boxes, colors, name)
@tf_export(v1=['image.draw_bounding_boxes'])
def draw_bounding_boxes(images, boxes, name=None, colors=None):
"""Draw bounding boxes on a batch of images.
Outputs a copy of `images` but draws on top of the pixels zero or more
bounding boxes specified by the locations in `boxes`. The coordinates of the
each bounding box in `boxes` are encoded as `[y_min, x_min, y_max, x_max]`.
The bounding box coordinates are floats in `[0.0, 1.0]` relative to the width
and height of the underlying image.
For example, if an image is 100 x 200 pixels (height x width) and the bounding
box is `[0.1, 0.2, 0.5, 0.9]`, the upper-left and bottom-right coordinates of
the bounding box will be `(40, 10)` to `(180, 50)` (in (x,y) coordinates).
Parts of the bounding box may fall outside the image.
Args:
images: A `Tensor`. Must be one of the following types: `float32`, `half`.
4-D with shape `[batch, height, width, depth]`. A batch of images.
boxes: A `Tensor` of type `float32`. 3-D with shape `[batch,
num_bounding_boxes, 4]` containing bounding boxes.
name: A name for the operation (optional).
Returns:
A `Tensor`. Has the same type as `images`.
"""
return draw_bounding_boxes_v2(images, boxes, colors, name)

2
tensorflow/tools/api/golden/v1/tensorflow.image.pbtxt
View File

@ -74,7 +74,7 @@ tf_module {
}
member_method {
name: "draw_bounding_boxes"
argspec: "args=[\'images\', \'boxes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
argspec: "args=[\'images\', \'boxes\', \'name\', \'colors\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
}
member_method {
name: "encode_jpeg"

4
tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
View File

@ -936,6 +936,10 @@ tf_module {
name: "DrawBoundingBoxes"
argspec: "args=[\'images\', \'boxes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "DrawBoundingBoxesV2"
argspec: "args=[\'images\', \'boxes\', \'colors\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "DynamicPartition"
argspec: "args=[\'data\', \'partitions\', \'num_partitions\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "

2
tensorflow/tools/api/golden/v2/tensorflow.image.pbtxt
View File

@ -74,7 +74,7 @@ tf_module {
}
member_method {
name: "draw_bounding_boxes"
argspec: "args=[\'images\', \'boxes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
argspec: "args=[\'images\', \'boxes\', \'colors\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "encode_jpeg"

4
tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
View File

@ -936,6 +936,10 @@ tf_module {
name: "DrawBoundingBoxes"
argspec: "args=[\'images\', \'boxes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "DrawBoundingBoxesV2"
argspec: "args=[\'images\', \'boxes\', \'colors\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "DynamicPartition"
argspec: "args=[\'data\', \'partitions\', \'num_partitions\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "