diff --git a/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator.cc b/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator.cc index bae926b7..6e06c4e3 100644 --- a/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator.cc +++ b/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator.cc @@ -120,7 +120,9 @@ absl::Status EmbeddingAggregationCalculator::Process(CalculatorContext* cc) { } kTimestampedEmbeddingsOut(cc).Send(std::move(results)); } else { - kEmbeddingsOut(cc).Send(kEmbeddingsIn(cc)); + auto result = kEmbeddingsIn(cc).Get(); + result.set_timestamp_ms(cc->InputTimestamp().Value() / 1000); + kEmbeddingsOut(cc).Send(result); } RET_CHECK(cached_embeddings_.empty()); return absl::OkStatus(); diff --git a/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator_test.cc b/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator_test.cc index ebb4d888..f2b2fa1d 100644 --- a/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator_test.cc +++ b/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator_test.cc @@ -120,7 +120,7 @@ class EmbeddingAggregationCalculatorTest : public tflite_shims::testing::Test { CalculatorGraph calculator_graph_; }; -TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithoutTimestamps) { +TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithoutAggregation) { EmbeddingResult embedding = ParseTextProtoOrDie( R"pb(embeddings { head_index: 0 })pb"); @@ -129,10 +129,12 @@ TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithoutTimestamps) { MP_ASSERT_OK(Send(embedding)); MP_ASSERT_OK_AND_ASSIGN(auto result, GetResult(poller)); - EXPECT_THAT(result, EqualsProto(embedding)); + EXPECT_THAT(result, EqualsProto(ParseTextProtoOrDie( + R"pb(timestamp_ms: 0 + embeddings { head_index: 0 })pb"))); } -TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithTimestamps) { +TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithAggregation) { MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph(/*connect_timestamps=*/true)); MP_ASSERT_OK(Send(ParseTextProtoOrDie(R"pb(embeddings { head_index: 0 diff --git a/mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph_test.cc b/mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph_test.cc index 163e46ee..809268a6 100644 --- a/mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph_test.cc +++ b/mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph_test.cc @@ -246,7 +246,7 @@ class PostprocessingTest : public tflite_shims::testing::Test { absl::make_unique>(); }; -TEST_F(PostprocessingTest, SucceedsWithoutTimestamps) { +TEST_F(PostprocessingTest, SucceedsWithoutAggregation) { // Build graph. proto::EmbedderOptions options; MP_ASSERT_OK_AND_ASSIGN(auto poller, @@ -261,7 +261,8 @@ TEST_F(PostprocessingTest, SucceedsWithoutTimestamps) { MP_ASSERT_OK_AND_ASSIGN(auto results, GetResult(poller)); // Validate results. - EXPECT_FALSE(results.has_timestamp_ms()); + EXPECT_TRUE(results.has_timestamp_ms()); + EXPECT_EQ(results.timestamp_ms(), 0); EXPECT_EQ(results.embeddings_size(), 1); EXPECT_EQ(results.embeddings(0).head_index(), 0); EXPECT_EQ(results.embeddings(0).head_name(), "feature"); @@ -273,7 +274,7 @@ TEST_F(PostprocessingTest, SucceedsWithoutTimestamps) { } } -TEST_F(PostprocessingTest, SucceedsWithTimestamps) { +TEST_F(PostprocessingTest, SucceedsWithAggregation) { // Build graph. proto::EmbedderOptions options; MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph(kMobileNetV3Embedder, options,