Add CombineTransforms operator

dali/operators/geometry/affine_transforms/combine_transforms.cc

@@ -0,0 +1,133 @@
+// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <string>
+#include <utility>
+#include <vector>
+#include "dali/core/format.h"
+#include "dali/core/geom/mat.h"
+#include "dali/core/static_switch.h"
+#include "dali/kernels/kernel_manager.h"
+#include "dali/pipeline/data/types.h"
+#include "dali/pipeline/operator/op_spec.h"
+#include "dali/pipeline/workspace/workspace.h"
+#include "dali/pipeline/operator/operator.h"
+
+#define TRANSFORM_INPUT_TYPES (float)
+
+namespace dali {
+
+template <int... values>
+using dims = std::integer_sequence<int, values...>;
+
+template <typename T, int mat_dim>
+using affine_mat_t = mat<mat_dim, mat_dim, T>;
+
+DALI_SCHEMA(transforms__Combine)
+  .DocStr(R"code(Combines two or more affine transforms.)code")
+  .NumInput(2, 99)
+  .NumOutput(1)
+  .AddParent("TransformAttr");
+
+class CombineTransformsCPU : public Operator<CPUBackend> {
+ public:
+  explicit CombineTransformsCPU(const OpSpec &spec) :
+      Operator<CPUBackend>(spec),
+      reverse_order_(spec.GetArgument<bool>("reverse_order")) {
+  }
+
+  bool CanInferOutputs() const override { return true; }
+
+ protected:
+  bool SetupImpl(std::vector<OutputDesc> &output_descs,
+                 const workspace_t<CPUBackend> &ws) override {
+    assert(ws.NumInput() > 1);
+    TensorListShape<> in0_shape = ws.template InputRef<CPUBackend>(0).shape();
+    ndim_ = in0_shape[0][0];
+    nsamples_ = in0_shape.size();
+
+    DALI_ENFORCE(in0_shape.sample_dim() == 2 &&
+                 in0_shape.size() > 0 &&
+                 in0_shape[0][1] == (in0_shape[0][0] + 1),
+      make_string(
+        "The input, if provided, is expected to be a 2D tensor with dimensions "
+        "(ndim, ndim+1) representing an affine transform. Got: ", in0_shape));
+
+    for (int i = 0; i < ws.NumInput(); i++) {
+      const auto &shape = ws.template InputRef<CPUBackend>(i).shape();
+      DALI_ENFORCE(shape == in0_shape,
+        make_string("All input transforms are expected to have the same shape. Got: ",
+                    in0_shape, " and ", shape, " for the ", i, "-th input."));
+    }
+
+    output_descs.resize(1);  // only one output
+    output_descs[0].type = TypeTable::GetTypeInfo(dtype_);
+    output_descs[0].shape = uniform_list_shape(nsamples_, {ndim_, ndim_+1});
+    return true;
+  }
+
+  template <typename T>
+  void RunImplTyped(workspace_t<CPUBackend> &ws, dims<>) {
+    DALI_FAIL(make_string("Unsupported number of dimensions ", ndim_));
+  }
+
+  template <typename T, int ndim, int... ndims>
+  void RunImplTyped(workspace_t<CPUBackend> &ws, dims<ndim, ndims...>) {
+    if (ndim_ != ndim) {
+      RunImplTyped<T>(ws, dims<ndims...>());
+      return;
+    }
+
+    constexpr int mat_dim = ndim + 1;
+    auto &out = ws.template OutputRef<CPUBackend>(0);
+    out.SetLayout({});  // no layout
+    auto out_view = view<T>(out);
+
+    for (int sample_idx = 0; sample_idx < nsamples_; sample_idx++) {
+      auto mat = affine_mat_t<T, mat_dim>::identity();
+      for (int input_idx = 0; input_idx < ws.NumInput(); input_idx++) {
+        auto &in = ws.template InputRef<CPUBackend>(input_idx);
+        auto in_view = view<T>(in);
+        auto next_mat = affine_mat_t<T, mat_dim>::identity();
+        for (int i = 0, k = 0; i < ndim; i++)
+          for (int j = 0; j < ndim + 1; j++, k++)
+            next_mat(i, j) = in_view[sample_idx].data[k];
+        mat = reverse_order_ ? mat * next_mat : next_mat * mat;  // mat mul
+      }
+
+      for (int i = 0, k = 0; i < ndim; i++) {
+        for (int j = 0; j < ndim + 1; j++, k++) {
+          out_view[sample_idx].data[k] = mat(i, j);
+        }
+      }
+    }
+  }
+
+  void RunImpl(workspace_t<CPUBackend> &ws) override {
+    TYPE_SWITCH(dtype_, type2id, T, TRANSFORM_INPUT_TYPES, (
+      RunImplTyped<T>(ws, SupportedDims());
+    ), DALI_FAIL(make_string("Unsupported data type: ", dtype_)));  // NOLINT
+  }
+
+ private:
+  using SupportedDims = dims<1, 2, 3, 4, 5, 6>;
+  DALIDataType dtype_ = DALI_FLOAT;
+  int ndim_ = -1;  // will be inferred from the arguments or the input
+  int nsamples_ = -1;
+  bool reverse_order_ = false;
+};
+
+DALI_REGISTER_OPERATOR(transforms__Combine, CombineTransformsCPU, CPU);
+
+}  // namespace dali

dali/operators/geometry/affine_transforms/transform_crop.cc

@@ -17,7 +17,7 @@
 
 namespace dali {
 
-DALI_SCHEMA(TransformCrop)
+DALI_SCHEMA(transforms__Crop)
   .DocStr(R"code(Produces an affine transform matrix that maps a reference coordinate space to another one.
 
 This transform can be used to adjust coordinates after a crop operation so that a ``from_start`` point will
@@ -164,6 +164,6 @@ class TransformCropCPU
   bool absolute_ = false;
 };
 
-DALI_REGISTER_OPERATOR(TransformCrop, TransformCropCPU, CPU);
+DALI_REGISTER_OPERATOR(transforms__Crop, TransformCropCPU, CPU);
 
 }  // namespace dali

dali/operators/geometry/affine_transforms/transform_rotation.cc

@@ -19,7 +19,7 @@
 
 namespace dali {
 
-DALI_SCHEMA(TransformRotation)
+DALI_SCHEMA(transforms__Rotation)
   .DocStr(R"code(Produces a rotation affine transform matrix.
 
 If another transform matrix is passed as an input, the operator applies rotation to the matrix provided.
@@ -137,6 +137,6 @@ class TransformRotationCPU
   Argument<std::vector<float>> center_;
 };
 
-DALI_REGISTER_OPERATOR(TransformRotation, TransformRotationCPU, CPU);
+DALI_REGISTER_OPERATOR(transforms__Rotation, TransformRotationCPU, CPU);
 
 }  // namespace dali

dali/operators/geometry/affine_transforms/transform_scale.cc

@@ -17,7 +17,7 @@
 
 namespace dali {
 
-DALI_SCHEMA(TransformScale)
+DALI_SCHEMA(transforms__Scale)
   .DocStr(R"code(Produces a scale affine transform matrix.
 
 If another transform matrix is passed as an input, the operator applies scaling to the matrix provided.
@@ -101,6 +101,6 @@ class TransformScaleCPU
   Argument<std::vector<float>> center_;
 };
 
-DALI_REGISTER_OPERATOR(TransformScale, TransformScaleCPU, CPU);
+DALI_REGISTER_OPERATOR(transforms__Scale, TransformScaleCPU, CPU);
 
 }  // namespace dali

dali/operators/geometry/affine_transforms/transform_shear.cc

@@ -19,7 +19,7 @@
 
 namespace dali {
 
-DALI_SCHEMA(TransformShear)
+DALI_SCHEMA(transforms__Shear)
   .DocStr(R"code(Produces a shear affine transform matrix.
 
 If another transform matrix is passed as an input, the operator applies the shear mapping to the matrix provided.
@@ -175,6 +175,6 @@ class TransformShearCPU
   Argument<std::vector<float>> center_;
 };
 
-DALI_REGISTER_OPERATOR(TransformShear, TransformShearCPU, CPU);
+DALI_REGISTER_OPERATOR(transforms__Shear, TransformShearCPU, CPU);
 
 }  // namespace dali

dali/operators/geometry/affine_transforms/transform_translation.cc

@@ -17,7 +17,7 @@
 
 namespace dali {
 
-DALI_SCHEMA(TransformTranslation)
+DALI_SCHEMA(transforms__Translation)
   .DocStr(R"code(Produces a translation affine transform matrix.
 
 If another transform matrix is passed as an input, the operator applies translation to the matrix provided.
@@ -75,6 +75,6 @@ class TransformTranslationCPU
   Argument<std::vector<float>> offset_;
 };
 
-DALI_REGISTER_OPERATOR(TransformTranslation, TransformTranslationCPU, CPU);
+DALI_REGISTER_OPERATOR(transforms__Translation, TransformTranslationCPU, CPU);
 
 }  // namespace dali

dali/pipeline/operator/operator_factory.h

@@ -60,7 +60,7 @@ class OperatorRegistry {
     for (const auto &pair : registry_) {
       auto& schema = SchemaRegistry::GetSchema(pair.first);
       if (internal_ops || !schema.IsInternal())
-        names.push_back(pair.first);
+        names.push_back(schema.name().length() ? schema.name() : pair.first);
     }
     return names;
   }

dali/python/nvidia/dali/fn.py

@@ -15,6 +15,7 @@
 #pylint: disable=no-member
 import sys
 from nvidia.dali.data_node import DataNode as _DataNode
+from nvidia.dali import internal as _internal
 
 _special_case_mapping = {
     "b_box" : "bbox",
@@ -85,13 +86,19 @@ def to_scalar(scalar):
         return op_class(**scalar_args)(*inputs, **tensor_args)
 
     op_wrapper.__name__ = wrapper_name
-    op_wrapper.__doc__ = "see :class:`nvidia.dali.ops.{0}`".format(op_class.__name__)
+    op_wrapper.__doc__ = "see :class:`{0}.{1}`".format(op_class.__module__, op_class.__name__)
     return op_wrapper
 
-def _wrap_op(op_class):
+def _wrap_op(op_class, submodule):
     wrapper_name = _to_snake_case(op_class.__name__)
-    if not hasattr(sys.modules[__name__], wrapper_name):
-        setattr(sys.modules[__name__], wrapper_name, _wrap_op_fn(op_class, wrapper_name))
+    fn_module = sys.modules[__name__]
+    module = _internal.get_submodule(fn_module, submodule)
+    if not hasattr(module, wrapper_name):
+        wrap_func = _wrap_op_fn(op_class, wrapper_name)
+        setattr(module, wrapper_name, wrap_func)
+        if submodule:
+            wrap_func.__module__ = module.__name__
+
 
 from nvidia.dali.external_source import external_source
 external_source.__module__ = __name__

dali/python/nvidia/dali/internal.py

@@ -0,0 +1,37 @@
+import sys
+import types
+
+def get_submodule(root, path):
+    """Gets or creates sumbodule(s) of `root`.
+If the module path contains multiple parts, multiple modules are traversed or created
+
+Parameters
+----------
+    `root`
+        module object or name of the root module
+    `path`
+        period-separated path of the submodule or a list/tuple of submodule names"""
+
+    if isinstance(root, str):
+        root = sys.modules[root]
+
+    if not path:
+        return root
+
+    if isinstance(path, str):
+        if str == '':
+            return root
+        path = path.split('.')
+
+    module_name = root.__name__
+    for part in path:
+        m = getattr(root, part, None)
+        module_name += '.' + part
+        if m is None:
+            m = sys.modules[module_name] = types.ModuleType(module_name)
+            setattr(root, part, m)
+        elif not isinstance(m, types.ModuleType):
+            raise RuntimeError("The module {} already contains an attribute \"{}\", which is not a module, but {}".format(
+                root, part, m))
+        root = m
+    return root