correct imports

Author: PatWie

tensorflow_addons/custom_ops/opticalflow/cc/kernels/correlation_cost_op.cc

@@ -13,8 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/contrib/correlation_cost/kernels/correlation_cost_op.h"
+#define EIGEN_USE_THREADS
 
+#if GOOGLE_CUDA
+#define EIGEN_USE_GPU
+#endif  // GOOGLE_CUDA
+
+#include "tensorflow_addons/custom_ops/opticalflow/cc/kernels/correlation_cost_op.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/types.h"

tensorflow_addons/custom_ops/opticalflow/cc/kernels/correlation_cost_op_gpu.cu.cc

@@ -17,8 +17,7 @@ limitations under the License.
 
 #define EIGEN_USE_GPU
 
-#include "tensorflow/contrib/correlation_cost/kernels/correlation_cost_op.h"
-
+#include "tensorflow_addons/custom_ops/opticalflow/cc/kernels/correlation_cost_op.h"
 #include "external/cub_archive/cub/device/device_reduce.cuh"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"

tensorflow_addons/custom_ops/opticalflow/cc/ops/correlation_cost_op.cc

@@ -90,7 +90,7 @@ REGISTER_OP("CorrelationCost")
 Compute Correlation costs.
 
 This layer implements the correlation operation from
-FlowNet: Learning Optical Flow with Convolutional Networks (Fischer et al.)
+FlowNet Learning Optical Flow with Convolutional Networks (Fischer et al.)
 
 input_a: A `Tensor` of the format specified by `data_format`.
 input_b: A `Tensor` of the format specified by `data_format`.

tensorflow_addons/opticalflow/correlation_cost.py

@@ -18,14 +18,13 @@
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow.contrib.correlation_cost.ops import gen_correlation_cost_op
-from tensorflow.contrib.util import loader
+import tensorflow as tf
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
-from tensorflow.python.platform import resource_loader
+from tensorflow_addons.utils.resource_loader import get_path_to_datafile
 
-_correlation_cost_op_so = loader.load_op_library(
-    resource_loader.get_path_to_datafile("_correlation_cost_op.so"))
+_correlation_cost_op_so = tf.load_op_library(
+    get_path_to_datafile("custom_ops/opticalflow/_correlation_cost_ops.so"))
 
 # pylint: disable=redefined-builtin
 
@@ -81,7 +80,7 @@ def correlation_cost(input_a,
     """
 
     with ops.name_scope(name, "correlation_cost"):
-        op_call = gen_correlation_cost_op.correlation_cost
+        op_call = _correlation_cost_op_so.correlation_cost
         ret = op_call(
             input_a,
             input_b,
@@ -98,7 +97,7 @@ def correlation_cost(input_a,
         return ret
 
 
-correlation_cost_grad = gen_correlation_cost_op.correlation_cost_grad
+correlation_cost_grad = _correlation_cost_op_so.correlation_cost_grad
 
 
 @ops.RegisterGradient("CorrelationCost")
@@ -114,7 +113,7 @@ def _correlation_cost_grad(op, grad_output):
     input_b = ops.convert_to_tensor(op.inputs[1], name="input_b")
     grad_output_tensor = ops.convert_to_tensor(grad_output, name="grad_output")
 
-    op_call = gen_correlation_cost_op.correlation_cost_grad
+    op_call = _correlation_cost_op_so.correlation_cost_grad
     grads = op_call(
         input_a,
         input_b,

tensorflow_addons/opticalflow/correlation_cost_test.py

@@ -19,7 +19,7 @@
 
 import numpy as np
 
-from tensorflow.contrib.correlation_cost.python.ops import correlation_cost_op
+from tensorflow_addons.opticalflow import correlation_cost
 from tensorflow.python.framework import dtypes
 from tensorflow.python.ops import array_ops
 from tensorflow.python.framework import ops
@@ -49,7 +49,7 @@ def _forward(self,
             stride_2 = 2
             pad = 4
 
-            call_op = correlation_cost_op.correlation_cost
+            call_op = correlation_cost
             actual_op = call_op(
                 input_a_op,
                 input_b_op,
@@ -179,7 +179,7 @@ def _gradients(self, data_format='NCHW', use_gpu=False):
             input_a_op = ops.convert_to_tensor(input_a, dtype=dtypes.float32)
             input_b_op = ops.convert_to_tensor(input_b, dtype=dtypes.float32)
 
-            call_op = correlation_cost_op.correlation_cost
+            call_op = correlation_cost
             actual_op = call_op(
                 input_a_op,
                 input_b_op,