1 /* Copyright 2020 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 #ifndef TENSORFLOW_CORE_KERNELS_IMAGE_IMAGE_OPS_H_ 17 #define TENSORFLOW_CORE_KERNELS_IMAGE_IMAGE_OPS_H_ 18 19 // See docs in ../ops/image_ops.cc. 20 21 #define EIGEN_USE_THREADS 22 23 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 24 25 #include "tensorflow/core/framework/tensor_types.h" 26 #include "tensorflow/core/platform/types.h" 27 28 namespace tensorflow { 29 30 namespace generator { 31 32 enum Interpolation { NEAREST, BILINEAR }; 33 enum Mode { FILL_REFLECT, FILL_WRAP, FILL_CONSTANT, FILL_NEAREST }; 34 35 using Eigen::array; 36 using Eigen::DenseIndex; 37 38 // Follow scipy's implementation 39 // https://github.com/scipy/scipy/blob/master/scipy/ndimage/src/ni_interpolation.c 40 template <typename Device, Mode M> 41 struct MapCoordinate { 42 float operator()(const float out_coord, const DenseIndex len); 43 }; 44 45 template <typename Device> 46 struct MapCoordinate<Device, Mode::FILL_REFLECT> { 47 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE float operator()(const float out_coord, 48 const DenseIndex len) { 49 // Reflect [abcd] to [dcba|abcd|dcba]. 50 float in_coord = out_coord; 51 if (in_coord < 0) { 52 if (len <= 1) { 53 in_coord = 0; 54 } else { 55 const DenseIndex sz2 = 2 * len; 56 if (in_coord < sz2) { 57 in_coord = sz2 * static_cast<DenseIndex>(-in_coord / sz2) + in_coord; 58 } 59 in_coord = (in_coord < -len) ? in_coord + sz2 : -in_coord - 1; 60 } 61 } else if (in_coord > len - 1) { 62 if (len <= 1) { 63 in_coord = 0; 64 } else { 65 const DenseIndex sz2 = 2 * len; 66 in_coord -= sz2 * static_cast<DenseIndex>(in_coord / sz2); 67 if (in_coord >= len) { 68 in_coord = sz2 - in_coord - 1; 69 } 70 } 71 } 72 // clamp is necessary because when out_coord = 3.5 and len = 4, 73 // in_coord = 3.5 and will be rounded to 4 in nearest interpolation. 74 return Eigen::internal::scalar_clamp_op<float>(0.0f, len - 1)(in_coord); 75 } 76 }; 77 78 template <typename Device> 79 struct MapCoordinate<Device, Mode::FILL_WRAP> { 80 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE float operator()(const float out_coord, 81 const DenseIndex len) { 82 // Wrap [abcd] to [abcd|abcd|abcd]. 83 float in_coord = out_coord; 84 if (in_coord < 0) { 85 if (len <= 1) { 86 in_coord = 0; 87 } else { 88 const DenseIndex sz = len - 1; 89 in_coord += len * (static_cast<DenseIndex>(-in_coord / sz) + 1); 90 } 91 } else if (in_coord > len - 1) { 92 if (len <= 1) { 93 in_coord = 0; 94 } else { 95 const DenseIndex sz = len - 1; 96 in_coord -= len * static_cast<DenseIndex>(in_coord / sz); 97 } 98 } 99 // clamp is necessary because when out_coord = -0.5 and len = 4, 100 // in_coord = 3.5 and will be rounded to 4 in nearest interpolation. 101 return Eigen::internal::scalar_clamp_op<float>(0.0f, len - 1)(in_coord); 102 } 103 }; 104 105 template <typename Device> 106 struct MapCoordinate<Device, Mode::FILL_CONSTANT> { 107 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE float operator()(const float out_coord, 108 const DenseIndex len) { 109 return out_coord; 110 } 111 }; 112 113 template <typename Device> 114 struct MapCoordinate<Device, Mode::FILL_NEAREST> { 115 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE float operator()(const float out_coord, 116 const DenseIndex len) { 117 return Eigen::internal::scalar_clamp_op<float>(0.0f, len - 1)(out_coord); 118 } 119 }; 120 121 template <typename Device, typename T, Mode M> 122 class ProjectiveGenerator { 123 private: 124 typename TTypes<T, 4>::ConstTensor input_; 125 typename TTypes<float>::ConstMatrix transforms_; 126 const Interpolation interpolation_; 127 const T fill_value_; 128 129 public: 130 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE 131 ProjectiveGenerator(typename TTypes<T, 4>::ConstTensor input, 132 typename TTypes<float>::ConstMatrix transforms, 133 const Interpolation interpolation, const T fill_value) 134 : input_(input), 135 transforms_(transforms), 136 interpolation_(interpolation), 137 fill_value_(fill_value) {} 138 139 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE T 140 operator()(const array<DenseIndex, 4>& coords) const { 141 const int64 output_y = coords[1]; 142 const int64 output_x = coords[2]; 143 const float* transform = 144 transforms_.dimension(0) == 1 145 ? transforms_.data() 146 : &transforms_.data()[transforms_.dimension(1) * coords[0]]; 147 float projection = transform[6] * output_x + transform[7] * output_y + 1.f; 148 if (projection == 0) { 149 // Return the fill value for infinite coordinates, 150 // which are outside the input image 151 return fill_value_; 152 } 153 const float input_x = 154 (transform[0] * output_x + transform[1] * output_y + transform[2]) / 155 projection; 156 const float input_y = 157 (transform[3] * output_x + transform[4] * output_y + transform[5]) / 158 projection; 159 160 // Map out-of-boundary input coordinates to in-boundary based on fill_mode. 161 auto map_functor = MapCoordinate<Device, M>(); 162 const float x = map_functor(input_x, input_.dimension(2)); 163 const float y = map_functor(input_y, input_.dimension(1)); 164 165 const DenseIndex batch = coords[0]; 166 const DenseIndex channels = coords[3]; 167 switch (interpolation_) { 168 case NEAREST: 169 return nearest_interpolation(batch, y, x, channels, fill_value_); 170 case BILINEAR: 171 return bilinear_interpolation(batch, y, x, channels, fill_value_); 172 } 173 // Unreachable; ImageProjectiveTransform only uses INTERPOLATION_NEAREST 174 // or INTERPOLATION_BILINEAR. 175 return fill_value_; 176 } 177 178 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE T 179 nearest_interpolation(const DenseIndex batch, const float y, const float x, 180 const DenseIndex channel, const T fill_value) const { 181 return read_with_fill_value(batch, DenseIndex(std::round(y)), 182 DenseIndex(std::round(x)), channel, fill_value); 183 } 184 185 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE T 186 bilinear_interpolation(const DenseIndex batch, const float y, const float x, 187 const DenseIndex channel, const T fill_value) const { 188 const float y_floor = std::floor(y); 189 const float x_floor = std::floor(x); 190 const float y_ceil = y_floor + 1; 191 const float x_ceil = x_floor + 1; 192 // f(x, y_floor) = (x_ceil - x) / (x_ceil - x_floor) * f(x_floor, y_floor) 193 // + (x - x_floor) / (x_ceil - x_floor) * f(x_ceil, y_floor) 194 const float value_yfloor = 195 (x_ceil - x) * static_cast<float>(read_with_fill_value( 196 batch, DenseIndex(y_floor), DenseIndex(x_floor), 197 channel, fill_value)) + 198 (x - x_floor) * static_cast<float>(read_with_fill_value( 199 batch, DenseIndex(y_floor), DenseIndex(x_ceil), 200 channel, fill_value)); 201 // f(x, y_ceil) = (x_ceil - x) / (x_ceil - x_floor) * f(x_floor, y_ceil) 202 // + (x - x_floor) / (x_ceil - x_floor) * f(x_ceil, y_ceil) 203 const float value_yceil = 204 (x_ceil - x) * static_cast<float>(read_with_fill_value( 205 batch, DenseIndex(y_ceil), DenseIndex(x_floor), 206 channel, fill_value)) + 207 (x - x_floor) * static_cast<float>(read_with_fill_value( 208 batch, DenseIndex(y_ceil), DenseIndex(x_ceil), 209 channel, fill_value)); 210 // f(x, y) = (y_ceil - y) / (y_ceil - y_floor) * f(x, y_floor) 211 // + (y - y_floor) / (y_ceil - y_floor) * f(x, y_ceil) 212 return T((y_ceil - y) * value_yfloor + (y - y_floor) * value_yceil); 213 } 214 215 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE T read_with_fill_value( 216 const DenseIndex batch, const DenseIndex y, const DenseIndex x, 217 const DenseIndex channel, const T fill_value) const { 218 // batch and channel must be correct, because they are passed unchanged from 219 // the input. 220 return (0 <= y && y < input_.dimension(1) && 0 <= x && 221 x < input_.dimension(2)) 222 ? input_(array<DenseIndex, 4>{batch, y, x, channel}) 223 : fill_value; 224 } 225 }; 226 227 } // end namespace generator 228 229 namespace functor { 230 231 using generator::Interpolation; 232 using generator::Mode; 233 using generator::ProjectiveGenerator; 234 235 template <typename Device, typename T> 236 struct FillProjectiveTransform { 237 typedef typename TTypes<T, 4>::Tensor OutputType; 238 typedef typename TTypes<T, 4>::ConstTensor InputType; 239 typedef typename TTypes<float, 2>::ConstTensor TransformsType; 240 const Interpolation interpolation; 241 242 explicit FillProjectiveTransform(Interpolation interpolation) 243 : interpolation(interpolation) {} 244 245 EIGEN_ALWAYS_INLINE 246 void operator()(const Device& device, OutputType* output, 247 const InputType& images, const TransformsType& transform, 248 const Mode fill_mode, const T fill_value) const { 249 switch (fill_mode) { 250 case Mode::FILL_REFLECT: 251 output->device(device) = 252 output->generate(ProjectiveGenerator<Device, T, Mode::FILL_REFLECT>( 253 images, transform, interpolation, fill_value)); 254 break; 255 case Mode::FILL_WRAP: 256 output->device(device) = 257 output->generate(ProjectiveGenerator<Device, T, Mode::FILL_WRAP>( 258 images, transform, interpolation, fill_value)); 259 break; 260 case Mode::FILL_CONSTANT: 261 output->device(device) = output->generate( 262 ProjectiveGenerator<Device, T, Mode::FILL_CONSTANT>( 263 images, transform, interpolation, fill_value)); 264 break; 265 case Mode::FILL_NEAREST: 266 output->device(device) = 267 output->generate(ProjectiveGenerator<Device, T, Mode::FILL_NEAREST>( 268 images, transform, interpolation, fill_value)); 269 break; 270 } 271 } 272 }; 273 274 } // end namespace functor 275 276 } // end namespace tensorflow 277 278 #endif // TENSORFLOW_CORE_KERNELS_IMAGE_IMAGE_OPS_H_ 279