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

Use DefaultColorTable utility function

Browse Source 
Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
This commit is contained in:
Yong Tang 2019-01-25 23:09:39 +00:00
parent ade7691974
commit d43bda0968
1 changed files with 30 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

65
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,47 +76,15 @@ 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 default_color_table_length = 10;
// 0: yellow
// 1: blue
// 2: red
// 3: lime
// 4: purple
// 5: olive
// 6: maroon
// 7: navy blue
// 8: aqua
// 9: fuchsia
float default_color_table[default_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},
};
std::vector<std::vector<float>> color_table;
for (int64 i = 0; i < default_color_table_length; i++) {
std::vector<float> color_value(4);
for (int64 j = 0; j < 4; j++) {
color_value[j] = default_color_table[i][j];
}
color_table.emplace_back(color_value);
}
// 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.size(); i++) {
color_table[i][0] = 1;
}
}
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) == 4,
errors::InvalidArgument("colors must be n x 4 (RGBA)",
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();
@ -107,6 +99,9 @@ class DrawBoundingBoxesOp : public OpKernel {
}
}
}
if (color_table.empty()) {
color_table = DefaultColorTable(depth);
}
Tensor* output;
OP_REQUIRES_OK(
context,