From d6119957a497d5c381966e3764416950042a0db0 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Fri, 1 Sep 2023 16:29:24 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 562075110 --- .../pack_media_sequence_calculator.cc | 24 ++++ .../pack_media_sequence_calculator_test.cc | 121 ++++++++++++++++++ 2 files changed, 145 insertions(+) diff --git a/mediapipe/calculators/tensorflow/pack_media_sequence_calculator.cc b/mediapipe/calculators/tensorflow/pack_media_sequence_calculator.cc index c8bdfd9e..fcb640ff 100644 --- a/mediapipe/calculators/tensorflow/pack_media_sequence_calculator.cc +++ b/mediapipe/calculators/tensorflow/pack_media_sequence_calculator.cc @@ -42,6 +42,7 @@ const char kImageLabelPrefixTag[] = "IMAGE_LABEL_"; const char kClipLabelPrefixTag[] = "CLIP_LABEL_"; const char kFloatContextFeaturePrefixTag[] = "FLOAT_CONTEXT_FEATURE_"; const char kIntsContextFeaturePrefixTag[] = "INTS_CONTEXT_FEATURE_"; +const char kBytesContextFeaturePrefixTag[] = "BYTES_CONTEXT_FEATURE_"; const char kFloatFeaturePrefixTag[] = "FLOAT_FEATURE_"; const char kIntFeaturePrefixTag[] = "INT_FEATURE_"; const char kBytesFeaturePrefixTag[] = "BYTES_FEATURE_"; @@ -179,6 +180,9 @@ class PackMediaSequenceCalculator : public CalculatorBase { if (absl::StartsWith(tag, kIntsContextFeaturePrefixTag)) { cc->Inputs().Tag(tag).Set>(); } + if (absl::StartsWith(tag, kBytesContextFeaturePrefixTag)) { + cc->Inputs().Tag(tag).Set>(); + } if (absl::StartsWith(tag, kFloatFeaturePrefixTag)) { cc->Inputs().Tag(tag).Set>(); } @@ -290,6 +294,13 @@ class PackMediaSequenceCalculator : public CalculatorBase { 1); mpms::ClearContextFeatureInts(key, sequence_.get()); } + if (absl::StartsWith(tag, kBytesContextFeaturePrefixTag)) { + const std::string& key = + tag.substr(sizeof(kBytesContextFeaturePrefixTag) / + sizeof(*kBytesContextFeaturePrefixTag) - + 1); + mpms::ClearContextFeatureBytes(key, sequence_.get()); + } if (absl::StartsWith(tag, kFloatFeaturePrefixTag)) { std::string key = tag.substr(sizeof(kFloatFeaturePrefixTag) / sizeof(*kFloatFeaturePrefixTag) - @@ -542,6 +553,19 @@ class PackMediaSequenceCalculator : public CalculatorBase { mpms::AddContextFeatureInts(key, value, sequence_.get()); } } + if (absl::StartsWith(tag, kBytesContextFeaturePrefixTag) && + !cc->Inputs().Tag(tag).IsEmpty()) { + const std::string& key = + tag.substr(sizeof(kBytesContextFeaturePrefixTag) / + sizeof(*kBytesContextFeaturePrefixTag) - + 1); + // To ensure only one packet is provided for this tag. + RET_CHECK_EQ(cc->InputTimestamp(), Timestamp::PostStream()); + for (const auto& value : + cc->Inputs().Tag(tag).Get>()) { + mpms::AddContextFeatureBytes(key, value, sequence_.get()); + } + } if (absl::StartsWith(tag, kFloatFeaturePrefixTag) && !cc->Inputs().Tag(tag).IsEmpty()) { std::string key = tag.substr(sizeof(kFloatFeaturePrefixTag) / diff --git a/mediapipe/calculators/tensorflow/pack_media_sequence_calculator_test.cc b/mediapipe/calculators/tensorflow/pack_media_sequence_calculator_test.cc index 3644505d..277e2d63 100644 --- a/mediapipe/calculators/tensorflow/pack_media_sequence_calculator_test.cc +++ b/mediapipe/calculators/tensorflow/pack_media_sequence_calculator_test.cc @@ -60,6 +60,8 @@ constexpr char kFloatContextFeatureOtherTag[] = "FLOAT_CONTEXT_FEATURE_OTHER"; constexpr char kFloatContextFeatureTestTag[] = "FLOAT_CONTEXT_FEATURE_TEST"; constexpr char kIntsContextFeatureTestTag[] = "INTS_CONTEXT_FEATURE_TEST"; constexpr char kIntsContextFeatureOtherTag[] = "INTS_CONTEXT_FEATURE_OTHER"; +constexpr char kBytesContextFeatureTestTag[] = "BYTES_CONTEXT_FEATURE_TEST"; +constexpr char kBytesContextFeatureOtherTag[] = "BYTES_CONTEXT_FEATURE_OTHER"; constexpr char kFloatFeatureOtherTag[] = "FLOAT_FEATURE_OTHER"; constexpr char kFloatFeatureTestTag[] = "FLOAT_FEATURE_TEST"; constexpr char kIntFeatureOtherTag[] = "INT_FEATURE_OTHER"; @@ -566,6 +568,125 @@ TEST_F(PackMediaSequenceCalculatorTest, AppendTwoContextIntLists) { testing::ElementsAre(2, 4, 7, 8)); } +TEST_F(PackMediaSequenceCalculatorTest, PackTwoContextByteLists) { + SetUpCalculator( + /*input_streams=*/{"BYTES_CONTEXT_FEATURE_TEST:test", + "BYTES_CONTEXT_FEATURE_OTHER:test2"}, + /*features=*/{}, + /*output_only_if_all_present=*/false, /*replace_instead_of_append=*/true); + auto input_sequence = absl::make_unique(); + + const std::vector vb_1 = {"value_1", "value_2"}; + runner_->MutableInputs() + ->Tag(kBytesContextFeatureTestTag) + .packets.push_back(MakePacket>(vb_1).At( + Timestamp::PostStream())); + const std::vector vb_2 = {"value_3", "value_4"}; + runner_->MutableInputs() + ->Tag(kBytesContextFeatureOtherTag) + .packets.push_back(MakePacket>(vb_2).At( + Timestamp::PostStream())); + + runner_->MutableSidePackets()->Tag(kSequenceExampleTag) = + Adopt(input_sequence.release()); + + MP_ASSERT_OK(runner_->Run()); + + const std::vector& output_packets = + runner_->Outputs().Tag(kSequenceExampleTag).packets; + ASSERT_EQ(1, output_packets.size()); + const tf::SequenceExample& output_sequence = + output_packets[0].Get(); + + ASSERT_THAT(mpms::GetContextFeatureBytes("TEST", output_sequence), + testing::ElementsAre("value_1", "value_2")); + ASSERT_THAT(mpms::GetContextFeatureBytes("OTHER", output_sequence), + testing::ElementsAre("value_3", "value_4")); +} + +TEST_F(PackMediaSequenceCalculatorTest, ReplaceTwoContextByteLists) { + SetUpCalculator( + /*input_streams=*/{"BYTES_CONTEXT_FEATURE_TEST:test", + "BYTES_CONTEXT_FEATURE_OTHER:test2"}, + /*features=*/{}, + /*output_only_if_all_present=*/false, /*replace_instead_of_append=*/true); + auto input_sequence = absl::make_unique(); + mpms::SetContextFeatureBytes("TEST", {"existing_value_1", "existing_value_2"}, + input_sequence.get()); + mpms::SetContextFeatureBytes( + "OTHER", {"existing_value_3", "existing_value_4"}, input_sequence.get()); + + const std::vector vb_1 = {"value_1", "value_2"}; + runner_->MutableInputs() + ->Tag(kBytesContextFeatureTestTag) + .packets.push_back(MakePacket>(vb_1).At( + Timestamp::PostStream())); + const std::vector vb_2 = {"value_3", "value_4"}; + runner_->MutableInputs() + ->Tag(kBytesContextFeatureOtherTag) + .packets.push_back(MakePacket>(vb_2).At( + Timestamp::PostStream())); + + runner_->MutableSidePackets()->Tag(kSequenceExampleTag) = + Adopt(input_sequence.release()); + + MP_ASSERT_OK(runner_->Run()); + + const std::vector& output_packets = + runner_->Outputs().Tag(kSequenceExampleTag).packets; + ASSERT_EQ(1, output_packets.size()); + const tf::SequenceExample& output_sequence = + output_packets[0].Get(); + + ASSERT_THAT(mpms::GetContextFeatureBytes("TEST", output_sequence), + testing::ElementsAre("value_1", "value_2")); + ASSERT_THAT(mpms::GetContextFeatureBytes("OTHER", output_sequence), + testing::ElementsAre("value_3", "value_4")); +} + +TEST_F(PackMediaSequenceCalculatorTest, AppendTwoContextByteLists) { + SetUpCalculator( + /*input_streams=*/{"BYTES_CONTEXT_FEATURE_TEST:test", + "BYTES_CONTEXT_FEATURE_OTHER:test2"}, + /*features=*/{}, + /*output_only_if_all_present=*/false, + /*replace_instead_of_append=*/false); + auto input_sequence = absl::make_unique(); + mpms::SetContextFeatureBytes("TEST", {"existing_value_1", "existing_value_2"}, + input_sequence.get()); + mpms::SetContextFeatureBytes( + "OTHER", {"existing_value_3", "existing_value_4"}, input_sequence.get()); + + const std::vector vb_1 = {"value_1", "value_2"}; + runner_->MutableInputs() + ->Tag(kBytesContextFeatureTestTag) + .packets.push_back(MakePacket>(vb_1).At( + Timestamp::PostStream())); + const std::vector vb_2 = {"value_3", "value_4"}; + runner_->MutableInputs() + ->Tag(kBytesContextFeatureOtherTag) + .packets.push_back(MakePacket>(vb_2).At( + Timestamp::PostStream())); + + runner_->MutableSidePackets()->Tag(kSequenceExampleTag) = + Adopt(input_sequence.release()); + + MP_ASSERT_OK(runner_->Run()); + + const std::vector& output_packets = + runner_->Outputs().Tag(kSequenceExampleTag).packets; + ASSERT_EQ(1, output_packets.size()); + const tf::SequenceExample& output_sequence = + output_packets[0].Get(); + + ASSERT_THAT(mpms::GetContextFeatureBytes("TEST", output_sequence), + testing::ElementsAre("existing_value_1", "existing_value_2", + "value_1", "value_2")); + ASSERT_THAT(mpms::GetContextFeatureBytes("OTHER", output_sequence), + testing::ElementsAre("existing_value_3", "existing_value_4", + "value_3", "value_4")); +} + TEST_F(PackMediaSequenceCalculatorTest, PacksAdditionalContext) { tf::Features context; (*context.mutable_feature())["TEST"].mutable_bytes_list()->add_value("YES");