Project import generated by Copybara.

PiperOrigin-RevId: 253489161
2019-06-16 16:06:57 -07:00
commit d68f5e4169
844 changed files with 134997 additions and 0 deletions
@@ -0,0 +1,392 @@
+// Copyright 2019 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/framework/input_stream_manager.h"
+
+#include <type_traits>
+#include <utility>
+
+#include "absl/strings/str_cat.h"
+#include "absl/synchronization/mutex.h"
+#include "mediapipe/framework/packet.h"
+#include "mediapipe/framework/port/logging.h"
+#include "mediapipe/framework/port/source_location.h"
+#include "mediapipe/framework/port/status_builder.h"
+#include "mediapipe/framework/tool/status_util.h"
+
+namespace mediapipe {
+
+::mediapipe::Status InputStreamManager::Initialize(
+    const std::string& name, const PacketType* packet_type, bool back_edge) {
+  name_ = name;
+  packet_type_ = packet_type;
+  back_edge_ = back_edge;
+  PrepareForRun();
+  return ::mediapipe::OkStatus();
+}
+
+const std::string& InputStreamManager::Name() const { return name_; }
+
+void InputStreamManager::SetQueueSizeCallbacks(
+    QueueSizeCallback becomes_full_callback,
+    QueueSizeCallback becomes_not_full_callback) {
+  becomes_full_callback_ = becomes_full_callback;
+  becomes_not_full_callback_ = becomes_not_full_callback;
+}
+
+void InputStreamManager::PrepareForRun() {
+  absl::MutexLock stream_lock(&stream_mutex_);
+  queue_.clear();
+  last_reported_stream_full_ = false;
+  num_packets_added_ = 0;
+  next_timestamp_bound_ = Timestamp::PreStream();
+  last_select_timestamp_ = Timestamp::Unstarted();
+  closed_ = false;
+  header_ = Packet();
+}
+
+bool InputStreamManager::IsEmpty() const {
+  absl::MutexLock stream_lock(&stream_mutex_);
+  return queue_.empty();
+}
+
+Packet InputStreamManager::QueueHead() const {
+  absl::MutexLock stream_lock(&stream_mutex_);
+  if (queue_.empty()) {
+    return Packet();
+  }
+  return queue_.front();
+}
+
+::mediapipe::Status InputStreamManager::SetHeader(const Packet& header) {
+  if (header.Timestamp() != Timestamp::Unset()) {
+    return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
+           << "Headers must not have a timestamp.  Stream: \"" << name_
+           << "\".";
+  }
+  header_ = header;
+  return ::mediapipe::OkStatus();
+}
+
+::mediapipe::Status InputStreamManager::AddPackets(
+    const std::list<Packet>& container, bool* notify) {
+  return AddOrMovePacketsInternal<const std::list<Packet>&>(container, notify);
+}
+
+::mediapipe::Status InputStreamManager::MovePackets(
+    std::list<Packet>* container, bool* notify) {
+  return AddOrMovePacketsInternal<std::list<Packet>&>(*container, notify);
+}
+
+template <typename Container>
+::mediapipe::Status InputStreamManager::AddOrMovePacketsInternal(
+    Container container, bool* notify) {
+  *notify = false;
+  bool queue_became_non_empty = false;
+  bool queue_became_full = false;
+  {
+    // Scope to prevent locking the stream when notification is called.
+    absl::MutexLock stream_lock(&stream_mutex_);
+    if (closed_) {
+      return ::mediapipe::OkStatus();
+    }
+    // Check if the queue was full before packets came in.
+    bool was_queue_full =
+        (max_queue_size_ != -1 && queue_.size() >= max_queue_size_);
+    // Check if the queue becomes non-empty.
+    queue_became_non_empty = queue_.empty() && !container.empty();
+    for (auto& packet : container) {
+      ::mediapipe::Status result = packet_type_->Validate(packet);
+      if (!result.ok()) {
+        return tool::AddStatusPrefix(
+            absl::StrCat(
+                "Packet type mismatch on a calculator receiving from stream \"",
+                name_, "\": "),
+            result);
+      }
+
+      const Timestamp timestamp = packet.Timestamp();
+      if (!timestamp.IsAllowedInStream()) {
+        return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
+               << "In stream \"" << name_
+               << "\", timestamp not specified or set to illegal value: "
+               << timestamp.DebugString();
+      }
+      if (enable_timestamps_) {
+        // Check that PostStream(), if used, is the only timestamp used.  This
+        // is also true for PreStream() but doesn't need to be checked because
+        // Timestamp::PreStream().NextAllowedInStream() is
+        // Timestamp::OneOverPostStream().
+        if (timestamp == Timestamp::PostStream() && num_packets_added_ > 0) {
+          return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
+                 << "In stream \"" << name_
+                 << "\", a packet at Timestamp::PostStream() must be the only "
+                    "Packet in an InputStream.";
+        }
+        if (timestamp < next_timestamp_bound_) {
+          return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
+                 << "Packet timestamp mismatch on a calculator receiving from "
+                    "stream \""
+                 << name_ << "\". Current minimum expected timestamp is "
+                 << next_timestamp_bound_.DebugString() << " but received "
+                 << timestamp.DebugString()
+                 << ". Are you using a custom InputStreamHandler? Note that "
+                    "some InputStreamHandlers allow timestamps that are not "
+                    "strictly monotonically increasing. See for example the "
+                    "ImmediateInputStreamHandler class comment.";
+        }
+      }
+      next_timestamp_bound_ = timestamp.NextAllowedInStream();
+
+      // If the caller is MovePackets(), packet's underlying holder should be
+      // transferred into queue_. Otherwise, queue_ keeps a copy of the packet.
+      if (std::is_const<
+              typename std::remove_reference<Container>::type>::value) {
+        queue_.emplace_back(packet);
+      } else {
+        queue_.emplace_back(std::move(packet));
+      }
+      ++num_packets_added_;
+      VLOG(2) << "Input stream:" << name_
+              << " has added packet at time: " << packet.Timestamp();
+    }
+    queue_became_full = (!was_queue_full && max_queue_size_ != -1 &&
+                         queue_.size() >= max_queue_size_);
+    VLOG_IF(2, queue_.size() > 1)
+        << "Queue size greater than 1: stream name: " << name_
+        << " queue_size: " << queue_.size();
+    VLOG(2) << "Input stream:" << name_
+            << " becomes non-empty status:" << queue_became_non_empty
+            << " Size: " << queue_.size();
+  }
+  if (queue_became_full) {
+    VLOG(2) << "Queue became full: " << Name();
+    becomes_full_callback_(this, &last_reported_stream_full_);
+  }
+  *notify = queue_became_non_empty;
+  return ::mediapipe::OkStatus();
+}
+
+::mediapipe::Status InputStreamManager::SetNextTimestampBound(
+    const Timestamp bound, bool* notify) {
+  *notify = false;
+  {
+    // Scope to prevent locking the stream when notification is called.
+    absl::MutexLock stream_lock(&stream_mutex_);
+    if (closed_) {
+      return ::mediapipe::OkStatus();
+    }
+
+    if (enable_timestamps_ && bound < next_timestamp_bound_) {
+      return ::mediapipe::UnknownErrorBuilder(MEDIAPIPE_LOC)
+             << "SetNextTimestampBound must be called with a timestamp greater "
+                "than or equal to the current bound. In stream \""
+             << name_ << "\". Current minimum expected timestamp is "
+             << next_timestamp_bound_.DebugString() << " but received "
+             << bound.DebugString();
+    }
+
+    // Even if enable_timestamps_ is false, Timestamp::Done() is used to
+    // indicate the end of stream. So this code is common to both timed and
+    // untimed scheduling policies.
+    if (bound > next_timestamp_bound_) {
+      next_timestamp_bound_ = bound;
+      if (queue_.empty()) {
+        // If the queue was not empty then a change to the next_timestamp_bound_
+        // is not detectable by the consumer.
+        *notify = true;
+      }
+    }
+  }
+  return ::mediapipe::OkStatus();
+}
+
+void InputStreamManager::DisableTimestamps() { enable_timestamps_ = false; }
+
+void InputStreamManager::Close() {
+  absl::MutexLock stream_lock(&stream_mutex_);
+  if (closed_) {
+    return;
+  }
+  next_timestamp_bound_ = Timestamp::Done();
+  last_select_timestamp_ = Timestamp::Done();
+  closed_ = true;
+}
+
+Timestamp InputStreamManager::MinTimestampOrBound(bool* is_empty) const {
+  absl::MutexLock stream_lock(&stream_mutex_);
+  if (is_empty) {
+    *is_empty = queue_.empty();
+  }
+  return queue_.empty() ? next_timestamp_bound_ : queue_.front().Timestamp();
+}
+
+Packet InputStreamManager::PopPacketAtTimestamp(Timestamp timestamp,
+                                                int* num_packets_dropped,
+                                                bool* stream_is_done) {
+  CHECK(enable_timestamps_);
+  *num_packets_dropped = -1;
+  *stream_is_done = false;
+  bool queue_became_non_full = false;
+  Packet packet;
+  {
+    absl::MutexLock stream_lock(&stream_mutex_);
+    // Make sure timestamp didn't decrease from last time.
+    CHECK_LE(last_select_timestamp_, timestamp);
+    last_select_timestamp_ = timestamp;
+
+    // Make sure AddPacket and SetNextTimestampBound are not called with
+    // timestamps we have already passed.
+    if (next_timestamp_bound_ <= timestamp) {
+      next_timestamp_bound_ = timestamp.NextAllowedInStream();
+    }
+
+    VLOG(2) << "Input stream " << name_
+            << " selecting at timestamp:" << timestamp.Value()
+            << " next timestamp bound: " << next_timestamp_bound_;
+
+    // Advances time to timestamp.
+    Timestamp current_timestamp = Timestamp::Unset();
+
+    // Checks if queue is full.
+    bool was_queue_full =
+        (max_queue_size_ != -1 && queue_.size() >= max_queue_size_);
+
+    while (!queue_.empty() && queue_.front().Timestamp() <= timestamp) {
+      packet = std::move(queue_.front());
+      queue_.pop_front();
+      current_timestamp = packet.Timestamp();
+      ++(*num_packets_dropped);
+    }
+    // Clear value_ if it doesn't have exactly the right timestamp.
+    if (current_timestamp != timestamp) {
+      packet = Packet();
+      ++(*num_packets_dropped);
+    }
+
+    VLOG(2) << "Input stream removed packets:" << name_
+            << " Size:" << queue_.size();
+    queue_became_non_full = (was_queue_full && queue_.size() < max_queue_size_);
+    *stream_is_done = IsDone();
+  }
+  if (queue_became_non_full) {
+    VLOG(2) << "Queue became non-full: " << Name();
+    becomes_not_full_callback_(this, &last_reported_stream_full_);
+  }
+  return packet;
+}
+
+Packet InputStreamManager::PopQueueHead(bool* stream_is_done) {
+  CHECK(!enable_timestamps_);
+  *stream_is_done = false;
+  bool queue_became_non_full = false;
+  Packet packet;
+  {
+    absl::MutexLock stream_lock(&stream_mutex_);
+
+    VLOG(2) << "Input stream " << name_ << " selecting at queue head";
+
+    // Check if queue is full.
+    bool was_queue_full =
+        (max_queue_size_ != -1 && queue_.size() >= max_queue_size_);
+
+    if (!queue_.empty()) {
+      packet = std::move(queue_.front());
+      queue_.pop_front();
+    } else {
+      packet = Packet();
+    }
+
+    VLOG(2) << "Input stream removed a packet:" << name_
+            << " Size:" << queue_.size();
+    queue_became_non_full = (was_queue_full && queue_.size() < max_queue_size_);
+    *stream_is_done = IsDone();
+  }
+  if (queue_became_non_full) {
+    VLOG(2) << "Queue became non-full: " << Name();
+    becomes_not_full_callback_(this, &last_reported_stream_full_);
+  }
+  return packet;
+}
+
+int InputStreamManager::QueueSize() const {
+  absl::MutexLock lock(&stream_mutex_);
+  return static_cast<int>(queue_.size());
+}
+
+int InputStreamManager::MaxQueueSize() const {
+  absl::MutexLock lock(&stream_mutex_);
+  return max_queue_size_;
+}
+
+void InputStreamManager::SetMaxQueueSize(int max_queue_size) {
+  bool was_full;
+  bool is_full;
+  {
+    absl::MutexLock lock(&stream_mutex_);
+    was_full = (max_queue_size_ != -1 && queue_.size() >= max_queue_size_);
+    max_queue_size_ = max_queue_size;
+    is_full = (max_queue_size_ != -1 && queue_.size() >= max_queue_size_);
+  }
+
+  // QueueSizeCallback is called with no mutexes held.
+  if (!was_full && is_full) {
+    VLOG(2) << "Queue became full: " << Name();
+    becomes_full_callback_(this, &last_reported_stream_full_);
+  } else if (was_full && !is_full) {
+    VLOG(2) << "Queue became non-full: " << Name();
+    becomes_not_full_callback_(this, &last_reported_stream_full_);
+  }
+}
+
+bool InputStreamManager::IsFull() const {
+  absl::MutexLock lock(&stream_mutex_);
+  return max_queue_size_ != -1 && queue_.size() >= max_queue_size_;
+}
+
+Timestamp InputStreamManager::GetMinTimestampAmongNLatest(int n) const {
+  absl::MutexLock lock(&stream_mutex_);
+  if (queue_.empty()) {
+    return Timestamp::Unset();
+  }
+  return (queue_.cend() - std::min((size_t)n, queue_.size()))->Timestamp();
+}
+
+void InputStreamManager::ErasePacketsEarlierThan(Timestamp timestamp) {
+  bool queue_became_non_full = false;
+  {
+    absl::MutexLock lock(&stream_mutex_);
+    // Checks if queue is full.
+    bool was_queue_full =
+        (max_queue_size_ != -1 && queue_.size() >= max_queue_size_);
+
+    while (!queue_.empty() && queue_.front().Timestamp() < timestamp) {
+      queue_.pop_front();
+    }
+
+    VLOG(2) << "Input stream removed packets:" << name_
+            << " Size:" << queue_.size();
+    queue_became_non_full = (was_queue_full && queue_.size() < max_queue_size_);
+  }
+  if (queue_became_non_full) {
+    VLOG(2) << "Queue became non-full: " << Name();
+    becomes_not_full_callback_(this, &last_reported_stream_full_);
+  }
+}
+
+bool InputStreamManager::IsDone() const {
+  return queue_.empty() && next_timestamp_bound_ == Timestamp::Done();
+}
+
+}  // namespace mediapipe