mediapipe_extension.diff

diff --git a/mediapipe/framework/calculator_graph.cc b/mediapipe/framework/calculator_graph.cc
index 2c34112..3b8ab93 100644
--- a/mediapipe/framework/calculator_graph.cc
+++ b/mediapipe/framework/calculator_graph.cc
@@ -432,7 +432,7 @@ absl::Status CalculatorGraph::ObserveOutputStream(
 }
 
 absl::StatusOr<OutputStreamPoller> CalculatorGraph::AddOutputStreamPoller(
-    const std::string& stream_name) {
+    const std::string& stream_name, bool observe_timestamp_bounds) {
   RET_CHECK(initialized_).SetNoLogging()
       << "CalculatorGraph is not initialized.";
   int output_stream_index = validated_graph_->OutputStreamIndex(stream_name);
@@ -446,7 +446,7 @@ absl::StatusOr<OutputStreamPoller> CalculatorGraph::AddOutputStreamPoller(
       stream_name, &any_packet_type_,
       std::bind(&CalculatorGraph::UpdateThrottledNodes, this,
                 std::placeholders::_1, std::placeholders::_2),
-      &output_stream_managers_[output_stream_index]));
+      &output_stream_managers_[output_stream_index], observe_timestamp_bounds));
   OutputStreamPoller poller(internal_poller);
   graph_output_streams_.push_back(std::move(internal_poller));
   return std::move(poller);
diff --git a/mediapipe/framework/calculator_graph.h b/mediapipe/framework/calculator_graph.h
index 77a1ee5..81beec2 100644
--- a/mediapipe/framework/calculator_graph.h
+++ b/mediapipe/framework/calculator_graph.h
@@ -164,7 +164,8 @@ class CalculatorGraph {
   // polling API for accessing a stream's output. Should only be called before
   // Run() or StartRun(). For asynchronous output, use ObserveOutputStream. See
   // also the helpers in tool/sink.h.
-  StatusOrPoller AddOutputStreamPoller(const std::string& stream_name);
+  StatusOrPoller AddOutputStreamPoller(const std::string& stream_name,
+                                       bool observe_timestamp_bounds = false);
 
   // Gets output side packet by name after the graph is done. However, base
   // packets (generated by PacketGenerators) can be retrieved before
diff --git a/mediapipe/framework/graph_output_stream.cc b/mediapipe/framework/graph_output_stream.cc
index 6639bb8..8235d3c 100644
--- a/mediapipe/framework/graph_output_stream.cc
+++ b/mediapipe/framework/graph_output_stream.cc
@@ -125,9 +125,10 @@ absl::Status OutputStreamObserver::Notify() {
 absl::Status OutputStreamPollerImpl::Initialize(
     const std::string& stream_name, const PacketType* packet_type,
     std::function<void(InputStreamManager*, bool*)> queue_size_callback,
-    OutputStreamManager* output_stream_manager) {
+    OutputStreamManager* output_stream_manager, bool observe_timestamp_bounds) {
   MP_RETURN_IF_ERROR(GraphOutputStream::Initialize(stream_name, packet_type,
-                                                   output_stream_manager));
+                                                   output_stream_manager,
+                                                   observe_timestamp_bounds));
   input_stream_handler_->SetQueueSizeCallbacks(queue_size_callback,
                                                queue_size_callback);
   return absl::OkStatus();
@@ -176,12 +177,16 @@ void OutputStreamPollerImpl::NotifyError() {
 bool OutputStreamPollerImpl::Next(Packet* packet) {
   CHECK(packet);
   bool empty_queue = true;
+  bool observed_timestamp_bound_change = false;
   Timestamp min_timestamp = Timestamp::Unset();
   mutex_.Lock();
   while (true) {
     min_timestamp = input_stream_->MinTimestampOrBound(&empty_queue);
+    observed_timestamp_bound_change =
+        input_stream_handler_->ProcessTimestampBounds() &&
+        prev_output_ts_ < min_timestamp.PreviousAllowedInStream();
     if (graph_has_error_ || !empty_queue ||
-        min_timestamp == Timestamp::Done()) {
+        min_timestamp == Timestamp::Done() || observed_timestamp_bound_change) {
       break;
     } else {
       handler_condvar_.Wait(&mutex_);
@@ -191,17 +196,26 @@ bool OutputStreamPollerImpl::Next(Packet* packet) {
     mutex_.Unlock();
     return false;
   }
+  if (empty_queue) {
+    prev_output_ts_ = min_timestamp.PreviousAllowedInStream();
+  } else {
+    prev_output_ts_ = min_timestamp;
+  }
   mutex_.Unlock();
   if (min_timestamp == Timestamp::Done()) {
     return false;
   }
-  int num_packets_dropped = 0;
-  bool stream_is_done = false;
-  *packet = input_stream_->PopPacketAtTimestamp(
-      min_timestamp, &num_packets_dropped, &stream_is_done);
-  CHECK_EQ(num_packets_dropped, 0)
-      << absl::Substitute("Dropped $0 packet(s) on input stream \"$1\".",
-                          num_packets_dropped, input_stream_->Name());
+  if (!empty_queue) {
+    int num_packets_dropped = 0;
+    bool stream_is_done = false;
+    *packet = input_stream_->PopPacketAtTimestamp(
+        min_timestamp, &num_packets_dropped, &stream_is_done);
+    CHECK_EQ(num_packets_dropped, 0)
+        << absl::Substitute("Dropped $0 packet(s) on input stream \"$1\".",
+                            num_packets_dropped, input_stream_->Name());
+  } else if (observed_timestamp_bound_change) {
+    *packet = Packet().At(Timestamp(min_timestamp.PreviousAllowedInStream()));
+  }
   return true;
 }
 
diff --git a/mediapipe/framework/graph_output_stream.h b/mediapipe/framework/graph_output_stream.h
index 393407a..1110a10 100644
--- a/mediapipe/framework/graph_output_stream.h
+++ b/mediapipe/framework/graph_output_stream.h
@@ -143,7 +143,8 @@ class OutputStreamPollerImpl : public GraphOutputStream {
   absl::Status Initialize(
       const std::string& stream_name, const PacketType* packet_type,
       std::function<void(InputStreamManager*, bool*)> queue_size_callback,
-      OutputStreamManager* output_stream_manager);
+      OutputStreamManager* output_stream_manager,
+      bool observe_timestamp_bounds = false);
 
   void PrepareForRun(std::function<void()> notification_callback,
                      std::function<void(absl::Status)> error_callback) override;
@@ -170,6 +171,7 @@ class OutputStreamPollerImpl : public GraphOutputStream {
   absl::Mutex mutex_;
   absl::CondVar handler_condvar_ ABSL_GUARDED_BY(mutex_);
   bool graph_has_error_ ABSL_GUARDED_BY(mutex_);
+  Timestamp prev_output_ts_ ABSL_GUARDED_BY(mutex_) = Timestamp::Min();
 };
 
 }  // namespace internal
diff --git a/mediapipe/gpu/gl_calculator_helper.cc b/mediapipe/gpu/gl_calculator_helper.cc
index aa708e7..c146f56 100644
--- a/mediapipe/gpu/gl_calculator_helper.cc
+++ b/mediapipe/gpu/gl_calculator_helper.cc
@@ -137,6 +137,10 @@ GlTexture GlCalculatorHelper::CreateDestinationTexture(int output_width,
   return impl_->CreateDestinationTexture(output_width, output_height, format);
 }
 
+GlTexture GlCalculatorHelper::CreateDestinationTexture(const GpuBuffer& pixel_buffer) {
+  return impl_->CreateDestinationTexture(pixel_buffer);
+}
+
 GlContext& GlCalculatorHelper::GetGlContext() const {
   return impl_->GetGlContext();
 }
diff --git a/mediapipe/gpu/gl_calculator_helper.h b/mediapipe/gpu/gl_calculator_helper.h
index b5cc699..e0aa081 100644
--- a/mediapipe/gpu/gl_calculator_helper.h
+++ b/mediapipe/gpu/gl_calculator_helper.h
@@ -144,6 +144,8 @@ class GlCalculatorHelper {
       int output_width, int output_height,
       GpuBufferFormat format = GpuBufferFormat::kBGRA32);
 
+  GlTexture CreateDestinationTexture(const GpuBuffer& pixel_buffer);
+
   // The OpenGL name of the output framebuffer.
   GLuint framebuffer() const;
 
diff --git a/mediapipe/gpu/gl_calculator_helper_impl.h b/mediapipe/gpu/gl_calculator_helper_impl.h
index 4381111..cdabc45 100644
--- a/mediapipe/gpu/gl_calculator_helper_impl.h
+++ b/mediapipe/gpu/gl_calculator_helper_impl.h
@@ -49,6 +49,7 @@ class GlCalculatorHelperImpl {
   // Creates a framebuffer and returns the texture that it is bound to.
   GlTexture CreateDestinationTexture(int output_width, int output_height,
                                      GpuBufferFormat format);
+  GlTexture CreateDestinationTexture(const GpuBuffer& gpu_buffer);
 
   GLuint framebuffer() const { return framebuffer_; }
   void BindFramebuffer(const GlTexture& dst);
diff --git a/mediapipe/gpu/gl_calculator_helper_impl_common.cc b/mediapipe/gpu/gl_calculator_helper_impl_common.cc
index fb2685b..af7ed96 100644
--- a/mediapipe/gpu/gl_calculator_helper_impl_common.cc
+++ b/mediapipe/gpu/gl_calculator_helper_impl_common.cc
@@ -196,14 +196,18 @@ GlTextureBufferSharedPtr GlCalculatorHelperImpl::MakeGlTextureBuffer(
 
 GlTexture GlCalculatorHelperImpl::CreateDestinationTexture(
     int width, int height, GpuBufferFormat format) {
+  GpuBuffer buffer =
+      gpu_resources_.gpu_buffer_pool().GetBuffer(width, height, format);
+
+  return CreateDestinationTexture(buffer);
+}
+
+GlTexture GlCalculatorHelperImpl::CreateDestinationTexture(const GpuBuffer& gpu_buffer) {
   if (!framebuffer_) {
     CreateFramebuffer();
   }
 
-  GpuBuffer buffer =
-      gpu_resources_.gpu_buffer_pool().GetBuffer(width, height, format);
-  GlTexture texture = MapGpuBuffer(buffer, 0);
-
+  GlTexture texture = MapGpuBuffer(gpu_buffer, 0);
   return texture;
 }
 
diff --git a/mediapipe/gpu/gl_scaler_calculator.cc b/mediapipe/gpu/gl_scaler_calculator.cc
index 6191876..19a8cf8 100644
--- a/mediapipe/gpu/gl_scaler_calculator.cc
+++ b/mediapipe/gpu/gl_scaler_calculator.cc
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// Modified to enable to specify the target GpuBuffer
+
 #include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/formats/image.h"
 #include "mediapipe/framework/port/ret_check.h"
@@ -61,6 +63,7 @@ using Image = mediapipe::Image;
 //   existing calculator options, depending on field merge_fields.
 //   OUTPUT_DIMENSIONS: the output width and height in pixels.
 //   ROTATION: the counterclockwise rotation angle in degrees.
+//   DESTINATION: the target GpuBuffer
 // These can also be specified as options.
 // To enable horizontal or vertical flip, specify them in options.
 // The flipping is applied after rotation.
@@ -85,6 +88,7 @@ class GlScalerCalculator : public CalculatorBase {
 
  private:
   GlCalculatorHelper helper_;
+  GpuBuffer dst_buffer_;
   int dst_width_ = 0;
   int dst_height_ = 0;
   float dst_scale_ = -1.f;
@@ -121,6 +125,9 @@ absl::Status GlScalerCalculator::GetContract(CalculatorContract* cc) {
   }
   MP_RETURN_IF_ERROR(GlCalculatorHelper::UpdateContract(cc));
 
+  if (cc->InputSidePackets().HasTag("DESTINATION")) {
+    cc->InputSidePackets().Tag("DESTINATION").Set<GpuBuffer>();
+  }
   if (cc->InputSidePackets().HasTag("OPTIONS")) {
     cc->InputSidePackets().Tag("OPTIONS").Set<GlScalerCalculatorOptions>();
   }
@@ -187,6 +194,11 @@ absl::Status GlScalerCalculator::Open(CalculatorContext* cc) {
     dst_width_ = dimensions[0];
     dst_height_ = dimensions[1];
   }
+  if (HasTagOrIndex(cc->InputSidePackets(), "DESTINATION", 1)) {
+    dst_buffer_ = cc->InputSidePackets().Tag("DESTINATION").Get<GpuBuffer>();
+    dst_width_ = dst_buffer_.width();
+    dst_height_ = dst_buffer_.height();
+  }
   if (cc->InputSidePackets().HasTag("ROTATION")) {
     rotation_ccw = cc->InputSidePackets().Tag("ROTATION").Get<int>();
   }
@@ -197,7 +209,7 @@ absl::Status GlScalerCalculator::Open(CalculatorContext* cc) {
 }
 
 absl::Status GlScalerCalculator::Process(CalculatorContext* cc) {
-  if (cc->Inputs().HasTag("OUTPUT_DIMENSIONS")) {
+  if (!dst_buffer_ && cc->Inputs().HasTag("OUTPUT_DIMENSIONS")) {
     if (cc->Inputs().Tag("OUTPUT_DIMENSIONS").IsEmpty()) {
       // OUTPUT_DIMENSIONS input stream is specified, but value is missing.
       return absl::OkStatus();
@@ -279,9 +291,18 @@ absl::Status GlScalerCalculator::Process(CalculatorContext* cc) {
               MakePacket<float>(left_right_padding).At(cc->InputTimestamp()));
     }
 
-    auto dst = helper_.CreateDestinationTexture(dst_width, dst_height,
-                                                GetOutputFormat());
-
+    GlTexture dst;
+#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+    // for iOS
+    dst = helper_.CreateDestinationTexture(dst_width, dst_height, GetOutputFormat());
+#else
+    if (dst_buffer_) {
+      dst_buffer_.GetGlTextureBufferSharedPtr()->Reuse();
+      dst = helper_.CreateDestinationTexture(dst_buffer_);
+    } else {
+      dst = helper_.CreateDestinationTexture(dst_width, dst_height, GetOutputFormat());
+    }
+#endif
     helper_.BindFramebuffer(dst);
     glActiveTexture(GL_TEXTURE1);
     glBindTexture(src1.target(), src1.name());
diff --git a/mediapipe/modules/objectron/objectron_gpu.pbtxt b/mediapipe/modules/objectron/objectron_gpu.pbtxt
index 16187de..a7703f5 100644
--- a/mediapipe/modules/objectron/objectron_gpu.pbtxt
+++ b/mediapipe/modules/objectron/objectron_gpu.pbtxt
@@ -11,6 +11,10 @@ input_side_packet: "MAX_NUM_OBJECTS:max_num_objects"
 
 # Collection of detected 3D objects, represented as a FrameAnnotation.
 output_stream: "FRAME_ANNOTATION:detected_objects"
+# Collection of box landmarks. (NormalizedLandmarkList)
+output_stream: "MULTI_LANDMARKS:multi_box_landmarks"
+# Crop rectangles derived from bounding box landmarks.
+output_stream: "NORM_RECTS:box_rects"
 
 # Defines whether landmarks from the previous video frame should be used to help
 # predict landmarks on the current video frame.