1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 // See docs for ImageConnectedComponents in ../ops/image_ops.cc, and description 17 // of the algorithm in segmentation_ops.h. 18 19 #define EIGEN_USE_THREADS 20 21 #include "tensorflow/contrib/image/kernels/segmentation_ops.h" 22 #include "tensorflow/core/framework/op_kernel.h" 23 #include "tensorflow/core/framework/register_types.h" 24 #include "tensorflow/core/framework/types.h" 25 #include "tensorflow/core/platform/types.h" 26 27 namespace tensorflow { 28 29 using tensorflow::functor::BlockedImageUnionFindFunctor; 30 using tensorflow::functor::FindRootFunctor; 31 using tensorflow::functor::ImageConnectedComponentsFunctor; 32 using tensorflow::functor::TensorRangeFunctor; 33 34 using OutputType = typename BlockedImageUnionFindFunctor<bool>::OutputType; 35 36 // Computes connected components on batches of 2D images. 37 template <typename Device, typename T> 38 class ImageConnectedComponents : public OpKernel { 39 public: ImageConnectedComponents(OpKernelConstruction * ctx)40 explicit ImageConnectedComponents(OpKernelConstruction* ctx) 41 : OpKernel(ctx) {} 42 Compute(OpKernelContext * ctx)43 void Compute(OpKernelContext* ctx) override { 44 const Tensor& images_t = ctx->input(0); 45 OP_REQUIRES(ctx, images_t.shape().dims() == 3, 46 errors::InvalidArgument("Input images must have rank 3")); 47 Tensor forest_t, rank_t; 48 OP_REQUIRES_OK(ctx, ctx->allocate_temp(tensorflow::DT_INT64, 49 images_t.shape(), &forest_t)); 50 OP_REQUIRES_OK(ctx, ctx->allocate_temp(tensorflow::DT_INT64, 51 images_t.shape(), &rank_t)); 52 Tensor* output_t; 53 OP_REQUIRES_OK(ctx, ctx->allocate_output(0, images_t.shape(), &output_t)); 54 55 // Fill forest with values from 0 to n - 1, so that each node points to 56 // itself. 57 TensorRangeFunctor<Device>()(ctx->eigen_device<Device>(), 58 forest_t.flat<OutputType>()); 59 auto rank = rank_t.tensor<OutputType, 3>(); 60 rank.device(ctx->eigen_device<Device>()) = rank.constant(OutputType(0)); 61 62 const auto images = images_t.tensor<T, 3>(); 63 auto forest = forest_t.tensor<OutputType, 3>(); 64 ImageConnectedComponentsFunctor<Device, T>()( 65 ctx, output_t->flat<OutputType>(), images, forest, rank); 66 } 67 }; 68 69 using CPUDevice = Eigen::ThreadPoolDevice; 70 71 namespace functor { 72 73 // Connected components CPU implementation. See `segmentation_ops.h` for a 74 // description of the algorithm. 75 template <typename T> 76 struct ImageConnectedComponentsFunctor<CPUDevice, T> { operator ()tensorflow::functor::ImageConnectedComponentsFunctor77 void operator()(OpKernelContext* ctx, 78 typename TTypes<OutputType>::Flat output, 79 typename TTypes<T, 3>::ConstTensor images, 80 typename TTypes<OutputType, 3>::Tensor forest, 81 typename TTypes<OutputType, 3>::Tensor rank) { 82 const int64 num_images = images.dimension(0), 83 num_rows = images.dimension(1), num_cols = images.dimension(2), 84 num_elements = images.size(); 85 // Bail out early for an empty image--no work to do. 86 if (num_elements == 0) { 87 return; 88 } 89 auto worker_threads = ctx->device()->tensorflow_cpu_worker_threads(); 90 BlockedImageUnionFindFunctor<T> union_find( 91 images.data(), num_rows, num_cols, forest.data(), rank.data()); 92 while (union_find.can_merge()) { 93 union_find.merge_blocks(); 94 int64 num_blocks_vertically = union_find.num_blocks_vertically(); 95 int64 num_blocks_horizontally = union_find.num_blocks_horizontally(); 96 // Merging each block calls union_down for each pixel in a row of the 97 // block, and union_right for each pixel in a column of the block. Assume 98 // 20 instructions for each call to union_down or union_right. find() may 99 // loop more while searching for the root, but this should not be very 100 // significant. 101 int cost = (union_find.block_height() + union_find.block_width()) * 20; 102 Shard(worker_threads->num_threads, worker_threads->workers, 103 num_images * num_blocks_vertically * num_blocks_horizontally, cost, 104 [&union_find, num_blocks_vertically, num_blocks_horizontally]( 105 int64 start_block, int64 limit_block) { 106 for (int64 i = start_block; i < limit_block; i++) { 107 int64 block_x = i % num_blocks_horizontally; 108 int64 block_y = 109 (i / num_blocks_horizontally) % num_blocks_vertically; 110 int64 image = 111 i / (num_blocks_horizontally * num_blocks_vertically); 112 union_find.merge_internal_block_edges(image, block_y, block_x); 113 } 114 }); 115 } 116 FindRootFunctor<CPUDevice, T>()(ctx->eigen_device<CPUDevice>(), output, 117 images.data(), union_find); 118 } 119 }; 120 121 } // end namespace functor 122 123 #define REGISTER_IMAGE_CONNECTED_COMPONENTS(TYPE) \ 124 REGISTER_KERNEL_BUILDER(Name("ImageConnectedComponents") \ 125 .Device(DEVICE_CPU) \ 126 .TypeConstraint<TYPE>("dtype"), \ 127 ImageConnectedComponents<CPUDevice, TYPE>) 128 // Connected components (arguably) make sense for number, bool, and string types 129 TF_CALL_NUMBER_TYPES(REGISTER_IMAGE_CONNECTED_COMPONENTS); 130 TF_CALL_bool(REGISTER_IMAGE_CONNECTED_COMPONENTS); 131 TF_CALL_string(REGISTER_IMAGE_CONNECTED_COMPONENTS); 132 #undef REGISTER_IMAGE_CONNECTED_COMPONENTS 133 134 // TODO(ringwalt): Implement on GPU. We probably want to stick to the original 135 // algorithm by Stava and Benes there for efficiency (computing small blocks in 136 // shared memory in CUDA thread blocks, instead of starting with single-pixel 137 // blocks). 138 139 } // end namespace tensorflow 140