Project import generated by Copybara.

GitOrigin-RevId: 73d686c40057684f8bfaca285368bf1813f9fc26

mediapipe/framework/tool/switch_mux_calculator.cc

@@ -28,6 +28,7 @@
 #include "mediapipe/framework/port/status.h"
 #include "mediapipe/framework/port/status_macros.h"
 #include "mediapipe/framework/tool/container_util.h"
+#include "mediapipe/framework/tool/switch_container.pb.h"
 
 namespace mediapipe {
 
@@ -68,6 +69,17 @@ class SwitchMuxCalculator : public CalculatorBase {
  private:
   int channel_index_;
   std::set<std::string> channel_tags_;
+  mediapipe::SwitchContainerOptions options_;
+  // This is used to keep around packets that we've received but not
+  // relayed yet (because we may not know which channel we should yet be using
+  // when synchronized_io flag is set).
+  std::map<Timestamp, std::map<CollectionItemId, Packet>> packet_history_;
+  // Historical channel index values for timestamps where we don't have all
+  // packets available yet (when synchronized_io flag is set).
+  std::map<Timestamp, int> channel_history_;
+  // Number of output steams that we already processed for the current output
+  // timestamp.
+  int current_processed_stream_count_ = 0;
 };
 REGISTER_CALCULATOR(SwitchMuxCalculator);
 
@@ -122,6 +134,7 @@ absl::Status SwitchMuxCalculator::GetContract(CalculatorContract* cc) {
 }
 
 absl::Status SwitchMuxCalculator::Open(CalculatorContext* cc) {
+  options_ = cc->Options<mediapipe::SwitchContainerOptions>();
   channel_index_ = tool::GetChannelIndex(*cc, channel_index_);
   channel_tags_ = ChannelTags(cc->Inputs().TagMap());
 
@@ -141,13 +154,79 @@ absl::Status SwitchMuxCalculator::Process(CalculatorContext* cc) {
   // Update the input channel index if specified.
   channel_index_ = tool::GetChannelIndex(*cc, channel_index_);
 
-  // Relay packets and timestamps only from channel_index_.
-  for (const std::string& tag : channel_tags_) {
-    for (int index = 0; index < cc->Outputs().NumEntries(tag); ++index) {
-      auto& output = cc->Outputs().Get(tag, index);
-      std::string input_tag = tool::ChannelTag(tag, channel_index_);
-      auto& input = cc->Inputs().Get(input_tag, index);
-      tool::Relay(input, &output);
+  if (options_.synchronize_io()) {
+    // Start with adding input signals into channel_history_ and packet_history_
+    if (cc->Inputs().HasTag("ENABLE") &&
+        !cc->Inputs().Tag("ENABLE").IsEmpty()) {
+      channel_history_[cc->Inputs().Tag("ENABLE").Value().Timestamp()] =
+          channel_index_;
+    }
+    if (cc->Inputs().HasTag("SELECT") &&
+        !cc->Inputs().Tag("SELECT").IsEmpty()) {
+      channel_history_[cc->Inputs().Tag("SELECT").Value().Timestamp()] =
+          channel_index_;
+    }
+    for (auto input_id = cc->Inputs().BeginId();
+         input_id < cc->Inputs().EndId(); ++input_id) {
+      auto& entry = cc->Inputs().Get(input_id);
+      if (entry.IsEmpty()) {
+        continue;
+      }
+      packet_history_[entry.Value().Timestamp()][input_id] = entry.Value();
+    }
+    // Now check if we have enough information to produce any outputs.
+    while (!channel_history_.empty()) {
+      // Look at the oldest unprocessed timestamp.
+      auto it = channel_history_.begin();
+      auto& packets = packet_history_[it->first];
+      int total_streams = 0;
+      // Loop over all outputs to see if we have anything new that we can relay.
+      for (const std::string& tag : channel_tags_) {
+        for (int index = 0; index < cc->Outputs().NumEntries(tag); ++index) {
+          ++total_streams;
+          auto input_id =
+              cc->Inputs().GetId(tool::ChannelTag(tag, it->second), index);
+          auto packet_it = packets.find(input_id);
+          if (packet_it != packets.end()) {
+            cc->Outputs().Get(tag, index).AddPacket(packet_it->second);
+            ++current_processed_stream_count_;
+          } else if (it->first <
+                     cc->Inputs().Get(input_id).Value().Timestamp()) {
+            // Getting here means that input stream that corresponds to this
+            // output at the timestamp we're trying to process right now has
+            // already advanced beyond this timestamp. This means that we will
+            // shouldn't expect a packet for this timestamp anymore, and we can
+            // safely advance timestamp on the output.
+            cc->Outputs()
+                .Get(tag, index)
+                .SetNextTimestampBound(it->first.NextAllowedInStream());
+            ++current_processed_stream_count_;
+          }
+        }
+      }
+      if (current_processed_stream_count_ == total_streams) {
+        // There's nothing else to wait for at the current timestamp, do the
+        // cleanup and move on to the next one.
+        packet_history_.erase(it->first);
+        channel_history_.erase(it);
+        current_processed_stream_count_ = 0;
+      } else {
+        // We're still missing some packets for the current timestamp. Clean up
+        // those that we just relayed and let the rest wait until the next
+        // Process() call.
+        packets.clear();
+        break;
+      }
+    }
+  } else {
+    // Relay packets and timestamps only from channel_index_.
+    for (const std::string& tag : channel_tags_) {
+      for (int index = 0; index < cc->Outputs().NumEntries(tag); ++index) {
+        auto& output = cc->Outputs().Get(tag, index);
+        std::string input_tag = tool::ChannelTag(tag, channel_index_);
+        auto& input = cc->Inputs().Get(input_tag, index);
+        tool::Relay(input, &output);
+      }
     }
   }
   return absl::OkStatus();