Project import generated by Copybara.

GitOrigin-RevId: 016275ca4057540b2370ed4531dbc81eb92caae2

docs/solutions/pose.md

@@ -399,4 +399,4 @@ on how to build MediaPipe examples.
 *   [Python Colab](https://mediapipe.page.link/pose_py_colab)
 
 [`mAP`]: https://cocodataset.org/#keypoints-eval
-[`PCK@0.2`]: https\://github.com/cbsudux/Human-Pose-Estimation-101
+[`PCK@0.2`]: https://github.com/cbsudux/Human-Pose-Estimation-101

mediapipe/calculators/image/BUILD

@@ -405,6 +405,36 @@ cc_library(
     alwayslink = 1,
 )
 
+mediapipe_proto_library(
+    name = "image_clone_calculator_proto",
+    srcs = ["image_clone_calculator.proto"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//mediapipe/framework:calculator_options_proto",
+        "//mediapipe/framework:calculator_proto",
+    ],
+)
+
+cc_library(
+    name = "image_clone_calculator",
+    srcs = ["image_clone_calculator.cc"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":image_clone_calculator_cc_proto",
+        "//mediapipe/framework/api2:node",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework/formats:image",
+        "//mediapipe/framework/port:ret_check",
+        "//mediapipe/framework/port:status",
+    ] + select({
+        "//mediapipe/gpu:disable_gpu": [],
+        "//conditions:default": [
+            "//mediapipe/gpu:gl_calculator_helper",
+        ],
+    }),
+    alwayslink = 1,
+)
+
 cc_library(
     name = "image_properties_calculator",
     srcs = ["image_properties_calculator.cc"],

mediapipe/calculators/image/image_clone_calculator.cc

@@ -0,0 +1,125 @@
+// 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/image_clone_calculator.pb.h"
+#include "mediapipe/framework/api2/node.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/formats/image.h"
+#include "mediapipe/framework/port/status.h"
+
+#if !MEDIAPIPE_DISABLE_GPU
+#include "mediapipe/gpu/gl_calculator_helper.h"
+#endif  // !MEDIAPIPE_DISABLE_GPU
+
+namespace mediapipe {
+namespace api2 {
+
+#if MEDIAPIPE_DISABLE_GPU
+// Just a placeholder to not have to depend on mediapipe::GpuBuffer.
+using GpuBuffer = AnyType;
+#else
+using GpuBuffer = mediapipe::GpuBuffer;
+#endif  // MEDIAPIPE_DISABLE_GPU
+
+// Clones an input image and makes sure in the output clone the pixel data are
+// stored on the target storage (CPU vs GPU) specified in the calculator option.
+//
+// The clone shares ownership of the input pixel data on the existing storage.
+// If the target storage is diffrent from the existing one, then the data is
+// further copied there.
+//
+// Example usage:
+// node {
+//   calculator: "ImageCloneCalculator"
+//   input_stream: "input"
+//   output_stream: "output"
+//   options: {
+//     [mediapipe.ImageCloneCalculatorOptions.ext] {
+//       output_on_gpu: true
+//     }
+//   }
+// }
+class ImageCloneCalculator : public Node {
+ public:
+  static constexpr Input<Image> kIn{""};
+  static constexpr Output<Image> kOut{""};
+
+  MEDIAPIPE_NODE_CONTRACT(kIn, kOut);
+
+  static absl::Status UpdateContract(CalculatorContract* cc) {
+#if MEDIAPIPE_DISABLE_GPU
+    if (cc->Options<mediapipe::ImageCloneCalculatorOptions>().output_on_gpu()) {
+      return absl::UnimplementedError(
+          "GPU processing is disabled in build flags");
+    }
+#else
+    MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
+#endif  // MEDIAPIPE_DISABLE_GPU
+    return absl::OkStatus();
+  }
+
+  absl::Status Open(CalculatorContext* cc) override {
+    const auto& options = cc->Options<mediapipe::ImageCloneCalculatorOptions>();
+    output_on_gpu_ = options.output_on_gpu();
+#if !MEDIAPIPE_DISABLE_GPU
+    MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
+#endif  // !MEDIAPIPE_DISABLE_GPU
+    return absl::OkStatus();
+  }
+
+  absl::Status Process(CalculatorContext* cc) override {
+    std::unique_ptr<Image> output;
+    const auto& input = *kIn(cc);
+    if (input.UsesGpu()) {
+#if !MEDIAPIPE_DISABLE_GPU
+      // Create an output Image that co-owns the underlying texture buffer as
+      // the input Image.
+      output = std::make_unique<Image>(input.GetGpuBuffer());
+#endif  // !MEDIAPIPE_DISABLE_GPU
+    } else {
+      // Make a copy of the input packet to co-own the input Image.
+      mediapipe::Packet* packet_copy_ptr =
+          new mediapipe::Packet(kIn(cc).packet());
+      // Create an output Image that (co-)owns a new ImageFrame that points to
+      // the same pixel data as the input Image and also owns the packet
+      // copy. As a result, the output Image indirectly co-owns the input
+      // Image. This ensures a correct life span of the shared pixel data.
+      output = std::make_unique<Image>(std::make_unique<mediapipe::ImageFrame>(
+          input.image_format(), input.width(), input.height(), input.step(),
+          const_cast<uint8*>(input.GetImageFrameSharedPtr()->PixelData()),
+          [packet_copy_ptr](uint8*) { delete packet_copy_ptr; }));
+    }
+
+    if (output_on_gpu_) {
+#if !MEDIAPIPE_DISABLE_GPU
+      gpu_helper_.RunInGlContext([&output]() { output->ConvertToGpu(); });
+#endif  // !MEDIAPIPE_DISABLE_GPU
+    } else {
+      output->ConvertToCpu();
+    }
+    kOut(cc).Send(std::move(output));
+
+    return absl::OkStatus();
+  }
+
+ private:
+  bool output_on_gpu_;
+#if !MEDIAPIPE_DISABLE_GPU
+  mediapipe::GlCalculatorHelper gpu_helper_;
+#endif  // !MEDIAPIPE_DISABLE_GPU
+};
+MEDIAPIPE_REGISTER_NODE(ImageCloneCalculator);
+
+}  // namespace api2
+}  // namespace mediapipe

mediapipe/calculators/image/image_clone_calculator.proto

@@ -0,0 +1,28 @@
+// 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.
+
+syntax = "proto2";
+
+package mediapipe;
+
+import "mediapipe/framework/calculator.proto";
+
+message ImageCloneCalculatorOptions {
+  extend CalculatorOptions {
+    optional ImageCloneCalculatorOptions ext = 372781894;
+  }
+
+  // Whether the output clone should have pixel data already available on GPU.
+  optional bool output_on_gpu = 1 [default = false];
+}

mediapipe/gpu/BUILD

@@ -753,6 +753,7 @@ cc_library(
         ":gl_simple_shaders",
         ":shader_util",
         "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework/formats:image",
         "//mediapipe/framework/port:ret_check",
         "//mediapipe/framework/port:status",
         "//mediapipe/framework/tool:options_util",

mediapipe/gpu/gl_scaler_calculator.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/formats/image.h"
 #include "mediapipe/framework/port/ret_check.h"
 #include "mediapipe/framework/port/status.h"
 #include "mediapipe/framework/tool/options_util.h"
@@ -33,6 +34,8 @@ typedef int DimensionsPacketType[2];
 
 namespace mediapipe {
 
+using Image = mediapipe::Image;
+
 // Scales, rotates, horizontal or vertical flips the image.
 // See GlSimpleCalculatorBase for inputs, outputs and input side packets.
 // Additional input streams:
@@ -99,8 +102,17 @@ REGISTER_CALCULATOR(GlScalerCalculator);
 
 // static
 absl::Status GlScalerCalculator::GetContract(CalculatorContract* cc) {
-  TagOrIndex(&cc->Inputs(), "VIDEO", 0).Set<GpuBuffer>();
-  TagOrIndex(&cc->Outputs(), "VIDEO", 0).Set<GpuBuffer>();
+  if (cc->Inputs().HasTag("IMAGE")) {
+    cc->Inputs().Tag("IMAGE").Set<Image>();
+  } else {
+    TagOrIndex(&cc->Inputs(), "VIDEO", 0).Set<GpuBuffer>();
+  }
+  if (cc->Outputs().HasTag("IMAGE")) {
+    cc->Outputs().Tag("IMAGE").Set<Image>();
+  } else {
+    TagOrIndex(&cc->Outputs(), "VIDEO", 0).Set<GpuBuffer>();
+  }
+
   if (cc->Inputs().HasTag("ROTATION")) {
     cc->Inputs().Tag("ROTATION").Set<int>();
   }
@@ -198,7 +210,10 @@ absl::Status GlScalerCalculator::Process(CalculatorContext* cc) {
   }
 
   return helper_.RunInGlContext([this, cc]() -> absl::Status {
-    const auto& input = TagOrIndex(cc->Inputs(), "VIDEO", 0).Get<GpuBuffer>();
+    const auto& input =
+        cc->Inputs().HasTag("IMAGE")
+            ? cc->Inputs().Tag("IMAGE").Get<Image>().GetGpuBuffer()
+            : TagOrIndex(cc->Inputs(), "VIDEO", 0).Get<GpuBuffer>();
     QuadRenderer* renderer = nullptr;
     GlTexture src1;
     GlTexture src2;
@@ -289,10 +304,14 @@ absl::Status GlScalerCalculator::Process(CalculatorContext* cc) {
 
     glFlush();
 
-    auto output = dst.GetFrame<GpuBuffer>();
-
-    TagOrIndex(&cc->Outputs(), "VIDEO", 0)
-        .Add(output.release(), cc->InputTimestamp());
+    if (cc->Outputs().HasTag("IMAGE")) {
+      auto output = dst.GetFrame<Image>();
+      cc->Outputs().Tag("IMAGE").Add(output.release(), cc->InputTimestamp());
+    } else {
+      auto output = dst.GetFrame<GpuBuffer>();
+      TagOrIndex(&cc->Outputs(), "VIDEO", 0)
+          .Add(output.release(), cc->InputTimestamp());
+    }
 
     return absl::OkStatus();
   });