Project import generated by Copybara.

GitOrigin-RevId: 1610e588e497817fae2d9a458093ab6a370e2972
2021-08-18 17:45:46 -07:00
parent b899d17f18
commit 710fb3de58
158 changed files with 10104 additions and 1568 deletions
@@ -0,0 +1,211 @@
+// Copyright 2021 The MediaPipe Authors.
+//
+// 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 "mediapipe/calculators/image/warp_affine_calculator.h"
+
+#include <array>
+#include <cstdint>
+#include <memory>
+
+#include "mediapipe/calculators/image/affine_transformation.h"
+#if !MEDIAPIPE_DISABLE_GPU
+#include "mediapipe/calculators/image/affine_transformation_runner_gl.h"
+#endif  // !MEDIAPIPE_DISABLE_GPU
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "mediapipe/calculators/image/affine_transformation_runner_opencv.h"
+#include "mediapipe/calculators/image/warp_affine_calculator.pb.h"
+#include "mediapipe/framework/api2/node.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/formats/image.h"
+#include "mediapipe/framework/formats/image_frame.h"
+#include "mediapipe/framework/port/ret_check.h"
+#if !MEDIAPIPE_DISABLE_GPU
+#include "mediapipe/gpu/gl_calculator_helper.h"
+#include "mediapipe/gpu/gpu_buffer.h"
+#endif  // !MEDIAPIPE_DISABLE_GPU
+
+namespace mediapipe {
+
+namespace {
+
+AffineTransformation::BorderMode GetBorderMode(
+    mediapipe::WarpAffineCalculatorOptions::BorderMode border_mode) {
+  switch (border_mode) {
+    case mediapipe::WarpAffineCalculatorOptions::BORDER_ZERO:
+      return AffineTransformation::BorderMode::kZero;
+    case mediapipe::WarpAffineCalculatorOptions::BORDER_UNSPECIFIED:
+    case mediapipe::WarpAffineCalculatorOptions::BORDER_REPLICATE:
+      return AffineTransformation::BorderMode::kReplicate;
+  }
+}
+
+template <typename ImageT>
+class WarpAffineRunnerHolder {};
+
+template <>
+class WarpAffineRunnerHolder<ImageFrame> {
+ public:
+  using RunnerType = AffineTransformation::Runner<ImageFrame, ImageFrame>;
+  absl::Status Open(CalculatorContext* cc) { return absl::OkStatus(); }
+  absl::StatusOr<RunnerType*> GetRunner() {
+    if (!runner_) {
+      ASSIGN_OR_RETURN(runner_, CreateAffineTransformationOpenCvRunner());
+    }
+    return runner_.get();
+  }
+
+ private:
+  std::unique_ptr<RunnerType> runner_;
+};
+
+#if !MEDIAPIPE_DISABLE_GPU
+template <>
+class WarpAffineRunnerHolder<mediapipe::GpuBuffer> {
+ public:
+  using RunnerType =
+      AffineTransformation::Runner<mediapipe::GpuBuffer,
+                                   std::unique_ptr<mediapipe::GpuBuffer>>;
+  absl::Status Open(CalculatorContext* cc) {
+    gpu_origin_ =
+        cc->Options<mediapipe::WarpAffineCalculatorOptions>().gpu_origin();
+    gl_helper_ = std::make_shared<mediapipe::GlCalculatorHelper>();
+    return gl_helper_->Open(cc);
+  }
+  absl::StatusOr<RunnerType*> GetRunner() {
+    if (!runner_) {
+      ASSIGN_OR_RETURN(
+          runner_, CreateAffineTransformationGlRunner(gl_helper_, gpu_origin_));
+    }
+    return runner_.get();
+  }
+
+ private:
+  mediapipe::GpuOrigin::Mode gpu_origin_;
+  std::shared_ptr<mediapipe::GlCalculatorHelper> gl_helper_;
+  std::unique_ptr<RunnerType> runner_;
+};
+#endif  // !MEDIAPIPE_DISABLE_GPU
+
+template <>
+class WarpAffineRunnerHolder<mediapipe::Image> {
+ public:
+  absl::Status Open(CalculatorContext* cc) { return runner_.Open(cc); }
+  absl::StatusOr<
+      AffineTransformation::Runner<mediapipe::Image, mediapipe::Image>*>
+  GetRunner() {
+    return &runner_;
+  }
+
+ private:
+  class Runner : public AffineTransformation::Runner<mediapipe::Image,
+                                                     mediapipe::Image> {
+   public:
+    absl::Status Open(CalculatorContext* cc) {
+      MP_RETURN_IF_ERROR(cpu_holder_.Open(cc));
+#if !MEDIAPIPE_DISABLE_GPU
+      MP_RETURN_IF_ERROR(gpu_holder_.Open(cc));
+#endif  // !MEDIAPIPE_DISABLE_GPU
+      return absl::OkStatus();
+    }
+    absl::StatusOr<mediapipe::Image> Run(
+        const mediapipe::Image& input, const std::array<float, 16>& matrix,
+        const AffineTransformation::Size& size,
+        AffineTransformation::BorderMode border_mode) override {
+      if (input.UsesGpu()) {
+#if !MEDIAPIPE_DISABLE_GPU
+        ASSIGN_OR_RETURN(auto* runner, gpu_holder_.GetRunner());
+        ASSIGN_OR_RETURN(auto result, runner->Run(input.GetGpuBuffer(), matrix,
+                                                  size, border_mode));
+        return mediapipe::Image(*result);
+#else
+        return absl::UnavailableError("GPU support is disabled");
+#endif  // !MEDIAPIPE_DISABLE_GPU
+      }
+      ASSIGN_OR_RETURN(auto* runner, cpu_holder_.GetRunner());
+      const auto& frame_ptr = input.GetImageFrameSharedPtr();
+      // Wrap image into image frame.
+      const ImageFrame image_frame(frame_ptr->Format(), frame_ptr->Width(),
+                                   frame_ptr->Height(), frame_ptr->WidthStep(),
+                                   const_cast<uint8_t*>(frame_ptr->PixelData()),
+                                   [](uint8* data) {});
+      ASSIGN_OR_RETURN(auto result,
+                       runner->Run(image_frame, matrix, size, border_mode));
+      return mediapipe::Image(std::make_shared<ImageFrame>(std::move(result)));
+    }
+
+   private:
+    WarpAffineRunnerHolder<ImageFrame> cpu_holder_;
+#if !MEDIAPIPE_DISABLE_GPU
+    WarpAffineRunnerHolder<mediapipe::GpuBuffer> gpu_holder_;
+#endif  // !MEDIAPIPE_DISABLE_GPU
+  };
+
+  Runner runner_;
+};
+
+template <typename InterfaceT>
+class WarpAffineCalculatorImpl : public mediapipe::api2::NodeImpl<InterfaceT> {
+ public:
+#if !MEDIAPIPE_DISABLE_GPU
+  static absl::Status UpdateContract(CalculatorContract* cc) {
+    if constexpr (std::is_same_v<InterfaceT, WarpAffineCalculatorGpu> ||
+                  std::is_same_v<InterfaceT, WarpAffineCalculator>) {
+      MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
+    }
+    return absl::OkStatus();
+  }
+#endif  // !MEDIAPIPE_DISABLE_GPU
+
+  absl::Status Open(CalculatorContext* cc) override { return holder_.Open(cc); }
+
+  absl::Status Process(CalculatorContext* cc) override {
+    if (InterfaceT::kInImage(cc).IsEmpty() ||
+        InterfaceT::kMatrix(cc).IsEmpty() ||
+        InterfaceT::kOutputSize(cc).IsEmpty()) {
+      return absl::OkStatus();
+    }
+    const std::array<float, 16>& transform = *InterfaceT::kMatrix(cc);
+    auto [out_width, out_height] = *InterfaceT::kOutputSize(cc);
+    AffineTransformation::Size output_size;
+    output_size.width = out_width;
+    output_size.height = out_height;
+    ASSIGN_OR_RETURN(auto* runner, holder_.GetRunner());
+    ASSIGN_OR_RETURN(
+        auto result,
+        runner->Run(
+            *InterfaceT::kInImage(cc), transform, output_size,
+            GetBorderMode(cc->Options<mediapipe::WarpAffineCalculatorOptions>()
+                              .border_mode())));
+    InterfaceT::kOutImage(cc).Send(std::move(result));
+
+    return absl::OkStatus();
+  }
+
+ private:
+  WarpAffineRunnerHolder<typename decltype(InterfaceT::kInImage)::PayloadT>
+      holder_;
+};
+
+}  // namespace
+
+MEDIAPIPE_NODE_IMPLEMENTATION(
+    WarpAffineCalculatorImpl<WarpAffineCalculatorCpu>);
+#if !MEDIAPIPE_DISABLE_GPU
+MEDIAPIPE_NODE_IMPLEMENTATION(
+    WarpAffineCalculatorImpl<WarpAffineCalculatorGpu>);
+#endif  // !MEDIAPIPE_DISABLE_GPU
+MEDIAPIPE_NODE_IMPLEMENTATION(WarpAffineCalculatorImpl<WarpAffineCalculator>);
+
+}  // namespace mediapipe