diff --git a/mediapipe/framework/BUILD b/mediapipe/framework/BUILD index e5e72cfbe..38812b39e 100644 --- a/mediapipe/framework/BUILD +++ b/mediapipe/framework/BUILD @@ -11,7 +11,6 @@ # 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. -# load("//mediapipe/framework/port:build_config.bzl", "mediapipe_proto_library") load("@bazel_skylib//:bzl_library.bzl", "bzl_library") @@ -368,6 +367,7 @@ cc_library( "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log:absl_check", "@com_google_absl//absl/log:absl_log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/memory", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", diff --git a/mediapipe/framework/calculator_graph.cc b/mediapipe/framework/calculator_graph.cc index 1890d799c..1bd356eac 100644 --- a/mediapipe/framework/calculator_graph.cc +++ b/mediapipe/framework/calculator_graph.cc @@ -28,6 +28,7 @@ #include "absl/container/flat_hash_set.h" #include "absl/log/absl_check.h" #include "absl/log/absl_log.h" +#include "absl/log/check.h" #include "absl/memory/memory.h" #include "absl/status/status.h" #include "absl/strings/str_cat.h" @@ -890,12 +891,12 @@ absl::Status CalculatorGraph::WaitForObservedOutput() { } absl::Status CalculatorGraph::AddPacketToInputStream( - const std::string& stream_name, const Packet& packet) { + absl::string_view stream_name, const Packet& packet) { return AddPacketToInputStreamInternal(stream_name, packet); } absl::Status CalculatorGraph::AddPacketToInputStream( - const std::string& stream_name, Packet&& packet) { + absl::string_view stream_name, Packet&& packet) { return AddPacketToInputStreamInternal(stream_name, std::move(packet)); } @@ -918,14 +919,18 @@ absl::Status CalculatorGraph::SetInputStreamTimestampBound( // std::forward will deduce the correct type as we pass along packet. template absl::Status CalculatorGraph::AddPacketToInputStreamInternal( - const std::string& stream_name, T&& packet) { + absl::string_view stream_name, T&& packet) { + auto stream_it = graph_input_streams_.find(stream_name); std::unique_ptr* stream = - mediapipe::FindOrNull(graph_input_streams_, stream_name); + stream_it == graph_input_streams_.end() ? nullptr : &stream_it->second; RET_CHECK(stream).SetNoLogging() << absl::Substitute( "AddPacketToInputStream called on input stream \"$0\" which is not a " "graph input stream.", stream_name); - int node_id = mediapipe::FindOrDie(graph_input_stream_node_ids_, stream_name); + auto node_id_it = graph_input_stream_node_ids_.find(stream_name); + ABSL_CHECK(node_id_it != graph_input_stream_node_ids_.end()) + << "Map key not found: " << stream_name; + int node_id = node_id_it->second; ABSL_CHECK_GE(node_id, validated_graph_->CalculatorInfos().size()); { absl::MutexLock lock(&full_input_streams_mutex_); diff --git a/mediapipe/framework/calculator_graph.h b/mediapipe/framework/calculator_graph.h index 4284beb7c..80af72650 100644 --- a/mediapipe/framework/calculator_graph.h +++ b/mediapipe/framework/calculator_graph.h @@ -32,6 +32,7 @@ #include "absl/container/flat_hash_set.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" #include "mediapipe/framework/calculator.pb.h" #include "mediapipe/framework/calculator_base.h" @@ -255,7 +256,7 @@ class CalculatorGraph { // sizes of the queues in the graph. The input stream must have been specified // in the configuration as a graph level input_stream. On error, nothing is // added. - absl::Status AddPacketToInputStream(const std::string& stream_name, + absl::Status AddPacketToInputStream(absl::string_view stream_name, const Packet& packet); // Same as the l-value version of this function by the same name, but moves @@ -265,7 +266,7 @@ class CalculatorGraph { // packet may remain valid. In particular, when using the ADD_IF_NOT_FULL // mode with a full queue, this will return StatusUnavailable and the caller // may try adding the packet again later. - absl::Status AddPacketToInputStream(const std::string& stream_name, + absl::Status AddPacketToInputStream(absl::string_view stream_name, Packet&& packet); // Indicates that input will arrive no earlier than a certain timestamp. @@ -509,7 +510,7 @@ class CalculatorGraph { // AddPacketToInputStream(Packet&& packet) or // AddPacketToInputStream(const Packet& packet). template - absl::Status AddPacketToInputStreamInternal(const std::string& stream_name, + absl::Status AddPacketToInputStreamInternal(absl::string_view stream_name, T&& packet); // Sets the executor that will run the nodes assigned to the executor