Project import generated by Copybara.

GitOrigin-RevId: ff83882955f1a1e2a043ff4e71278be9d7217bbe

mediapipe/gpu/gl_surface_sink_calculator.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "absl/synchronization/mutex.h"
+#include "mediapipe/framework/api2/node.h"
 #include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/port/ret_check.h"
 #include "mediapipe/framework/port/status.h"
@@ -21,9 +22,11 @@
 #include "mediapipe/gpu/gl_calculator_helper.h"
 #include "mediapipe/gpu/gl_quad_renderer.h"
 #include "mediapipe/gpu/gl_surface_sink_calculator.pb.h"
+#include "mediapipe/gpu/gpu_buffer.h"
 #include "mediapipe/gpu/shader_util.h"
 
 namespace mediapipe {
+namespace api2 {
 
 enum { kAttribVertex, kAttribTexturePosition, kNumberOfAttributes };
 
@@ -37,45 +40,52 @@ enum { kAttribVertex, kAttribTexturePosition, kNumberOfAttributes };
 //   GPU_SHARED: shared GPU resources.
 //
 // See GlSurfaceSinkCalculatorOptions for options.
-class GlSurfaceSinkCalculator : public CalculatorBase {
+class GlSurfaceSinkCalculator : public Node {
  public:
-  GlSurfaceSinkCalculator() : initialized_(false) {}
-  ~GlSurfaceSinkCalculator() override;
+  static constexpr Input<
+      OneOf<mediapipe::Image, mediapipe::GpuBuffer>>::Optional kInVideo{
+      "VIDEO"};
+  static constexpr Input<
+      OneOf<mediapipe::Image, mediapipe::GpuBuffer>>::Optional kIn{""};
+  static constexpr SideInput<std::unique_ptr<mediapipe::EglSurfaceHolder>>
+      kSurface{"SURFACE"};
 
-  static absl::Status GetContract(CalculatorContract* cc);
+  MEDIAPIPE_NODE_INTERFACE(GlSurfaceSinkCalculator, kInVideo, kIn, kSurface);
 
-  absl::Status Open(CalculatorContext* cc) override;
-  absl::Status Process(CalculatorContext* cc) override;
+  ~GlSurfaceSinkCalculator();
+
+  static absl::Status UpdateContract(CalculatorContract* cc);
+
+  absl::Status Open(CalculatorContext* cc) final;
+  absl::Status Process(CalculatorContext* cc) final;
 
  private:
-  GlCalculatorHelper helper_;
-  EglSurfaceHolder* surface_holder_;
-  bool initialized_;
-  std::unique_ptr<QuadRenderer> renderer_;
-  FrameScaleMode scale_mode_ = FrameScaleMode::kFillAndCrop;
+  mediapipe::GlCalculatorHelper helper_;
+  mediapipe::EglSurfaceHolder* surface_holder_;
+  bool initialized_ = false;
+  std::unique_ptr<mediapipe::QuadRenderer> renderer_;
+  mediapipe::FrameScaleMode scale_mode_ =
+      mediapipe::FrameScaleMode::kFillAndCrop;
 };
-REGISTER_CALCULATOR(GlSurfaceSinkCalculator);
+MEDIAPIPE_REGISTER_NODE(GlSurfaceSinkCalculator);
 
 // static
-absl::Status GlSurfaceSinkCalculator::GetContract(CalculatorContract* cc) {
-  TagOrIndex(&(cc->Inputs()), "VIDEO", 0).Set<GpuBuffer>();
-  cc->InputSidePackets()
-      .Tag("SURFACE")
-      .Set<std::unique_ptr<EglSurfaceHolder>>();
+absl::Status GlSurfaceSinkCalculator::UpdateContract(CalculatorContract* cc) {
+  RET_CHECK(kInVideo(cc).IsConnected() ^ kIn(cc).IsConnected())
+      << "Only one of VIDEO or index 0 input is expected.";
+
   // Currently we pass GL context information and other stuff as external
   // inputs, which are handled by the helper.
-  return GlCalculatorHelper::UpdateContract(cc);
+  return mediapipe::GlCalculatorHelper::UpdateContract(cc);
 }
 
 absl::Status GlSurfaceSinkCalculator::Open(CalculatorContext* cc) {
-  surface_holder_ = cc->InputSidePackets()
-                        .Tag("SURFACE")
-                        .Get<std::unique_ptr<EglSurfaceHolder>>()
-                        .get();
+  surface_holder_ = kSurface(cc).Get().get();
 
   scale_mode_ = FrameScaleModeFromProto(
-      cc->Options<GlSurfaceSinkCalculatorOptions>().frame_scale_mode(),
-      FrameScaleMode::kFillAndCrop);
+      cc->Options<mediapipe::GlSurfaceSinkCalculatorOptions>()
+          .frame_scale_mode(),
+      mediapipe::FrameScaleMode::kFillAndCrop);
 
   // Let the helper access the GL context information.
   return helper_.Open(cc);
@@ -90,9 +100,20 @@ absl::Status GlSurfaceSinkCalculator::Process(CalculatorContext* cc) {
       return absl::OkStatus();
     }
 
-    const auto& input = TagOrIndex(cc->Inputs(), "VIDEO", 0).Get<GpuBuffer>();
+    mediapipe::Packet packet;
+    if (kInVideo(cc).IsConnected())
+      packet = kInVideo(cc).packet();
+    else
+      packet = kIn(cc).packet();
+
+    mediapipe::GpuBuffer input;
+    if (packet.ValidateAsType<mediapipe::GpuBuffer>().ok())
+      input = packet.Get<mediapipe::GpuBuffer>();
+    if (packet.ValidateAsType<mediapipe::Image>().ok())
+      input = packet.Get<mediapipe::Image>().GetGpuBuffer();
+
     if (!initialized_) {
-      renderer_ = absl::make_unique<QuadRenderer>();
+      renderer_ = absl::make_unique<mediapipe::QuadRenderer>();
       MP_RETURN_IF_ERROR(renderer_->GlSetup());
       initialized_ = true;
     }
@@ -125,7 +146,7 @@ absl::Status GlSurfaceSinkCalculator::Process(CalculatorContext* cc) {
 
     MP_RETURN_IF_ERROR(
         renderer_->GlRender(src.width(), src.height(), dst_width, dst_height,
-                            scale_mode_, FrameRotation::kNone,
+                            scale_mode_, mediapipe::FrameRotation::kNone,
                             /*flip_horizontal=*/false, /*flip_vertical=*/false,
                             /*flip_texture=*/surface_holder_->flip_y));
 
@@ -145,7 +166,7 @@ absl::Status GlSurfaceSinkCalculator::Process(CalculatorContext* cc) {
 GlSurfaceSinkCalculator::~GlSurfaceSinkCalculator() {
   if (renderer_) {
     // TODO: use move capture when we have C++14 or better.
-    QuadRenderer* renderer = renderer_.release();
+    mediapipe::QuadRenderer* renderer = renderer_.release();
     helper_.RunInGlContext([renderer] {
       renderer->GlTeardown();
       delete renderer;
@@ -153,4 +174,5 @@ GlSurfaceSinkCalculator::~GlSurfaceSinkCalculator() {
   }
 }
 
+}  // namespace api2
 }  // namespace mediapipe